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

data/README.md

@@ -25,7 +25,7 @@ Polars.read_csv("iris.csv")
   .collect
 ```
 
-You can follow [Polars tutorials](https://pola-rs.github.io/polars-book/user-guide/introduction.html) and convert the code to Ruby in many cases. Feel free to open an issue if you run into problems.
+You can follow [Polars tutorials](https://pola-rs.github.io/polars-book/user-guide/) and convert the code to Ruby in many cases. Feel free to open an issue if you run into problems.
 
 ## Reference
 
@@ -348,7 +348,7 @@ df.to_numo
 You can specify column types when creating a data frame
 
 ```ruby
-Polars::DataFrame.new(data, columns: {"a" => Polars::Int32, "b" => Polars::Float32})
+Polars::DataFrame.new(data, schema: {"a" => Polars::Int32, "b" => Polars::Float32})
 ```
 
 Supported types are:
@@ -357,8 +357,9 @@ Supported types are:
 - float - `Float64`, `Float32`
 - integer - `Int64`, `Int32`, `Int16`, `Int8`
 - unsigned integer - `UInt64`, `UInt32`, `UInt16`, `UInt8`
-- string - `Utf8`, `Categorical`
+- string - `Utf8`, `Binary`, `Categorical`
 - temporal - `Date`, `Datetime`, `Time`, `Duration`
+- other - `Object`, `List`, `Struct`, `Array` [unreleased]
 
 Get column types
 

data/lib/polars/array_expr.rb

@@ -0,0 +1,84 @@
+module Polars
+  # Namespace for array related expressions.
+  class ArrayExpr
+    # @private
+    attr_accessor :_rbexpr
+
+    # @private
+    def initialize(expr)
+      self._rbexpr = expr._rbexpr
+    end
+
+    # Compute the min values of the sub-arrays.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(2, Polars::Int64)}
+    #   )
+    #   df.select(Polars.col("a").arr.min)
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ i64 │
+    #   # ╞═════╡
+    #   # │ 1   │
+    #   # │ 3   │
+    #   # └─────┘
+    def min
+      Utils.wrap_expr(_rbexpr.array_min)
+    end
+
+    # Compute the max values of the sub-arrays.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(2, Polars::Int64)}
+    #   )
+    #   df.select(Polars.col("a").arr.max)
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ i64 │
+    #   # ╞═════╡
+    #   # │ 2   │
+    #   # │ 4   │
+    #   # └─────┘
+    def max
+      Utils.wrap_expr(_rbexpr.array_max)
+    end
+
+    # Compute the sum values of the sub-arrays.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(2, Polars::Int64)}
+    #   )
+    #   df.select(Polars.col("a").arr.sum)
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ i64 │
+    #   # ╞═════╡
+    #   # │ 3   │
+    #   # │ 7   │
+    #   # └─────┘
+    def sum
+      Utils.wrap_expr(_rbexpr.array_sum)
+    end
+  end
+end

data/lib/polars/array_name_space.rb

@@ -0,0 +1,77 @@
+module Polars
+  # Series.arr namespace.
+  class ArrayNameSpace
+    include ExprDispatch
+
+    self._accessor = "arr"
+
+    # @private
+    def initialize(series)
+      self._s = series._s
+    end
+
+    # Compute the min values of the sub-arrays.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new(
+    #     "a", [[1, 2], [4, 3]], dtype: Polars::Array.new(2, Polars::Int64)
+    #   )
+    #   s.arr.min
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: 'a' [i64]
+    #   # [
+    #   #         1
+    #   #         3
+    #   # ]
+    def min
+      super
+    end
+
+    # Compute the max values of the sub-arrays.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new(
+    #     "a", [[1, 2], [4, 3]], dtype: Polars::Array.new(2, Polars::Int64)
+    #   )
+    #   s.arr.max
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: 'a' [i64]
+    #   # [
+    #   #         2
+    #   #         4
+    #   # ]
+    def max
+      super
+    end
+
+    # Compute the sum values of the sub-arrays.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(2, Polars::Int64)}
+    #   )
+    #   df.select(Polars.col("a").arr.sum)
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ i64 │
+    #   # ╞═════╡
+    #   # │ 3   │
+    #   # │ 7   │
+    #   # └─────┘
+    def sum
+      super
+    end
+  end
+end

data/lib/polars/batched_csv_reader.rb

@@ -41,7 +41,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"

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