polars-df 0.5.0 → 0.6.0

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)
 

data/lib/polars/lazy_frame.rb

@@ -4,6 +4,22 @@ module Polars
     # @private
     attr_accessor :_ldf
 
+    # Create a new LazyFrame.
+    def initialize(data = nil, schema: nil, schema_overrides: nil, orient: nil, infer_schema_length: 100, nan_to_null: false)
+      self._ldf = (
+        DataFrame.new(
+          data,
+          schema: schema,
+          schema_overrides: schema_overrides,
+          orient: orient,
+          infer_schema_length: infer_schema_length,
+          nan_to_null: nan_to_null
+        )
+        .lazy
+        ._ldf
+      )
+    end
+
     # @private
     def self._from_rbldf(rb_ldf)
       ldf = LazyFrame.allocate
@@ -379,16 +395,16 @@ module Polars
     #   # │ 2   ┆ 7.0 ┆ b   │
     #   # │ 1   ┆ 6.0 ┆ a   │
     #   # └─────┴─────┴─────┘
-    def sort(by, reverse: false, nulls_last: false)
+    def sort(by, reverse: false, nulls_last: false, maintain_order: false)
       if by.is_a?(String)
-        _from_rbldf(_ldf.sort(by, reverse, nulls_last))
+        return _from_rbldf(_ldf.sort(by, reverse, nulls_last, maintain_order))
       end
       if Utils.bool?(reverse)
         reverse = [reverse]
       end
 
       by = Utils.selection_to_rbexpr_list(by)
-      _from_rbldf(_ldf.sort_by_exprs(by, reverse, nulls_last))
+      _from_rbldf(_ldf.sort_by_exprs(by, reverse, nulls_last, maintain_order))
     end
 
     # def profile
@@ -921,6 +937,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 [LazyFrame]
     #
@@ -933,8 +955,8 @@ module Polars
     #     "2020-01-03 19:45:32",
     #     "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)
+    #   df = Polars::LazyFrame.new({"dt" => dates, "a" => [3, 7, 5, 9, 2, 1]}).with_column(
+    #     Polars.col("dt").str.strptime(Polars::Datetime).set_sorted
     #   )
     #   df.groupby_rolling(index_column: "dt", period: "2d").agg(
     #     [
@@ -942,7 +964,7 @@ module Polars
     #       Polars.min("a").alias("min_a"),
     #       Polars.max("a").alias("max_a")
     #     ]
-    #   )
+    #   ).collect
     #   # =>
     #   # shape: (6, 4)
     #   # ┌─────────────────────┬───────┬───────┬───────┐
@@ -962,7 +984,8 @@ module Polars
       period:,
       offset: nil,
       closed: "right",
-      by: nil
+      by: nil,
+      check_sorted: true
     )
       index_column = Utils.expr_to_lit_or_expr(index_column, str_to_lit: false)
       if offset.nil?
@@ -974,7 +997,7 @@ module Polars
       offset = Utils._timedelta_to_pl_duration(offset)
 
       lgb = _ldf.groupby_rolling(
-        index_column._rbexpr, period, offset, closed, rbexprs_by
+        index_column._rbexpr, period, offset, closed, rbexprs_by, check_sorted
       )
       LazyGroupBy.new(lgb, self.class)
     end
@@ -1112,21 +1135,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(
@@ -1193,7 +1216,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)
     #   # ┌─────────────────┬─────────────────┬─────┬─────────────────┐
@@ -1216,12 +1239,9 @@ module Polars
       by: nil,
       start_by: "window"
     )
+      index_column = Utils.expr_to_lit_or_expr(index_column, str_to_lit: false)
       if offset.nil?
-        if period.nil?
-          offset = "-#{every}"
-        else
-          offset = "0ns"
-        end
+        offset = period.nil? ? "-#{every}" : "0ns"
       end
 
       if period.nil?
@@ -1234,7 +1254,7 @@ module Polars
 
       rbexprs_by = by.nil? ? [] : Utils.selection_to_rbexpr_list(by)
       lgb = _ldf.groupby_dynamic(
-        index_column,
+        index_column._rbexpr,
         every,
         period,
         offset,
@@ -1351,7 +1371,7 @@ module Polars
       if by.is_a?(String)
         by_left_ = [by]
         by_right_ = [by]
-      elsif by.is_a?(Array)
+      elsif by.is_a?(::Array)
         by_left_ = by
         by_right_ = by
       end
@@ -1619,7 +1639,7 @@ module Polars
     #   # │ null │
     #   # └──────┘
     def with_context(other)
-      if !other.is_a?(Array)
+      if !other.is_a?(::Array)
         other = [other]
       end
 
@@ -2228,7 +2248,7 @@ module Polars
     #
     # @return [LazyFrame]
     def unique(maintain_order: true, subset: nil, keep: "first")
-      if !subset.nil? && !subset.is_a?(Array)
+      if !subset.nil? && !subset.is_a?(::Array)
         subset = [subset]
       end
       _from_rbldf(_ldf.unique(maintain_order, subset, keep))
@@ -2261,7 +2281,7 @@ module Polars
     #   # │ 3   ┆ 8   ┆ c   │
     #   # └─────┴─────┴─────┘
     def drop_nulls(subset: nil)
-      if !subset.nil? && !subset.is_a?(Array)
+      if !subset.nil? && !subset.is_a?(::Array)
         subset = [subset]
       end
       _from_rbldf(_ldf.drop_nulls(subset))
@@ -2423,6 +2443,38 @@ module Polars
       _from_rbldf(_ldf.unnest(names))
     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 [LazyFrame]
+    def set_sorted(
+      column,
+      *more_columns,
+      descending: false
+    )
+      columns = Utils.selection_to_rbexpr_list(column)
+      if more_columns.any?
+        columns.concat(Utils.selection_to_rbexpr_list(more_columns))
+      end
+      with_columns(
+        columns.map { |e| Utils.wrap_expr(e).set_sorted(descending: descending) }
+      )
+    end
+
+    # TODO
+    # def update
+    # end
+
     private
 
     def initialize_copy(other)