polars-df 0.14.0 → 0.16.0

data/lib/polars/functions/as_datatype.rb

@@ -86,8 +86,57 @@ module Polars
     # Concat the arrays in a Series dtype List in linear time.
     #
     # @return [Expr]
-    def concat_list(exprs)
-      exprs = Utils.parse_into_list_of_expressions(exprs)
+    #
+    # @example Concatenate two existing list columns. Null values are propagated.
+    #   df = Polars::DataFrame.new({"a" => [[1, 2], [3], [4, 5]], "b" => [[4], [], nil]})
+    #   df.with_columns(concat_list: Polars.concat_list("a", "b"))
+    #   # =>
+    #   # shape: (3, 3)
+    #   # ┌───────────┬───────────┬─────────────┐
+    #   # │ a         ┆ b         ┆ concat_list │
+    #   # │ ---       ┆ ---       ┆ ---         │
+    #   # │ list[i64] ┆ list[i64] ┆ list[i64]   │
+    #   # ╞═══════════╪═══════════╪═════════════╡
+    #   # │ [1, 2]    ┆ [4]       ┆ [1, 2, 4]   │
+    #   # │ [3]       ┆ []        ┆ [3]         │
+    #   # │ [4, 5]    ┆ null      ┆ null        │
+    #   # └───────────┴───────────┴─────────────┘
+    #
+    # @example Non-list columns are cast to a list before concatenation. The output data type is the supertype of the concatenated columns.
+    #   df.select("a", concat_list: Polars.concat_list("a", Polars.lit("x")))
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌───────────┬─────────────────┐
+    #   # │ a         ┆ concat_list     │
+    #   # │ ---       ┆ ---             │
+    #   # │ list[i64] ┆ list[str]       │
+    #   # ╞═══════════╪═════════════════╡
+    #   # │ [1, 2]    ┆ ["1", "2", "x"] │
+    #   # │ [3]       ┆ ["3", "x"]      │
+    #   # │ [4, 5]    ┆ ["4", "5", "x"] │
+    #   # └───────────┴─────────────────┘
+    #
+    # @example Create lagged columns and collect them into a list. This mimics a rolling window.
+    #   df = Polars::DataFrame.new({"A" => [1.0, 2.0, 9.0, 2.0, 13.0]})
+    #   df = df.select(3.times.map { |i| Polars.col("A").shift(i).alias("A_lag_#{i}") })
+    #   df.select(
+    #     Polars.concat_list(3.times.map { |i| "A_lag_#{i}" }.reverse).alias("A_rolling")
+    #   )
+    #   # =>
+    #   # shape: (5, 1)
+    #   # ┌───────────────────┐
+    #   # │ A_rolling         │
+    #   # │ ---               │
+    #   # │ list[f64]         │
+    #   # ╞═══════════════════╡
+    #   # │ [null, null, 1.0] │
+    #   # │ [null, 1.0, 2.0]  │
+    #   # │ [1.0, 2.0, 9.0]   │
+    #   # │ [2.0, 9.0, 2.0]   │
+    #   # │ [9.0, 2.0, 13.0]  │
+    #   # └───────────────────┘
+    def concat_list(exprs, *more_exprs)
+      exprs = Utils.parse_into_list_of_expressions(exprs, *more_exprs)
       Utils.wrap_expr(Plr.concat_list(exprs))
     end
 

data/lib/polars/functions/col.rb

@@ -23,7 +23,7 @@ module Polars
         Utils.wrap_expr(Plr.col(name.to_s))
       elsif Utils.is_polars_dtype(name)
         Utils.wrap_expr(Plr.dtype_cols([name]))
-      elsif name.is_a?(::Array)
+      elsif name.is_a?(::Array) || name.is_a?(::Set)
         names = Array(name)
         if names.empty?
           return Utils.wrap_expr(Plr.cols(names))

data/lib/polars/functions/eager.rb

