polars-df 0.23.0 → 0.24.0

data/lib/polars/list_name_space.rb

@@ -77,7 +77,6 @@ module Polars
     def len
       super
     end
-    alias_method :lengths, :len
 
     # Drop all null values in the list.
     #
@@ -259,6 +258,11 @@ module Polars
 
     # Sort the arrays in the list.
     #
+    # @param descending [Boolean]
+    #   Sort in descending order.
+    # @param nulls_last [Boolean]
+    #   Place null values last.
+    #
     # @return [Series]
     #
     # @example
@@ -273,7 +277,7 @@ module Polars
     #   # ]
     #
     # @example
-    #   s.list.sort(reverse: true)
+    #   s.list.sort(descending: true)
     #   # =>
     #   # shape: (2,)
     #   # Series: 'a' [list[i64]]
@@ -281,7 +285,7 @@ module Polars
     #   #         [3, 2, 1]
     #   #         [9, 2, 1]
     #   # ]
-    def sort(reverse: false)
+    def sort(descending: false, nulls_last: false)
       super
     end
 
@@ -305,6 +309,9 @@ module Polars
 
     # Get the unique/distinct values in the list.
     #
+    # @param maintain_order [Boolean]
+    #   Maintain order of data. This requires more work.
+    #
     # @return [Series]
     #
     # @example
@@ -317,7 +324,7 @@ module Polars
     #   #         [1, 2]
     #   #         [2, 3]
     #   # ]
-    def unique
+    def unique(maintain_order: false)
       super
     end
 
@@ -400,8 +407,8 @@ module Polars
     #   Indices to return per sublist
     # @param null_on_oob [Boolean]
     #   Behavior if an index is out of bounds:
-    #   True -> set as null
-    #   False -> raise an error
+    #   true -> set as null
+    #   false -> raise an error
     #   Note that defaulting to raising an error is much cheaper
     #
     # @return [Series]
@@ -461,6 +468,11 @@ module Polars
     #
     # @param separator [String]
     #   string to separate the items with
+    # @param ignore_nulls [Boolean]
+    #   Ignore null values (default).
+    #
+    #   If set to `false`, null values will be propagated.
+    #   If the sub-list contains any null values, the output is `nil`.
     #
     # @return [Series]
     #
@@ -474,7 +486,7 @@ module Polars
     #   #         "foo-bar"
     #   #         "hello-world"
     #   # ]
-    def join(separator)
+    def join(separator, ignore_nulls: true)
       super
     end
 
@@ -516,10 +528,33 @@ module Polars
       super
     end
 
+    # Get the single value of the sublists.
+    #
+    # This errors if the sublist length is not exactly one.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("a", [[1], [4], [6]])
+    #   s.list.item
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'a' [i64]
+    #   # [
+    #   #         1
+    #   #         4
+    #   #         6
+    #   # ]
+    def item
+      super
+    end
+
     # Check if sublists contain the given item.
     #
     # @param item [Object]
     #   Item that will be checked for membership.
+    # @param nulls_equal [Boolean]
+    #   If true, treat null as a distinct value. Null values will not propagate.
     #
     # @return [Series]
     #
@@ -534,7 +569,7 @@ module Polars
     #   #         false
     #   #         true
     #   # ]
-    def contains(item)
+    def contains(item, nulls_equal: true)
       super
     end
 
@@ -599,7 +634,7 @@ module Polars
 
     # Shift values by the given period.
     #
-    # @param periods [Integer]
+    # @param n [Integer]
     #   Number of places to shift (may be negative).
     #
     # @return [Series]
@@ -614,7 +649,7 @@ module Polars
     #   #         [null, 1, … 3]
     #   #         [null, 10, 2]
     #   # ]
-    def shift(periods = 1)
+    def shift(n = 1)
       super
     end
 
@@ -799,22 +834,51 @@ module Polars
     # @return [Series]
     #
     # @example
-    #   df = Polars::DataFrame.new({"a" => [1, 8, 3], "b" => [4, 5, 2]})
-    #   df.with_column(
-    #     Polars.concat_list(["a", "b"]).list.eval(Polars.element.rank).alias("rank")
-    #   )
+    #   s = Polars::Series.new("a", [[1, 4], [8, 5], [3, 2]])
+    #   s.list.eval(Polars.element.rank)
     #   # =>
