polars-df 0.25.0 → 0.26.0

data/lib/polars/lazy_frame.rb

@@ -262,6 +262,59 @@ module Polars
       function.(self, *args, **kwargs, &block)
     end
 
+    # Allows to alter the lazy frame during the plan stage with the resolved schema.
+    #
+    # In contrast to `pipe`, this method does not execute `function` immediately but
+    # only during the plan stage. This allows to use the resolved schema of the input
+    # to dynamically alter the lazy frame. This also means that any exceptions raised
+    # by `function` will only be emitted during the plan stage.
+    #
+    # @note
+    #   This functionality is considered **unstable**. It may be changed at any
+    #   point without it being considered a breaking change.
+    #
+    # @param function [Object]
+    #   Callable; will receive the frame as the first parameter and the resolved
+    #   schema as the second parameter.
+    #
+    # @return [LazyFrame]
+    #
+    # @example
+    #   cast_to_float_if_necessary = lambda do |lf, schema|
+    #     required_casts =
+    #       schema.filter_map do |name, dtype|
+    #         Polars.col(name).cast(Polars::Float64) if !dtype.float?
+    #       end
+    #     lf.with_columns(required_casts)
+    #   end
+    #   lf = Polars::LazyFrame.new(
+    #     {"a" => [1.0, 2.0], "b" => ["1.0", "2.5"], "c" => [2.0, 3.0]},
+    #     schema: {"a" => Polars::Float64, "b" => Polars::String, "c" => Polars::Float32}
+    #   )
+    #   lf.pipe_with_schema(cast_to_float_if_necessary).collect
+    #   # =>
+    #   # shape: (2, 3)
+    #   # ┌─────┬─────┬─────┐
+    #   # │ a   ┆ b   ┆ c   │
+    #   # │ --- ┆ --- ┆ --- │
+    #   # │ f64 ┆ f64 ┆ f32 │
+    #   # ╞═════╪═════╪═════╡
+    #   # │ 1.0 ┆ 1.0 ┆ 2.0 │
+    #   # │ 2.0 ┆ 2.5 ┆ 3.0 │
+    #   # └─────┴─────┴─────┘
+    def pipe_with_schema(function)
+      wrapper = lambda do |lf_and_schema|
+        # The last index is because we return a list for multiple inputs
+        # to make `pipe_with_schemas` (plural) work, but we don't use that
+        function.(
+          _from_rbldf(lf_and_schema[0][0]),
+          lf_and_schema[1][0]
+        )._ldf
+      end
+
+      _from_rbldf(_ldf.pipe_with_schema(wrapper))
+    end
+
     # Creates a summary of statistics for a LazyFrame, returning a DataFrame.
     #
     # @param percentiles [Array]
@@ -1109,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)
 
