polars-df 0.15.0 → 0.26.0

data/lib/polars/scan_cast_options.rb

@@ -0,0 +1,86 @@
+module Polars
+  # Options for scanning files.
+  class ScanCastOptions
+    attr_reader :integer_cast, :float_cast, :datetime_cast, :missing_struct_fields, :extra_struct_fields, :categorical_to_string
+
+    # Common configuration for scanning files.
+    #
+    # @note
+    #   This functionality is considered **unstable**. It may be changed
+    #   at any point without it being considered a breaking change.
+    #
+    # @param integer_cast ['upcast', 'forbid']
+    #   Configuration for casting from integer types:
+    #
+    #   * `upcast`: Allow lossless casting to wider integer types.
+    #   * `forbid`: Raises an error if dtypes do not match.
+    #
+    # @param float_cast ['upcast', 'downcast', 'forbid']
+    #   Configuration for casting from float types:
+    #
+    #   * `upcast`: Allow casting to higher precision float types.
+    #   * `downcast`: Allow casting to lower precision float types.
+    #   * `forbid`: Raises an error if dtypes do not match.
+    #
+    # @param datetime_cast ['nanosecond-downcast', 'convert-timezone', 'forbid']
+    #   Configuration for casting from datetime types:
+    #
+    #   * `nanosecond-downcast`: Allow nanosecond precision datetime to be
+    #     downcasted to any lower precision. This has a similar effect to
+    #     PyArrow's `coerce_int96_timestamp_unit`.
+    #   * `convert-timezone`: Allow casting to a different timezone.
+    #   * `forbid`: Raises an error if dtypes do not match.
+    #
+    # @param missing_struct_fields ['insert', 'raise']
+    #   Configuration for behavior when struct fields defined in the schema
+    #   are missing from the data:
+    #
+    #   * `insert`: Inserts the missing fields.
+    #   * `raise`: Raises an error.
+    #
+    # @param extra_struct_fields ['ignore', 'raise']
+    #   Configuration for behavior when extra struct fields outside of the
+    #   defined schema are encountered in the data:
+    #
+    #   * `ignore`: Silently ignores.
+    #   * `raise`: Raises an error.
+    def initialize(
+      integer_cast: "forbid",
+      float_cast: "forbid",
+      datetime_cast: "forbid",
+      missing_struct_fields: "raise",
+      extra_struct_fields: "raise",
+      categorical_to_string: "forbid",
+      _internal_call: false
+    )
+      if !_internal_call
+        warn "ScanCastOptions is considered unstable."
+      end
+
+      @integer_cast = integer_cast
+      @float_cast = float_cast
+      @datetime_cast = datetime_cast
+      @missing_struct_fields = missing_struct_fields
+      @extra_struct_fields = extra_struct_fields
+      @categorical_to_string = categorical_to_string
+    end
+
+    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",
+          categorical_to_string: "allow",
+          _internal_call: true
+        )
+      end
+    end
+  end
+end

data/lib/polars/schema.rb

@@ -0,0 +1,128 @@
+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)
+      @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
+        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)
+      _check_dtype(dtype)
+      @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]
+    #
+    # @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 each(&block)
+      @schema.each(&block)
+    end
+
+    # @private
+    def include?(name)
+      @schema.include?(name)
+    end
+
+    # @private
+    def to_h(&block)
+      @schema.to_h(&block)
+    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

@@ -0,0 +1,245 @@
+module Polars
+  # Base column selector expression/proxy.
+  class Selector < Expr
+    # @private
+    attr_accessor :_rbselector
+
+    # @private
+    def self._from_rbselector(rbselector)
+      slf = new
+      slf._rbselector = rbselector
+      slf._rbexpr = RbExpr.new_selector(rbselector)
+      slf
+    end
+
+    # Returns a string representing the Selector.
+    #
+    # @return [String]
+    def inspect
+      Expr._from_rbexpr(_rbexpr).to_s
+    end
+
+    # @private
+    def self._by_dtype(dtypes)
+      selectors = []
+      concrete_dtypes = []
+      dtypes.each do |dt|
+        if Utils.is_polars_dtype(dt)
+          concrete_dtypes += [dt]
+        else
+          raise Todo
+        end
+      end
+
+      dtype_selector = _from_rbselector(RbSelector.by_dtype(concrete_dtypes))
+
+      if selectors.length == 0
+        return dtype_selector
+      end
+
+      selector = selectors[0]
+      selectors[1..].each do |s|
+        selector = selector | s
+      end
+      if concrete_dtypes.length == 0
+        selector
+      else
+        dtype_selector | selector
+      end
+    end
+
+    # @private
+    def self._by_name(names, strict:, expand_patterns:)
+      _from_rbselector(RbSelector.by_name(names, strict, expand_patterns))
+    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
+        other = by_name(colname)
+      end
+      if Utils.is_selector(other)
+        Selector._from_rbselector(
+          _rbselector.intersect(other._rbselector)
+        )
+      else
+        as_expr & other
+      end
+    end
+
+    # OR.
+    #
+    # @return [Selector]
+    def |(other)
+      if Utils.is_column(other)
+        other = by_name(other.meta.output_name)
+      end
+      if Utils.is_selector(other)
+        Selector._from_rbselector(
+          _rbselector.union(other._rbselector)
+        )
+      else
+        as_expr | other
+      end
+    end
+
+    # Difference.
+    #
+    # @return [Selector]
+    def -(other)
+      if Utils.is_selector(other)
+        Selector._from_rbselector(
+          _rbselector.difference(other._rbselector)
+        )
+      else
+        as_expr - other
+      end
+    end
+
+    # XOR.
+    #
+    # @return [Selector]
+    def ^(other)
+      if Utils.is_column(other)
+        other = by_name(other.meta.output_name)
+      end
+      if Utils.is_selector(other)
+        Selector._from_rbselector(
+          _rbselector.exclusive_or(other._rbselector)
+        )
+      else
+        as_expr ^ other
+      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).
+    #
+    # @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 = []
+      ((columns.is_a?(::Array) ? columns : [columns]) + more_columns).each do |item|
+        if item.is_a?(::String)
+          exclude_cols << item
+        elsif Utils.is_polars_dtype(item)
+          exclude_dtypes << item
+        else
+          msg = (
+            "invalid input for `exclude`" +
+            "\n\nExpected one or more `str` or `DataType`; found #{item.inspect} instead."
+          )
+          raise TypeError, msg
+        end
+      end
+
+      if exclude_cols.any? && exclude_dtypes.any?
+        msg = "cannot exclude by both column name and dtype; use a selector instead"
+        raise TypeError, msg
+      elsif exclude_dtypes.any?
+        self - Selectors.by_dtype(exclude_dtypes)
+      else
+        self - Selector._by_name(exclude_cols, strict: false, expand_patterns: true)
+      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
+  end
+end