polars-df 0.6.0 → 0.7.0

data/lib/polars/data_frame.rb

@@ -20,15 +20,9 @@ module Polars
     #   this does not yield conclusive results, column orientation is used.
     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)
-        data = {}
-        result.columns.each_with_index do |k, i|
-          data[k] = result.rows.map { |r| r[i] }
-        end
+        raise ArgumentError, "Use read_database instead"
       end
 
       if data.nil?
@@ -905,6 +899,7 @@ module Polars
     def write_csv(
       file = nil,
       has_header: true,
+      include_header: nil,
       sep: ",",
       quote: '"',
       batch_size: 1024,
@@ -914,6 +909,8 @@ module Polars
       float_precision: nil,
       null_value: nil
     )
+      include_header = has_header if include_header.nil?
+
       if sep.length > 1
         raise ArgumentError, "only single byte separator is allowed"
       elsif quote.length > 1
@@ -927,7 +924,7 @@ module Polars
         buffer.set_encoding(Encoding::BINARY)
         _df.write_csv(
           buffer,
-          has_header,
+          include_header,
           sep.ord,
           quote.ord,
           batch_size,
@@ -946,7 +943,7 @@ module Polars
 
       _df.write_csv(
         file,
-        has_header,
+        include_header,
         sep.ord,
         quote.ord,
         batch_size,
@@ -1151,22 +1148,8 @@ module Polars
     #   # │ b   ┆ 1   ┆ 2   ┆ 3   │
     #   # └─────┴─────┴─────┴─────┘
     def transpose(include_header: false, header_name: "column", column_names: nil)
-      df = _from_rbdf(_df.transpose(include_header, header_name))
-      if !column_names.nil?
-        names = []
-        n = df.width
-        if include_header
-          names << header_name
-          n -= 1
-        end
-
-        column_names = column_names.each
-        n.times do
-          names << column_names.next
-        end
-        df.columns = names
-      end
-      df
+      keep_names_as = include_header ? header_name : nil
+      _from_rbdf(_df.transpose(keep_names_as, column_names))
     end
 
     # Reverse the DataFrame.
@@ -1811,13 +1794,13 @@ module Polars
       _from_rbdf(_df.with_row_count(name, offset))
     end
 
-    # Start a groupby operation.
+    # Start a group by operation.
     #
     # @param by [Object]
     #   Column(s) to group by.
     # @param maintain_order [Boolean]
     #   Make sure that the order of the groups remain consistent. This is more
-    #   expensive than a default groupby. Note that this only works in expression
+    #   expensive than a default group by. Note that this only works in expression
     #   aggregations.
     #
     # @return [GroupBy]
@@ -1830,7 +1813,7 @@ module Polars
     #       "c" => [6, 5, 4, 3, 2, 1]
     #     }
     #   )
-    #   df.groupby("a").agg(Polars.col("b").sum).sort("a")
+    #   df.group_by("a").agg(Polars.col("b").sum).sort("a")
     #   # =>
     #   # shape: (3, 2)
     #   # ┌─────┬─────┐
@@ -1842,25 +1825,26 @@ module Polars
     #   # │ b   ┆ 11  │
     #   # │ c   ┆ 6   │
     #   # └─────┴─────┘
-    def groupby(by, maintain_order: false)
+    def group_by(by, maintain_order: false)
       if !Utils.bool?(maintain_order)
-        raise TypeError, "invalid input for groupby arg `maintain_order`: #{maintain_order}."
+        raise TypeError, "invalid input for group_by arg `maintain_order`: #{maintain_order}."
       end
       GroupBy.new(
-        _df,
+        self,
         by,
-        self.class,
         maintain_order: maintain_order
       )
     end
+    alias_method :groupby, :group_by
+    alias_method :group, :group_by
 
     # Create rolling groups based on a time column.
     #
     # Also works for index values of type `:i32` or `:i64`.
     #
-    # Different from a `dynamic_groupby` the windows are now determined by the
+    # Different from a `dynamic_group_by` the windows are now determined by the
     # individual values and are not of constant intervals. For constant intervals use
-    # *groupby_dynamic*
+    # *group_by_dynamic*
     #
     # The `period` and `offset` arguments are created either from a timedelta, or
     # by using the following string language:
@@ -1880,7 +1864,7 @@ module Polars
     # Or combine them:
     # "3d12h4m25s" # 3 days, 12 hours, 4 minutes, and 25 seconds
     #
-    # In case of a groupby_rolling on an integer column, the windows are defined by:
+    # In case of a group_by_rolling on an integer column, the windows are defined by:
     #
     # - **"1i"      # length 1**
     # - **"10i"     # length 10**
@@ -1891,7 +1875,7 @@ module Polars
     #   This column must be sorted in ascending order. If not the output will not
     #   make sense.
     #
-    #   In case of a rolling groupby on indices, dtype needs to be one of
+    #   In case of a rolling group by on indices, dtype needs to be one of
     #   `:i32`, `:i64`. Note that `:i32` gets temporarily cast to `:i64`, so if
     #   performance matters use an `:i64` column.
     # @param period [Object]
@@ -1923,7 +1907,7 @@ module Polars
     #   df = Polars::DataFrame.new({"dt" => dates, "a" => [3, 7, 5, 9, 2, 1]}).with_column(
     #     Polars.col("dt").str.strptime(Polars::Datetime).set_sorted
     #   )
-    #   df.groupby_rolling(index_column: "dt", period: "2d").agg(
+    #   df.group_by_rolling(index_column: "dt", period: "2d").agg(
     #     [
     #       Polars.sum("a").alias("sum_a"),
     #       Polars.min("a").alias("min_a"),
@@ -1944,7 +1928,7 @@ module Polars
     #   # │ 2020-01-03 19:45:32 ┆ 11    ┆ 2     ┆ 9     │
     #   # │ 2020-01-08 23:16:43 ┆ 1     ┆ 1     ┆ 1     │
     #   # └─────────────────────┴───────┴───────┴───────┘
-    def groupby_rolling(
+    def group_by_rolling(
       index_column:,
       period:,
       offset: nil,
@@ -1954,11 +1938,12 @@ module Polars
     )
       RollingGroupBy.new(self, index_column, period, offset, closed, by, check_sorted)
     end
+    alias_method :groupby_rolling, :group_by_rolling
 
     # Group based on a time value (or index value of type `:i32`, `:i64`).
     #
     # Time windows are calculated and rows are assigned to windows. Different from a
-    # normal groupby is that a row can be member of multiple groups. The time/index
+    # normal group by is that a row can be member of multiple groups. The time/index
     # window could be seen as a rolling window, with a window size determined by
     # dates/times/values instead of slots in the DataFrame.
     #
@@ -1986,7 +1971,7 @@ module Polars
     # Or combine them:
     # "3d12h4m25s" # 3 days, 12 hours, 4 minutes, and 25 seconds
     #
-    # In case of a groupby_dynamic on an integer column, the windows are defined by:
+    # In case of a group_by_dynamic on an integer column, the windows are defined by:
     #
     # - "1i"      # length 1
     # - "10i"     # length 10
@@ -1997,7 +1982,7 @@ module Polars
     #   This column must be sorted in ascending order. If not the output will not
     #   make sense.
     #
-    #   In case of a dynamic groupby on indices, dtype needs to be one of
+    #   In case of a dynamic group by on indices, dtype needs to be one of
     #   `:i32`, `:i64`. Note that `:i32` gets temporarily cast to `:i64`, so if
     #   performance matters use an `:i64` column.
     # @param every
@@ -2048,7 +2033,7 @@ module Polars
     #   # └─────────────────────┴─────┘
     #
     # @example Group by windows of 1 hour starting at 2021-12-16 00:00:00.
-    #   df.groupby_dynamic("time", every: "1h", closed: "right").agg(
+    #   df.group_by_dynamic("time", every: "1h", closed: "right").agg(
     #     [
     #       Polars.col("time").min.alias("time_min"),
     #       Polars.col("time").max.alias("time_max")
@@ -2068,7 +2053,7 @@ module Polars
     #   # └─────────────────────┴─────────────────────┴─────────────────────┘
     #
     # @example The window boundaries can also be added to the aggregation result.
-    #   df.groupby_dynamic(
+    #   df.group_by_dynamic(
     #     "time", every: "1h", include_boundaries: true, closed: "right"
     #   ).agg([Polars.col("time").count.alias("time_count")])
     #   # =>
@@ -2085,7 +2070,7 @@ module Polars
     #   # └─────────────────────┴─────────────────────┴─────────────────────┴────────────┘
     #
     # @example When closed="left", should not include right end of interval.
-    #   df.groupby_dynamic("time", every: "1h", closed: "left").agg(
+    #   df.group_by_dynamic("time", every: "1h", closed: "left").agg(
     #     [
     #       Polars.col("time").count.alias("time_count"),
     #       Polars.col("time").alias("time_agg_list")
@@ -2105,7 +2090,7 @@ module Polars
     #   # └─────────────────────┴────────────┴───────────────────────────────────┘
     #
     # @example When closed="both" the time values at the window boundaries belong to 2 groups.
-    #   df.groupby_dynamic("time", every: "1h", closed: "both").agg(
+    #   df.group_by_dynamic("time", every: "1h", closed: "both").agg(
     #     [Polars.col("time").count.alias("time_count")]
     #   )
     #   # =>
@@ -2122,7 +2107,7 @@ module Polars
     #   # │ 2021-12-16 03:00:00 ┆ 1          │
     #   # └─────────────────────┴────────────┘
     #
-    # @example Dynamic groupbys can also be combined with grouping on normal keys.
+    # @example Dynamic group bys can also be combined with grouping on normal keys.
     #   df = Polars::DataFrame.new(
     #     {
     #       "time" => Polars.date_range(
@@ -2133,7 +2118,7 @@ module Polars
     #       "groups" => ["a", "a", "a", "b", "b", "a", "a"]
     #     }
     #   )
-    #   df.groupby_dynamic(
+    #   df.group_by_dynamic(
     #     "time",
     #     every: "1h",
     #     closed: "both",
@@ -2156,14 +2141,14 @@ module Polars
     #   # │ b      ┆ 2021-12-16 02:00:00 ┆ 2021-12-16 03:00:00 ┆ 2021-12-16 02:00:00 ┆ 1          │
     #   # └────────┴─────────────────────┴─────────────────────┴─────────────────────┴────────────┘
     #
-    # @example Dynamic groupby on an index column.
+    # @example Dynamic group by on an index column.
     #   df = Polars::DataFrame.new(
     #     {
     #       "idx" => Polars.arange(0, 6, eager: true),
     #       "A" => ["A", "A", "B", "B", "B", "C"]
     #     }
     #   )
-    #   df.groupby_dynamic(
+    #   df.group_by_dynamic(
     #     "idx",
     #     every: "2i",
     #     period: "3i",
@@ -2181,7 +2166,7 @@ module Polars
     #   # │ 2               ┆ 5               ┆ 2   ┆ ["B", "B", "C"] │
     #   # │ 4               ┆ 7               ┆ 4   ┆ ["C"]           │
     #   # └─────────────────┴─────────────────┴─────┴─────────────────┘
-    def groupby_dynamic(
+    def group_by_dynamic(
       index_column,
       every:,
       period: nil,
@@ -2205,6 +2190,7 @@ module Polars
         start_by
       )
     end
+    alias_method :groupby_dynamic, :group_by_dynamic
 
     # Upsample a DataFrame at a regular frequency.
     #
@@ -3464,8 +3450,10 @@ module Polars
 
     # Shift values by the given period.
     #
-    # @param periods [Integer]
+    # @param n [Integer]
     #   Number of places to shift (may be negative).
+    # @param fill_value [Object]
+    #  Fill the resulting null values with this value.
     #
     # @return [DataFrame]
     #
@@ -3503,8 +3491,8 @@ module Polars
     #   # │ 3    ┆ 8    ┆ c    │
     #   # │ null ┆ null ┆ null │
     #   # └──────┴──────┴──────┘
-    def shift(periods)
-      _from_rbdf(_df.shift(periods))
+    def shift(n, fill_value: nil)
+      lazy.shift(n, fill_value: fill_value).collect(_eager: true)
     end
 
     # Shift the values by a given period and fill the resulting null values.
@@ -3537,9 +3525,7 @@ module Polars
     #   # │ 2   ┆ 7   ┆ b   │
     #   # └─────┴─────┴─────┘
     def shift_and_fill(periods, fill_value)
-      lazy
-        .shift_and_fill(periods, fill_value)
-        .collect(no_optimization: true, string_cache: false)
+      shift(periods, fill_value: fill_value)
     end
 
     # Get a mask of all duplicated rows in this DataFrame.
@@ -3790,7 +3776,7 @@ module Polars
       if axis == 0
         _from_rbdf(_df.max)
       elsif axis == 1
-        Utils.wrap_s(_df.hmax)
+        Utils.wrap_s(_df.max_horizontal)
       else
         raise ArgumentError, "Axis should be 0 or 1."
       end
@@ -3822,7 +3808,7 @@ module Polars
       if axis == 0
         _from_rbdf(_df.min)
       elsif axis == 1
-        Utils.wrap_s(_df.hmin)
+        Utils.wrap_s(_df.min_horizontal)
       else
         raise ArgumentError, "Axis should be 0 or 1."
       end
@@ -3871,7 +3857,7 @@ module Polars
       when 0
         _from_rbdf(_df.sum)
       when 1
-        Utils.wrap_s(_df.hsum(null_strategy))
+        Utils.wrap_s(_df.sum_horizontal(null_strategy))
       else
         raise ArgumentError, "Axis should be 0 or 1."
       end
@@ -3909,7 +3895,7 @@ module Polars
       when 0
         _from_rbdf(_df.mean)
       when 1
-        Utils.wrap_s(_df.hmean(null_strategy))
+        Utils.wrap_s(_df.mean_horizontal(null_strategy))
       else
         raise ArgumentError, "Axis should be 0 or 1."
       end
@@ -4294,15 +4280,20 @@ module Polars
       end
 
       if n.nil? && !frac.nil?
+        frac = Series.new("frac", [frac]) unless frac.is_a?(Series)
+
         _from_rbdf(
-          _df.sample_frac(frac, with_replacement, shuffle, seed)
+          _df.sample_frac(frac._s, with_replacement, shuffle, seed)
         )
       end
 
       if n.nil?
         n = 1
       end
-      _from_rbdf(_df.sample_n(n, with_replacement, shuffle, seed))
+
+      n = Series.new("", [n]) unless n.is_a?(Series)
+
+      _from_rbdf(_df.sample_n(n._s, with_replacement, shuffle, seed))
     end
 
     # Apply a horizontal reduction on a DataFrame.
@@ -4601,7 +4592,7 @@ module Polars
     #
     # @example
     #   s = Polars::DataFrame.new({"a" => [1, 2, 3, 4], "b" => [5, 6, 7, 8]})
-    #   s.take_every(2)
+    #   s.gather_every(2)
     #   # =>
     #   # shape: (2, 2)
     #   # ┌─────┬─────┐
@@ -4612,9 +4603,10 @@ module Polars
     #   # │ 1   ┆ 5   │
     #   # │ 3   ┆ 7   │
     #   # └─────┴─────┘
-    def take_every(n)
-      select(Utils.col("*").take_every(n))
+    def gather_every(n)
+      select(Utils.col("*").gather_every(n))
     end
+    alias_method :take_every, :gather_every
 
     # Hash and combine the rows in this DataFrame.
     #
@@ -4671,16 +4663,16 @@ module Polars
     #   df.interpolate
     #   # =>
     #   # shape: (4, 3)
-    #   # ┌─────┬──────┬─────┐
-    #   # │ foo ┆ bar  ┆ baz │
-    #   # │ --- ┆ ---  ┆ --- │
-    #   # │ i64 ┆ i64  ┆ i64 │
-    #   # ╞═════╪══════╪═════╡
-    #   # │ 1   ┆ 6    ┆ 1   │
-    #   # │ 5   ┆ 7    ┆ 3   │
-    #   # │ 9   ┆ 9    ┆ 6   │
-    #   # │ 10  ┆ null ┆ 9   │
-    #   # └─────┴──────┴─────┘
+    #   # ┌──────┬──────┬──────────┐
+    #   # │ foo  ┆ bar  ┆ baz      │
+    #   # │ ---  ┆ ---  ┆ ---      │
+    #   # │ f64  ┆ f64  ┆ f64      │
+    #   # ╞══════╪══════╪══════════╡
+    #   # │ 1.0  ┆ 6.0  ┆ 1.0      │
+    #   # │ 5.0  ┆ 7.0  ┆ 3.666667 │
+    #   # │ 9.0  ┆ 9.0  ┆ 6.333333 │
+    #   # │ 10.0 ┆ null ┆ 9.0      │
+    #   # └──────┴──────┴──────────┘
     def interpolate
       select(Utils.col("*").interpolate)
     end
@@ -4952,8 +4944,8 @@ module Polars
           [lookup[col[0]] || col[0], col[1]]
         end
 
-      if schema_overrides
-        raise Todo
+      if schema_overrides && schema_overrides.any?
+        column_dtypes.merge!(schema_overrides)
       end
 
       column_dtypes.each do |col, dtype|
@@ -5056,13 +5048,54 @@ module Polars
         return rbdf
       elsif data[0].is_a?(::Array)
         if orient.nil? && !columns.nil?
-          orient = columns.length == data.length ? "col" : "row"
+          first_element = data[0]
+          row_types = first_element.filter_map { |value| value.class }.uniq
+          if row_types.include?(Integer) && row_types.include?(Float)
+            row_types.delete(Integer)
+          end
+          orient = row_types.length == 1 ? "col" : "row"
         end
 
         if orient == "row"
-          raise Todo
+          column_names, schema_overrides = _unpack_schema(
+            schema, schema_overrides: schema_overrides, n_expected: first_element.length
+          )
+          local_schema_override = (
+            schema_overrides.any? ? (raise Todo) : {}
+          )
+          if column_names.any? && first_element.length > 0 && first_element.length != column_names.length
+            raise ArgumentError, "the row data does not match the number of columns"
+          end
+
+          unpack_nested = false
+          local_schema_override.each do |col, tp|
+            raise Todo
+          end
+
+          if unpack_nested
+            raise Todo
+          else
+            rbdf = RbDataFrame.read_rows(
+              data,
+              infer_schema_length,
+              local_schema_override.any? ? local_schema_override : nil
+            )
+          end
+          if column_names.any? || schema_overrides.any?
+            rbdf = _post_apply_columns(
+              rbdf, column_names, schema_overrides: schema_overrides
+            )
+          end
+          return rbdf
         elsif orient == "col" || orient.nil?
-          raise Todo
+          column_names, schema_overrides = _unpack_schema(
+            schema, schema_overrides: schema_overrides, n_expected: data.length
+          )
+          data_series =
+            data.map.with_index do |element, i|
+              Series.new(column_names[i], element, dtype: schema_overrides[column_names[i]])._s
+            end
+          return RbDataFrame.new(data_series)
         else
           raise ArgumentError, "orient must be one of {{'col', 'row', nil}}, got #{orient} instead."
         end
@@ -5108,10 +5141,10 @@ module Polars
 
     def _compare_to_other_df(other, op)
       if columns != other.columns
-        raise ArgmentError, "DataFrame columns do not match"
+        raise ArgumentError, "DataFrame columns do not match"
       end
       if shape != other.shape
-        raise ArgmentError, "DataFrame dimensions do not match"
+        raise ArgumentError, "DataFrame dimensions do not match"
       end
 
       suffix = "__POLARS_CMP_OTHER"

data/lib/polars/date_time_expr.rb

@@ -97,15 +97,20 @@ module Polars
     #   # │ 2001-01-01 00:50:00 ┆ 2001-01-01 00:30:00 │
     #   # │ 2001-01-01 01:00:00 ┆ 2001-01-01 01:00:00 │
     #   # └─────────────────────┴─────────────────────┘
-    def truncate(every, offset: nil)
+    def truncate(every, offset: nil, use_earliest: nil)
       if offset.nil?
         offset = "0ns"
       end
 
+      if !every.is_a?(Expr)
+        every = Utils._timedelta_to_pl_duration(every)
+      end
+      every = Utils.parse_as_expression(every, str_as_lit: true)
+
       Utils.wrap_expr(
         _rbexpr.dt_truncate(
-          Utils._timedelta_to_pl_duration(every),
-          Utils._timedelta_to_pl_duration(offset)
+          every,
+          Utils._timedelta_to_pl_duration(offset),
         )
       )
     end
@@ -1026,21 +1031,10 @@ module Polars
     #   Time zone for the `Datetime` Series.
     #
     # @return [Expr]
-    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.
-    #
-    # This method takes a naive Datetime Series and makes this time zone aware.
-    # It does not move the time to another time zone.
-    #
-    # @param tz [String]
-    #   Time zone for the `Datetime` Series.
-    #
-    # @return [Expr]
-    def tz_localize(tz)
-      Utils.wrap_expr(_rbexpr.dt_tz_localize(tz))
+    def replace_time_zone(tz, use_earliest: nil, ambiguous: "raise")
+      ambiguous = Utils.rename_use_earliest_to_ambiguous(use_earliest, ambiguous)
+      ambiguous = Polars.lit(ambiguous) unless ambiguous.is_a?(Expr)
+      Utils.wrap_expr(_rbexpr.dt_replace_time_zone(tz, ambiguous._rbexpr))
     end
 
     # Extract the days from a Duration type.
@@ -1348,6 +1342,7 @@ module Polars
     #   # │ 2006-01-01 00:00:00 ┆ 2003-11-01 00:00:00 │
     #   # └─────────────────────┴─────────────────────┘
     def offset_by(by)
+      by = Utils.parse_as_expression(by, str_as_lit: true)
       Utils.wrap_expr(_rbexpr.dt_offset_by(by))
     end
 

data/lib/polars/date_time_name_space.rb

@@ -23,18 +23,8 @@ module Polars
     # @return [Object]
     #
     # @example
-    #   date = Polars.date_range(DateTime.new(2001, 1, 1), DateTime.new(2001, 1, 3), "1d")
-    #   # =>
-    #   # shape: (3,)
-    #   # Series: '' [datetime[μs]]
-    #   # [
-    #   #         2001-01-01 00:00:00
-    #   #         2001-01-02 00:00:00
-    #   #         2001-01-03 00:00:00
-    #   # ]
-    #
-    # @example
-    #   date.dt.min
+    #   s = Polars.date_range(DateTime.new(2001, 1, 1), DateTime.new(2001, 1, 3), "1d")
+    #   s.dt.min
     #   # => 2001-01-01 00:00:00 UTC
     def min
       Utils.wrap_s(_s).min
@@ -45,18 +35,8 @@ module Polars
     # @return [Object]
     #
     # @example
-    #   date = Polars.date_range(DateTime.new(2001, 1, 1), DateTime.new(2001, 1, 3), "1d")
-    #   # =>
-    #   # shape: (3,)
-    #   # Series: '' [datetime[μs]]
-    #   # [
-    #   #         2001-01-01 00:00:00
-    #   #         2001-01-02 00:00:00
-    #   #         2001-01-03 00:00:00
-    #   # ]
-    #
-    # @example
-    #   date.dt.max
+    #   s = Polars.date_range(DateTime.new(2001, 1, 1), DateTime.new(2001, 1, 3), "1d")
+    #   s.dt.max
     #   # => 2001-01-03 00:00:00 UTC
     def max
       Utils.wrap_s(_s).max
@@ -1400,7 +1380,7 @@ module Polars
     #   #         2001-01-01 00:30:00
     #   #         2001-01-01 01:00:00
     #   # ]
-    def truncate(every, offset: nil)
+    def truncate(every, offset: nil, use_earliest: nil)
       super
     end
 

data/lib/polars/dynamic_group_by.rb

@@ -2,7 +2,7 @@ module Polars
   # A dynamic grouper.
   #
   # This has an `.agg` method which allows you to run all polars expressions in a
-  # groupby context.
+  # group by context.
   class DynamicGroupBy
     def initialize(
       df,
@@ -34,7 +34,7 @@ module Polars
 
     def agg(aggs)
       @df.lazy
-        .groupby_dynamic(
+        .group_by_dynamic(
           @time_column,
           every: @every,
           period: @period,