@@ -127,7 +127,7 @@ module Polars
     #   af1, af2, af3 = Polars.align_frames(
     #     df1, df2, df3, on: "dt", select: ["x", "y"]
     #   )
-    #   (af1 * af2 * af3).fill_null(0).select(Polars.sum(Polars.col("*")).alias("dot"))
+    #   (af1 * af2 * af3).fill_null(0).select(Polars.sum_horizontal("*").alias("dot"))
     #   # =>
     #   # shape: (3, 1)
     #   # ┌───────┐
@@ -136,9 +136,7 @@ module Polars
     #   # │ f64   │
     #   # ╞═══════╡
     #   # │ 0.0   │
-    #   # ├╌╌╌╌╌╌╌┤
     #   # │ 167.5 │
-    #   # ├╌╌╌╌╌╌╌┤
     #   # │ 47.0  │
     #   # └───────┘
     def align_frames(

data/lib/polars/functions/lazy.rb

@@ -729,16 +729,20 @@ module Polars
       a,
       b,
       method: "pearson",
-      ddof: 1,
+      ddof: nil,
       propagate_nans: false
     )
+      if !ddof.nil?
+        warn "The `ddof` parameter has no effect. Do not use it."
+      end
+
       a = Utils.parse_into_expression(a)
       b = Utils.parse_into_expression(b)
 
       if method == "pearson"
-        Utils.wrap_expr(Plr.pearson_corr(a, b, ddof))
+        Utils.wrap_expr(Plr.pearson_corr(a, b))
       elsif method == "spearman"
-        Utils.wrap_expr(Plr.spearman_rank_corr(a, b, ddof, propagate_nans))
+        Utils.wrap_expr(Plr.spearman_rank_corr(a, b, propagate_nans))
       else
         msg = "method must be one of {{'pearson', 'spearman'}}, got #{method}"
         raise ArgumentError, msg
@@ -824,6 +828,29 @@ module Polars
     # @note
     #   If you simply want the first encountered expression as accumulator,
     #   consider using `cumreduce`.
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "a" => [1, 2, 3],
+    #       "b" => [3, 4, 5],
+    #       "c" => [5, 6, 7]
+    #     }
+    #   )
+    #   df.with_columns(
+    #     Polars.cum_fold(Polars.lit(1), ->(acc, x) { acc + x }, Polars.all)
+    #   )
+    #   # =>
+    #   # shape: (3, 4)
+    #   # ┌─────┬─────┬─────┬───────────┐
+    #   # │ a   ┆ b   ┆ c   ┆ cum_fold  │
+    #   # │ --- ┆ --- ┆ --- ┆ ---       │
+    #   # │ i64 ┆ i64 ┆ i64 ┆ struct[3] │
+    #   # ╞═════╪═════╪═════╪═══════════╡
+    #   # │ 1   ┆ 3   ┆ 5   ┆ {2,5,10}  │
+    #   # │ 2   ┆ 4   ┆ 6   ┆ {3,7,13}  │
+    #   # │ 3   ┆ 5   ┆ 7   ┆ {4,9,16}  │
+    #   # └─────┴─────┴─────┴───────────┘
     def cum_fold(acc, f, exprs, include_init: false)
       acc = Utils.parse_into_expression(acc, str_as_lit: true)
       if exprs.is_a?(Expr)
@@ -831,7 +858,7 @@ module Polars
       end
 
       exprs = Utils.parse_into_list_of_expressions(exprs)
-      Utils.wrap_expr(Plr.cum_fold(acc, f, exprs, include_init))
+      Utils.wrap_expr(Plr.cum_fold(acc, f, exprs, include_init)._alias("cum_fold"))
     end
     alias_method :cumfold, :cum_fold
 
@@ -1024,15 +1051,70 @@ module Polars
     #   Default is ascending.
     #
     # @return [Expr]