@@ -1151,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(
@@ -1175,6 +1230,133 @@ module Polars
       LazyFrame._from_rbldf(ldf_rb)
     end
 
+    # Sink DataFrame as delta table.
+    #
+    # @note
+    #   This functionality is considered **unstable**. It may be changed
+    #   at any point without it being considered a breaking change.
+    #
+    # @param target [Object]
+    #   URI of a table or a DeltaTable object.
+    # @param mode ['error', 'append', 'overwrite', 'ignore', 'merge']
+    #   How to handle existing data.
+    #
+    #   - If 'error', throw an error if the table already exists (default).
+    #   - If 'append', will add new data.
+    #   - If 'overwrite', will replace table with new data.
+    #   - If 'ignore', will not write anything if table already exists.
+    #   - If 'merge', return a `TableMerger` object to merge data from the DataFrame
+    #     with the existing data.
+    # @param storage_options [Object]
+    #     Extra options for the storage backends supported by `deltalake`.
+    #     For cloud storages, this may include configurations for authentication etc.
+    #
+    #     - See a list of supported storage options for S3 [here](https://docs.rs/object_store/latest/object_store/aws/enum.AmazonS3ConfigKey.html#variants).
+    #     - See a list of supported storage options for GCS [here](https://docs.rs/object_store/latest/object_store/gcp/enum.GoogleConfigKey.html#variants).
+    #     - See a list of supported storage options for Azure [here](https://docs.rs/object_store/latest/object_store/azure/enum.AzureConfigKey.html#variants).
+    # @param credential_provider [Object]
+    #   Provide a function that can be called to provide cloud storage
+    #   credentials. The function is expected to return a dictionary of
+    #   credential keys along with an optional credential expiry time.
+    # @param delta_write_options [Hash]
+    #   Additional keyword arguments while writing a Delta lake Table.
+    #   See a list of supported write options [here](https://delta-io.github.io/delta-rs/api/delta_writer/#deltalake.write_deltalake).
+    # @param delta_merge_options [Hash]
+    #   Keyword arguments which are required to `MERGE` a Delta lake Table.
+    #   See a list of supported merge options [here](https://delta-io.github.io/delta-rs/api/delta_table/#deltalake.DeltaTable.merge).
+    # @param optimizations [Object]
+    #   The optimization passes done during query optimization.
+    #
+    # @return [Object]
+    #
+    # @example Sink a large than fits into memory dataset to a Delta Lake table.
+    #   lf = Polars.scan_parquet(
+    #     "/path/to/my_larger_than_ram_file.parquet"
+    #   )
+    #   table_path = "/path/to/delta-table/"
+    #   lf.sink_delta(table_path)
+    def sink_delta(
+      target,
+      mode: "error",
+      storage_options: nil,
+      credential_provider: "auto",
+      delta_write_options: nil,
+      delta_merge_options: nil,
+      optimizations: DEFAULT_QUERY_OPT_FLAGS
+    )
+      Polars.send(:_check_if_delta_available)
+
+      # TODO
+      # _check_for_unsupported_types(collect_schema.dtypes)
+
+      if Utils.pathlike?(target)
+        target = Polars.send(:_resolve_delta_lake_uri, target.to_s, strict: false)
+      end
+
+      _init_credential_provider_builder = Polars.method(:_init_credential_provider_builder)
+
+      if !target.is_a?(DeltaLake::Table)
+        credential_provider_builder = _init_credential_provider_builder.(
+          credential_provider, target, storage_options, "sink_delta"
+        )
+      elsif !credential_provider.nil? && credential_provider != "auto"
+        msg = "cannot use credential_provider when passing a DeltaTable object"
+        raise ArgumentError, msg
+      else
+        credential_provider_builder = nil
+      end
+
+      credential_provider_creds = {}
+
+      if credential_provider_builder
+        raise Todo
+      end
+
+      # We aren't calling into polars-native write functions so we just update
+      # the storage_options here.
+      storage_options =
+        if !storage_options.nil? || !credential_provider_builder.nil?
+          (storage_options || {}).merge(credential_provider_creds)
+        else
+          nil
+        end
+
+      stream = collect_batches(
+        engine: "streaming",
+        maintain_order: true,
+        chunk_size: nil,
+        lazy: true,
+        optimizations: optimizations
+      )
+
+      if mode == "merge"
+        if delta_merge_options.nil?
+          msg = "you need to pass delta_merge_options with at least a given predicate for `MERGE` to work."
+          raise ArgumentError, msg
+        end
+        if target.is_a?(::String)
+          dt = DeltaLake::Table.new(target, storage_options: storage_options)
+        else
+          dt = target
+        end
+
+        dt.merge(stream, **delta_merge_options)
+      else
+        if delta_write_options.nil?
+          delta_write_options = {}
+        end
+
+        DeltaLake.write(
+          target,
+          stream,
+          mode: mode,
+          storage_options: storage_options,
+          **delta_write_options
+        )
+        nil
+      end
+    end
+
     # Evaluate the query in streaming mode and write to an IPC file.
     #
     # This allows streaming results that are larger than RAM to be written to disk.
@@ -1591,6 +1773,71 @@ module Polars
       LazyFrame._from_rbldf(ldf_rb)
     end
 
+    # Evaluate the query and call a user-defined function for every ready batch.
+    #
+    # This allows streaming results that are larger than RAM in certain cases.
+    #
+    # @note
+    #   This functionality is considered **unstable**. It may be changed
+    #   at any point without it being considered a breaking change.
+    #
+    # @note
+    #   This method is much slower than native sinks. Only use it if you cannot
+    #   implement your logic otherwise.
+    #
+    # @param chunk_size [Integer]
+    #   The number of rows that are buffered before the callback is called.
+    # @param maintain_order [Boolean]
+    #   Maintain the order in which data is processed.
+    #   Setting this to `false` will be slightly faster.
+    # @param lazy [Boolean]
+    #   Wait to start execution until `collect` is called.
+    # @param engine [String]
+    #   Select the engine used to process the query, optional.
+    #   At the moment, if set to `"auto"` (default), the query is run
+    #   using the polars streaming engine. Polars will also
+    #   attempt to use the engine set by the `POLARS_ENGINE_AFFINITY`
+    #   environment variable. If it cannot run the query using the
+    #   selected engine, the query is run using the polars streaming
+    #   engine.
+    # @param optimizations [Object]
+    #   The optimization passes done during query optimization.
+    #
+    #   This has no effect if `lazy` is set to `True`.
+    #
+    # @return [Object]
+    #
+    # @example
+    #   lf = Polars.scan_csv("/path/to/my_larger_than_ram_file.csv")
+    #   lf.sink_batches { |df| p df }
+    def sink_batches(
+      chunk_size: nil,
+      maintain_order: true,
+      lazy: false,
+      engine: "auto",
+      optimizations: DEFAULT_QUERY_OPT_FLAGS,
+      &function
+    )
+      _wrap = lambda do |rbdf|
+        df = Utils.wrap_df(rbdf)
+        !!function.(df)
+      end
+
+      ldf = _ldf.sink_batches(
+        _wrap,
+        maintain_order,
+        chunk_size
+      )
+
+      if !lazy
+        ldf = ldf.with_optimizations(optimizations._rboptflags)
+        lf = LazyFrame._from_rbldf(ldf)
+        lf.collect(engine: engine)
+        return nil
+      end
+      LazyFrame._from_rbldf(ldf)
+    end
+
     # Evaluate the query in streaming mode and get a generator that returns chunks.
     #
     # This allows streaming results that are larger than RAM to be written to disk.
@@ -3234,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]
@@ -4398,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]
     #
