polars-df 0.6.0 → 0.7.0

data/lib/polars/functions.rb

@@ -19,8 +19,8 @@ module Polars
     #   DataFrames/Series/LazyFrames to concatenate.
     # @param rechunk [Boolean]
     #   Make sure that all data is in contiguous memory.
-    # @param how ["vertical", "diagonal", "horizontal"]
-    #   Lazy only supports the 'vertical' strategy.
+    # @param how ["vertical", "vertical_relaxed", "diagonal", "horizontal"]
+    #   LazyFrames do not support the `horizontal` strategy.
     #
     #   - Vertical: applies multiple `vstack` operations.
     #   - Diagonal: finds a union between the column schemas and fills missing column values with null.
@@ -55,18 +55,21 @@ module Polars
         if how == "vertical"
           out = Utils.wrap_df(_concat_df(items))
         elsif how == "diagonal"
-          out = Utils.wrap_df(_diag_concat_df(items))
+          out = Utils.wrap_df(_concat_df_diagonal(items))
         elsif how == "horizontal"
-          out = Utils.wrap_df(_hor_concat_df(items))
+          out = Utils.wrap_df(_concat_df_horizontal(items))
         else
           raise ArgumentError, "how must be one of {{'vertical', 'diagonal', 'horizontal'}}, got #{how}"
         end
       elsif first.is_a?(LazyFrame)
         if how == "vertical"
-          # TODO
-          return Utils.wrap_ldf(_concat_lf(items, rechunk, parallel))
+          return Utils.wrap_ldf(_concat_lf(items, rechunk, parallel, false))
+        elsif how == "vertical_relaxed"
+          return Utils.wrap_ldf(_concat_lf(items, rechunk, parallel, true))
+        elsif how == "diagonal"
+          return Utils.wrap_ldf(_concat_lf_diagonal(items, rechunk, parallel, false))
         else
-          raise ArgumentError, "Lazy only allows 'vertical' concat strategy."
+          raise ArgumentError, "Lazy only allows 'vertical', 'vertical_relaxed', and 'diagonal' concat strategy."
         end
       elsif first.is_a?(Series)
         # TODO
@@ -89,9 +92,9 @@ module Polars
 
     # Create a range of type `Datetime` (or `Date`).
     #
-    # @param low [Object]
+    # @param start [Object]
     #   Lower bound of the date range.
-    # @param high [Object]
+    # @param stop [Object]
     #   Upper bound of the date range.
     # @param interval [Object]
     #   Interval periods. It can be a polars duration string, such as `3d12h4m25s`