-    def arg_sort_by(exprs, reverse: false)
-      if !exprs.is_a?(::Array)
-        exprs = [exprs]
-      end
-      if reverse == true || reverse == false
-        reverse = [reverse] * exprs.length
-      end
-      exprs = Utils.parse_into_list_of_expressions(exprs)
-      Utils.wrap_expr(Plr.arg_sort_by(exprs, reverse))
+    #
+    # @example Pass a single column name to compute the arg sort by that column.
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "a" => [0, 1, 1, 0],
+    #       "b" => [3, 2, 3, 2],
+    #       "c" => [1, 2, 3, 4]
+    #     }
+    #   )
+    #   df.select(Polars.arg_sort_by("a"))
+    #   # =>
+    #   # shape: (4, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ u32 │
+    #   # ╞═════╡
+    #   # │ 0   │
+    #   # │ 3   │
+    #   # │ 1   │
+    #   # │ 2   │
+    #   # └─────┘
+    #
+    # @example Compute the arg sort by multiple columns by either passing a list of columns, or by specifying each column as a positional argument.
+    #   df.select(Polars.arg_sort_by(["a", "b"], reverse: true))
+    #   # =>
+    #   # shape: (4, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ u32 │
+    #   # ╞═════╡
+    #   # │ 2   │
+    #   # │ 1   │
+    #   # │ 0   │
+    #   # │ 3   │
+    #   # └─────┘
+    #
+    # @example Use gather to apply the arg sort to other columns.
+    #   df.select(Polars.col("c").gather(Polars.arg_sort_by("a")))
+    #   # =>
+    #   # shape: (4, 1)
+    #   # ┌─────┐
+    #   # │ c   │
+    #   # │ --- │
+    #   # │ i64 │
+    #   # ╞═════╡
+    #   # │ 1   │
+    #   # │ 4   │
+    #   # │ 2   │
+    #   # │ 3   │
+    #   # └─────┘
+    def arg_sort_by(
+      exprs,
+      *more_exprs,
+      reverse: false,
+      nulls_last: false,
+      multithreaded: true,
+      maintain_order: false
+    )
+      exprs = Utils.parse_into_list_of_expressions(exprs, *more_exprs)
+      reverse = Utils.extend_bool(reverse, exprs.length, "reverse", "exprs")
+      nulls_last = Utils.extend_bool(nulls_last, exprs.length, "nulls_last", "exprs")
+      Utils.wrap_expr(Plr.arg_sort_by(exprs, reverse, nulls_last, multithreaded, maintain_order))
     end
     alias_method :argsort_by, :arg_sort_by
 

data/lib/polars/functions/range/time_range.rb

@@ -18,7 +18,7 @@ module Polars
     #
     # @example
     #   Polars.time_range(
-    #     time(14, 0),
+    #     Time.utc(2000, 1, 1, 14, 0),
     #     nil,
     #     "3h15m",
     #     eager: true
@@ -48,12 +48,12 @@ module Polars
       end
 
       if start.nil?
-        # start = time(0, 0, 0)
-        raise Todo
+        # date part is ignored
+        start = ::Time.utc(2000, 1, 1, 0, 0, 0)
       end
       if stop.nil?
-        # stop = time(23, 59, 59, 999999)
-        raise Todo
+        # date part is ignored
+        stop = ::Time.utc(2000, 1, 1, 23, 59, 59, 999999)
       end
 
       start_rbexpr = Utils.parse_into_expression(start)
@@ -87,21 +87,21 @@ module Polars
     # @example
     #   df = Polars::DataFrame.new(
     #     {
-    #       "start" => [time(9, 0), time(10, 0)],
-    #       "end" => time(11, 0)
+    #       "start" => [Time.utc(2000, 1, 1, 9, 0), Time.utc(2000, 1, 1, 10, 0)],
+    #       "end" => Time.utc(2000, 1, 1, 11, 0)
     #     }
     #   )
