polars-df 0.21.0 → 0.22.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]
@@ -184,6 +248,31 @@ module Polars
       Selector._from_rbselector(_rbexpr.into_selector)
     end
 
+    # Serialize this expression to a file or string.
+    #
+    # @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.
     #
     # @param return_as_string [Boolean]

data/lib/polars/name_expr.rb

@@ -194,5 +194,41 @@ module Polars
     def to_uppercase
       Utils.wrap_expr(_rbexpr.name_to_uppercase)
     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
+    #   # => {"x"=>Polars::Struct({"prefix_a"=>Polars::Int64, "prefix_b"=>Polars::Int64})}
+    def prefix_fields(prefix)
+      Utils.wrap_expr(_rbexpr.name_prefix_fields(prefix))
+    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
+    #   # => {"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,50 @@
+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
+    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
+    )
+      @_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
+      )
+    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
+    )
+      raise Todo
+    end
+  end
+end

data/lib/polars/scan_cast_options.rb

@@ -1,6 +1,8 @@
 module Polars
   # Options for scanning files.
   class ScanCastOptions
+    attr_reader :integer_cast, :float_cast, :datetime_cast, :missing_struct_fields, :extra_struct_fields
+
     # Common configuration for scanning files.
     #
     # @note
@@ -50,6 +52,10 @@ module Polars
       extra_struct_fields: "raise",
       _internal_call: false
     )
+      if !_internal_call
+        warn "ScanCastOptions is considered unstable."
+      end
+
       @integer_cast = integer_cast
       @float_cast = float_cast
       @datetime_cast = datetime_cast
@@ -57,8 +63,21 @@ module Polars
       @extra_struct_fields = extra_struct_fields
     end
 
-    def self.default
+    def self._default
       new(_internal_call: true)
     end
+
+    def self._default_iceberg
+      @_default_cast_options_iceberg ||= begin
+        ScanCastOptions.new(
+          integer_cast: "upcast",
+          float_cast: ["upcast", "downcast"],
+          datetime_cast: ["nanosecond-downcast", "convert-timezone"],
+          missing_struct_fields: "insert",
+          extra_struct_fields: "ignore",
+          _internal_call: true
+        )
+      end
+    end
   end
 end

data/lib/polars/schema.rb

@@ -1,34 +1,110 @@
 module Polars
   class Schema
+    # 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)
-      raise Todo if check_dtypes
-      @schema = schema.to_h
+      input = schema || {}
+      @schema = {}
+      input.each do |name, tp|
+        if !check_dtypes
+          @schema[name] = tp
+        elsif Utils.is_polars_dtype(tp)
+          @schema[name] = _check_dtype(tp)
+        else
+          self[name] = tp
+        end
+      end
     end
 
+    # Returns the data type of the column.
+    #
+    # @return [Object]
     def [](key)
       @schema[key]
     end
 
+    # Sets the data type of the column.
+    #
+    # @return [Object]
     def []=(name, dtype)
-      # TODO check dtype if needed
+      _check_dtype(dtype)
       @schema[name] = dtype
     end
 
+    # Get the column names of the schema.
+    #
+    # @return [Array]
+    #
+    # @example
+    #   s = Polars::Schema.new({"x" => Polars::Float64.new, "y" => Polars::Datetime.new(time_zone: "UTC")})
+    #   s.names
+    #   # => ["x", "y"]
     def names
       @schema.keys
     end
 
+    # Get the data types of the schema.
+    #
+    # @return [Array]
+    #
+    # @example
+    #   s = Polars::Schema.new({"x" => Polars::UInt8.new, "y" => Polars::List.new(Polars::UInt8)})
+    #   s.dtypes
+    #   # => [Polars::UInt8, Polars::List(Polars::UInt8)]
     def dtypes
       @schema.values
     end
 
+    # Get the number of schema entries.
+    #
+    # @return [Integer]
+    #
+    # @example
+    #   s = Polars::Schema.new({"x" => Polars::Int32.new, "y" => Polars::List.new(Polars::String)})
+    #   s.length
+    #   # => 2
     def length
       @schema.length
     end
 
+    # Returns a string representing the Schema.
+    #
+    # @return [String]
     def to_s
       "#{self.class.name}(#{@schema})"
     end
     alias_method :inspect, :to_s
