polars-df 0.5.0-x86_64-linux → 0.6.0-x86_64-linux

data/lib/polars/data_types.rb

@@ -1,6 +1,25 @@
 module Polars
   # Base class for all Polars data types.
   class DataType
+    def self.base_type
+      self
+    end
+
+    def base_type
+      is_a?(DataType) ? self.class : self
+    end
+
+    def self.nested?
+      false
+    end
+
+    def nested?
+      self.class.nested?
+    end
+
+    def self.==(other)
+      eql?(other) || other.is_a?(self)
+    end
   end
 
   # Base class for numeric data types.
@@ -15,12 +34,19 @@ module Polars
   class FractionalType < NumericType
   end
 
+  # Base class for float data types.
+  class FloatType < FractionalType
+  end
+
   # Base class for temporal data types.
   class TemporalType < DataType
   end
 
   # Base class for nested data types.
   class NestedType < DataType
+    def self.nested?
+      true
+    end
   end
 
   # 8-bit signed integer type.
@@ -56,11 +82,37 @@ module Polars
   end
 
   # 32-bit floating point type.
-  class Float32 < FractionalType
+  class Float32 < FloatType
   end
 
   # 64-bit floating point type.
-  class Float64 < FractionalType
+  class Float64 < FloatType
+  end
+
+  # Decimal 128-bit type with an optional precision and non-negative scale.
+  #
+  # NOTE: this is an experimental work-in-progress feature and may not work as expected.
+  class Decimal < FractionalType
+    attr_reader :precision, :scale
+
+    def initialize(precision, scale)
+      @precision = precision
+      @scale = scale
+    end
+
+    def ==(other)
+      if other.eql?(Decimal)
+        true
+      elsif other.is_a?(Decimal)
+        precision == other.precision && scale == other.scale
+      else
+        false
+      end
+    end
+
+    def to_s
+      "#{self.class.name}(precision: #{precision.inspect}, scale: #{scale.inspect})"
+    end
   end
 
   # Boolean type.
@@ -71,17 +123,18 @@ module Polars
   class Utf8 < DataType
   end
 
-  # Nested list/array type.
-  class List < NestedType
-    def initialize(inner)
-      @inner = Utils.rb_type_to_dtype(inner)
-    end
+  # Binary type.
+  class Binary < DataType
   end
 
   # Calendar date type.
   class Date < TemporalType
   end
 
+  # Time of day type.
+  class Time < TemporalType
+  end
+
   # Calendar date and time type.
   class Datetime < TemporalType
     attr_reader :time_unit, :time_zone
@@ -91,6 +144,20 @@ module Polars
       @time_unit = time_unit || "us"
       @time_zone = time_zone
     end
+
+    def ==(other)
+      if other.eql?(Datetime)
+        true
+      elsif other.is_a?(Datetime)
+        time_unit == other.time_unit && time_zone == other.time_zone
+      else
+        false
+      end
+    end
+
+    def to_s
+      "#{self.class.name}(time_unit: #{time_unit.inspect}, time_zone: #{time_zone.inspect})"
+    end
   end
 
   # Time duration/delta type.
@@ -101,18 +168,85 @@ module Polars
     def initialize(time_unit = "us")
       @time_unit = time_unit
     end
+
+    def ==(other)
+      if other.eql?(Duration)
+        true
+      elsif other.is_a?(Duration)
+        time_unit == other.time_unit
+      else
+        false
+      end
+    end
+
+    def to_s
+      "#{self.class.name}(time_unit: #{time_unit.inspect})"
+    end
   end
 
-  # Time of day type.
-  class Time < TemporalType
+  # A categorical encoding of a set of strings.
+  class Categorical < DataType
   end
 
   # Type for wrapping arbitrary Ruby objects.
   class Object < DataType
   end
 
