polars-df 0.4.0 → 0.5.0

data/lib/polars/convert.rb

@@ -26,8 +26,8 @@ module Polars
     #   # │ 1   ┆ 3   │
     #   # │ 2   ┆ 4   │
     #   # └─────┴─────┘
-    def from_hash(data, columns: nil)
-      DataFrame._from_hash(data, columns: columns)
+    def from_hash(data, schema: nil, columns: nil)
+      DataFrame._from_hash(data, schema: schema || columns)
     end
 
     # Construct a DataFrame from a sequence of dictionaries. This operation clones data.

data/lib/polars/data_frame.rb

@@ -18,7 +18,10 @@ module Polars
     #   Whether to interpret two-dimensional data as columns or as rows. If `nil`,
     #   the orientation is inferred by matching the columns and data dimensions. If
     #   this does not yield conclusive results, column orientation is used.
-    def initialize(data = nil, columns: nil, orient: nil)
+    def initialize(data = nil, schema: nil, columns: nil, schema_overrides: nil, orient: nil, infer_schema_length: 100, nan_to_null: false)
+      schema ||= columns
+      raise Todo if schema_overrides
+
       # TODO deprecate in favor of read_sql
       if defined?(ActiveRecord) && (data.is_a?(ActiveRecord::Relation) || data.is_a?(ActiveRecord::Result))
         result = data.is_a?(ActiveRecord::Result) ? data : data.connection.select_all(data.to_sql)
@@ -29,14 +32,14 @@ module Polars
       end
 
       if data.nil?
-        self._df = self.class.hash_to_rbdf({}, columns: columns)
+        self._df = self.class.hash_to_rbdf({}, schema: schema, schema_overrides: schema_overrides)
       elsif data.is_a?(Hash)
         data = data.transform_keys { |v| v.is_a?(Symbol) ? v.to_s : v }
-        self._df = self.class.hash_to_rbdf(data, columns: columns)
+        self._df = self.class.hash_to_rbdf(data, schema: schema, schema_overrides: schema_overrides, nan_to_null: nan_to_null)
       elsif data.is_a?(Array)
-        self._df = self.class.sequence_to_rbdf(data, columns: columns, orient: orient)
+        self._df = self.class.sequence_to_rbdf(data, schema: schema, schema_overrides: schema_overrides, orient: orient, infer_schema_length: infer_schema_length)
       elsif data.is_a?(Series)
-        self._df = self.class.series_to_rbdf(data, columns: columns)
+        self._df = self.class.series_to_rbdf(data, schema: schema, schema_overrides: schema_overrides)
       else
         raise ArgumentError, "DataFrame constructor called with unsupported type; got #{data.class.name}"
       end
@@ -56,8 +59,8 @@ module Polars
     end
 
     # @private
-    def self._from_hash(data, columns: nil)
-      _from_rbdf(hash_to_rbdf(data, columns: columns))
+    def self._from_hash(data, schema: nil, schema_overrides: nil)
+      _from_rbdf(hash_to_rbdf(data, schema: schema, schema_overrides: schema_overrides))
     end
 
     # def self._from_records
@@ -336,6 +339,7 @@ module Polars
     end
     alias_method :count, :height
     alias_method :length, :height
+    alias_method :size, :height
 
     # Get the width of the DataFrame.
     #
@@ -546,6 +550,13 @@ module Polars
     end
     alias_method :inspect, :to_s
 
+    # Returns an array representing the DataFrame
+    #
+    # @return [Array]
+    def to_a
+      rows(named: true)
+    end
+
     # Check if DataFrame includes column.
     #
     # @return [Boolean]
@@ -655,7 +666,7 @@ module Polars
       end
 
       # Ruby-specific
-      if item.is_a?(Expr)
+      if item.is_a?(Expr) || item.is_a?(Series)
         return filter(item)
       end
 
@@ -665,15 +676,42 @@ module Polars
     # Set item.
     #
     # @return [Object]