@@ -4437,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"
@@ -4511,7 +4833,8 @@ module Polars
           values_selector._rbselector,
           agg._rbexpr,
           maintain_order,
-          separator
+          separator,
+          column_naming
         )
       )
     end
@@ -4527,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]
@@ -4575,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
@@ -4660,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,
@@ -4775,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
       )
@@ -4794,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]
     #
@@ -4820,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.
@@ -5126,6 +5454,181 @@ module Polars
       _from_rbldf(_ldf.count)
     end
 
+    # Match or evolve the schema of a LazyFrame into a specific schema.
+    #
+    # By default, match_to_schema returns an error if the input schema does not
+    # exactly match the target schema. It also allows columns to be freely reordered,
+    # with additional coercion rules available through optional parameters.
+    #
+    # @note
+    #   This functionality is considered **unstable**. It may be changed
+    #   at any point without it being considered a breaking change.
+    #
+    # @param schema [Object]
+    #   Target schema to match or evolve to.
+    # @param missing_columns [Object]
+    #   Raise of insert missing columns from the input with respect to the `schema`.
+    #
+    #   This can also be an expression per column with what to insert if it is
+    #   missing.
+    # @param missing_struct_fields [Object]
+    #   Raise of insert missing struct fields from the input with respect to the
+    #   `schema`.
+    # @param extra_columns [Object]
+    #   Raise of ignore extra columns from the input with respect to the `schema`.
+    # @param extra_struct_fields [Object]
+    #   Raise of ignore extra struct fields from the input with respect to the
+    #   `schema`.
+    # @param integer_cast [Object]
+    #   Forbid of upcast for integer columns from the input to the respective column
+    #   in `schema`.
+    # @param float_cast [Object]
+    #   Forbid of upcast for float columns from the input to the respective column
+    #   in `schema`.
+    #
+    # @return [LazyFrame]
+    #
+    # @example Ensuring the schema matches
+    #   lf = Polars::LazyFrame.new({"a" => [1, 2, 3], "b" => ["A", "B", "C"]})
+    #   lf.match_to_schema({"a" => Polars::Int64, "b" => Polars::String}).collect
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌─────┬─────┐
+    #   # │ a   ┆ b   │
+    #   # │ --- ┆ --- │
+    #   # │ i64 ┆ str │
+    #   # ╞═════╪═════╡
+    #   # │ 1   ┆ A   │
+    #   # │ 2   ┆ B   │
+    #   # │ 3   ┆ C   │
+    #   # └─────┴─────┘
+    #
+    # @example Adding missing columns
+    #   Polars::LazyFrame.new({"a" => [1, 2, 3]})
+    #   .match_to_schema(
+    #       {"a" => Polars::Int64, "b" => Polars::String},
+    #       missing_columns: "insert"
+    #   )
+    #   .collect
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌─────┬──────┐
+    #   # │ a   ┆ b    │
+    #   # │ --- ┆ ---  │
+    #   # │ i64 ┆ str  │
+    #   # ╞═════╪══════╡
+    #   # │ 1   ┆ null │
+    #   # │ 2   ┆ null │
+    #   # │ 3   ┆ null │
+    #   # └─────┴──────┘
+    #
+    # @example
+    #   Polars::LazyFrame.new({"a" => [1, 2, 3]})
+    #   .match_to_schema(
+    #     {"a" => Polars::Int64, "b" => Polars::String},
+    #     missing_columns: {"b" => Polars.col("a").cast(Polars::String)}
+    #   )
+    #   .collect
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌─────┬─────┐
+    #   # │ a   ┆ b   │
+    #   # │ --- ┆ --- │
+    #   # │ i64 ┆ str │
+    #   # ╞═════╪═════╡
+    #   # │ 1   ┆ 1   │
+    #   # │ 2   ┆ 2   │
+    #   # │ 3   ┆ 3   │
+    #   # └─────┴─────┘
+    #
+    # @example Removing extra columns
+    #   Polars::LazyFrame.new({"a" => [1, 2, 3], "b" => ["A", "B", "C"]})
+    #   .match_to_schema(
+    #     {"a" => Polars::Int64},
+    #     extra_columns: "ignore"
+    #   )
+    #   .collect
+    #   # =>
+    #   # shape: (3, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ i64 │
+    #   # ╞═════╡
+    #   # │ 1   │
+    #   # │ 2   │
+    #   # │ 3   │
+    #   # └─────┘
+    #
+    # @example Upcasting integers and floats
+    #   Polars::LazyFrame.new(
+    #     {"a" => [1, 2, 3], "b" => [1.0, 2.0, 3.0]},
+    #     schema: {"a" => Polars::Int32, "b" => Polars::Float32}
+    #   )
+    #   .match_to_schema(
+    #     {"a" => Polars::Int64, "b" => Polars::Float64},
+    #     integer_cast: "upcast",
+    #     float_cast: "upcast"
+    #   )
+    #   .collect
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌─────┬─────┐
+    #   # │ a   ┆ b   │
+    #   # │ --- ┆ --- │
+    #   # │ i64 ┆ f64 │
+    #   # ╞═════╪═════╡
+    #   # │ 1   ┆ 1.0 │
+    #   # │ 2   ┆ 2.0 │
+    #   # │ 3   ┆ 3.0 │
+    #   # └─────┴─────┘
+    def match_to_schema(
+      schema,
+      missing_columns: "raise",
+      missing_struct_fields: "raise",
+      extra_columns: "raise",
+      extra_struct_fields: "raise",
+      integer_cast: "forbid",
+      float_cast: "forbid"
+    )
+      prepare_missing_columns = lambda do |value|
+        if value.is_a?(Expr)
+          value._rbexpr
+        else
+          value
+        end
+      end
+
+      if schema.is_a?(Hash)
+        schema_prep = Schema.new(schema)
+      else
+        schema_prep = schema
+      end
+
+      if missing_columns.is_a?(Hash)
+        missing_columns_rbexpr =
+          missing_columns.to_h do |key, value|
+            [key.to_s, prepare_missing_columns.(value)]
+          end
+      elsif missing_columns.is_a?(Expr)
+        missing_columns_rbexpr = prepare_missing_columns.(missing_columns)
+      else
+        missing_columns_rbexpr = missing_columns
+      end
+
+      LazyFrame._from_rbldf(
+        _ldf.match_to_schema(
+          schema_prep,
+          missing_columns_rbexpr,
+          missing_struct_fields,
+          extra_columns,
+          extra_struct_fields,
+          integer_cast,
+          float_cast
+        )
+      )
+    end
+
     private
 
     def initialize_copy(other)