-  # A categorical encoding of a set of strings.
-  class Categorical < DataType
+  # Type representing Null / None values.
+  class Null < DataType
+  end
+
+  # Type representing Datatype values that could not be determined statically.
+  class Unknown < DataType
+  end
+
+  # Nested list/array type.
+  class List < NestedType
+    attr_reader :inner
+
+    def initialize(inner)
+      @inner = Utils.rb_type_to_dtype(inner)
+    end
+
+    def ==(other)
+      if other.eql?(List)
+        true
+      elsif other.is_a?(List)
+        @inner.nil? || other.inner.nil? || @inner == other.inner
+      else
+        false
+      end
+    end
+
+    def to_s
+      "#{self.class.name}(#{inner})"
+    end
+  end
+
+  # Nested list/array type.
+  class Array < NestedType
+    attr_reader :width, :inner
+
+    def initialize(width, inner = nil)
+      @width = width
+      @inner = Utils.rb_type_to_dtype(inner) if inner
+    end
+
+    # TODO check width?
+    def ==(other)
+      if other.eql?(Array)
+        true
+      elsif other.is_a?(Array)
+        @inner.nil? || other.inner.nil? || @inner == other.inner
+      else
+        false
+      end
+    end
+
+    # TODO add width?
+    def to_s
+      "#{self.class.name}(#{inner})"
+    end
   end
 
   # Definition of a single field within a `Struct` DataType.
@@ -124,9 +258,12 @@ module Polars
       @dtype = Utils.rb_type_to_dtype(dtype)
     end
 
-    def inspect
-      class_name = self.class.name
-      "#{class_name}(#{@name}: #{@dtype})"
+    def ==(other)
+      name == other.name && dtype == other.dtype
+    end
+
+    def to_s
+      "#{self.class.name}(#{name.inspect}, #{dtype})"
     end
   end
 
@@ -142,25 +279,22 @@ module Polars
       end
     end
 
-    def inspect
-      class_name = self.class.name
-      "#{class_name}(#{@fields})"
+    def ==(other)
+      if other.eql?(Struct)
+        true
+      elsif other.is_a?(Struct)
+        fields == other.fields
+      else
+        false
+      end
+    end
+
+    def to_s
+      "#{self.class.name}([#{fields.map(&:to_s).join("\n")}])"
     end
 
     def to_schema
       @fields.to_h { |f| [f.name, f.dtype] }
     end
   end
-
-  # Binary type.
-  class Binary < DataType
-  end
-
-  # Type representing Null / None values.
-  class Null < DataType
-  end
-
-  # Type representing Datatype values that could not be determined statically.
-  class Unknown < DataType
-  end
 end

data/lib/polars/date_time_name_space.rb

@@ -82,8 +82,15 @@ module Polars
     #   # => 2001-01-02 00:00:00 UTC
     def median
       s = Utils.wrap_s(_s)
-      out = s.median.to_i
-      Utils._to_ruby_datetime(out, s.dtype, tu: s.time_unit)
+      out = s.median
+      if !out.nil?
+        if s.dtype == Date
+          return Utils._to_ruby_date(out.to_i)
+        else
+          return Utils._to_ruby_datetime(out.to_i, s.time_unit)
+        end
+      end
+      nil
     end
 
     # Return mean as Ruby object.
@@ -107,7 +114,14 @@ module Polars
     def mean
       s = Utils.wrap_s(_s)
       out = s.mean.to_i
-      Utils._to_ruby_datetime(out, s.dtype, tu: s.time_unit)
+      if !out.nil?
+        if s.dtype == Date
+          return Utils._to_ruby_date(out.to_i)
+        else
+          return Utils._to_ruby_datetime(out.to_i, s.time_unit)
+        end
+      end
+      nil
     end
 
     # Format Date/datetime with a formatting rule.

data/lib/polars/expr.rb

@@ -362,7 +362,7 @@ module Polars
       if columns.is_a?(String)
         columns = [columns]
         return wrap_expr(_rbexpr.exclude(columns))
-      elsif !columns.is_a?(Array)
+      elsif !columns.is_a?(::Array)
         columns = [columns]
         return wrap_expr(_rbexpr.exclude_dtype(columns))
       end
