polars-df 0.25.1 → 0.26.1

data/lib/polars/functions/aggregation/horizontal.rb

@@ -242,10 +242,10 @@ module Polars
       rbexprs = Utils.parse_into_list_of_expressions(*exprs)
       exprs_wrapped = rbexprs.map { |e| Utils.wrap_expr(e) }
 
-      # (Expr): use u32 as that will not cast to float as eagerly
-      Polars.cum_fold(Polars.lit(0).cast(UInt32), exprs_wrapped) { |a, b| a + b }.alias(
-        "cum_sum"
-      )
+      Polars.cum_fold(
+        Polars.lit(0).cast(Polars.dtype_of(Polars.sum_horizontal(exprs))),
+        exprs_wrapped
+      ) { |a, b| a + b }.alias("cum_sum")
     end
   end
 end

data/lib/polars/functions/business.rb

@@ -81,13 +81,13 @@ module Polars
     )
       start_rbexpr = Utils.parse_into_expression(start)
       end_rbexpr = Utils.parse_into_expression(stop)
-      unix_epoch = ::Date.new(1970, 1, 1)
+      holidays_rbexpr = Utils._holidays_to_expr(holidays)
       Utils.wrap_expr(
         Plr.business_day_count(
           start_rbexpr,
           end_rbexpr,
           week_mask,
-          holidays.map { |holiday| holiday - unix_epoch }
+          holidays_rbexpr
         )
       )
     end

data/lib/polars/functions/eager.rb

@@ -399,6 +399,85 @@ module Polars
       out
     end
 
+    # Merge multiple sorted DataFrames or LazyFrames by the sorted key.
+    #
+    # The output of this operation will also be sorted.
+    # It is the callers responsibility that the frames
+    # are sorted in ascending order by that key otherwise
+    # the output will not make sense.
+    #
+    # @note
+    #   This functionality is considered **unstable**. It may be changed
+    #   at any point without it being considered a breaking change.
+    #
+    # @param items [Array]
+    #   DataFrames or LazyFrames to merge.
+    # @param key [String]
+    #   Key that is sorted.
+    # @param maintain_order [Boolean]
+    #   If `true`, the output is guaranteed to have left-biased ordering
+    #   for equal keys: rows from the left frame appear before rows from
+    #   the right frame when their keys are equal.
+    #
+    # @return [Object]
+    #
+    # @example
+    #   df0 = Polars::DataFrame.new(
+    #     {"name" => ["steve", "elise", "bob"], "age" => [42, 44, 18]}
+    #   ).sort("age")
+    #   df1 = Polars::DataFrame.new(
+    #     {"name" => ["anna", "megan", "steve", "thomas"], "age" => [21, 33, 17, 20]}
+    #   ).sort("age")
+    #   df2 = Polars::DataFrame.new({"name" => ["ida", "maya"], "age" => [37, 27]}).sort("age")
+    #   Polars.merge_sorted([df0, df1, df2], "age")
+    #   # =>
+    #   # shape: (9, 2)
+    #   # ┌────────┬─────┐
+    #   # │ name   ┆ age │
+    #   # │ ---    ┆ --- │
+    #   # │ str    ┆ i64 │
+    #   # ╞════════╪═════╡
+    #   # │ steve  ┆ 17  │
+    #   # │ bob    ┆ 18  │
+    #   # │ thomas ┆ 20  │
+    #   # │ anna   ┆ 21  │
+    #   # │ maya   ┆ 27  │
+    #   # │ megan  ┆ 33  │
+    #   # │ ida    ┆ 37  │
+    #   # │ steve  ┆ 42  │
+    #   # │ elise  ┆ 44  │
+    #   # └────────┴─────┘
+    def merge_sorted(
+      items,
+      key,
+      maintain_order: false
+    )
+      elems = items.to_a
+
+      if elems.empty?
+        msg = "cannot merge_sort empty list"
+        raise ArgumentError, msg
+      end
+      if elems.length == 1 && (elems[0].is_a?(DataFrame) || elems[0].is_a?(LazyFrame))
+        return elems[0]
+      end
+
+      if !Utils.is_non_empty_sequence_of(elems, DataFrame) && !Utils.is_non_empty_sequence_of(elems, LazyFrame)
+        msg = "merge_sorted is not supported for #{elems[0].inspect}"
+        raise TypeError, msg
+      end
+
+      frames = elems.map { |df| df.lazy }
+
+      reduce_fn = lambda do |x, y|
+        x.merge_sorted(y, key, maintain_order: maintain_order)
+      end
+
+      lf = Utils.reduce_balanced(reduce_fn, frames)
+      eager = elems[0].is_a?(DataFrame)
+      eager ? lf.collect : lf
+    end
+
     # Align an array of frames using the unique values from one or more columns as a key.
     #
     # Frames that do not contain the given key values have rows injected (with nulls
