polars-df 0.6.0-x86_64-linux → 0.8.0-x86_64-linux

data/lib/polars/lazy_frame.rb

@@ -97,7 +97,8 @@ module Polars
       row_count_offset: 0,
       storage_options: nil,
       low_memory: false,
-      use_statistics: true
+      use_statistics: true,
+      hive_partitioning: true
     )
       _from_rbldf(
         RbLazyFrame.new_from_parquet(
@@ -108,7 +109,8 @@ module Polars
           rechunk,
           Utils._prepare_row_count_args(row_count_name, row_count_offset),
           low_memory,
-          use_statistics
+          use_statistics,
+          hive_partitioning
         )
       )
     end
@@ -216,7 +218,7 @@ module Polars
     #     }
     #   ).lazy
     #   lf.dtypes
-    #   # => [Polars::Int64, Polars::Float64, Polars::Utf8]
+    #   # => [Polars::Int64, Polars::Float64, Polars::String]
     def dtypes
       _ldf.dtypes
     end
@@ -234,7 +236,7 @@ module Polars
     #     }
     #   ).lazy
     #   lf.schema
-    #   # => {"foo"=>Polars::Int64, "bar"=>Polars::Float64, "ham"=>Polars::Utf8}
+    #   # => {"foo"=>Polars::Int64, "bar"=>Polars::Float64, "ham"=>Polars::String}
     def schema
       _ldf.schema
     end
@@ -350,6 +352,7 @@ module Polars
         slice_pushdown,
         common_subplan_elimination,
         allow_streaming,
+        false
       )
 
       ldf.describe_optimized_plan
@@ -396,7 +399,7 @@ module Polars
     #   # │ 1   ┆ 6.0 ┆ a   │
     #   # └─────┴─────┴─────┘
     def sort(by, reverse: false, nulls_last: false, maintain_order: false)
-      if by.is_a?(String)
+      if by.is_a?(::String)
         return _from_rbldf(_ldf.sort(by, reverse, nulls_last, maintain_order))
       end
       if Utils.bool?(reverse)
@@ -445,7 +448,7 @@ module Polars
     #       "c" => [6, 5, 4, 3, 2, 1]
     #     }
     #   ).lazy
-    #   df.groupby("a", maintain_order: true).agg(Polars.all.sum).collect
+    #   df.group_by("a", maintain_order: true).agg(Polars.all.sum).collect
     #   # =>
     #   # shape: (3, 3)
     #   # ┌─────┬─────┬─────┐
@@ -466,7 +469,8 @@ module Polars
       no_optimization: false,
       slice_pushdown: true,
       common_subplan_elimination: true,
-      allow_streaming: false
+      allow_streaming: false,
+      _eager: false
     )
       if no_optimization
         predicate_pushdown = false
@@ -486,7 +490,8 @@ module Polars
         simplify_expression,
         slice_pushdown,
         common_subplan_elimination,
-        allow_streaming
+        allow_streaming,
+        _eager
       )
       Utils.wrap_df(ldf.collect)
     end
@@ -568,7 +573,8 @@ module Polars
         simplify_expression,
         slice_pushdown,
         false,
-        true
+        true,
+        false
       )
       lf.sink_parquet(
         path,
@@ -623,7 +629,7 @@ module Polars
     #       "c" => [6, 5, 4, 3, 2, 1]
     #     }
     #   ).lazy
-    #   df.groupby("a", maintain_order: true).agg(Polars.all.sum).fetch(2)
+    #   df.group_by("a", maintain_order: true).agg(Polars.all.sum).fetch(2)
     #   # =>
     #   # shape: (2, 3)
     #   # ┌─────┬─────┬─────┐
@@ -660,7 +666,8 @@ module Polars
         simplify_expression,
         slice_pushdown,
         common_subplan_elimination,
-        allow_streaming
+        allow_streaming,
+        false
       )
       Utils.wrap_df(ldf.fetch(n_rows))
     end
@@ -853,13 +860,13 @@ module Polars
       _from_rbldf(_ldf.select(exprs))
     end
 
-    # Start a groupby operation.
+    # Start a group by operation.
     #
     # @param by [Object]
     #   Column(s) to group by.
     # @param maintain_order [Boolean]
     #   Make sure that the order of the groups remain consistent. This is more
-    #   expensive than a default groupby.
+    #   expensive than a default group by.
     #
     # @return [LazyGroupBy]
     #
@@ -871,7 +878,7 @@ module Polars
     #       "c" => [6, 5, 4, 3, 2, 1]
     #     }
     #   ).lazy