@@ -820,18 +820,18 @@ module Polars
     #   df.select(Polars.repeat(nil, 3).append(Polars.col("a")).rechunk)
     #   # =>
     #   # shape: (6, 1)
-    #   # ┌─────────┐
-    #   # │ literal │
-    #   # │ ---     │
-    #   # │ i64     │
-    #   # ╞═════════╡
-    #   # │ null    │
-    #   # │ null    │
-    #   # │ null    │
-    #   # │ 1       │
-    #   # │ 1       │
-    #   # │ 2       │
-    #   # └─────────┘
+    #   # ┌────────┐
+    #   # │ repeat │
+    #   # │ ---    │
+    #   # │ i64    │
+    #   # ╞════════╡
+    #   # │ null   │
+    #   # │ null   │
+    #   # │ null   │
+    #   # │ 1      │
+    #   # │ 1      │
+    #   # │ 2      │
+    #   # └────────┘
     def rechunk
       wrap_expr(_rbexpr.rechunk)
     end
@@ -1534,10 +1534,10 @@ module Polars
     #   # │ two   │
     #   # └───────┘
     def sort_by(by, reverse: false)
-      if !by.is_a?(Array)
+      if !by.is_a?(::Array)
         by = [by]
       end
-      if !reverse.is_a?(Array)
+      if !reverse.is_a?(::Array)
         reverse = [reverse]
       end
       by = Utils.selection_to_rbexpr_list(by)
@@ -1578,7 +1578,7 @@ module Polars
     #   # │ two   ┆ 99    │
     #   # └───────┴───────┘
     def take(indices)
-      if indices.is_a?(Array)
+      if indices.is_a?(::Array)
         indices_lit = Polars.lit(Series.new("", indices, dtype: :u32))
       else
         indices_lit = Utils.expr_to_lit_or_expr(indices, str_to_lit: false)
@@ -2436,14 +2436,14 @@ module Polars
     #   ).sort("group_col")
     #   # =>
     #   # shape: (2, 3)
-    #   # ┌───────────┬──────┬─────┐
-    #   # │ group_col ┆ lt   ┆ gte │
-    #   # │ ---       ┆ ---  ┆ --- │
-    #   # │ str       ┆ i64  ┆ i64 │
-    #   # ╞═══════════╪══════╪═════╡
-    #   # │ g1        ┆ 1    ┆ 2   │
-    #   # │ g2        ┆ null ┆ 3   │
-    #   # └───────────┴──────┴─────┘
+    #   # ┌───────────┬─────┬─────┐
+    #   # │ group_col ┆ lt  ┆ gte │
+    #   # │ ---       ┆ --- ┆ --- │
+    #   # │ str       ┆ i64 ┆ i64 │
+    #   # ╞═══════════╪═════╪═════╡
+    #   # │ g1        ┆ 1   ┆ 2   │
+    #   # │ g2        ┆ 0   ┆ 3   │
+    #   # └───────────┴─────┴─────┘
     def filter(predicate)
       wrap_expr(_rbexpr.filter(predicate._rbexpr))
     end
@@ -2474,14 +2474,14 @@ module Polars
     #   ).sort("group_col")
     #   # =>
     #   # shape: (2, 3)
-    #   # ┌───────────┬──────┬─────┐
-    #   # │ group_col ┆ lt   ┆ gte │
-    #   # │ ---       ┆ ---  ┆ --- │
-    #   # │ str       ┆ i64  ┆ i64 │
-    #   # ╞═══════════╪══════╪═════╡
-    #   # │ g1        ┆ 1    ┆ 2   │
-    #   # │ g2        ┆ null ┆ 3   │
-    #   # └───────────┴──────┴─────┘
+    #   # ┌───────────┬─────┬─────┐
+    #   # │ group_col ┆ lt  ┆ gte │
+    #   # │ ---       ┆ --- ┆ --- │
+    #   # │ str       ┆ i64 ┆ i64 │
+    #   # ╞═══════════╪═════╪═════╡
+    #   # │ g1        ┆ 1   ┆ 2   │
+    #   # │ g2        ┆ 0   ┆ 3   │
+    #   # └───────────┴─────┴─────┘
     def where(predicate)
       filter(predicate)
     end