@@ -464,16 +543,10 @@ module Polars
     def align_frames(
       *frames,
       on:,
-      how: nil,
+      how: "full",
       select: nil,
       descending: false
     )
-      # TODO update
-      if how.nil?
-        warn "The default `how` for `align_frames` method will change from `left` to `full` in a future version"
-        how = "left"
-      end
-
       if frames.empty?
         return []
       elsif frames.map(&:class).uniq.length != 1

data/lib/polars/functions/lazy.rb

@@ -141,6 +141,9 @@ module Polars
     #
     # @param columns [Array]
     #   One or more column names.
+    # @param maintain_order [Boolean]
+    #   Whether to preserve the order of elements in the list. Setting this
+    #   to `false` can improve performance, especially within `group_by`.
     #
     # @return [Expr]
     #
@@ -174,8 +177,8 @@ module Polars
     #   # ╞═══════════╪═══════════════════════╡
     #   # │ [9, 8, 7] ┆ ["foo", "bar", "foo"] │
     #   # └───────────┴───────────────────────┘
-    def implode(*columns)
-      col(*columns).implode
+    def implode(*columns, maintain_order: true)
+      col(*columns).implode(maintain_order: maintain_order)
     end
 
     # Get the standard deviation.

data/lib/polars/iceberg_dataset.rb

@@ -17,24 +17,48 @@ module Polars
 
       table = scan.table
       snapshot = scan.snapshot
-      schema = snapshot ? table.schema_by_id(snapshot[:schema_id]) : table.current_schema
+      schema = snapshot ? table.schema_by_id(snapshot.is_a?(Hash) ? snapshot[:schema_id] : snapshot.schema_id) : table.current_schema
 
-      if files.empty?
-        # TODO improve
+      if files.empty? && !schema.respond_to?(:arrow_c_schema)
+        # TODO remove in 0.27.0
         schema =
           schema.fields.to_h do |field|
             dtype =
               case field[:type]
+              when "unknown"
+                Unknown
+              when "boolean"
+                Boolean
               when "int"
-                Polars::Int32
+                Int32
               when "long"
-                Polars::Int64
+                Int64
+              when "float"
+                Float32
               when "double"
-                Polars::Float64
+                Float64
+              when "decimal"
+                Decimal.new(field[:precision], field[:scale])
               when "string"
-                Polars::String
+                String
+              when "uuid"
+                Binary
+              when "fixed"
+                Binary
+              when "binary"
+                Binary
+              when "date"
+                Date
+              when "time"
+                Time # ns instead of us since does not support time unit
               when "timestamp"
-                Polars::Datetime
+                Datetime.new("us")
+              when "timestamp_ns"
+                Datetime.new("ns")
+              when "timestamptz"
+                Datetime.new("us", "+00:00")
+              when "timestamptz_ns"
+                Datetime.new("ns", "+00:00")
               else
                 raise Todo
               end
@@ -44,25 +68,26 @@ module Polars
 
         LazyFrame.new(schema: schema)
       else
