polars-df 0.4.0-x86_64-linux → 0.6.0-x86_64-linux

data/lib/polars/lazy_frame.rb

@@ -1,9 +1,25 @@
 module Polars
-  # Representation of a Lazy computation graph/query againat a DataFrame.
+  # Representation of a Lazy computation graph/query against a DataFrame.
   class LazyFrame
     # @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(: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,8 +984,10 @@ 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?
         offset = "-#{period}"
       end
@@ -973,7 +997,7 @@ module Polars
       offset = Utils._timedelta_to_pl_duration(offset)
 
       lgb = _ldf.groupby_rolling(
-        index_column, period, offset, closed, rbexprs_by
+        index_column._rbexpr, period, offset, closed, rbexprs_by, check_sorted
       )
       LazyGroupBy.new(lgb, self.class)
     end
@@ -1111,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(
@@ -1192,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)
     #   # ┌─────────────────┬─────────────────┬─────┬─────────────────┐
@@ -1215,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?
@@ -1233,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,
@@ -1350,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
@@ -1618,7 +1639,7 @@ module Polars
     #   # │ null │
     #   # └──────┘
     def with_context(other)
-      if !other.is_a?(Array)
+      if !other.is_a?(::Array)
         other = [other]
       end
 
@@ -2227,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))
@@ -2260,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))
@@ -2422,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)

data/lib/polars/lazy_functions.rb

@@ -14,7 +14,7 @@ module Polars
 
       if name.is_a?(DataType)
         Utils.wrap_expr(_dtype_cols([name]))
-      elsif name.is_a?(Array)
+      elsif name.is_a?(::Array)
         if name.length == 0 || Utils.strlike?(name[0])
           name = name.map { |v| v.is_a?(Symbol) ? v.to_s : v }
           Utils.wrap_expr(RbExpr.cols(name))
@@ -36,7 +36,7 @@ module Polars
     # @example A horizontal rank computation by taking the elements of a list
     #   df = Polars::DataFrame.new({"a" => [1, 8, 3], "b" => [4, 5, 2]})
     #   df.with_column(
-    #     Polars.concat_list(["a", "b"]).arr.eval(Polars.element.rank).alias("rank")
+    #     Polars.concat_list(["a", "b"]).list.eval(Polars.element.rank).alias("rank")
     #   )
     #   # =>
     #   # shape: (3, 3)
@@ -156,9 +156,8 @@ module Polars
         column.sum
       elsif Utils.strlike?(column)
         col(column.to_s).sum
-      elsif column.is_a?(Array)
+      elsif column.is_a?(::Array)
         exprs = Utils.selection_to_rbexpr_list(column)
-        # TODO
         Utils.wrap_expr(_sum_exprs(exprs))
       else
         fold(lit(0).cast(:u32), ->(a, b) { a + b }, column).alias("sum")
@@ -283,18 +282,33 @@ module Polars
     # Return an expression representing a literal value.
     #
     # @return [Expr]
-    def lit(value)
-      if value.is_a?(Polars::Series)
+    def lit(value, dtype: nil, allow_object: nil)
+      if value.is_a?(::Time) || value.is_a?(::DateTime)
+        time_unit = dtype&.time_unit || "ns"
+        time_zone = dtype.&time_zone
+        e = lit(Utils._datetime_to_pl_timestamp(value, time_unit)).cast(Datetime.new(time_unit))
+        if time_zone
+          return e.dt.replace_time_zone(time_zone.to_s)
+        else
+          return e
+        end
+      elsif value.is_a?(::Date)
+        return lit(::Time.utc(value.year, value.month, value.day)).cast(Date)
+      elsif value.is_a?(Polars::Series)
         name = value.name
         value = value._s
-        e = Utils.wrap_expr(RbExpr.lit(value))
+        e = Utils.wrap_expr(RbExpr.lit(value, allow_object))
         if name == ""
           return e
         end
         return e.alias(name)
+      elsif (defined?(Numo::NArray) && value.is_a?(Numo::NArray)) || value.is_a?(::Array)
+        return lit(Series.new("", value))
+      elsif dtype
+        return Utils.wrap_expr(RbExpr.lit(value, allow_object)).cast(dtype)
       end
 
-      Utils.wrap_expr(RbExpr.lit(value))
+      Utils.wrap_expr(RbExpr.lit(value, allow_object))
     end
 
     # Cumulatively sum values in a column/Series, or horizontally across list of columns/expressions.