-    #
-    # def []=(key, value)
-    #   if key.is_a?(String)
-    #     raise TypeError, "'DataFrame' object does not support 'Series' assignment by index. Use 'DataFrame.with_columns'"
-    #   end
+    def []=(*key, value)
+      if key.length == 1
+        key = key.first
+      elsif key.length != 2
+        raise ArgumentError, "wrong number of arguments (given #{key.length + 1}, expected 2..3)"
+      end
 
-    #   raise Todo
-    # end
+      if Utils.strlike?(key)
+        if value.is_a?(Array) || (defined?(Numo::NArray) && value.is_a?(Numo::NArray))
+          value = Series.new(value)
+        elsif !value.is_a?(Series)
+          value = Polars.lit(value)
+        end
+        self._df = with_column(value.alias(key.to_s))._df
+      elsif key.is_a?(Array)
+        row_selection, col_selection = key
 
+        if Utils.strlike?(col_selection)
+          s = self[col_selection]
+        elsif col_selection.is_a?(Integer)
+          raise Todo
+        else
+          raise ArgumentError, "column selection not understood: #{col_selection}"
+        end
+
+        s[row_selection] = value
+
+        if col_selection.is_a?(Integer)
+          replace_at_idx(col_selection, s)
+        elsif Utils.strlike?(col_selection)
+          replace(col_selection, s)
+        end
+      else
+        raise Todo
+      end
+    end
 
     # Return the dataframe as a scalar.
     #
@@ -1462,6 +1500,20 @@ module Polars
       end
     end
 
+    # Sort the DataFrame by column in-place.
+    #
+    # @param by [String]
+    #   By which column to sort.
+    # @param reverse [Boolean]
+    #   Reverse/descending sort.
+    # @param nulls_last [Boolean]
+    #   Place null values last. Can only be used if sorted by a single column.
+    #
+    # @return [DataFrame]
+    def sort!(by, reverse: false, nulls_last: false)
+      self._df = sort(by, reverse: reverse, nulls_last: nulls_last)._df
+    end
+
     # Check if DataFrame is equal to other.
     #
     # @param other [DataFrame]
@@ -1519,7 +1571,7 @@ module Polars
     #   # │ 30  ┆ 6   │
     #   # └─────┴─────┘
     def replace(column, new_col)
-      _df.replace(column, new_col._s)
+      _df.replace(column.to_s, new_col._s)
       self
     end
 
@@ -1860,7 +1912,7 @@ module Polars
     #     "2020-01-08 23:16:43"
     #   ]
     #   df = Polars::DataFrame.new({"dt" => dates, "a" => [3, 7, 5, 9, 2, 1]}).with_column(
-    #     Polars.col("dt").str.strptime(:datetime)
+    #     Polars.col("dt").str.strptime(Polars::Datetime)
     #   )
     #   df.groupby_rolling(index_column: "dt", period: "2d").agg(
     #     [
@@ -2791,6 +2843,16 @@ module Polars
       Utils.wrap_s(_df.drop_in_place(name))
     end
 
+    # Drop in place if exists.
+    #
+    # @param name [Object]
+    #   Column to drop.
+    #
+    # @return [Series]
+    def delete(name)
+      drop_in_place(name) if include?(name)
+    end
+
     # Create an empty copy of the current DataFrame.
     #
     # Returns a DataFrame with identical schema but no data.
@@ -3202,7 +3264,7 @@ module Polars
     #   # │ B    ┆ 1    │
     #   # │ C    ┆ 2    │
     #   # │ D    ┆ 3    │
-    #   # │ …    ┆ …    │
+    #   # │ E    ┆ 4    │
     #   # │ F    ┆ 5    │
     #   # │ G    ┆ 6    │
     #   # │ H    ┆ 7    │
@@ -4428,7 +4490,7 @@ module Polars
       end
     end
 
-    # Returns an iterator over the DataFrame of rows of python-native values.
+    # Returns an iterator over the DataFrame of rows of Ruby-native values.
     #
     # @param named [Boolean]
     #   Return hashes instead of arrays. The hashes are a mapping of