-    #   # shape: (3, 3)
-    #   # ┌─────┬─────┬────────────┐
-    #   # │ a   ┆ b   ┆ rank       │
-    #   # │ --- ┆ --- ┆ ---        │
-    #   # │ i64 ┆ i64 ┆ list[f64]  │
-    #   # ╞═════╪═════╪════════════╡
-    #   # │ 1   ┆ 4   ┆ [1.0, 2.0] │
-    #   # │ 8   ┆ 5   ┆ [2.0, 1.0] │
-    #   # │ 3   ┆ 2   ┆ [2.0, 1.0] │
-    #   # └─────┴─────┴────────────┘
+    #   # shape: (3,)
+    #   # Series: 'a' [list[f64]]
+    #   # [
+    #   #         [1.0, 2.0]
+    #   #         [2.0, 1.0]
+    #   #         [2.0, 1.0]
+    #   # ]
     def eval(expr)
+      s = Utils.wrap_s(_s)
+      s.to_frame.select(F.col(s.name).list.eval(expr)).to_series
+    end
+
+    # Run any polars aggregation expression against the list' elements.
+    #
+    # @param expr [Expr]
+    #   Expression to run. Note that you can select an element with `Polars.element`.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("a", [[1, nil], [42, 13], [nil, nil]])
+    #   s.list.agg(Polars.element.null_count)
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'a' [u32]
+    #   # [
+    #   #         1
+    #   #         0
+    #   #         2
+    #   # ]
+    #
+    # @example
+    #   s.list.agg(Polars.element.drop_nulls)
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'a' [list[i64]]
+    #   # [
+    #   #         [1]
+    #   #         [42, 13]
+    #   #         []
+    #   # ]
+    def agg(expr)
       super
     end
 

data/lib/polars/meta_expr.rb

@@ -180,8 +180,13 @@ module Polars
     #   # => "foo"
     #   Polars.len.meta.output_name
     #   # => "len"
-    def output_name
+    def output_name(raise_if_undetermined: true)
       _rbexpr.meta_output_name
+    rescue Polars::ComputeError
+      if !raise_if_undetermined
+        return nil
+      end
+      raise
     end
 
     # Pop the latest expression and return the input(s) of the popped expression.
@@ -292,5 +297,41 @@ module Polars
         nil
       end
     end
+
+    # Format the expression as a Graphviz graph.
+    #
+    # Note that Graphviz must be installed to render the visualization (if not
+    # already present, you can download it here: https://graphviz.org/download.
+    #
+    # @param show [Boolean]
+    #   Show the figure.
+    # @param output_path [String]
+    #   Write the figure to disk.
+    # @param raw_output [Boolean]
+    #   Return dot syntax. This cannot be combined with `show` and/or `output_path`.
+    # @param figsize [Array]
+    #   Passed to matplotlib if `show == true`.
+    #
+    # @return [Object]
+    #
+    # @example
+    #   e = (Polars.col("foo") * Polars.col("bar")).sum.over(Polars.col("ham")) / 2
+    #   e.meta.show_graph
+    def show_graph(
+      show: true,
+      output_path: nil,
+      raw_output: false,
+      figsize: [16.0, 12.0],
+      schema: nil
+    )
+      dot = _rbexpr.meta_show_graph(schema)
+      Utils.display_dot_graph(
+        dot: dot,
+        show: show,
+        output_path: output_path,
+        raw_output: raw_output,
+        figsize: figsize
+      )
+    end
   end
 end

data/lib/polars/name_expr.rb

@@ -77,8 +77,11 @@ module Polars
     #   # │ 2         ┆ y         ┆ 2   ┆ y   │
     #   # │ 1         ┆ x         ┆ 3   ┆ z   │
     #   # └───────────┴───────────┴─────┴─────┘
-    def map(&f)
-      Utils.wrap_expr(_rbexpr.name_map(f))
+    def map(&function)
+      # need to mark function for GC
+      raise Todo
+
+      Utils.wrap_expr(_rbexpr.name_map(function))
     end
 
     # Add a prefix to the root column name of the expression.
@@ -208,11 +211,45 @@ module Polars
     # @example
     #   df = Polars::DataFrame.new({"x" => {"a" => 1, "b" => 2}})
     #   df.select(Polars.col("x").name.prefix_fields("prefix_")).schema
