polars-df 0.5.0 → 0.6.0

data/lib/polars/data_frame.rb

@@ -36,7 +36,7 @@ module Polars
       elsif data.is_a?(Hash)
         data = data.transform_keys { |v| v.is_a?(Symbol) ? v.to_s : v }
         self._df = self.class.hash_to_rbdf(data, schema: schema, schema_overrides: schema_overrides, nan_to_null: nan_to_null)
-      elsif data.is_a?(Array)
+      elsif data.is_a?(::Array)
         self._df = self.class.sequence_to_rbdf(data, schema: schema, schema_overrides: schema_overrides, orient: orient, infer_schema_length: infer_schema_length)
       elsif data.is_a?(Series)
         self._df = self.class.series_to_rbdf(data, schema: schema, schema_overrides: schema_overrides)
@@ -116,7 +116,7 @@ module Polars
           dtypes.each do|k, v|
             dtype_list << [k, Utils.rb_type_to_dtype(v)]
           end
-        elsif dtypes.is_a?(Array)
+        elsif dtypes.is_a?(::Array)
           dtype_slice = dtypes
         else
           raise ArgumentError, "dtype arg should be list or dict"
@@ -590,7 +590,7 @@ module Polars
 
         # df[2, ..] (select row as df)
         if row_selection.is_a?(Integer)
-          if col_selection.is_a?(Array)
+          if col_selection.is_a?(::Array)
             df = self[0.., col_selection]
             return df.slice(row_selection, 1)
           end
@@ -611,7 +611,7 @@ module Polars
           return series[row_selection]
         end
 
-        if col_selection.is_a?(Array)
+        if col_selection.is_a?(::Array)
           # df[.., [1, 2]]
           if Utils.is_int_sequence(col_selection)
             series_list = col_selection.map { |i| to_series(i) }
@@ -641,7 +641,7 @@ module Polars
           return Slice.new(self).apply(item)
         end
 
-        if item.is_a?(Array) && item.all? { |v| Utils.strlike?(v) }
+        if item.is_a?(::Array) && item.all? { |v| Utils.strlike?(v) }
           # select multiple columns
           # df[["foo", "bar"]]
           return _from_rbdf(_df.select(item.map(&:to_s)))
@@ -684,13 +684,13 @@ module Polars
       end
 
       if Utils.strlike?(key)
-        if value.is_a?(Array) || (defined?(Numo::NArray) && value.is_a?(Numo::NArray))
+        if value.is_a?(::Array) || (defined?(Numo::NArray) && value.is_a?(Numo::NArray))
           value = Series.new(value)
         elsif !value.is_a?(Series)
           value = Polars.lit(value)
         end
         self._df = with_column(value.alias(key.to_s))._df
-      elsif key.is_a?(Array)
+      elsif key.is_a?(::Array)
         row_selection, col_selection = key
 
         if Utils.strlike?(col_selection)
@@ -994,14 +994,21 @@ module Polars
     #
     # @return [nil]
     def write_ipc(file, compression: "uncompressed")
-      if compression.nil?
-        compression = "uncompressed"
+      return_bytes = file.nil?
+      if return_bytes
+        file = StringIO.new
+        file.set_encoding(Encoding::BINARY)
       end
       if Utils.pathlike?(file)
         file = Utils.normalise_filepath(file)
       end
 
+      if compression.nil?
+        compression = "uncompressed"
+      end
+
       _df.write_ipc(file, compression)
+      return_bytes ? file.string : nil
     end
 
     # Write to Apache Parquet file.
@@ -1491,13 +1498,9 @@ module Polars
     #   # │ 1   ┆ 6.0 ┆ a   │
     #   # └─────┴─────┴─────┘
     def sort(by, reverse: false, nulls_last: false)
-      if by.is_a?(Array) || by.is_a?(Expr)
-        lazy
-          .sort(by, reverse: reverse, nulls_last: nulls_last)
-          .collect(no_optimization: true, string_cache: false)
-      else
-        _from_rbdf(_df.sort(by, reverse, nulls_last))
-      end
+      lazy
+        .sort(by, reverse: reverse, nulls_last: nulls_last)
+        .collect(no_optimization: true)
     end
 
     # Sort the DataFrame by column in-place.