@@ -4489,6 +4551,24 @@ module Polars
       end
     end
 
+    # Returns an iterator over the DataFrame of rows of Ruby-native values.
+    #
+    # @param named [Boolean]
+    #   Return hashes instead of arrays. The hashes are a mapping of
+    #   column name to row value. This is more expensive than returning an
+    #   array, but allows for accessing values by column name.
+    # @param buffer_size [Integer]
+    #   Determines the number of rows that are buffered internally while iterating
+    #   over the data; you should only modify this in very specific cases where the
+    #   default value is determined not to be a good fit to your access pattern, as
+    #   the speedup from using the buffer is significant (~2-4x). Setting this
+    #   value to zero disables row buffering.
+    #
+    # @return [Object]
+    def each_row(named: true, buffer_size: 500, &block)
+      iter_rows(named: named, buffer_size: buffer_size, &block)
+    end
+
     # Shrink DataFrame memory usage.
     #
     # Shrinks to fit the exact capacity needed to hold the data.
@@ -4742,20 +4822,63 @@ module Polars
     end
 
     # @private
-    def self.hash_to_rbdf(data, columns: nil)
-      if !columns.nil?
-        columns, dtypes = _unpack_columns(columns, lookup_names: data.keys)
+    def self.expand_hash_scalars(data, schema_overrides: nil, order: nil, nan_to_null: false)
+      updated_data = {}
+      unless data.empty?
+        dtypes = schema_overrides || {}
+        array_len = data.values.map { |val| Utils.arrlen(val) || 0 }.max
+        if array_len > 0
+          data.each do |name, val|
+            dtype = dtypes[name]
+            if val.is_a?(Hash) && dtype != Struct
+              updated_data[name] = DataFrame.new(val).to_struct(name)
+            elsif !Utils.arrlen(val).nil?
+              updated_data[name] = Series.new(String.new(name), val, dtype: dtype)
+            elsif val.nil? || [Integer, Float, TrueClass, FalseClass, String, ::Date, ::DateTime, ::Time].any? { |cls| val.is_a?(cls) }
+              dtype = Polars::Float64 if val.nil? && dtype.nil?
+              updated_data[name] = Series.new(String.new(name), [val], dtype: dtype).extend_constant(val, array_len - 1)
+            else
+              raise Todo
+            end
+          end
+        elsif data.values.all? { |val| Utils.arrlen(val) == 0 }
+          data.each do |name, val|
+            updated_data[name] = Series.new(name, val, dtype: dtypes[name])
+          end
+        elsif data.values.all? { |val| Utils.arrlen(val).nil? }
+          data.each do |name, val|
+            updated_data[name] = Series.new(name, [val], dtype: dtypes[name])
+          end
+        end
+      end
+      updated_data
+    end
 
-        if data.empty? && dtypes
-          data_series = columns.map { |name| Series.new(name, [], dtype: dtypes[name])._s }
-        else
-          data_series = data.map { |name, values| Series.new(name, values, dtype: dtypes[name])._s }
+    # @private
+    def self.hash_to_rbdf(data, schema: nil, schema_overrides: nil, nan_to_null: nil)
+      if schema.is_a?(Hash) && !data.empty?
+        if !data.all? { |col, _| schema[col] }
+          raise ArgumentError, "The given column-schema names do not match the data dictionary"
         end
-        data_series = _handle_columns_arg(data_series, columns: columns)
-        return RbDataFrame.new(data_series)
+
+        data = schema.to_h { |col| [col, data[col]] }
+      end
+
+      column_names, schema_overrides = _unpack_schema(
+        schema, lookup_names: data.keys, schema_overrides: schema_overrides
+      )
+      if column_names.empty?
+        column_names = data.keys
       end
 
