polars-df 0.21.0 → 0.21.1

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
+      # 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

@@ -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).