@@ -1899,6 +1902,12 @@ module Polars
     #   Define whether the temporal window interval is closed or not.
     # @param by [Object]
     #   Also group by this column/these columns.
+    # @param check_sorted [Boolean]
+    #   When the `by` argument is given, polars can not check sortedness
+    #   by the metadata and has to do a full scan on the index column to
+    #   verify data is sorted. This is expensive. If you are sure the
+    #   data within the by groups is sorted, you can set this to `false`.
+    #   Doing so incorrectly will lead to incorrect output
     #
     # @return [RollingGroupBy]
     #
@@ -1912,7 +1921,7 @@ module Polars
     #     "2020-01-08 23:16:43"
     #   ]
     #   df = Polars::DataFrame.new({"dt" => dates, "a" => [3, 7, 5, 9, 2, 1]}).with_column(
-    #     Polars.col("dt").str.strptime(Polars::Datetime)
+    #     Polars.col("dt").str.strptime(Polars::Datetime).set_sorted
     #   )
     #   df.groupby_rolling(index_column: "dt", period: "2d").agg(
     #     [
@@ -1940,9 +1949,10 @@ module Polars
       period:,
       offset: nil,
       closed: "right",
-      by: nil
+      by: nil,
+      check_sorted: true
     )
-      RollingGroupBy.new(self, index_column, period, offset, closed, by)
+      RollingGroupBy.new(self, index_column, period, offset, closed, by, check_sorted)
     end
 
     # Group based on a time value (or index value of type `:i32`, `:i64`).
@@ -2078,21 +2088,21 @@ module Polars
     #   df.groupby_dynamic("time", every: "1h", closed: "left").agg(
     #     [
     #       Polars.col("time").count.alias("time_count"),
-    #       Polars.col("time").list.alias("time_agg_list")
+    #       Polars.col("time").alias("time_agg_list")
     #     ]
     #   )
     #   # =>
     #   # shape: (4, 3)
-    #   # ┌─────────────────────┬────────────┬─────────────────────────────────────┐
-    #   # │ time                ┆ time_count ┆ time_agg_list                       │
-    #   # │ ---                 ┆ ---        ┆ ---                                 │
-    #   # │ datetime[μs]        ┆ u32        ┆ list[datetime[μs]]                  │
-    #   # ╞═════════════════════╪════════════╪═════════════════════════════════════╡
-    #   # │ 2021-12-16 00:00:00 ┆ 2          ┆ [2021-12-16 00:00:00, 2021-12-16... │
-    #   # │ 2021-12-16 01:00:00 ┆ 2          ┆ [2021-12-16 01:00:00, 2021-12-16... │
-    #   # │ 2021-12-16 02:00:00 ┆ 2          ┆ [2021-12-16 02:00:00, 2021-12-16... │
-    #   # │ 2021-12-16 03:00:00 ┆ 1          ┆ [2021-12-16 03:00:00]               │
-    #   # └─────────────────────┴────────────┴─────────────────────────────────────┘
+    #   # ┌─────────────────────┬────────────┬───────────────────────────────────┐
+    #   # │ time                ┆ time_count ┆ time_agg_list                     │
+    #   # │ ---                 ┆ ---        ┆ ---                               │
+    #   # │ datetime[μs]        ┆ u32        ┆ list[datetime[μs]]                │
+    #   # ╞═════════════════════╪════════════╪═══════════════════════════════════╡
+    #   # │ 2021-12-16 00:00:00 ┆ 2          ┆ [2021-12-16 00:00:00, 2021-12-16… │
+    #   # │ 2021-12-16 01:00:00 ┆ 2          ┆ [2021-12-16 01:00:00, 2021-12-16… │
+    #   # │ 2021-12-16 02:00:00 ┆ 2          ┆ [2021-12-16 02:00:00, 2021-12-16… │
+    #   # │ 2021-12-16 03:00:00 ┆ 1          ┆ [2021-12-16 03:00:00]             │
+    #   # └─────────────────────┴────────────┴───────────────────────────────────┘
     #
     # @example When closed="both" the time values at the window boundaries belong to 2 groups.
     #   df.groupby_dynamic("time", every: "1h", closed: "both").agg(