-      RbDataFrame.read_hash(data)
+      if data.empty? && !schema_overrides.empty?
+        data_series = column_names.map { |name| Series.new(name, [], dtype: schema_overrides[name], nan_to_null: nan_to_null)._s }
+      else
+        data_series = expand_hash_scalars(data, schema_overrides: schema_overrides, nan_to_null: nan_to_null).values.map(&:_s)
+      end
+
+      data_series = _handle_columns_arg(data_series, columns: column_names, from_hash: true)
+      RbDataFrame.new(data_series)
     end
 
     # @private
@@ -4764,14 +4887,12 @@ module Polars
     end
 
     # @private
-    def self._unpack_columns(columns, schema_overrides: nil, lookup_names: nil, n_expected: nil)
-      raise Todo if schema_overrides
-
-      if columns.is_a?(Hash)
-        columns = columns.to_a
+    def self._unpack_schema(schema, schema_overrides: nil, n_expected: nil, lookup_names: nil, include_overrides_in_columns: false)
+      if schema.is_a?(Hash)
+        schema = schema.to_a
       end
       column_names =
-        (columns || []).map.with_index do |col, i|
+        (schema || []).map.with_index do |col, i|
           if col.is_a?(String)
             col || "column_#{i}"
           else
@@ -4784,21 +4905,38 @@ module Polars
       # TODO zip_longest
       lookup = column_names.zip(lookup_names || []).to_h
 
-      [
-        column_names,
-        (columns || []).select { |col| !col.is_a?(String) && col[1] }.to_h do |col|
+      column_dtypes =
+        (schema || []).select { |col| !col.is_a?(String) && col[1] }.to_h do |col|
           [lookup[col[0]] || col[0], col[1]]
         end
-      ]
+
+      if schema_overrides
+        raise Todo
+      end
+
+      column_dtypes.each do |col, dtype|
+        if !Utils.is_polars_dtype(dtype, include_unknown: true) && !dtype.nil?
+          column_dtypes[col] = Utils.rb_type_to_dtype(dtype)
+        end
+      end
+
+      [column_names, column_dtypes]
     end
 
-    def self._handle_columns_arg(data, columns: nil)
-      if columns.nil?
+    def self._handle_columns_arg(data, columns: nil, from_hash: false)
+      if columns.nil? || columns.empty?
         data
       else
         if data.empty?
           columns.map { |c| Series.new(c, nil)._s }
         elsif data.length == columns.length
+          if from_hash
+            series_map = data.to_h { |s| [s.name, s] }
+            if columns.all? { |col| series_map.key?(col) }
+              return columns.map { |col| series_map[col] }
+            end
+          end
+
           columns.each_with_index do |c, i|
             # not in-place?
             data[i].rename(c)
@@ -4813,7 +4951,7 @@ module Polars
     def self._post_apply_columns(rbdf, columns, structs: nil, schema_overrides: nil)
       rbdf_columns = rbdf.columns
       rbdf_dtypes = rbdf.dtypes