+
+    # @private
+    def include?(name)
+      @schema.include?(name)
+    end
+
+    # @private
+    def to_h
+      @schema.to_h
+    end
+
+    private
+
+    def _check_dtype(tp)
+      if !tp.is_a?(DataType)
+        # note: if nested/decimal, or has signature params, this implies required args
+        if tp.nested? || tp.decimal? || _required_init_args(tp)
+          msg = "dtypes must be fully-specified, got: #{tp.inspect}"
+          raise TypeError, msg
+        end
+        tp = tp.new
+      end
+      tp
+    end
+
+    def _required_init_args(tp)
+      arity = tp.method(:new).arity
+      arity > 0 || arity < -1
+    end
   end
 end

data/lib/polars/selector.rb

@@ -12,6 +12,9 @@ module Polars
       slf
     end
 
+    # Returns a string representing the Selector.
+    #
+    # @return [String]
     def inspect
       Expr._from_rbexpr(_rbexpr).to_s
     end
@@ -50,10 +53,16 @@ module Polars
       _from_rbselector(RbSelector.by_name(names, strict))
     end
 
+    # Invert the selector.
+    #
+    # @return [Selector]
     def ~
       Selectors.all - self
     end
 
+    # AND.
+    #
+    # @return [Selector]
     def &(other)
       if Utils.is_column(other)
         colname = other.meta.output_name
@@ -68,6 +77,9 @@ module Polars
       end
     end
 
+    # OR.
+    #
+    # @return [Selector]
     def |(other)
       if Utils.is_column(other)
         other = by_name(other.meta.output_name)
@@ -81,6 +93,9 @@ module Polars
       end
     end
 
+    # Difference.
+    #
+    # @return [Selector]
     def -(other)
       if Utils.is_selector(other)
         Selector._from_rbselector(
@@ -91,6 +106,9 @@ module Polars
       end
     end
 
+    # XOR.
+    #
+    # @return [Selector]
     def ^(other)
       if Utils.is_column(other)
         other = by_name(other.meta.output_name)
@@ -104,6 +122,19 @@ module Polars
       end
     end
 
+    # Exclude columns from a multi-column expression.
+    #
+    # 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.
     def exclude(columns, *more_columns)
       exclude_cols = []
       exclude_dtypes = []
@@ -131,6 +162,47 @@ module Polars
       end
     end
 
+    # Materialize the `selector` as a normal expression.
+    #
+    # This ensures that the operators `|`, `&`, `~` and `-`
+    # are applied on the data and not on the selector sets.
+    #
+    # @return [Expr]
+    #
+    # @example Inverting the boolean selector will choose the non-boolean columns:
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "colx" => ["aa", "bb", "cc"],
+    #       "coly" => [true, false, true],
+    #       "colz" => [1, 2, 3]
+    #     }
+    #   )
+    #   df.select(~Polars.cs.boolean)
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌──────┬──────┐
+    #   # │ colx ┆ colz │
+    #   # │ ---  ┆ ---  │
+    #   # │ str  ┆ i64  │
+    #   # ╞══════╪══════╡
+    #   # │ aa   ┆ 1    │
+    #   # │ bb   ┆ 2    │
+    #   # │ cc   ┆ 3    │
+    #   # └──────┴──────┘
+    #
+    # @example To invert the *values* in the selected boolean columns, we need to materialize the selector as a standard expression instead:
+    #   df.select(~Polars.cs.boolean.as_expr)
+    #   # =>
+    #   # shape: (3, 1)
+    #   # ┌───────┐
+    #   # │ coly  │
+    #   # │ ---   │
+    #   # │ bool  │
+    #   # ╞═══════╡
+    #   # │ false │
+    #   # │ true  │
+    #   # │ false │
+    #   # └───────┘
     def as_expr
       Expr._from_rbexpr(_rbexpr)
     end

data/lib/polars/selectors.rb

@@ -287,7 +287,7 @@ module Polars
     #   # │ b"hello" ┆ world ┆ b"!"   ┆ :)  │
     #   # └──────────┴───────┴────────┴─────┘
     #
-    # @example Select binary columns and export as a dict:
+    # @example Select binary columns and export as a hash:
     #   df.select(Polars.cs.binary).to_h(as_series: false)
     #   # => {"a"=>["hello"], "c"=>["!"]}
     #
@@ -628,7 +628,7 @@ module Polars
     #   # └──────┘
     #
     # @example Select all columns *except* for those that are enum:
-    #   df.select(~Polars.cs.enum())
+    #   df.select(~Polars.cs.enum)
     #   # =>
     #   # shape: (2, 2)
     #   # ┌─────┬─────┐
@@ -928,7 +928,7 @@ module Polars
     #   # │ 456 ┆ 5.5 │
     #   # └─────┴─────┘
     def self.categorical
-      Selector._from_rbselector(RbSelector.categorical())
+      Selector._from_rbselector(RbSelector.categorical)
     end
 
     # Select columns whose names contain the given literal substring(s).