-        sources = files.map { |v| v[:data_file_path] }
+        sources = files.map { |v| v.is_a?(Hash) ? v[:data_file_path] : v.file.file_path }
 
         column_mapping = [
           "iceberg-column-mapping",
-          arrow_schema(schema)
+          schema.respond_to?(:arrow_c_schema) ? schema : arrow_schema(schema)
         ]
 
         deletion_files = [
           "iceberg-position-delete",
           files.map.with_index
-            .select { |v, i| v[:deletes].any? }
-            .to_h { |v, i| [i, v[:deletes].map { |d| d[:file_path] }] }
+            .select { |v, i| (v.is_a?(Hash) ? v[:deletes] : v.delete_files).any? }
+            .to_h { |v, i| [i, v.is_a?(Hash) ? v[:deletes].map { |d| d[:file_path] } : v.delete_files.map(&:file_path)] }
         ]
 
         scan_options = {
-          storage_options: @storage_options,
+          schema: schema.respond_to?(:arrow_c_schema) ? Schema.new(schema) : nil,
           cast_options: Polars::ScanCastOptions._default_iceberg,
           missing_columns: "insert",
           extra_columns: "ignore",
+          storage_options: @storage_options,
           _column_mapping: column_mapping,
           _deletion_files: deletion_files
         }
@@ -73,11 +98,14 @@ module Polars
 
     private
 
+    # TODO remove in 0.27.0
     def arrow_schema(schema)
       fields =
         schema.fields.map do |field|
           type =
             case field[:type]
+            when "unknown"
+              "unknown"
             when "boolean"
               "boolean"
             when "int"
@@ -88,11 +116,42 @@ module Polars
               "float32"
             when "double"
               "float64"
+            when "decimal"
+              "decimal"
+            when "string"
+              "string"
+            when "uuid"
+              limit = 16
+              "fixed_size_binary"
+            when "fixed"
+              limit = field[:limit]
+              "fixed_size_binary"
+            when "binary"
+              "large_binary"
+            when "date"
+              "date32"
+            when "time"
+              time_unit = "us"
+              "time64"
+            when "timestamp"
+              time_unit = "us"
+              "timestamp"
+            when "timestamp_ns"
+              time_unit = "ns"
+              "timestamp"
+            when "timestamptz"
+              time_unit = "us"
+              time_zone = "+00:00"
+              "timestamp"
+            when "timestamptz_ns"
+              time_unit = "ns"
+              time_zone = "+00:00"
+              "timestamp"
             else
               raise Todo
             end
 
-          {
+          arrow_field = {
             name: field[:name],
             type: type,
             nullable: !field[:required],
@@ -100,6 +159,14 @@ module Polars
               "PARQUET:field_id" => field[:id].to_s
             }
           }
+          arrow_field[:limit] = limit if limit
+          if type == "decimal"
+            arrow_field[:precision] = field[:precision]
+            arrow_field[:scale] = field[:scale]
+          end
+          arrow_field[:time_unit] = time_unit if time_unit
+          arrow_field[:time_zone] = time_zone if time_zone
+          arrow_field
         end
 
       {fields: fields}

data/lib/polars/io/csv.rb

@@ -648,6 +648,12 @@ module Polars
     #   (which defaults to 1 hour) if not given.
     # @param include_file_paths [String]
     #   Include the path of the source file(s) as a column with this name.
+    # @param missing_columns ['insert', 'raise']
+    #   Configuration for behavior when columns defined in the schema are
+    #   missing from the data:
+    #
+    #   * `"insert"`: Insert the missing columns with NULL values.
+    #   * `"raise"`: Raise an error.
     #
     # @return [LazyFrame]
     def scan_csv(
@@ -663,7 +669,7 @@ module Polars
       null_values: nil,
       missing_utf8_is_empty_string: false,
       ignore_errors: false,
-      cache: true,
+      cache: nil,
       with_column_names: nil,
       infer_schema: true,
       infer_schema_length: N_INFER_DEFAULT,
@@ -685,7 +691,8 @@ module Polars
       credential_provider: "auto",
       retries: nil,
       file_cache_ttl: nil,
-      include_file_paths: nil
+      include_file_paths: nil,
+      missing_columns: nil
     )
       if new_columns&.any? && schema_overrides.is_a?(::Array)
         msg = "expected 'schema_overrides' hash, found #{schema_overrides.inspect}"