-    #   # => {"x"=>Polars::Struct({"prefix_a"=>Polars::Int64, "prefix_b"=>Polars::Int64})}
+    #   # => Polars::Schema({"x"=>Polars::Struct({"prefix_a"=>Polars::Int64, "prefix_b"=>Polars::Int64})})
     def prefix_fields(prefix)
       Utils.wrap_expr(_rbexpr.name_prefix_fields(prefix))
     end
 
+    # Replace matching regex/literal substring in the name with a new value.
+    #
+    # @param pattern [String]
+    #   A valid regular expression pattern, compatible with the [regex crate](https://docs.rs/regex/latest/regex/).
+    # @param value [String]
+    #   String that will replace the matched substring.
+    # @param literal [Boolean]
+    #   Treat `pattern` as a literal string, not a regex.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "n_foo" => [1, 2, 3],
+    #       "n_bar" => ["x", "y", "z"]
+    #     }
+    #   )
+    #   df.select(Polars.all.name.replace("^n_", "col_"))
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌─────────┬─────────┐
+    #   # │ col_foo ┆ col_bar │
+    #   # │ ---     ┆ ---     │
+    #   # │ i64     ┆ str     │
+    #   # ╞═════════╪═════════╡
+    #   # │ 1       ┆ x       │
+    #   # │ 2       ┆ y       │
+    #   # │ 3       ┆ z       │
+    #   # └─────────┴─────────┘
+    def replace(pattern, value, literal: false)
+      Utils.wrap_expr(_rbexpr.name_replace(pattern, value, literal))
+    end
+
     # Add a suffix to all field names of a struct.
     #
     # @note
@@ -226,7 +263,7 @@ module Polars
     # @example
     #   df = Polars::DataFrame.new({"x" => {"a" => 1, "b" => 2}})
     #   df.select(Polars.col("x").name.suffix_fields("_suffix")).schema
-    #   # => {"x"=>Polars::Struct({"a_suffix"=>Polars::Int64, "b_suffix"=>Polars::Int64})}
+    #   # => Polars::Schema({"x"=>Polars::Struct({"a_suffix"=>Polars::Int64, "b_suffix"=>Polars::Int64})})
     def suffix_fields(suffix)
       Utils.wrap_expr(_rbexpr.name_suffix_fields(suffix))
     end

data/lib/polars/query_opt_flags.rb

@@ -5,6 +5,9 @@ module Polars
   #   This functionality is considered **unstable**. It may be changed
   #   at any point without it being considered a breaking change.
   class QueryOptFlags
+    # @private
+    attr_accessor :_rboptflags
+
     def initialize(
       predicate_pushdown: nil,
       projection_pushdown: nil,
@@ -17,7 +20,7 @@ module Polars
       check_order_observe: nil,
       fast_projection: nil
     )
-      @_rboptflags = RbOptFlags.default
+      self._rboptflags = RbOptFlags.default
       update(
         predicate_pushdown: predicate_pushdown,
         projection_pushdown: projection_pushdown,
@@ -32,6 +35,35 @@ module Polars
       )
     end
 
+    # Create new empty set off optimizations.
+    def self.none(
+      predicate_pushdown: nil,
+      projection_pushdown: nil,
+      simplify_expression: nil,
+      slice_pushdown: nil,
+      comm_subplan_elim: nil,
+      comm_subexpr_elim: nil,
+      cluster_with_columns: nil,
+      collapse_joins: nil,
+      check_order_observe: nil,
+      fast_projection: nil
+    )
+      optflags = QueryOptFlags.new
+      optflags.no_optimizations
+      optflags.update(
+        predicate_pushdown: predicate_pushdown,
+        projection_pushdown: projection_pushdown,
+        simplify_expression: simplify_expression,
+        slice_pushdown: slice_pushdown,
+        comm_subplan_elim: comm_subplan_elim,
+        comm_subexpr_elim: comm_subexpr_elim,
+        cluster_with_columns: cluster_with_columns,
+        collapse_joins: collapse_joins,
+        check_order_observe: check_order_observe,
+        fast_projection: fast_projection
+      )
+    end
+
     def update(
       predicate_pushdown: nil,
       projection_pushdown: nil,
@@ -44,7 +76,171 @@ module Polars
       check_order_observe: nil,
       fast_projection: nil
     )