@@ -2159,7 +2169,7 @@ module Polars
     #     period: "3i",
     #     include_boundaries: true,
     #     closed: "right"
-    #   ).agg(Polars.col("A").list.alias("A_agg_list"))
+    #   ).agg(Polars.col("A").alias("A_agg_list"))
     #   # =>
     #   # shape: (3, 4)
     #   # ┌─────────────────┬─────────────────┬─────┬─────────────────┐
@@ -2242,7 +2252,7 @@ module Polars
     #       "groups" => ["A", "B", "A", "B"],
     #       "values" => [0, 1, 2, 3]
     #     }
-    #   )
+    #   ).set_sorted("time")
     #   df.upsample(
     #     time_column: "time", every: "1mo", by: "groups", maintain_order: true
     #   ).select(Polars.all.forward_fill)
@@ -2360,7 +2370,7 @@ module Polars
     #       ],  # note record date: Jan 1st (sorted!)
     #       "gdp" => [4164, 4411, 4566, 4696]
     #     }
-    #   )
+    #   ).set_sorted("date")
     #   population = Polars::DataFrame.new(
     #     {
     #       "date" => [
@@ -2371,7 +2381,7 @@ module Polars
     #       ],  # note record date: May 12th (sorted!)
     #       "population" => [82.19, 82.66, 83.12, 83.52]
     #     }
-    #   )
+    #   ).set_sorted("date")
     #   population.join_asof(
     #     gdp, left_on: "date", right_on: "date", strategy: "backward"
     #   )
@@ -2674,7 +2684,7 @@ module Polars
     #   # │ 3   ┆ 8   ┆ c   ┆ 30    │
     #   # └─────┴─────┴─────┴───────┘
     def hstack(columns, in_place: false)
-      if !columns.is_a?(Array)
+      if !columns.is_a?(::Array)
         columns = columns.get_columns
       end
       if in_place
@@ -2804,7 +2814,7 @@ module Polars
     #   # │ 3   ┆ 8.0 │
     #   # └─────┴─────┘
     def drop(columns)
-      if columns.is_a?(Array)
+      if columns.is_a?(::Array)
         df = clone
         columns.each do |n|
           df._df.drop_in_place(n)
@@ -3317,7 +3327,7 @@ module Polars
       n_fill = n_cols * n_rows - height
 
       if n_fill > 0
-        if !fill_values.is_a?(Array)
+        if !fill_values.is_a?(::Array)
           fill_values = [fill_values] * df.width
         end
 
@@ -3426,29 +3436,29 @@ module Polars
     #   # ╞═════╪═════╪═════╡
     #   # │ C   ┆ 2   ┆ l   │
     #   # └─────┴─────┴─────┘}
-    def partition_by(groups, maintain_order: true, as_dict: false)
+    def partition_by(groups, maintain_order: true, include_key: true, as_dict: false)
       if groups.is_a?(String)
         groups = [groups]
-      elsif !groups.is_a?(Array)
+      elsif !groups.is_a?(::Array)
         groups = Array(groups)
       end
 
       if as_dict
         out = {}
         if groups.length == 1
-          _df.partition_by(groups, maintain_order).each do |df|
+          _df.partition_by(groups, maintain_order, include_key).each do |df|
             df = _from_rbdf(df)
             out[df[groups][0, 0]] = df
           end
         else
-          _df.partition_by(groups, maintain_order).each do |df|
+          _df.partition_by(groups, maintain_order, include_key).each do |df|
             df = _from_rbdf(df)
             out[df[groups].row(0)] = df
           end
         end
         out
       else