@@ -625,23 +639,42 @@ module Polars
     # @return [Expr, Series]
     #
     # @example
-    #   df.lazy.filter(Polars.col("foo") < Polars.arange(0, 100)).collect
+    #   Polars.arange(0, 3, eager: true)
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'arange' [i64]
+    #   # [
+    #   #         0
+    #   #         1
+    #   #         2
+    #   # ]
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"a" => [1, 2], "b" => [3, 4]})
+    #   df.select(Polars.arange(Polars.col("a"), Polars.col("b")))
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌───────────┐
+    #   # │ arange    │
+    #   # │ ---       │
+    #   # │ list[i64] │
+    #   # ╞═══════════╡
+    #   # │ [1, 2]    │
+    #   # │ [2, 3]    │
+    #   # └───────────┘
     def arange(low, high, step: 1, eager: false, dtype: nil)
       low = Utils.expr_to_lit_or_expr(low, str_to_lit: false)
       high = Utils.expr_to_lit_or_expr(high, str_to_lit: false)
       range_expr = Utils.wrap_expr(RbExpr.arange(low._rbexpr, high._rbexpr, step))
 
-      if !dtype.nil? && dtype != "i64"
+      if !dtype.nil? && !["i64", Int64].include?(dtype)
         range_expr = range_expr.cast(dtype)
       end
 
       if !eager
         range_expr
       else
-        DataFrame.new
-          .select(range_expr)
-          .to_series
-          .rename("arange", in_place: true)
+        DataFrame.new.select(range_expr.alias("arange")).to_series
       end
     end
 
@@ -658,7 +691,7 @@ module Polars
     #
     # @return [Expr]
     def arg_sort_by(exprs, reverse: false)
-      if !exprs.is_a?(Array)
+      if !exprs.is_a?(::Array)
         exprs = [exprs]
       end
       if reverse == true || reverse == false
@@ -997,19 +1030,24 @@ module Polars
     #   Only used in `eager` mode. As expression, use `alias`.
     #
     # @return [Expr]
-    def repeat(value, n, eager: false, name: nil)
+    def repeat(value, n, dtype: nil, eager: false, name: nil)
+      if !name.nil?
+        warn "the `name` argument is deprecated. Use the `alias` method instead."
+      end
+
+      if n.is_a?(Integer)
+        n = lit(n)
+      end
+
+      value = Utils.parse_as_expression(value, str_as_lit: true)
+      expr = Utils.wrap_expr(RbExpr.repeat(value, n._rbexpr, dtype))
+      if !name.nil?
+        expr = expr.alias(name)
+      end
       if eager
-        if name.nil?
-          name = ""
-        end
-        dtype = py_type_to_dtype(type(value))
-        Series._repeat(name, value, n, dtype)
-      else
-        if n.is_a?(Integer)
-          n = lit(n)
-        end
-        Utils.wrap_expr(RbExpr.repeat(value, n._rbexpr))
+        return select(expr).to_series
       end
+      expr
     end
 
     # Return indices where `condition` evaluates `true`.
@@ -1077,8 +1115,11 @@ module Polars
     #   # │ null ┆ null ┆ 3.0  ┆ 3.0  │
     #   # │ null ┆ null ┆ null ┆ 99.9 │
     #   # └──────┴──────┴──────┴──────┘
-    def coalesce(exprs)
+    def coalesce(exprs, *more_exprs)
       exprs = Utils.selection_to_rbexpr_list(exprs)
+      if more_exprs.any?
+        exprs.concat(Utils.selection_to_rbexpr_list(more_exprs))
+      end
       Utils.wrap_expr(_coalesce_exprs(exprs))
     end
 
@@ -1121,13 +1162,11 @@ module Polars
       end
 
       if unit == "d"
-        expr = column.cast(:date)
+        expr = column.cast(Date)
       elsif unit == "s"
-        raise Todo
-        # expr = (column.cast(:i64) * 1_000_000).cast(Datetime("us"))
+        expr = (column.cast(Int64) * 1_000_000).cast(Datetime.new("us"))
       elsif Utils::DTYPE_TEMPORAL_UNITS.include?(unit)
-        raise Todo
-        # expr = column.cast(Datetime(unit))
+        expr = column.cast(Datetime.new(unit))
       else
         raise ArgumentError, "'unit' must be one of {{'ns', 'us', 'ms', 's', 'd'}}, got '#{unit}'."
       end