-      raise Todo
+      if !predicate_pushdown.nil?
+        self.predicate_pushdown = predicate_pushdown
+      end
+      if !projection_pushdown.nil?
+        self.projection_pushdown = projection_pushdown
+      end
+      if !simplify_expression.nil?
+        self.simplify_expression = simplify_expression
+      end
+      if !slice_pushdown.nil?
+        self.slice_pushdown = slice_pushdown
+      end
+      if !comm_subplan_elim.nil?
+        self.comm_subplan_elim = comm_subplan_elim
+      end
+      if !comm_subexpr_elim.nil?
+        self.comm_subexpr_elim = comm_subexpr_elim
+      end
+      if !cluster_with_columns.nil?
+        self.cluster_with_columns = cluster_with_columns
+      end
+      if !collapse_joins.nil?
+        Utils.issue_deprecation_warning(
+          "the `collapse_joins` parameter for `QueryOptFlags` is deprecated. " +
+          "Use `predicate_pushdown` instead."
+        )
+        if !collapse_joins
+          self.predicate_pushdown = false
+        end
+      end
+      if !check_order_observe.nil?
+        self.check_order_observe = check_order_observe
+      end
+      if !fast_projection.nil?
+        self.fast_projection = fast_projection
+      end
+
+      self
+    end
+
+    # Create new empty set off optimizations.
+    def self._eager
+      optflags = QueryOptFlags.new
+      optflags.no_optimizations
+      optflags._rboptflags.eager = true
+      optflags.simplify_expression = true
+      optflags
+    end
+
+    # Remove selected optimizations.
+    def no_optimizations
+      _rboptflags.no_optimizations
+    end
+
+    # Only read columns that are used later in the query.
+    def projection_pushdown
+      _rboptflags.projection_pushdown
+    end
+
+    def projection_pushdown=(value)
+      _rboptflags.projection_pushdown = value
+    end
+
+    # Apply predicates/filters as early as possible.
+    def predicate_pushdown
+      _rboptflags.predicate_pushdown
+    end
+
+    def predicate_pushdown=(value)
+      _rboptflags.predicate_pushdown = value
+    end
+
+    # Cluster sequential `with_columns` calls to independent calls.
+    def cluster_with_columns
+      _rboptflags.cluster_with_columns
+    end
+
+    def cluster_with_columns=(value)
+      _rboptflags.cluster_with_columns = value
+    end
+
+    # Run many expression optimization rules until fixed point.
+    def simplify_expression
+      _rboptflags.simplify_expression
+    end
+
+    def simplify_expression=(value)
+      _rboptflags.simplify_expression = value
+    end
+
+    # Pushdown slices/limits.
+    def slice_pushdown
+      _rboptflags.slice_pushdown
+    end
+
+    def slice_pushdown=(value)
+      _rboptflags.slice_pushdown = value
+    end
+
+    # Elide duplicate plans and caches their outputs.
+    def comm_subplan_elim
+      _rboptflags.comm_subplan_elim
+    end
+
+    def comm_subplan_elim=(value)
+      _rboptflags.comm_subplan_elim = value
+    end
+
+    # Elide duplicate expressions and caches their outputs.
+    def comm_subexpr_elim
+      _rboptflags.comm_subexpr_elim
+    end
+
+    def comm_subexpr_elim=(value)
+      _rboptflags.comm_subexpr_elim = value
+    end
+
+    # Do not maintain order if the order would not be observed.
+    def check_order_observe
+      _rboptflags.check_order_observe
+    end
+
+    def check_order_observe=(value)
+      _rboptflags.check_order_observe = value
+    end
+
+    # Replace simple projections with a faster inlined projection that skips the expression engine.
+    def fast_projection
+      _rboptflags.fast_projection
+    end
+
+    def fast_projection=(value)
+      _rboptflags.fast_projection = value
+    end
+
+    def to_s
+      <<~STR
+        QueryOptFlags {
+            type_coercion: #{_rboptflags.type_coercion}
+            type_check: #{_rboptflags.type_check}
+
+            predicate_pushdown: #{predicate_pushdown}
+            projection_pushdown: #{projection_pushdown}
+            simplify_expression: #{simplify_expression}
+            slice_pushdown: #{slice_pushdown}
+            comm_subplan_elim: #{comm_subplan_elim}
+            comm_subexpr_elim: #{comm_subexpr_elim}
+            cluster_with_columns: #{cluster_with_columns}
+            check_order_observe: #{check_order_observe}
+            fast_projection: #{fast_projection}
+
+            eager: #{_rboptflags.eager}
+            streaming: #{_rboptflags.streaming}
+        }
+      STR
+    end
+    alias_method :inspect, :to_s
+
+    private
+
+    def initialize_copy(other)
+      super
+      self._rboptflags = _rboptflags.copy
     end
   end