-    #   df.with_columns(time_range: Polars.time_ranges("start", "end"))
+    #   df.select(time_range: Polars.time_ranges("start", "end"))
     #   # =>
-    #   # shape: (2, 3)
-    #   # ┌──────────┬──────────┬────────────────────────────────┐
-    #   # │ start    ┆ end      ┆ time_range                     │
-    #   # │ ---      ┆ ---      ┆ ---                            │
-    #   # │ time     ┆ time     ┆ list[time]                     │
-    #   # ╞══════════╪══════════╪════════════════════════════════╡
-    #   # │ 09:00:00 ┆ 11:00:00 ┆ [09:00:00, 10:00:00, 11:00:00] │
-    #   # │ 10:00:00 ┆ 11:00:00 ┆ [10:00:00, 11:00:00]           │
-    #   # └──────────┴──────────┴────────────────────────────────┘
+    #   # shape: (2, 1)
+    #   # ┌────────────────────────────────┐
+    #   # │ time_range                     │
+    #   # │ ---                            │
+    #   # │ list[time]                     │
+    #   # ╞════════════════════════════════╡
+    #   # │ [09:00:00, 10:00:00, 11:00:00] │
+    #   # │ [10:00:00, 11:00:00]           │
+    #   # └────────────────────────────────┘
     def time_ranges(
       start = nil,
       stop = nil,
@@ -118,12 +118,12 @@ module Polars
       end
 
       if start.nil?
-        # start = time(0, 0, 0)
-        raise Todo
+        # date part is ignored
+        start = ::Time.utc(2000, 1, 1, 0, 0, 0)
       end
       if stop.nil?
-        # stop = time(23, 59, 59, 999999)
-        raise Todo
+        # date part is ignored
+        stop = ::Time.utc(2000, 1, 1, 23, 59, 59, 999999)
       end
 
       start_rbexpr = Utils.parse_into_expression(start)

data/lib/polars/io/csv.rb

@@ -75,9 +75,6 @@ module Polars
     #   the DataFrame.
     # @param row_count_offset [Integer]
     #   Offset to start the row_count column (only used if the name is set).
-    # @param sample_size [Integer]
-    #   Set the sample size. This is used to sample statistics to estimate the
-    #   allocation needed.
     # @param eol_char [String]
     #   Single byte end of line character.
     # @param truncate_ragged_lines [Boolean]
@@ -114,7 +111,6 @@ module Polars
       skip_rows_after_header: 0,
       row_count_name: nil,
       row_count_offset: 0,
-      sample_size: 1024,
       eol_char: "\n",
       truncate_ragged_lines: false
     )
@@ -163,7 +159,6 @@ module Polars
           skip_rows_after_header: skip_rows_after_header,
           row_count_name: row_count_name,
           row_count_offset: row_count_offset,
-          sample_size: sample_size,
           eol_char: eol_char,
           truncate_ragged_lines: truncate_ragged_lines
         )
@@ -201,7 +196,6 @@ module Polars
       skip_rows_after_header: 0,
       row_count_name: nil,
       row_count_offset: 0,
-      sample_size: 1024,
       eol_char: "\n",
       raise_if_empty: true,
       truncate_ragged_lines: false,
@@ -305,7 +299,6 @@ module Polars
           parse_dates,
           skip_rows_after_header,
           Utils.parse_row_index_args(row_count_name, row_count_offset),
-          sample_size,
           eol_char,
           raise_if_empty,
           truncate_ragged_lines,
@@ -392,9 +385,6 @@ module Polars
     #   the DataFrame.
     # @param row_count_offset [Integer]
     #   Offset to start the row_count column (only used if the name is set).
-    # @param sample_size [Integer]
-    #   Set the sample size. This is used to sample statistics to estimate the
-    #   allocation needed.
     # @param eol_char [String]
     #   Single byte end of line character.
     # @param truncate_ragged_lines [Boolean]
