polars-df 0.20.0-x86_64-darwin → 0.21.1-x86_64-darwin

data/lib/polars/schema.rb

@@ -1,29 +1,110 @@
 module Polars
   class Schema
-    def initialize(schema, check_dtypes: true)
-      raise Todo if check_dtypes
-      @schema = schema.to_h
+    # 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 = {}
+      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
+      @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

@@ -0,0 +1,210 @@
+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:)
+      _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
+        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).
+    #
+    # @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 = []
+      ((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 - Selectors.by_name(exclude_cols, require_all: false)
+      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