-        _df.partition_by(groups, maintain_order).map { |df| _from_rbdf(df) }
+        _df.partition_by(groups, maintain_order, include_key).map { |df| _from_rbdf(df) }
       end
     end
 
@@ -3716,7 +3726,7 @@ module Polars
     #   # │ 4   ┆ 13.0 ┆ true  ┆ 16.0 ┆ 6.5  ┆ false │
     #   # └─────┴──────┴───────┴──────┴──────┴───────┘
     def with_columns(exprs)
-      if !exprs.nil? && !exprs.is_a?(Array)
+      if !exprs.nil? && !exprs.is_a?(::Array)
         exprs = [exprs]
       end
       lazy
@@ -4097,11 +4107,11 @@ module Polars
     #   # │ 1     ┆ 0     ┆ 1     ┆ 0     ┆ 1     ┆ 0     │
     #   # │ 0     ┆ 1     ┆ 0     ┆ 1     ┆ 0     ┆ 1     │
     #   # └───────┴───────┴───────┴───────┴───────┴───────┘
-    def to_dummies(columns: nil, separator: "_")
+    def to_dummies(columns: nil, separator: "_", drop_first: false)
       if columns.is_a?(String)
         columns = [columns]
       end
-      _from_rbdf(_df.to_dummies(columns, separator))
+      _from_rbdf(_df.to_dummies(columns, separator, drop_first))
     end
 
     # Drop duplicate rows from this DataFrame.
@@ -4189,7 +4199,7 @@ module Polars
         subset = [subset]
       end
 
-      if subset.is_a?(Array) && subset.length == 1
+      if subset.is_a?(::Array) && subset.length == 1
         expr = Utils.expr_to_lit_or_expr(subset[0], str_to_lit: false)
       else
         struct_fields = subset.nil? ? Polars.all : subset
@@ -4758,6 +4768,38 @@ module Polars
       _from_rbdf(_df.unnest(names))
     end
 
+    # TODO
+    # def corr
+    # end
+
+    # TODO
+    # def merge_sorted
+    # end
+
+    # Indicate that one or multiple columns are sorted.
+    #
+    # @param column [Object]
+    #   Columns that are sorted
+    # @param more_columns [Object]
+    #   Additional columns that are sorted, specified as positional arguments.
+    # @param descending [Boolean]
+    #   Whether the columns are sorted in descending order.
+    #
+    # @return [DataFrame]
+    def set_sorted(
+      column,
+      *more_columns,
+      descending: false
+    )
+      lazy
+        .set_sorted(column, *more_columns, descending: descending)
+        .collect(no_optimization: true)
+    end
+
+    # TODO
+    # def update
+    # end
+
     private
 
     def initialize_copy(other)
@@ -4967,7 +5009,7 @@ module Polars
       columns.each do |col, i|
         if dtypes[col] == Categorical # != rbdf_dtypes[i]
           column_casts << Polars.col(col).cast(Categorical)._rbexpr
-        elsif structs.any? && structs.include?(col) && structs[col] != rbdf_dtypes[i]
+        elsif structs&.any? && structs.include?(col) && structs[col] != rbdf_dtypes[i]
           column_casts << Polars.col(col).cast(structs[col])._rbexpr
         elsif dtypes.include?(col) && dtypes[col] != rbdf_dtypes[i]
           column_casts << Polars.col(col).cast(dtypes[col])._rbexpr
@@ -5012,7 +5054,7 @@ module Polars
           rbdf = _post_apply_columns(rbdf, column_names)
         end
         return rbdf
-      elsif data[0].is_a?(Array)
+      elsif data[0].is_a?(::Array)
         if orient.nil? && !columns.nil?
           orient = columns.length == data.length ? "col" : "row"
         end
@@ -5117,7 +5159,7 @@ module Polars
 
     def _prepare_other_arg(other)
       if !other.is_a?(Series)
-        if other.is_a?(Array)
+        if other.is_a?(::Array)
           raise ArgumentError, "Operation not supported."
         end
 

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.