-      columns, dtypes = _unpack_columns(
+      columns, dtypes = _unpack_schema(
         (columns || rbdf_columns), schema_overrides: schema_overrides
       )
       column_subset = []
@@ -4851,20 +4989,23 @@ module Polars
     end
 
     # @private
-    def self.sequence_to_rbdf(data, columns: nil, orient: nil, infer_schema_length: 50)
+    def self.sequence_to_rbdf(data, schema: nil, schema_overrides: nil, orient: nil, infer_schema_length: 50)
+      raise Todo if schema_overrides
+      columns = schema
+
       if data.length == 0
-        return hash_to_rbdf({}, columns: columns)
+        return hash_to_rbdf({}, schema: schema, schema_overrides: schema_overrides)
       end
 
       if data[0].is_a?(Series)
         # series_names = data.map(&:name)
-        # columns, dtypes = _unpack_columns(columns || series_names, n_expected: data.length)
+        # columns, dtypes = _unpack_schema(columns || series_names, n_expected: data.length)
         data_series = []
         data.each do |s|
           data_series << s._s
         end
       elsif data[0].is_a?(Hash)
-        column_names, dtypes = _unpack_columns(columns)
+        column_names, dtypes = _unpack_schema(columns)
         schema_overrides = dtypes ? include_unknowns(dtypes, column_names) : nil
         rbdf = RbDataFrame.read_hashes(data, infer_schema_length, schema_overrides)
         if column_names
@@ -4890,11 +5031,21 @@ module Polars
     end
 
     # @private
-    def self.series_to_rbdf(data, columns: nil)
-      if columns
-        raise Todo
+    def self.series_to_rbdf(data, schema: nil, schema_overrides: nil)
+      data_series = [data._s]
+      series_name = data_series.map(&:name)
+      column_names, schema_overrides = _unpack_schema(
+        schema || series_name, schema_overrides: schema_overrides, n_expected: 1
+      )
+      if schema_overrides.any?
+        new_dtype = schema_overrides.values[0]
+        if new_dtype != data.dtype
+          data_series[0] = data_series[0].cast(new_dtype, true)
+        end
       end
-      RbDataFrame.new([data._s])
+
+      data_series = _handle_columns_arg(data_series, columns: column_names)
+      RbDataFrame.new(data_series)
     end
 
     def wrap_ldf(ldf)

data/lib/polars/data_types.rb

@@ -84,20 +84,22 @@ module Polars
 
   # Calendar date and time type.
   class Datetime < TemporalType
-    attr_reader :tu
+    attr_reader :time_unit, :time_zone
+    alias_method :tu, :time_unit
 
     def initialize(time_unit = "us", time_zone = nil)
-      @tu = time_unit || "us"
+      @time_unit = time_unit || "us"
       @time_zone = time_zone
     end
   end
 
   # Time duration/delta type.
   class Duration < TemporalType
-    attr_reader :tu
+    attr_reader :time_unit
+    alias_method :tu, :time_unit
 
     def initialize(time_unit = "us")
-      @tu = time_unit
+      @time_unit = time_unit
     end
   end
 

data/lib/polars/date_time_expr.rb

@@ -218,6 +218,25 @@ module Polars
       )
     end
 
+    # Create a naive Datetime from an existing Date/Datetime expression and a Time.
+    #
+    # If the underlying expression is a Datetime then its time component is replaced,
+    # and if it is a Date then a new Datetime is created by combining the two values.
+    #
+    # @param time [Object]
+    #   A Ruby time literal or Polars expression/column that resolves to a time.
+    # @param time_unit ["ns", "us", "ms"]
+    #   Unit of time.
+    #
+    # @return [Expr]
+    def combine(time, time_unit: "us")
+      unless time.is_a?(Time) || time.is_a?(Expr)
+        raise TypeError, "expected 'time' to be a Ruby time or Polars expression, found #{time}"
+      end
+      time = Utils.expr_to_lit_or_expr(time)
+      Utils.wrap_expr(_rbexpr.dt_combine(time._rbexpr, time_unit))
+    end
+
     # Format Date/datetime with a formatting rule.
     #
     # See [chrono strftime/strptime](https://docs.rs/chrono/latest/chrono/format/strftime/index.html).
@@ -270,6 +289,34 @@ module Polars
       Utils.wrap_expr(_rbexpr.year)
     end
 
+    # Determine whether the year of the underlying date is a leap year.
+    #
+    # Applies to Date and Datetime columns.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   start = DateTime.new(2000, 1, 1)
+    #   stop = DateTime.new(2002, 1, 1)
+    #   df = Polars::DataFrame.new(
+    #     {"date" => Polars.date_range(start, stop, "1y")}
+    #   )
+    #   df.select(Polars.col("date").dt.is_leap_year)
+    #   # =>
+    #   # shape: (3, 1)
+    #   # ┌───────┐
+    #   # │ date  │
+    #   # │ ---   │
+    #   # │ bool  │
+    #   # ╞═══════╡
+    #   # │ true  │
+    #   # │ false │
+    #   # │ false │
+    #   # └───────┘
+    def is_leap_year
+      Utils.wrap_expr(_rbexpr.dt_is_leap_year)
+    end
+
     # Extract ISO year from underlying Date representation.
     #
     # Applies to Date and Datetime columns.