+
+  DEFAULT_QUERY_OPT_FLAGS = QueryOptFlags.new
 end

data/lib/polars/rolling_group_by.rb

@@ -25,11 +25,11 @@ module Polars
 
     def agg(*aggs, **named_aggs)
       @df.lazy
-        .group_by_rolling(
-          index_column: @time_column, period: @period, offset: @offset, closed: @closed, by: @group_by
+        .rolling(
+          index_column: @time_column, period: @period, offset: @offset, closed: @closed, group_by: @group_by
         )
         .agg(*aggs, **named_aggs)
-        .collect(no_optimization: true, string_cache: false)
+        .collect(optimizations: QueryOptFlags.none)
     end
   end
 end

data/lib/polars/schema.rb

@@ -1,13 +1,21 @@
 module Polars
   class Schema
+    include Enumerable
+
     # Ordered mapping of column names to their data type.
     #
     # @param schema [Object]
     #   The schema definition given by column names and their associated
     #   Polars data type. Accepts a mapping or an enumerable of arrays.
     def initialize(schema = nil, check_dtypes: true)
-      input = schema || {}
       @schema = {}
+
+      if schema.respond_to?(:arrow_c_schema) && !schema.is_a?(Schema)
+        Plr.init_polars_schema_from_arrow_c_schema(@schema, schema)
+        return
+      end
+
+      input = schema || {}
       input.each do |name, tp|
         if !check_dtypes
           @schema[name] = tp
@@ -34,6 +42,11 @@ module Polars
       @schema[name] = dtype
     end
 
+    # @private
+    def arrow_c_schema
+      Plr.polars_schema_to_rbcapsule(self)
+    end
+
     # Get the column names of the schema.
     #
     # @return [Array]
@@ -78,14 +91,19 @@ module Polars
     end
     alias_method :inspect, :to_s
 
+    # @private
+    def each(&block)
+      @schema.each(&block)
+    end
+
     # @private
     def include?(name)
       @schema.include?(name)
     end
 
     # @private
-    def to_h
-      @schema.to_h
+    def to_h(&block)
+      @schema.to_h(&block)
     end
 
     private

data/lib/polars/selector.rb

@@ -127,14 +127,49 @@ module Polars
     # Only works after a wildcard or regex column selection, and you cannot provide
     # both string column names *and* dtypes (you may prefer to use selectors instead).
     #
-    # @return [Selector]
-    #
     # @param columns [Object]
     #   The name or datatype of the column(s) to exclude. Accepts regular expression
     #   input. Regular expressions should start with `^` and end with `$`.
     # @param more_columns [Array]
     #   Additional names or datatypes of columns to exclude, specified as positional
     #   arguments.
+    #
+    # @return [Selector]
+    #
+    # @example Exclude by column name(s):
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "aa" => [1, 2, 3],
+    #       "ba" => ["a", "b", nil],
+    #       "cc" => [nil, 2.5, 1.5]
+    #     }
+    #   )
+    #   df.select(Polars.cs.exclude("ba", "xx"))
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌─────┬──────┐
+    #   # │ aa  ┆ cc   │
+    #   # │ --- ┆ ---  │
+    #   # │ i64 ┆ f64  │
+    #   # ╞═════╪══════╡
+    #   # │ 1   ┆ null │
+    #   # │ 2   ┆ 2.5  │
+    #   # │ 3   ┆ 1.5  │
+    #   # └─────┴──────┘
+    #
+    # @example Exclude using a column name, a selector, and a dtype:
+    #   df.select(Polars.cs.exclude("aa", Polars.cs.string, Polars::UInt32))
+    #   # =>
+    #   # shape: (3, 1)
+    #   # ┌──────┐
+    #   # │ cc   │
+    #   # │ ---  │
+    #   # │ f64  │
+    #   # ╞══════╡
+    #   # │ null │
+    #   # │ 2.5  │
+    #   # │ 1.5  │
+    #   # └──────┘
     def exclude(columns, *more_columns)
       exclude_cols = []
       exclude_dtypes = []