polars-df 0.15.0 → 0.26.0

data/lib/polars/meta_expr.rb

@@ -97,6 +97,70 @@ module Polars
       _rbexpr.meta_is_regex_projection
     end
 
+    # Indicate if this expression only selects columns (optionally with aliasing).
+    #
+    # This can include bare columns, columns matched by regex or dtype, selectors
+    # and exclude ops, and (optionally) column/expression aliasing.
+    #
+    # @param allow_aliasing [Boolean]
+    #   If false (default), any aliasing is not considered to be column selection.
+    #   Set true to allow for column selection that also includes aliasing.
+    #
+    # @return [Boolean]
+    #
+    # @example
+    #   e = Polars.col("foo")
+    #   e.meta.is_column_selection
+    #   # => true
+    #
+    # @example
+    #   e = Polars.col("foo").alias("bar")
+    #   e.meta.is_column_selection
+    #   # => false
+    #
+    # @example
+    #   e.meta.is_column_selection(allow_aliasing: true)
+    #   # => true
+    #
+    # @example
+    #   e = Polars.col("foo") * Polars.col("bar")
+    #   e.meta.is_column_selection
+    #   # => false
+    #
+    # @example
+    #   e = Polars.cs.starts_with("foo")
+    #   e.meta.is_column_selection
+    #   # => true
+    #
+    # @example
+    #   e = Polars.cs.starts_with("foo").exclude("foo!")
+    #   e.meta.is_column_selection
+    #   # => true
+    def is_column_selection(allow_aliasing: false)
+      _rbexpr.meta_is_column_selection(allow_aliasing)
+    end
+
+    # Indicate if this expression is a literal value (optionally aliased).
+    #
+    # @param allow_aliasing [Boolean]
+    #   If false (default), only a bare literal will match.
+    #   Set true to also allow for aliased literals.
+    #
+    # @return [Boolean]
+    #
+    # @example
+    #   e = Polars.lit(123)
+    #   e.meta.is_literal
+    #   # => true
+    #
+    # @example
+    #   e = Polars.lit(987.654321).alias("foo")
+    #   e.meta.is_literal
+    #   # => false
+    def is_literal(allow_aliasing: false)
+      _rbexpr.meta_is_literal(allow_aliasing)
+    end
+
     # Get the column name that this expression would produce.
     #
     # @return [String]
@@ -116,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.
@@ -125,14 +194,14 @@ module Polars
     # @return [Array]
     #
     # @example
-    #   e = Polars.col("foo").alias("bar")
+    #   e = Polars.col("foo") + Polars.col("bar")
     #   first = e.meta.pop[0]
-    #   _ = first.meta == Polars.col("foo")
-    #   # => true
     #   _ = first.meta == Polars.col("bar")
+    #   # => true
+    #   _ = first.meta == Polars.col("foo")
     #   # => false
-    def pop
-      _rbexpr.meta_pop.map { |e| Utils.wrap_expr(e) }
+    def pop(schema: nil)
+      _rbexpr.meta_pop(schema).map { |e| Utils.wrap_expr(e) }
     end
 
     # Get a list with the root column name.
@@ -171,32 +240,49 @@ module Polars
       Utils.wrap_expr(_rbexpr.meta_undo_aliases)
     end
 
-    # Turn this expression in a selector.
+    # Return the original expression.
     #
     # @return [Expr]
-    def _as_selector
-      Utils.wrap_expr(_rbexpr._meta_as_selector)
+    def as_expression
+      Utils.wrap_expr(_rbexpr)
     end
 
-    # Add selectors.
+    # Try to turn this expression in a selector.
     #
-    # @return [Expr]
-    def _selector_add(other)
-      Utils.wrap_expr(_rbexpr._meta_selector_add(other._rbexpr))
-    end
-
-    # Subtract selectors.
+    # Raises if the underlying expressions is not a column or selector.
     #
     # @return [Expr]