@@ -2616,25 +2616,23 @@ module Polars
     # @return [Expr]
     #
     # @example
-    #   df = Polars::DataFrame.new({"foo" => ["hello", "world"]})
-    #   df.select(Polars.col("foo").flatten)
-    #   # =>
-    #   # shape: (10, 1)
-    #   # ┌─────┐
-    #   # │ foo │
-    #   # │ --- │
-    #   # │ str │
-    #   # ╞═════╡
-    #   # │ h   │
-    #   # │ e   │
-    #   # │ l   │
-    #   # │ l   │
-    #   # │ …   │
-    #   # │ o   │
-    #   # │ r   │
-    #   # │ l   │
-    #   # │ d   │
-    #   # └─────┘
+    #  df = Polars::DataFrame.new(
+    #    {
+    #      "group" => ["a", "b", "b"],
+    #      "values" => [[1, 2], [2, 3], [4]]
+    #    }
+    #  )
+    #  df.groupby("group").agg(Polars.col("values").flatten)
+    #  # =>
+    #  # shape: (2, 2)
+    #  # ┌───────┬───────────┐
+    #  # │ group ┆ values    │
+    #  # │ ---   ┆ ---       │
+    #  # │ str   ┆ list[i64] │
+    #  # ╞═══════╪═══════════╡
+    #  # │ a     ┆ [1, 2]    │
+    #  # │ b     ┆ [2, 3, 4] │
+    #  # └───────┴───────────┘
     def flatten
       wrap_expr(_rbexpr.explode)
     end
@@ -2798,7 +2796,7 @@ module Polars
     #   # │ false    │
     #   # └──────────┘
     def is_in(other)
-      if other.is_a?(Array)
+      if other.is_a?(::Array)
         if other.length == 0
           other = Polars.lit(nil)
         else
@@ -3502,14 +3500,15 @@ module Polars
       min_periods: nil,
       center: false,
       by: nil,
-      closed: "left"
+      closed: "left",
+      ddof: 1
     )
       window_size, min_periods = _prepare_rolling_window_args(
         window_size, min_periods
       )
       wrap_expr(
         _rbexpr.rolling_std(
-          window_size, weights, min_periods, center, by, closed
+          window_size, weights, min_periods, center, by, closed, ddof
         )
       )
     end
@@ -3591,14 +3590,15 @@ module Polars
       min_periods: nil,
       center: false,
       by: nil,
-      closed: "left"
+      closed: "left",
+      ddof: 1
     )
       window_size, min_periods = _prepare_rolling_window_args(
         window_size, min_periods
       )
       wrap_expr(
         _rbexpr.rolling_var(
-          window_size, weights, min_periods, center, by, closed
+          window_size, weights, min_periods, center, by, closed, ddof
         )
       )
     end
@@ -4558,11 +4558,11 @@ module Polars
     #   # │ 1   │
     #   # │ 3   │
     #   # └─────┘
-    def shuffle(seed: nil)
+    def shuffle(seed: nil, fixed_seed: false)
       if seed.nil?
         seed = rand(10000)
       end
-      wrap_expr(_rbexpr.shuffle(seed))
+      wrap_expr(_rbexpr.shuffle(seed, fixed_seed))
     end
 
     # Sample from this expression.
@@ -4600,21 +4600,22 @@ module Polars
       with_replacement: true,
       shuffle: false,
       seed: nil,
-      n: nil
+      n: nil,
+      fixed_seed: false
     )
       if !n.nil? && !frac.nil?
         raise ArgumentError, "cannot specify both `n` and `frac`"
       end
 
       if !n.nil? && frac.nil?