@@ -145,8 +148,8 @@ module Polars
     #   #         1985-01-10 00:00:00
     #   # ]
     def date_range(
-      low,
-      high,
+      start,
+      stop,
       interval,
       lazy: false,
       closed: "both",
@@ -163,39 +166,28 @@ module Polars
         end
       end
 
-      if low.is_a?(Expr) || high.is_a?(Expr) || lazy
-        low = Utils.expr_to_lit_or_expr(low, str_to_lit: true)
-        high = Utils.expr_to_lit_or_expr(high, str_to_lit: true)
-        return Utils.wrap_expr(
-          _rb_date_range_lazy(low, high, interval, closed, name, time_zone)
-        )
+      if time_unit.nil?
+        if interval.include?("ns")
+          time_unit = "ns"
+        else
+          time_unit = "us"
+        end
       end
 
-      low, low_is_date = _ensure_datetime(low)
-      high, high_is_date = _ensure_datetime(high)
+      start_rbexpr = Utils.parse_as_expression(start)
+      stop_rbexpr = Utils.parse_as_expression(stop)
 
-      if !time_unit.nil?
-        tu = time_unit
-      elsif interval.include?("ns")
-        tu = "ns"
-      else
-        tu = "us"
-      end
+      result = Utils.wrap_expr(
+        _rb_date_range(start_rbexpr, stop_rbexpr, interval, closed, time_unit, time_zone)
+      )
 
-      start = Utils._datetime_to_pl_timestamp(low, tu)
-      stop = Utils._datetime_to_pl_timestamp(high, tu)
-      if name.nil?
-        name = ""
-      end
+      result = result.alias(name.to_s)
 
-      dt_range = Utils.wrap_s(
-        _rb_date_range(start, stop, interval, closed, name, tu, time_zone)
-      )
-      if low_is_date && high_is_date && !["h", "m", "s"].any? { |v| _interval_granularity(interval).end_with?(v) }
-        dt_range = dt_range.cast(Date)
+      if !lazy
+        return select(result).to_series
       end
 
-      dt_range
+      result
     end
 
     # Bin values into discrete values.

data/lib/polars/group_by.rb

@@ -2,23 +2,19 @@ module Polars
   # Starts a new GroupBy operation.
   class GroupBy
     # @private
-    attr_accessor :_df, :_dataframe_class, :by, :maintain_order
-
-    # @private
-    def initialize(df, by, dataframe_class, maintain_order: false)
-      self._df = df
-      self._dataframe_class = dataframe_class
-      self.by = by
-      self.maintain_order = maintain_order
+    def initialize(df, by, maintain_order: false)
+      @df = df
+      @by = by
+      @maintain_order = maintain_order
     end
 
-    # Allows iteration over the groups of the groupby operation.
+    # Allows iteration over the groups of the group by operation.
     #
     # @return [Object]
     #
     # @example
     #   df = Polars::DataFrame.new({"foo" => ["a", "a", "b"], "bar" => [1, 2, 3]})
-    #   df.groupby("foo", maintain_order: true).each.to_h
+    #   df.group_by("foo", maintain_order: true).each.to_h
     #   # =>
     #   # {"a"=>shape: (2, 2)
     #   # ┌─────┬─────┐
@@ -41,10 +37,9 @@ module Polars
 
       temp_col = "__POLARS_GB_GROUP_INDICES"
       groups_df =
-        Utils.wrap_df(_df)
-          .lazy
+        @df.lazy
           .with_row_count(name: temp_col)
-          .groupby(by, maintain_order: maintain_order)
+          .group_by(@by, maintain_order: @maintain_order)
           .agg(Polars.col(temp_col))
           .collect(no_optimization: true)
 
@@ -52,7 +47,7 @@ module Polars
 
       # When grouping by a single column, group name is a single value
       # When grouping by multiple columns, group name is a tuple of values
-      if by.is_a?(String) || by.is_a?(Expr)
+      if @by.is_a?(String) || @by.is_a?(Expr)
         _group_names = group_names.to_series.each
       else
         _group_names = group_names.iter_rows
@@ -62,10 +57,8 @@ module Polars
       _current_index = 0
 
       while _current_index < _group_indices.length
-        df = _dataframe_class._from_rbdf(_df)
-
         group_name = _group_names.next
-        group_data = df[_group_indices[_current_index]]
+        group_data = @df[_group_indices[_current_index]]
         _current_index += 1
 
         yield group_name, group_data
@@ -96,7 +89,7 @@ module Polars
     #       "shape" => ["square", "triangle", "square", "triangle", "square"]
     #     }
     #   )
-    #   df.groupby("color").apply { |group_df| group_df.sample(2) }
+    #   df.group_by("color").apply { |group_df| group_df.sample(2) }
     #   # =>
     #   # shape: (4, 3)
     #   # ┌─────┬───────┬──────────┐
@@ -110,7 +103,7 @@ module Polars
     #   # │ 3   ┆ red   ┆ triangle │
     #   # └─────┴───────┴──────────┘
     # def apply(&f)
-    #   _dataframe_class._from_rbdf(_df.groupby_apply(by, f))
+    #   _dataframe_class._from_rbdf(_df.group_by_apply(by, f))
     # end
 
     # Use multiple aggregations on columns.
@@ -126,7 +119,7 @@ module Polars
     #   df = Polars::DataFrame.new(
     #     {"foo" => ["one", "two", "two", "one", "two"], "bar" => [5, 3, 2, 4, 1]}
     #   )
-    #   df.groupby("foo", maintain_order: true).agg(
+    #   df.group_by("foo", maintain_order: true).agg(
     #     [
     #       Polars.sum("bar").suffix("_sum"),
     #       Polars.col("bar").sort.tail(2).sum.suffix("_tail_sum")
@@ -143,12 +136,10 @@ module Polars
     #   # │ two ┆ 6       ┆ 5            │
     #   # └─────┴─────────┴──────────────┘
     def agg(aggs)
-      df = Utils.wrap_df(_df)
-        .lazy
-        .groupby(by, maintain_order: maintain_order)
+      @df.lazy
+        .group_by(@by, maintain_order: @maintain_order)
         .agg(aggs)
-        .collect(no_optimization: true, string_cache: false)
-      _dataframe_class._from_rbdf(df._df)
+        .collect(no_optimization: true)
     end
 
     # Get the first `n` rows of each group.
@@ -181,7 +172,7 @@ module Polars
     #   # └─────────┴─────┘
     #
     # @example
-    #   df.groupby("letters").head(2).sort("letters")
+    #   df.group_by("letters").head(2).sort("letters")
     #   # =>
     #   # shape: (5, 2)
     #   # ┌─────────┬─────┐
@@ -196,14 +187,10 @@ module Polars
     #   # │ c       ┆ 2   │
     #   # └─────────┴─────┘
     def head(n = 5)
-      df = (
-        Utils.wrap_df(_df)
-          .lazy
-          .groupby(by, maintain_order: maintain_order)
-          .head(n)
-          .collect(no_optimization: true, string_cache: false)
-      )
-      _dataframe_class._from_rbdf(df._df)
+      @df.lazy
+        .group_by(@by, maintain_order: @maintain_order)
+        .head(n)
+        .collect(no_optimization: true)
     end
 
     # Get the last `n` rows of each group.
@@ -236,7 +223,7 @@ module Polars
     #   # └─────────┴─────┘
     #
     # @example
-    #   df.groupby("letters").tail(2).sort("letters")
+    #   df.group_by("letters").tail(2).sort("letters")
     #   # =>
     #   # shape: (5, 2)
     #   # ┌─────────┬─────┐
@@ -251,14 +238,10 @@ module Polars
     #   # │ c       ┆ 4   │
     #   # └─────────┴─────┘
     def tail(n = 5)
-      df = (
-        Utils.wrap_df(_df)
-          .lazy
-          .groupby(by, maintain_order: maintain_order)
-          .tail(n)
-          .collect(no_optimization: true, string_cache: false)
-      )
-      _dataframe_class._from_rbdf(df._df)
+      @df.lazy
+        .group_by(@by, maintain_order: @maintain_order)
+        .tail(n)
+        .collect(no_optimization: true)
     end
 
     # Aggregate the first values in the group.
@@ -274,7 +257,7 @@ module Polars
     #       "d" => ["Apple", "Orange", "Apple", "Apple", "Banana", "Banana"]
     #     }
     #   )
-    #   df.groupby("d", maintain_order: true).first
+    #   df.group_by("d", maintain_order: true).first
     #   # =>
     #   # shape: (3, 4)
     #   # ┌────────┬─────┬──────┬───────┐
@@ -303,7 +286,7 @@ module Polars
     #       "d" => ["Apple", "Orange", "Apple", "Apple", "Banana", "Banana"]
     #     }
     #   )
-    #   df.groupby("d", maintain_order: true).last
+    #   df.group_by("d", maintain_order: true).last
     #   # =>
     #   # shape: (3, 4)
     #   # ┌────────┬─────┬──────┬───────┐
@@ -332,7 +315,7 @@ module Polars
     #       "d" => ["Apple", "Orange", "Apple", "Apple", "Banana", "Banana"]
     #     }
     #   )
-    #   df.groupby("d", maintain_order: true).sum
+    #   df.group_by("d", maintain_order: true).sum
     #   # =>
     #   # shape: (3, 4)
     #   # ┌────────┬─────┬──────┬─────┐
@@ -361,7 +344,7 @@ module Polars
     #       "d" => ["Apple", "Orange", "Apple", "Apple", "Banana", "Banana"],
     #     }
     #   )
-    #   df.groupby("d", maintain_order: true).min
+    #   df.group_by("d", maintain_order: true).min
     #   # =>
     #   # shape: (3, 4)
     #   # ┌────────┬─────┬──────┬───────┐
@@ -390,7 +373,7 @@ module Polars
     #       "d" => ["Apple", "Orange", "Apple", "Apple", "Banana", "Banana"]
     #     }
     #   )
-    #   df.groupby("d", maintain_order: true).max
+    #   df.group_by("d", maintain_order: true).max
     #   # =>
     #   # shape: (3, 4)
     #   # ┌────────┬─────┬──────┬──────┐
@@ -419,7 +402,7 @@ module Polars
     #       "d" => ["Apple", "Orange", "Apple", "Apple", "Banana", "Banana"]
     #     }
     #   )
-    #   df.groupby("d", maintain_order: true).count
+    #   df.group_by("d", maintain_order: true).count
     #   # =>
     #   # shape: (3, 2)
     #   # ┌────────┬───────┐
@@ -448,7 +431,7 @@ module Polars
     #       "d" => ["Apple", "Orange", "Apple", "Apple", "Banana", "Banana"]
     #     }
     #   )
-    #   df.groupby("d", maintain_order: true).mean
+    #   df.group_by("d", maintain_order: true).mean
     #   # =>
     #   # shape: (3, 4)
     #   # ┌────────┬─────┬──────────┬──────────┐
@@ -476,7 +459,7 @@ module Polars
     #       "d" => ["Apple", "Banana", "Apple", "Apple", "Banana", "Banana"]
     #     }
     #   )
-    #   df.groupby("d", maintain_order: true).n_unique
+    #   df.group_by("d", maintain_order: true).n_unique
     #   # =>
     #   # shape: (2, 3)
     #   # ┌────────┬─────┬─────┐
@@ -508,7 +491,7 @@ module Polars
     #       "d" => ["Apple", "Orange", "Apple", "Apple", "Banana", "Banana"]
     #     }
     #   )
-    #   df.groupby("d", maintain_order: true).quantile(1)
+    #   df.group_by("d", maintain_order: true).quantile(1)
     #   # =>
     #   # shape: (3, 3)
     #   # ┌────────┬─────┬──────┐
@@ -536,7 +519,7 @@ module Polars
     #       "d" => ["Apple", "Banana", "Apple", "Apple", "Banana", "Banana"]
     #     }
     #   )
-    #   df.groupby("d", maintain_order: true).median
+    #   df.group_by("d", maintain_order: true).median
     #   # =>
     #   # shape: (2, 3)
     #   # ┌────────┬─────┬──────┐
@@ -555,11 +538,11 @@ module Polars
     #
     # @return [Vega::LiteChart]
     def plot(*args, **options)
-      raise ArgumentError, "Multiple groups not supported" if by.is_a?(::Array) && by.size > 1
+      raise ArgumentError, "Multiple groups not supported" if @by.is_a?(::Array) && @by.size > 1
       # same message as Ruby
       raise ArgumentError, "unknown keyword: :group" if options.key?(:group)
 
-      Utils.wrap_df(_df).plot(*args, **options, group: by)
+      @df.plot(*args, **options, group: @by)
     end
   end
 end

data/lib/polars/io.rb

@@ -621,11 +621,46 @@ module Polars
         else
           raise ArgumentError, "Expected ActiveRecord::Relation, ActiveRecord::Result, or String"
         end
+
       data = {}
+      schema_overrides = {}
+
       result.columns.each_with_index do |k, i|
-        data[k] = result.rows.map { |r| r[i] }
+        column_type = result.column_types[i]
+
+        data[k] =
+          if column_type
+            result.rows.map { |r| column_type.deserialize(r[i]) }
+          else
+            result.rows.map { |r| r[i] }
+          end
+
+        polars_type =
+          case column_type&.type
+          when :binary
+            Binary
+          when :boolean
+            Boolean
+          when :date
+            Date
+          when :datetime, :timestamp
+            Datetime
+          when :decimal
+            Decimal
+          when :float
+            Float64
+          when :integer
+            Int64
+          when :string, :text
+            Utf8
+          when :time
+            Time
+          end
+
+        schema_overrides[k] = polars_type if polars_type
       end
-      DataFrame.new(data)
+
+      DataFrame.new(data, schema_overrides: schema_overrides)
     end
     alias_method :read_sql, :read_database