-    def _selector_sub(other)
-      Utils.wrap_expr(_rbexpr._meta_selector_sub(other._rbexpr))
+    #
+    # @note
+    #   This functionality is considered **unstable**. It may be changed
+    #   at any point without it being considered a breaking change.
+    def as_selector
+      Selector._from_rbselector(_rbexpr.into_selector)
     end
 
-    # & selectors.
+    # Serialize this expression to a file or string.
     #
-    # @return [Expr]
-    def _selector_and(other)
-      Utils.wrap_expr(_rbexpr._meta_selector_and(other._rbexpr))
+    # @param file [Object]
+    #   File path to which the result should be written. If set to `nil`
+    #   (default), the output is returned as a string instead.
+    #
+    # @return [Object]
+    #
+    # @note
+    #   Serialization is not stable across Polars versions: a LazyFrame serialized
+    #   in one Polars version may not be deserializable in another Polars version.
+    #
+    # @example Serialize the expression into a binary representation.
+    #   expr = Polars.col("foo").sum.over("bar")
+    #   bytes = expr.meta.serialize
+    #   Polars::Expr.deserialize(StringIO.new(bytes))
+    #   # => col("foo").sum().over([col("bar")])
+    def serialize(file = nil)
+      raise Todo unless _rbexpr.respond_to?(:serialize_binary)
+
+      serializer = _rbexpr.method(:serialize_binary)
+
+      Utils.serialize_polars_object(serializer, file)
     end
 
     # Format the expression as a tree.
@@ -209,8 +295,8 @@ module Polars
     # @example
     #   e = (Polars.col("foo") * Polars.col("bar")).sum.over(Polars.col("ham")) / 2
     #   e.meta.tree_format(return_as_string: true)
-    def tree_format(return_as_string: false)
-      s = _rbexpr.meta_tree_format
+    def tree_format(return_as_string: false, schema: nil)
+      s = _rbexpr.meta_tree_format(schema)
       if return_as_string
         s
       else
@@ -218,5 +304,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,8 @@ module Polars
     #   # │ 2         ┆ y         ┆ 2   ┆ y   │
     #   # │ 1         ┆ x         ┆ 3   ┆ z   │
     #   # └───────────┴───────────┴─────┴─────┘
-    def map(&f)
-      Utils.wrap_expr(_rbexpr.name_map(f))
+    def map(&function)
+      Utils.wrap_expr(_rbexpr.name_map(function))
     end
 
     # Add a prefix to the root column name of the expression.
@@ -194,5 +194,90 @@ module Polars
     def to_uppercase
       Utils.wrap_expr(_rbexpr.name_to_uppercase)
     end
+
+    # Rename fields of a struct by mapping a function over the field name(s).
+    #
+    # @note
+    #   This only takes effect for struct columns.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"x" => {"a" => 1, "b" => 2}})
+    #   df.select(Polars.col("x").name.map_fields  { |x| x.upcase }).schema
+    #   # => Polars::Schema({"x"=>Polars::Struct({"A"=>Polars::Int64, "B"=>Polars::Int64})})
+    def map_fields(&function)
+      Utils.wrap_expr(_rbexpr.name_map_fields(function))
+    end
+
+    # Add a prefix to all field names of a struct.
+    #
+    # @note
+    #   This only takes effect for struct columns.
+    #
+    # @param prefix [String]
+    #   Prefix to add to the field name.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"x" => {"a" => 1, "b" => 2}})
+    #   df.select(Polars.col("x").name.prefix_fields("prefix_")).schema
+    #   # => 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
+    #   This only takes effect for struct columns.
+    #
+    # @param suffix [String]
+    #   Suffix to add to the field name.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"x" => {"a" => 1, "b" => 2}})
+    #   df.select(Polars.col("x").name.suffix_fields("_suffix")).schema
+    #   # => 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
   end
 end

data/lib/polars/query_opt_flags.rb