@@ -734,6 +741,11 @@ module Polars
         storage_options["file_cache_ttl"] = file_cache_ttl
       end
 
+      if !missing_columns.nil?
+        msg = "The `missing_columns` parameter of `scan_csv` is considered unstable."
+        Utils.issue_unstable_warning(msg)
+      end
+
       credential_provider_builder = _init_credential_provider_builder(
         credential_provider, source, storage_options, "scan_csv"
       )
@@ -768,7 +780,8 @@ module Polars
         glob: glob,
         storage_options: storage_options,
         credential_provider: credential_provider_builder,
-        include_file_paths: include_file_paths
+        include_file_paths: include_file_paths,
+        missing_columns: missing_columns
       )
     end
 
@@ -804,7 +817,8 @@ module Polars
       glob: true,
       storage_options: nil,
       credential_provider: nil,
-      include_file_paths: nil
+      include_file_paths: nil,
+      missing_columns: nil
     )
       dtype_list = nil
       if !schema_overrides.nil?
@@ -822,6 +836,13 @@ module Polars
         sources = []
       end
 
+      # TODO: This is a hack. We conditionally set `missing_columns` to mimic
+      # existing behavior. This should be removed once the workaround is no
+      # longer needed.
+      if missing_columns.nil? && !schema.nil? && has_header
+        missing_columns = "insert"
+      end
+
       rblf =
         RbLazyFrame.new_from_csv(
           source,
@@ -854,7 +875,8 @@ module Polars
           schema,
           storage_options,
           credential_provider,
-          include_file_paths
+          include_file_paths,
+          missing_columns
         )
       Utils.wrap_ldf(rblf)
     end

data/lib/polars/io/delta.rb

@@ -25,6 +25,8 @@ module Polars
       storage_options: nil,
       delta_table_options: nil
     )
+      require "deltalake-rb"
+
       df =
         scan_delta(
           source,
@@ -62,6 +64,8 @@ module Polars
       delta_table_options: nil,
       rechunk: nil
     )