-    #   df.groupby("a", maintain_order: true).agg(Polars.col("b").sum).collect
+    #   df.group_by("a", maintain_order: true).agg(Polars.col("b").sum).collect
     #   # =>
     #   # shape: (3, 2)
     #   # ┌─────┬─────┐
@@ -883,19 +890,21 @@ module Polars
     #   # │ b   ┆ 11  │
     #   # │ c   ┆ 6   │
     #   # └─────┴─────┘
-    def groupby(by, maintain_order: false)
+    def group_by(by, maintain_order: false)
       rbexprs_by = Utils.selection_to_rbexpr_list(by)
-      lgb = _ldf.groupby(rbexprs_by, maintain_order)
-      LazyGroupBy.new(lgb, self.class)
+      lgb = _ldf.group_by(rbexprs_by, maintain_order)
+      LazyGroupBy.new(lgb)
     end
+    alias_method :groupby, :group_by
+    alias_method :group, :group_by
 
     # Create rolling groups based on a time column.
     #
     # Also works for index values of type `:i32` or `:i64`.
     #
-    # Different from a `dynamic_groupby` the windows are now determined by the
+    # Different from a `dynamic_group_by` the windows are now determined by the
     # individual values and are not of constant intervals. For constant intervals
-    # use *groupby_dynamic*.
+    # use *group_by_dynamic*.
     #
     # The `period` and `offset` arguments are created either from a timedelta, or
     # by using the following string language:
@@ -915,7 +924,7 @@ module Polars
     # Or combine them:
     # "3d12h4m25s" # 3 days, 12 hours, 4 minutes, and 25 seconds
     #
-    # In case of a groupby_rolling on an integer column, the windows are defined by:
+    # In case of a group_by_rolling on an integer column, the windows are defined by:
     #
     # - "1i"      # length 1
     # - "10i"     # length 10
@@ -926,7 +935,7 @@ module Polars
     #   This column must be sorted in ascending order. If not the output will not
     #   make sense.
     #
-    #   In case of a rolling groupby on indices, dtype needs to be one of
+    #   In case of a rolling group by on indices, dtype needs to be one of
     #   `:i32`, `:i64`. Note that `:i32` gets temporarily cast to `:i64`, so if
     #   performance matters use an `:i64` column.
     # @param period [Object]
@@ -958,7 +967,7 @@ module Polars
     #   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(
+    #   df.group_by_rolling(index_column: "dt", period: "2d").agg(
     #     [
     #       Polars.sum("a").alias("sum_a"),
     #       Polars.min("a").alias("min_a"),
@@ -979,7 +988,7 @@ module Polars
     #   # │ 2020-01-03 19:45:32 ┆ 11    ┆ 2     ┆ 9     │
     #   # │ 2020-01-08 23:16:43 ┆ 1     ┆ 1     ┆ 1     │
     #   # └─────────────────────┴───────┴───────┴───────┘