@@ -0,0 +1,264 @@
+module Polars
+  # The set of the optimizations considered during query optimization.
+  #
+  # @note
+  #   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,
+      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,
+      sort_collapse: nil
+    )
+      self._rboptflags = RbOptFlags.default
+      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,
+        sort_collapse: sort_collapse
+      )
+    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,
+      sort_collapse: 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,
+        sort_collapse: sort_collapse
+      )
+    end
+
+    def update(
+      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,
+      sort_collapse: nil
+    )
+      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
+      if !sort_collapse.nil?
+        self.sort_collapse = sort_collapse
+      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
+
+    # Collapse sequential sort nodes into a single sort node.
+    def sort_collapse
+      _rboptflags.sort_collapse
+    end
+
+    def sort_collapse=(value)
+      _rboptflags.sort_collapse = 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
+
+  # @private
+  DEFAULT_QUERY_OPT_FLAGS = QueryOptFlags.new
+end

data/lib/polars/rolling_group_by.rb

@@ -4,13 +4,15 @@ module Polars
   # This has an `.agg` method which will allow you to run all polars expressions in a
   # group by context.
   class RollingGroupBy
+    # @private
     def initialize(
       df,
       index_column,
       period,
       offset,
       closed,
-      group_by
+      group_by,
+      predicates
     )
       period = Utils.parse_as_duration_string(period)
       offset = Utils.parse_as_duration_string(offset)
@@ -21,15 +23,98 @@ module Polars
       @offset = offset
       @closed = closed
       @group_by = group_by
+      @predicates = predicates
     end
 
+    # Filter groups with a list of predicates after aggregation.
+    #
+    # Using this method is equivalent to adding the predicates to the aggregation and
+    # filtering afterwards.
+    #
+    # This method can be chained and all conditions will be combined using `&`.
+    #
+    # @param predicates [Array]
+    #   Expressions that evaluate to a boolean value for each group. Typically, this
+    #   requires the use of an aggregation function. Multiple predicates are
+    #   combined using `&`.
+    #
+    # @return [RollingGroupBy]
+    def having(*predicates)
+      RollingGroupBy.new(
+        @df,
+        @time_column,
+        @period,
+        @offset,
+        @closed,
+        @group_by,
+        Utils._chain_predicates(@predicates, predicates)
+      )
+    end
+
+    # Compute aggregations for each group of a group by operation.
+    #
+    # @param aggs [Array]
+    #   Aggregations to compute for each group of the group by operation,
+    #   specified as positional arguments.
+    #   Accepts expression input. Strings are parsed as column names.
+    # @param named_aggs [Hash]
+    #   Additional aggregations, specified as keyword arguments.
+    #   The resulting columns will be renamed to the keyword used.
+    #
+    # @return [DataFrame]
     def agg(*aggs, **named_aggs)
+      group_by =
+        @df.lazy.rolling(
+          index_column: @time_column, period: @period, offset: @offset, closed: @closed, group_by: @group_by
+        )
+
+      if @predicates&.any?
+        group_by = group_by.having(@predicates)
+      end
+
+      group_by.agg(*aggs, **named_aggs).collect(
+        optimizations: QueryOptFlags.none
+      )
+    end
+
+    # Apply a custom/user-defined function (UDF) over the groups as a new DataFrame.
+    #
+    # Using this is considered an anti-pattern as it will be very slow because:
+    #
+    # - it forces the engine to materialize the whole `DataFrames` for the groups.
+    # - it is not parallelized.
+    # - it blocks optimizations as the passed python function is opaque to the
+    #   optimizer.
+    #
+    # The idiomatic way to apply custom functions over multiple columns is using:
+    #
+    # `Polars.struct([my_columns]).map_elements { |struct_series| ... }`
+    #
+    # @param schema [Object]
+    #   Schema of the output function. This has to be known statically. If the
+    #   given schema is incorrect, this is a bug in the caller's query and may
+    #   lead to errors. If set to None, polars assumes the schema is unchanged.
+    #
+    # @return [DataFrame]
+    def map_groups(
+      schema,
+      &function
+    )
+      if @predicates&.any?
+        msg = "cannot call `map_groups` when filtering groups with `having`"
+        raise TypeError, msg
+      end
+
       @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)
+        .map_groups(schema, &function)
+        .collect(optimizations: QueryOptFlags.none)
     end
   end
 end