@@ -550,6 +597,27 @@ module Polars
       Utils.wrap_expr(_rbexpr.ordinal_day)
     end
 
+    # Time
+    #
+    # @return [Expr]
+    def time
+      Utils.wrap_expr(_rbexpr.dt_time)
+    end
+
+    # Date
+    #
+    # @return [Expr]
+    def date
+      Utils.wrap_expr(_rbexpr.dt_date)
+    end
+
+    # Datetime
+    #
+    # @return [Expr]
+    def datetime
+      Utils.wrap_expr(_rbexpr.dt_datetime)
+    end
+
     # Extract hour from underlying DateTime representation.
     #
     # Applies to Datetime columns.
@@ -958,8 +1026,8 @@ module Polars
     #   Time zone for the `Datetime` Series.
     #
     # @return [Expr]
-    def replace_time_zone(tz)
-      Utils.wrap_expr(_rbexpr.dt_replace_time_zone(tz))
+    def replace_time_zone(tz, use_earliest: nil)
+      Utils.wrap_expr(_rbexpr.dt_replace_time_zone(tz, use_earliest))
     end
 
     # Localize tz-naive Datetime Series to tz-aware Datetime Series.
@@ -1282,5 +1350,77 @@ module Polars
     def offset_by(by)
       Utils.wrap_expr(_rbexpr.dt_offset_by(by))
     end
+
+    # Roll backward to the first day of the month.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "dates" => Polars.date_range(
+    #         DateTime.new(2000, 1, 15, 2),
+    #         DateTime.new(2000, 12, 15, 2),
+    #         "1mo"
+    #       )
+    #     }
+    #   )
+    #   df.select(Polars.col("dates").dt.month_start)
+    #   # =>
+    #   # shape: (12, 1)
+    #   # ┌─────────────────────┐
+    #   # │ dates               │
+    #   # │ ---                 │
+    #   # │ datetime[μs]        │
+    #   # ╞═════════════════════╡
+    #   # │ 2000-01-01 02:00:00 │
+    #   # │ 2000-02-01 02:00:00 │
+    #   # │ 2000-03-01 02:00:00 │
+    #   # │ 2000-04-01 02:00:00 │
+    #   # │ …                   │
+    #   # │ 2000-09-01 02:00:00 │
+    #   # │ 2000-10-01 02:00:00 │
+    #   # │ 2000-11-01 02:00:00 │
+    #   # │ 2000-12-01 02:00:00 │
+    #   # └─────────────────────┘
+    def month_start
+      Utils.wrap_expr(_rbexpr.dt_month_start)
+    end
+
+    # Roll forward to the last day of the month.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "dates" => Polars.date_range(
+    #         DateTime.new(2000, 1, 15, 2),
+    #         DateTime.new(2000, 12, 15, 2),
+    #         "1mo"
+    #       )
+    #     }
+    #   )
+    #   df.select(Polars.col("dates").dt.month_end)
+    #   # =>
+    #   # shape: (12, 1)
+    #   # ┌─────────────────────┐
+    #   # │ dates               │
+    #   # │ ---                 │
+    #   # │ datetime[μs]        │
+    #   # ╞═════════════════════╡
+    #   # │ 2000-01-31 02:00:00 │
+    #   # │ 2000-02-29 02:00:00 │
+    #   # │ 2000-03-31 02:00:00 │
+    #   # │ 2000-04-30 02:00:00 │
+    #   # │ …                   │
+    #   # │ 2000-09-30 02:00:00 │
+    #   # │ 2000-10-31 02:00:00 │
+    #   # │ 2000-11-30 02:00:00 │
+    #   # │ 2000-12-31 02:00:00 │
+    #   # └─────────────────────┘
+    def month_end
+      Utils.wrap_expr(_rbexpr.dt_month_end)
+    end
   end
 end