+      require "deltalake-rb"
+
       dl_tbl =
         _get_delta_lake_table(
           source,

data/lib/polars/io/ipc.rb

@@ -218,7 +218,7 @@ module Polars
     def scan_ipc(
       source,
       n_rows: nil,
-      cache: true,
+      cache: nil,
       rechunk: false,
       row_index_name: nil,
       row_index_offset: 0,

data/lib/polars/io/lines.rb

@@ -49,7 +49,7 @@ module Polars
     #   # =>
     #   # shape: (2, 1)
     #   # ┌───────┐
-    #   # │ lines │
+    #   # │ line  │
     #   # │ ---   │
     #   # │ str   │
     #   # ╞═══════╡
@@ -58,7 +58,7 @@ module Polars
     #   # └───────┘
     def read_lines(
       source,
-      name: "lines",
+      name: "line",
       n_rows: nil,
       row_index_name: nil,
       row_index_offset: 0,
@@ -129,7 +129,7 @@ module Polars
     #   # =>
     #   # shape: (2, 1)
     #   # ┌───────┐
-    #   # │ lines │
+    #   # │ line  │
     #   # │ ---   │
     #   # │ str   │
     #   # ╞═══════╡
@@ -138,7 +138,7 @@ module Polars
     #   # └───────┘
     def scan_lines(
       source,
-      name: "lines",
+      name: "line",
       n_rows: nil,
       row_index_name: nil,
       row_index_offset: 0,

data/lib/polars/io/sink_options.rb

@@ -2,20 +2,22 @@ module Polars
   module IO
     # @private
     class SinkOptions
-      attr_reader :mkdir, :maintain_order, :sync_on_close, :storage_options, :credential_provider
+      attr_reader :mkdir, :maintain_order, :sync_on_close, :storage_options, :credential_provider, :sinked_paths_callback
 
       def initialize(
         mkdir: nil,
         maintain_order: nil,
         sync_on_close: nil,
         storage_options: nil,
-        credential_provider: nil
+        credential_provider: nil,
+        sinked_paths_callback: nil
       )
         @mkdir = mkdir
         @maintain_order = maintain_order
         @sync_on_close = sync_on_close
         @storage_options = storage_options
         @credential_provider = credential_provider
+        @sinked_paths_callback = sinked_paths_callback
       end
     end
   end

data/lib/polars/lazy_frame.rb

@@ -1162,7 +1162,8 @@ module Polars
       lazy: false,
       arrow_schema: nil,
       engine: "auto",
-      optimizations: DEFAULT_QUERY_OPT_FLAGS
+      optimizations: DEFAULT_QUERY_OPT_FLAGS,
+      _sinked_paths_callback: nil
     )
       engine = _select_engine(engine)
 
@@ -1204,7 +1205,8 @@ module Polars
         maintain_order: maintain_order,
         sync_on_close: sync_on_close,
         storage_options: storage_options,
-        credential_provider: credential_provider_builder
+        credential_provider: credential_provider_builder,
+        sinked_paths_callback: _sinked_paths_callback
       )
 
       ldf_rb = _ldf.sink_parquet(
@@ -3479,6 +3481,78 @@ module Polars
       )
     end
 
+    # Selects rows from this LazyFrame at the given indices.
+    #
+    # @note
+    #   This functionality is experimental. It may be
+    #   changed at any point without it being considered a breaking change.
+    #
+    # @param indices [Object]
+    #   The indices of the rows to select.
+    #
+    #   Due to the lack of a `LazySeries` it's permitted to pass a single-width
+    #   `LazyFrame` as indices as well.
+    # @param null_on_oob [Boolean]
+    #   If true when an index is out-of-bounds a null row will be generated
+    #   instead of raising an error.
+    #
+    # @return [LazyFrame]
+    #
+    # @example
+    #   lf = Polars::LazyFrame.new({"x" => [2, 1, 0], "s" => ["foo", "bar", "baz"]})
+    #   lf.gather([2, 0, 0]).collect
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌─────┬─────┐
+    #   # │ x   ┆ s   │
+    #   # │ --- ┆ --- │
+    #   # │ i64 ┆ str │
+    #   # ╞═════╪═════╡
+    #   # │ 0   ┆ baz │
+    #   # │ 2   ┆ foo │
+    #   # │ 2   ┆ foo │
+    #   # └─────┴─────┘
+    #
+    # @example
+    #   lf.gather([0, 10, 1], null_on_oob: true).collect
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌──────┬──────┐
+    #   # │ x    ┆ s    │
+    #   # │ ---  ┆ ---  │
+    #   # │ i64  ┆ str  │
+    #   # ╞══════╪══════╡
+    #   # │ 2    ┆ foo  │
+    #   # │ null ┆ null │
+    #   # │ 1    ┆ bar  │
+    #   # └──────┴──────┘
+    #
+    # @example
+    #   idxs = Polars::LazyFrame.new({"i" => [1, 10, 0], "b" => [true, false, true]})
+    #   lf.gather(idxs.filter(Polars.col("b")).select(Polars.col("i"))).collect
+    #   # =>
+    #   # shape: (2, 2)
+    #   # ┌─────┬─────┐
+    #   # │ x   ┆ s   │
+    #   # │ --- ┆ --- │
+    #   # │ i64 ┆ str │
+    #   # ╞═════╪═════╡
+    #   # │ 1   ┆ bar │
+    #   # │ 2   ┆ foo │
+    #   # └─────┴─────┘
+    def gather(indices, null_on_oob: false)
+      if !indices.is_a?(LazyFrame)
+        if indices.is_a?(::Array)
+          indices_expr = F.lit(Series.new("", indices, dtype: Int64))
+        else
+          indices_expr = wrap_expr(Utils.parse_into_expression(indices))
+        end
+        indices = select(indices_expr)
+      end
+
+      _from_rbldf(_ldf.gather(indices._ldf, null_on_oob))
+    end
+
     # Add or overwrite multiple columns in a DataFrame.
     #
     # @param exprs [Object]
@@ -4643,6 +4717,8 @@ module Polars
     # @param separator [String]
     #   Used as separator/delimiter in generated column names in case of multiple
     #   `values` columns.
+    # @param column_naming ['auto', 'combine']
+    #   How resulting column names will be constructed.
     #
     # @return [LazyFrame]
     #
@@ -4682,7 +4758,8 @@ module Polars
       values: nil,
       aggregate_function: nil,
       maintain_order: false,
-      separator: "_"
+      separator: "_",
+      column_naming: "auto"
     )
       if index.nil? && values.nil?
         msg = "`pivot` needs either `index or `values` needs to be specified"
@@ -4756,7 +4833,8 @@ module Polars
           values_selector._rbselector,
           agg._rbexpr,
           maintain_order,
-          separator
+          separator,
+          column_naming
         )
       )
     end