-    def groupby_rolling(
+    def group_by_rolling(
       index_column:,
       period:,
       offset: nil,
@@ -987,7 +996,7 @@ module Polars
       by: nil,
       check_sorted: true
     )
-      index_column = Utils.expr_to_lit_or_expr(index_column, str_to_lit: false)
+      index_column = Utils.parse_as_expression(index_column)
       if offset.nil?
         offset = "-#{period}"
       end
@@ -996,16 +1005,17 @@ module Polars
       period = Utils._timedelta_to_pl_duration(period)
       offset = Utils._timedelta_to_pl_duration(offset)
 
-      lgb = _ldf.groupby_rolling(
-        index_column._rbexpr, period, offset, closed, rbexprs_by, check_sorted
+      lgb = _ldf.group_by_rolling(
+        index_column, period, offset, closed, rbexprs_by, check_sorted
       )
-      LazyGroupBy.new(lgb, self.class)
+      LazyGroupBy.new(lgb)
     end
+    alias_method :groupby_rolling, :group_by_rolling
 
     # Group based on a time value (or index value of type `:i32`, `:i64`).
     #
     # Time windows are calculated and rows are assigned to windows. Different from a
-    # normal groupby is that a row can be member of multiple groups. The time/index
+    # normal group by is that a row can be member of multiple groups. The time/index
     # window could be seen as a rolling window, with a window size determined by
     # dates/times/values instead of slots in the DataFrame.
     #
@@ -1033,37 +1043,43 @@ module Polars
     # Or combine them:
     # "3d12h4m25s" # 3 days, 12 hours, 4 minutes, and 25 seconds
     #
-    # In case of a groupby_dynamic on an integer column, the windows are defined by:
+    # In case of a group_by_dynamic on an integer column, the windows are defined by:
     #
     # - "1i"      # length 1
     # - "10i"     # length 10
     #
-    # @param index_column
+    # @param index_column [Object]
     #   Column used to group based on the time window.
     #   Often to type Date/Datetime
     #   This column must be sorted in ascending order. If not the output will not
     #   make sense.
     #
-    #   In case of a dynamic groupby on indices, dtype needs to be one of
+    #   In case of a dynamic group by on indices, dtype needs to be one of
     #   `:i32`, `:i64`. Note that `:i32` gets temporarily cast to `:i64`, so if
     #   performance matters use an `:i64` column.
-    # @param every
+    # @param every [Object]
     #   Interval of the window.
-    # @param period
+    # @param period [Object]
     #   Length of the window, if None it is equal to 'every'.
-    # @param offset
+    # @param offset [Object]
     #   Offset of the window if None and period is None it will be equal to negative
     #   `every`.
-    # @param truncate
+    # @param truncate [Boolean]
     #   Truncate the time value to the window lower bound.
-    # @param include_boundaries
+    # @param include_boundaries [Boolean]
     #   Add the lower and upper bound of the window to the "_lower_bound" and
     #   "_upper_bound" columns. This will impact performance because it's harder to
     #   parallelize
     # @param closed ["right", "left", "both", "none"]
     #   Define whether the temporal window interval is closed or not.
-    # @param by
+    # @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 [DataFrame]
     #
@@ -1095,7 +1111,7 @@ module Polars
     #   # └─────────────────────┴─────┘
     #
     # @example Group by windows of 1 hour starting at 2021-12-16 00:00:00.
-    #   df.groupby_dynamic("time", every: "1h", closed: "right").agg(
+    #   df.group_by_dynamic("time", every: "1h", closed: "right").agg(
     #     [
     #       Polars.col("time").min.alias("time_min"),
     #       Polars.col("time").max.alias("time_max")
@@ -1115,7 +1131,7 @@ module Polars
     #   # └─────────────────────┴─────────────────────┴─────────────────────┘
     #
     # @example The window boundaries can also be added to the aggregation result.
-    #   df.groupby_dynamic(
+    #   df.group_by_dynamic(
     #     "time", every: "1h", include_boundaries: true, closed: "right"
     #   ).agg([Polars.col("time").count.alias("time_count")])
     #   # =>
@@ -1132,7 +1148,7 @@ module Polars
     #   # └─────────────────────┴─────────────────────┴─────────────────────┴────────────┘
     #
     # @example When closed="left", should not include right end of interval.
-    #   df.groupby_dynamic("time", every: "1h", closed: "left").agg(
+    #   df.group_by_dynamic("time", every: "1h", closed: "left").agg(
     #     [
     #       Polars.col("time").count.alias("time_count"),
     #       Polars.col("time").alias("time_agg_list")
@@ -1152,7 +1168,7 @@ module Polars
     #   # └─────────────────────┴────────────┴───────────────────────────────────┘
     #
     # @example When closed="both" the time values at the window boundaries belong to 2 groups.
-    #   df.groupby_dynamic("time", every: "1h", closed: "both").agg(
+    #   df.group_by_dynamic("time", every: "1h", closed: "both").agg(
     #     [Polars.col("time").count.alias("time_count")]
     #   )
     #   # =>
@@ -1169,7 +1185,7 @@ module Polars
     #   # │ 2021-12-16 03:00:00 ┆ 1          │
     #   # └─────────────────────┴────────────┘
     #
-    # @example Dynamic groupbys can also be combined with grouping on normal keys.
+    # @example Dynamic group bys can also be combined with grouping on normal keys.
     #   df = Polars::DataFrame.new(
     #     {
     #       "time" => Polars.date_range(
@@ -1180,7 +1196,7 @@ module Polars
     #       "groups" => ["a", "a", "a", "b", "b", "a", "a"]
     #     }
     #   )
-    #   df.groupby_dynamic(
+    #   df.group_by_dynamic(
     #     "time",
     #     every: "1h",
     #     closed: "both",
@@ -1203,14 +1219,14 @@ module Polars
     #   # │ b      ┆ 2021-12-16 02:00:00 ┆ 2021-12-16 03:00:00 ┆ 2021-12-16 02:00:00 ┆ 1          │
     #   # └────────┴─────────────────────┴─────────────────────┴─────────────────────┴────────────┘
     #
-    # @example Dynamic groupby on an index column.
+    # @example Dynamic group by on an index column.
     #   df = Polars::DataFrame.new(
     #     {
     #       "idx" => Polars.arange(0, 6, eager: true),
     #       "A" => ["A", "A", "B", "B", "B", "C"]
     #     }
     #   )
-    #   df.groupby_dynamic(
+    #   df.group_by_dynamic(
     #     "idx",
     #     every: "2i",
     #     period: "3i",
@@ -1228,17 +1244,23 @@ module Polars
     #   # │ 2               ┆ 5               ┆ 2   ┆ ["B", "B", "C"] │
     #   # │ 4               ┆ 7               ┆ 4   ┆ ["C"]           │
     #   # └─────────────────┴─────────────────┴─────┴─────────────────┘
-    def groupby_dynamic(
+    def group_by_dynamic(
       index_column,
       every:,
       period: nil,
       offset: nil,
-      truncate: true,
+      truncate: nil,
       include_boundaries: false,
       closed: "left",
+      label: "left",
       by: nil,
-      start_by: "window"
+      start_by: "window",
+      check_sorted: true
     )
+      if !truncate.nil?
+        label = truncate ? "left" : "datapoint"
+      end
+
       index_column = Utils.expr_to_lit_or_expr(index_column, str_to_lit: false)
       if offset.nil?
         offset = period.nil? ? "-#{every}" : "0ns"
@@ -1253,19 +1275,21 @@ module Polars
       every = Utils._timedelta_to_pl_duration(every)
 
       rbexprs_by = by.nil? ? [] : Utils.selection_to_rbexpr_list(by)
-      lgb = _ldf.groupby_dynamic(
+      lgb = _ldf.group_by_dynamic(
         index_column._rbexpr,
         every,
         period,
         offset,
-        truncate,
+        label,
         include_boundaries,
         closed,
         rbexprs_by,
-        start_by
+        start_by,
+        check_sorted
       )
-      LazyGroupBy.new(lgb, self.class)
+      LazyGroupBy.new(lgb)
     end
+    alias_method :groupby_dynamic, :group_by_dynamic
 
     # Perform an asof join.
     #
@@ -1347,7 +1371,7 @@ module Polars
         raise ArgumentError, "Expected a `LazyFrame` as join table, got #{other.class.name}"
       end
 
-      if on.is_a?(String)
+      if on.is_a?(::String)
         left_on = on
         right_on = on
       end
@@ -1356,19 +1380,19 @@ module Polars
         raise ArgumentError, "You should pass the column to join on as an argument."
       end
 
-      if by_left.is_a?(String) || by_left.is_a?(Expr)
+      if by_left.is_a?(::String) || by_left.is_a?(Expr)
         by_left_ = [by_left]
       else
         by_left_ = by_left
       end
 
-      if by_right.is_a?(String) || by_right.is_a?(Expr)
+      if by_right.is_a?(::String) || by_right.is_a?(Expr)
         by_right_ = [by_right]
       else
         by_right_ = by_right
       end
 
-      if by.is_a?(String)
+      if by.is_a?(::String)
         by_left_ = [by]
         by_right_ = [by]
       elsif by.is_a?(::Array)
@@ -1378,7 +1402,7 @@ module Polars
 
       tolerance_str = nil
       tolerance_num = nil
-      if tolerance.is_a?(String)
+      if tolerance.is_a?(::String)
         tolerance_str = tolerance
       else
         tolerance_num = tolerance
@@ -1454,17 +1478,17 @@ module Polars
     # @example
     #   df.join(other_df, on: "ham", how: "outer").collect
     #   # =>
-    #   # shape: (4, 4)
-    #   # ┌──────┬──────┬─────┬───────┐
-    #   # │ foo  ┆ bar  ┆ ham ┆ apple │
-    #   # │ ---  ┆ ---  ┆ --- ┆ ---   │
-    #   # │ i64  ┆ f64  ┆ str ┆ str   │
-    #   # ╞══════╪══════╪═════╪═══════╡
-    #   # │ 1    ┆ 6.0  ┆ a   ┆ x     │
-    #   # │ 2    ┆ 7.0  ┆ b   ┆ y     │
-    #   # │ null ┆ null ┆ d   ┆ z     │
-    #   # │ 3    ┆ 8.0  ┆ c   ┆ null  │
-    #   # └──────┴──────┴─────┴───────┘
+    #   # shape: (4, 5)
+    #   # ┌──────┬──────┬──────┬───────┬───────────┐
+    #   # │ foo  ┆ bar  ┆ ham  ┆ apple ┆ ham_right │
+    #   # │ ---  ┆ ---  ┆ ---  ┆ ---   ┆ ---       │
+    #   # │ i64  ┆ f64  ┆ str  ┆ str   ┆ str       │
+    #   # ╞══════╪══════╪══════╪═══════╪═══════════╡
+    #   # │ 1    ┆ 6.0  ┆ a    ┆ x     ┆ a         │
+    #   # │ 2    ┆ 7.0  ┆ b    ┆ y     ┆ b         │
+    #   # │ null ┆ null ┆ null ┆ z     ┆ d         │
+    #   # │ 3    ┆ 8.0  ┆ c    ┆ null  ┆ null      │
+    #   # └──────┴──────┴──────┴───────┴───────────┘
     #
     # @example
     #   df.join(other_df, on: "ham", how: "left").collect
@@ -1698,7 +1722,7 @@ module Polars
     #
     # @return [LazyFrame]
     def drop(columns)
-      if columns.is_a?(String)
+      if columns.is_a?(::String)
         columns = [columns]
       end
       _from_rbldf(_ldf.drop_columns(columns))
@@ -1725,8 +1749,10 @@ module Polars
 
     # Shift the values by a given period.
     #
-    # @param periods [Integer]
+    # @param n [Integer]
     #   Number of places to shift (may be negative).
+    # @param fill_value [Object]
+    #   Fill the resulting null values with this value.
     #
     # @return [LazyFrame]
     #
@@ -1763,8 +1789,12 @@ module Polars
     #   # │ 5    ┆ 6    │
     #   # │ null ┆ null │
     #   # └──────┴──────┘
-    def shift(periods)
-      _from_rbldf(_ldf.shift(periods))
+    def shift(n, fill_value: nil)
+      if !fill_value.nil?
+        fill_value = Utils.parse_as_expression(fill_value, str_as_lit: true)
+      end
+      n = Utils.parse_as_expression(n)
+      _from_rbldf(_ldf.shift(n, fill_value))
     end
 
     # Shift the values by a given period and fill the resulting null values.
@@ -1810,10 +1840,7 @@ module Polars
     #   # │ 0   ┆ 0   │
     #   # └─────┴─────┘
     def shift_and_fill(periods, fill_value)
-      if !fill_value.is_a?(Expr)
-        fill_value = Polars.lit(fill_value)
-      end
-      _from_rbldf(_ldf.shift_and_fill(periods, fill_value._rbexpr))
+      shift(periods, fill_value: fill_value)
     end
 
     # Get a slice of this DataFrame.
@@ -2336,10 +2363,10 @@ module Polars
     #   # │ z   ┆ c        ┆ 6     │
     #   # └─────┴──────────┴───────┘
     def melt(id_vars: nil, value_vars: nil, variable_name: nil, value_name: nil, streamable: true)
-      if value_vars.is_a?(String)
+      if value_vars.is_a?(::String)
         value_vars = [value_vars]
       end
-      if id_vars.is_a?(String)
+      if id_vars.is_a?(::String)
         id_vars = [id_vars]
       end
       if value_vars.nil?
@@ -2371,16 +2398,16 @@ module Polars
     #   df.interpolate.collect
     #   # =>
     #   # shape: (4, 3)
-    #   # ┌─────┬──────┬─────┐
-    #   # │ foo ┆ bar  ┆ baz │
-    #   # │ --- ┆ ---  ┆ --- │
-    #   # │ i64 ┆ i64  ┆ i64 │
-    #   # ╞═════╪══════╪═════╡
-    #   # │ 1   ┆ 6    ┆ 1   │
-    #   # │ 5   ┆ 7    ┆ 3   │
-    #   # │ 9   ┆ 9    ┆ 6   │
-    #   # │ 10  ┆ null ┆ 9   │
-    #   # └─────┴──────┴─────┘
+    #   # ┌──────┬──────┬──────────┐
+    #   # │ foo  ┆ bar  ┆ baz      │
+    #   # │ ---  ┆ ---  ┆ ---      │
+    #   # │ f64  ┆ f64  ┆ f64      │
+    #   # ╞══════╪══════╪══════════╡
+    #   # │ 1.0  ┆ 6.0  ┆ 1.0      │
+    #   # │ 5.0  ┆ 7.0  ┆ 3.666667 │
+    #   # │ 9.0  ┆ 9.0  ┆ 6.333333 │
+    #   # │ 10.0 ┆ null ┆ 9.0      │
+    #   # └──────┴──────┴──────────┘
     def interpolate
       select(Utils.col("*").interpolate)
     end
@@ -2437,7 +2464,7 @@ module Polars
     #   # │ bar    ┆ 2   ┆ b   ┆ null ┆ [3]       ┆ womp  │
     #   # └────────┴─────┴─────┴──────┴───────────┴───────┘
     def unnest(names)
-      if names.is_a?(String)
+      if names.is_a?(::String)
         names = [names]
       end
       _from_rbldf(_ldf.unnest(names))