-        return wrap_expr(_rbexpr.sample_n(n, with_replacement, shuffle, seed))
+        return wrap_expr(_rbexpr.sample_n(n, with_replacement, shuffle, seed, fixed_seed))
       end
 
       if frac.nil?
         frac = 1.0
       end
       wrap_expr(
-        _rbexpr.sample_frac(frac, with_replacement, shuffle, seed)
+        _rbexpr.sample_frac(frac, with_replacement, shuffle, seed, fixed_seed)
       )
     end
 
@@ -4929,8 +4930,8 @@ module Polars
     #
     # Enables downstream code to user fast paths for sorted arrays.
     #
-    # @param reverse [Boolean]
-    #   If the `Series` order is reversed, e.g. descending.
+    # @param descending [Boolean]
+    #   Whether the `Series` order is descending.
     #
     # @return [Expr]
     #
@@ -4950,9 +4951,9 @@ module Polars
     #   # ╞════════╡
     #   # │ 3      │
     #   # └────────┘
-    # def set_sorted(reverse: false)
-    #   map { |s| s.set_sorted(reverse) }
-    # end
+    def set_sorted(descending: false)
+      wrap_expr(_rbexpr.set_sorted_flag(descending))
+    end
 
     # Aggregate to list.
     #
@@ -4965,7 +4966,7 @@ module Polars
     #       "b" => [4, 5, 6]
     #     }
     #   )
-    #   df.select(Polars.all.list)
+    #   df.select(Polars.all.implode)
     #   # =>
     #   # shape: (1, 2)
     #   # ┌───────────┬───────────┐
@@ -4978,7 +4979,6 @@ module Polars
     def implode
       wrap_expr(_rbexpr.implode)
     end
-    alias_method :list, :implode
 
     # Shrink numeric columns to the minimal required datatype.
     #
@@ -5018,10 +5018,17 @@ module Polars
     # Create an object namespace of all list related methods.
     #
     # @return [ListExpr]
-    def arr
+    def list
       ListExpr.new(self)
     end
 
+    # Create an object namespace of all array related methods.
+    #
+    # @return [ArrayExpr]
+    def arr
+      ArrayExpr.new(self)
+    end
+
     # Create an object namespace of all binary related methods.
     #
     # @return [BinaryExpr]

data/lib/polars/functions.rb

@@ -43,7 +43,6 @@ module Polars
     #   # │ i64 ┆ i64 │
     #   # ╞═════╪═════╡
     #   # │ 1   ┆ 3   │
-    #   # ├╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ 2   ┆ 4   │
     #   # └─────┴─────┘
     def concat(items, rechunk: true, how: "vertical", parallel: true)

data/lib/polars/group_by.rb

@@ -551,32 +551,11 @@ module Polars
       agg(Polars.all.median)
     end
 
-    # Aggregate the groups into Series.
-    #
-    # @return [DataFrame]
-    #
-    # @example
-    #   df = Polars::DataFrame.new({"a" => ["one", "two", "one", "two"], "b" => [1, 2, 3, 4]})
-    #   df.groupby("a", maintain_order: true).agg_list
-    #   # =>
-    #   # shape: (2, 2)
-    #   # ┌─────┬─────────────────┐
-    #   # │ a   ┆ b               │
-    #   # │ --- ┆ ---             │
-    #   # │ str ┆ list[list[i64]] │
-    #   # ╞═════╪═════════════════╡
-    #   # │ one ┆ [[1, 3]]        │
-    #   # │ two ┆ [[2, 4]]        │
-    #   # └─────┴─────────────────┘
-    def agg_list
-      agg(Polars.all.list)
-    end
-
     # Plot data.
     #
     # @return [Vega::LiteChart]
     def plot(*args, **options)
-      raise ArgumentError, "Multiple groups not supported" if by.is_a?(Array) && by.size > 1
+      raise ArgumentError, "Multiple groups not supported" if by.is_a?(::Array) && by.size > 1
       # same message as Ruby
       raise ArgumentError, "unknown keyword: :group" if options.key?(:group)