@@ -4772,7 +4850,8 @@ module Polars
     #
     # @param on [Object]
     #   Column(s) or selector(s) to use as values variables; if `on`
-    #   is empty all columns that are not in `index` will be used.
+    #   is empty no columns will be used. If set to `nil` (default)
+    #   all columns that are not in `index` will be used.
     # @param index [Object]
     #   Column(s) or selector(s) to use as identifier variables.
     # @param variable_name [String]
@@ -4820,12 +4899,12 @@ module Polars
         warn "The `streamable` parameter for `LazyFrame.unpivot` is deprecated"
       end
 
-      selector_on = on.nil? ? Selectors.empty : Utils.parse_list_into_selector(on)
+      selector_on = on.nil? ? nil : Utils.parse_list_into_selector(on)._rbselector
       selector_index = index.nil? ? Selectors.empty : Utils.parse_list_into_selector(index)
 
       _from_rbldf(
         _ldf.unpivot(
-          selector_on._rbselector,
+          selector_on,
           selector_index._rbselector,
           value_name,
           variable_name
@@ -4905,10 +4984,10 @@ module Polars
     #   # │ -2      ┆ 199998 │
     #   # └─────────┴────────┘
     def map_batches(
-      predicate_pushdown: true,
-      projection_pushdown: true,
-      slice_pushdown: true,
-      no_optimizations: false,
+      predicate_pushdown: false,
+      projection_pushdown: false,
+      slice_pushdown: false,
+      no_optimizations: nil,
       schema: nil,
       validate_output_schema: true,
       streamable: false,
@@ -5020,7 +5099,7 @@ module Polars
     #   # │ foo    ┆ 1   ┆ a   ┆ true ┆ [1, 2]    ┆ baz   │
     #   # │ bar    ┆ 2   ┆ b   ┆ null ┆ [3]       ┆ womp  │
     #   # └────────┴─────┴─────┴──────┴───────────┴───────┘
-    def unnest(columns, *more_columns, separator: nil)
+    def unnest(columns = nil, *more_columns, separator: nil)
       subset = Utils.parse_list_into_selector(columns) | Utils.parse_list_into_selector(
         more_columns
       )
@@ -5039,6 +5118,10 @@ module Polars
     #   Other DataFrame that must be merged
     # @param key [String]
     #   Key that is sorted.
+    # @param maintain_order [Boolean]
+    #   If `true`, the output is guaranteed to have left-biased ordering
+    #   for equal keys: rows from the left frame appear before rows from
+    #   the right frame when their keys are equal.
     #
     # @return [LazyFrame]
     #
@@ -5065,8 +5148,8 @@ module Polars
     #   # │ steve  ┆ 42  │
     #   # │ elise  ┆ 44  │
     #   # └────────┴─────┘
-    def merge_sorted(other, key)
-      _from_rbldf(_ldf.merge_sorted(other._ldf, key))
+    def merge_sorted(other, key, maintain_order: false)
+      _from_rbldf(_ldf.merge_sorted(other._ldf, key, maintain_order))
     end
 
     # Flag a column as sorted.