@@ -431,7 +421,6 @@ module Polars
       skip_rows_after_header: 0,
       row_count_name: nil,
       row_count_offset: 0,
-      sample_size: 1024,
       eol_char: "\n",
       raise_if_empty: true,
       truncate_ragged_lines: false,
@@ -474,7 +463,6 @@ module Polars
         skip_rows_after_header: skip_rows_after_header,
         row_count_name: row_count_name,
         row_count_offset: row_count_offset,
-        sample_size: sample_size,
         eol_char: eol_char,
         new_columns: new_columns,
         raise_if_empty: raise_if_empty,
@@ -618,7 +606,7 @@ module Polars
 
     # @private
     def _scan_csv_impl(
-      file,
+      source,
       has_header: true,
       sep: ",",
       comment_char: nil,
@@ -650,9 +638,16 @@ module Polars
       end
       processed_null_values = Utils._process_null_values(null_values)
 
+      if source.is_a?(::Array)
+        sources = source
+        source = nil
+      else
+        sources = []
+      end
+
       rblf =
         RbLazyFrame.new_from_csv(
-          file,
+          source,
           sep,
           has_header,
           ignore_errors,
@@ -672,7 +667,8 @@ module Polars
           Utils.parse_row_index_args(row_count_name, row_count_offset),
           parse_dates,
           eol_char,
-          truncate_ragged_lines
+          truncate_ragged_lines,
+          sources
         )
       Utils.wrap_ldf(rblf)
     end
@@ -681,7 +677,9 @@ module Polars
 
     def _prepare_file_arg(file)
       if file.is_a?(::String) && file =~ /\Ahttps?:\/\//
-        raise ArgumentError, "use URI(...) for remote files"
+        require "uri"
+
+        file = URI(file)
       end
 
       if defined?(URI) && file.is_a?(URI)

data/lib/polars/io/database.rb

@@ -18,9 +18,9 @@ module Polars
         if query.is_a?(ActiveRecord::Result)
           query
         elsif query.is_a?(ActiveRecord::Relation)
-          query.connection.select_all(query.to_sql)
+          query.connection_pool.with_connection { |c| c.select_all(query.to_sql) }
         elsif query.is_a?(::String)
-          ActiveRecord::Base.connection.select_all(query)
+          ActiveRecord::Base.connection_pool.with_connection { |c| c.select_all(query) }
         else
           raise ArgumentError, "Expected ActiveRecord::Relation, ActiveRecord::Result, or String"
         end

data/lib/polars/io/delta.rb

@@ -0,0 +1,126 @@
+module Polars
+  module IO
+    # Reads into a DataFrame from a Delta lake table.
+    #
+    # @param source [Object]
+    #   DeltaTable or a Path or URI to the root of the Delta lake table.
+    # @param version [Object]
+    #   Numerical version or timestamp version of the Delta lake table.
+    # @param columns [Array]
+    #   Columns to select. Accepts a list of column names.
+    # @param rechunk [Boolean]
+    #   Make sure that all columns are contiguous in memory by
+    #   aggregating the chunks into a single array.
+    # @param storage_options [Hash]
+    #   Extra options for the storage backends supported by `deltalake-rb`.
+    # @param delta_table_options [Hash]
+    #   Additional keyword arguments while reading a Delta lake Table.
+    #
+    # @return [DataFrame]
+    def read_delta(
+      source,
+      version: nil,
+      columns: nil,
+      rechunk: false,
+      storage_options: nil,
+      delta_table_options: nil
+    )
+      dl_tbl =
+        _get_delta_lake_table(
+          source,
+          version: version,
+          storage_options: storage_options,
+          delta_table_options: delta_table_options
+        )
+
+      dl_tbl.to_polars(columns: columns, rechunk: rechunk)
+    end
+
+    # Lazily read from a Delta lake table.
+    #
+    # @param source [Object]
+    #   DeltaTable or a Path or URI to the root of the Delta lake table.
+    # @param version [Object]
+    #   Numerical version or timestamp version of the Delta lake table.
+    # @param storage_options [Hash]
+    #   Extra options for the storage backends supported by `deltalake-rb`.
+    # @param delta_table_options [Hash]
+    #   Additional keyword arguments while reading a Delta lake Table.
+    #
+    # @return [LazyFrame]
+    def scan_delta(
+      source,
+      version: nil,
+      storage_options: nil,
+      delta_table_options: nil
+    )
+      dl_tbl =
+        _get_delta_lake_table(
+          source,
+          version: version,
+          storage_options: storage_options,
+          delta_table_options: delta_table_options
+        )
+
+      dl_tbl.to_polars(eager: false)
+    end
+
+    private
+
+    def _resolve_delta_lake_uri(table_uri, strict: true)
+      require "uri"
+
+      parsed_result = URI(table_uri)
+
+      resolved_uri =
+        if parsed_result.scheme == ""
+          Utils.normalize_filepath(table_uri)
+        else
+          table_uri
+        end
+
+      resolved_uri
+    end
+
+    def _get_delta_lake_table(
+      table_path,
+      version: nil,
+      storage_options: nil,
+      delta_table_options: nil
+    )
+      _check_if_delta_available
+
+      if table_path.is_a?(DeltaLake::Table)
+        return table_path
+      end
+      delta_table_options ||= {}
+      resolved_uri = _resolve_delta_lake_uri(table_path)
+      if !version.is_a?(::String) && !version.is_a?(::Time)
+        dl_tbl =
+          DeltaLake::Table.new(
+            resolved_uri,
+            version: version,
+            storage_options: storage_options,
+            **delta_table_options
+          )
+      else
+        dl_tbl =
+          DeltaLake::Table.new(
+            resolved_uri,
+            storage_options: storage_options,
+            **delta_table_options
+          )
+        dl_tbl.load_as_version(version)
+      end
+
+      dl_tbl = DeltaLake::Table.new(table_path)
+      dl_tbl
+    end
+
+    def _check_if_delta_available
+      if !defined?(DeltaLake)
+        raise Error, "Delta Lake not available"
+      end
+    end
+  end
+end

data/lib/polars/io/ipc.rb

@@ -233,7 +233,7 @@ module Polars
 
     # @private
     def _scan_ipc_impl(
-      file,
+      source,
       n_rows: nil,
       cache: true,
       rechunk: true,
@@ -245,13 +245,23 @@ module Polars
       try_parse_hive_dates: true,
       include_file_paths: nil
     )
-      if Utils.pathlike?(file)
-        file = Utils.normalize_filepath(file)
+      sources = []
+      if Utils.pathlike?(source)
+        source = Utils.normalize_filepath(source)
+      elsif source.is_a?(::Array)
+        if Utils.is_path_or_str_sequence(source)
+          sources = source.map { |s| Utils.normalize_filepath(s) }
+        else
+          sources = source
+        end
+
+        source = nil
       end
 
       rblf =
         RbLazyFrame.new_from_ipc(
-          file,
+          source,
+          sources,
           n_rows,
           cache,
           rechunk,

data/lib/polars/io/ndjson.rb

@@ -60,13 +60,23 @@ module Polars
       row_count_name: nil,
       row_count_offset: 0
     )
+      sources = []
       if Utils.pathlike?(source)
         source = Utils.normalize_filepath(source)
+      elsif source.is_a?(::Array)
+        if Utils.is_path_or_str_sequence(source)
+          sources = source.map { |s| Utils.normalize_filepath(s) }
+        else
+          sources = source
+        end
+
+        source = nil
       end
 
       rblf =
         RbLazyFrame.new_from_ndjson(
           source,
+          sources,
           infer_schema_length,
           batch_size,
           n_rows,