polars-df 0.25.0 → 0.26.0

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.
@@ -1568,6 +1571,25 @@ module Polars
       result
     end
 
+    # Explain multiple LazyFrames as if passed to `collect_all`.
+    #
+    # Common Subplan Elimination is applied on the combined plan, meaning
+    # that diverging queries will run only once.
+    #
+    # @param lazy_frames [Array]
+    #   A list of LazyFrames to collect.
+    # @param optimizations [Object]
+    #   The optimization passes done during query optimization.
+    #
+    # @return [String]
+    def explain_all(
+      lazy_frames,
+      optimizations: DEFAULT_QUERY_OPT_FLAGS
+    )
+      lfs = lazy_frames.map { |lf| lf._ldf }
+      Plr.explain_all(lfs, optimizations._rboptflags)
+    end
+
     # Run polars expressions without a context.
     #
     # This is syntactic sugar for running `df.select` on an empty DataFrame.
@@ -1890,6 +1912,79 @@ module Polars
       end
     end
 
+    # Generates a sequence of integers.
+    #
+    # The length of the returned sequence will match the context length, and the
+    # datatype will match the one returned by `get_index_dtype()`.
+    #
+    # If you would like to generate sequences with custom offsets / length /
+    # step size / datatypes, it is recommended to use `int_range` instead.
+    #
+    # @note
+    #   This functionality is considered **unstable**. It may be changed
+    #   at any point without it being considered a breaking change.
+    #
+    # @param name [String]
+    #   Name of the returned column.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"x" => ["A", "A", "B", "B", "B"]})
+    #   df.with_columns(Polars.row_index, Polars.row_index("another_index"))
+    #   # =>
+    #   # shape: (5, 3)
+    #   # ┌─────┬───────┬───────────────┐
+    #   # │ x   ┆ index ┆ another_index │
+    #   # │ --- ┆ ---   ┆ ---           │
+    #   # │ str ┆ u32   ┆ u32           │
+    #   # ╞═════╪═══════╪═══════════════╡
+    #   # │ A   ┆ 0     ┆ 0             │
+    #   # │ A   ┆ 1     ┆ 1             │
+    #   # │ B   ┆ 2     ┆ 2             │
+    #   # │ B   ┆ 3     ┆ 3             │
+    #   # │ B   ┆ 4     ┆ 4             │
+    #   # └─────┴───────┴───────────────┘
+    #
+    # @example
+    #   df.group_by("x").agg(Polars.row_index).sort("x")
+    #   # =>
+    #   # shape: (2, 2)
+    #   # ┌─────┬───────────┐
+    #   # │ x   ┆ index     │
+    #   # │ --- ┆ ---       │
+    #   # │ str ┆ list[u32] │
+    #   # ╞═════╪═══════════╡
+    #   # │ A   ┆ [0, 1]    │
+    #   # │ B   ┆ [0, 1, 2] │
+    #   # └─────┴───────────┘
+    #
+    # @example
+    #   df.select(Polars.row_index)
+    #   # =>
+    #   # shape: (5, 1)
+    #   # ┌───────┐
+    #   # │ index │
+    #   # │ ---   │
+    #   # │ u32   │
+    #   # ╞═══════╡
+    #   # │ 0     │
+    #   # │ 1     │
+    #   # │ 2     │
+    #   # │ 3     │
+    #   # │ 4     │
+    #   # └───────┘
+    def row_index(name = "index")
+      # Notes
+      # * Dispatching to `int_range` means that we cannot accept an offset
+      #   parameter, as unlike `DataFrame.with_row_index`, `int_range` will simply
+      #   truncate instead of raising an error.
+      F.int_range(
+        F.len,
+        dtype: get_index_type
+      ).alias(name)
+    end
+
     private
 
     def _wrap_acc_lambda(function)

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

@@ -1,5 +1,123 @@
 module Polars
   module Functions
+    # Create sequence of evenly-spaced points.
+    #
+    # @param start [Object]
+    #   Lower bound of the range.
+    # @param stop [Object]
+    #   Upper bound of the range.
+    # @param num_samples [Object]
+    #   Number of samples in the output sequence.
+    # @param closed ['both', 'left', 'right', 'none']
+    #   Define which sides of the interval are closed (inclusive).
+    # @param eager [Boolean]
+    #   Evaluate immediately and return a `Series`.
+    #   If set to `false` (default), return an expression instead.
+    #
+    # @return [Object]
+    #
+    # @note
+    #   This functionality is experimental. It may be changed at any point without it
+    #   being considered a breaking change.
+    #
+    # @example
+    #   Polars.linear_space(0, 1, 3, eager: true)
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'literal' [f64]
+    #   # [
+    #   #         0.0
+    #   #         0.5
+    #   #         1.0
+    #   # ]
+    #
+    # @example
+    #   Polars.linear_space(0, 1, 3, closed: "left", eager: true)
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'literal' [f64]
+    #   # [
+    #   #         0.0
+    #   #         0.333333
+    #   #         0.666667
+    #   # ]
+    #
+    # @example
+    #   Polars.linear_space(0, 1, 3, closed: "right", eager: true)
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'literal' [f64]
+    #   # [
+    #   #         0.333333
+    #   #         0.666667
+    #   #         1.0
+    #   # ]
+    #
+    # @example
+    #   Polars.linear_space(0, 1, 3, closed: "none", eager: true)
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'literal' [f64]
+    #   # [
+    #   #         0.25
+    #   #         0.5
+    #   #         0.75
+    #   # ]
+    #
+    # @example `Date` endpoints generate a sequence of `Datetime` values:
+    #   Polars.linear_space(
+    #     Date.new(2025, 1, 1),
+    #     Date.new(2025, 2, 1),
+    #     3,
+    #     closed: "right",
+    #     eager: true
+    #   )
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'literal' [datetime[μs]]
+    #   # [
+    #   #         2025-01-11 08:00:00
+    #   #         2025-01-21 16:00:00
+    #   #         2025-02-01 00:00:00
+    #   # ]
+    #
+    # @example When `eager: false` (default), an expression is produced. You can generate a sequence using the length of the dataframe:
+    #   df = Polars::DataFrame.new({"a" => [1, 2, 3, 4, 5]})
+    #   df.with_columns(Polars.linear_space(0, 1, Polars.len).alias("ls"))
+    #   # =>
+    #   # shape: (5, 2)
+    #   # ┌─────┬──────┐
+    #   # │ a   ┆ ls   │
+    #   # │ --- ┆ ---  │
+    #   # │ i64 ┆ f64  │
+    #   # ╞═════╪══════╡
+    #   # │ 1   ┆ 0.0  │
+    #   # │ 2   ┆ 0.25 │
+    #   # │ 3   ┆ 0.5  │
+    #   # │ 4   ┆ 0.75 │
+    #   # │ 5   ┆ 1.0  │
+    #   # └─────┴──────┘
+    def linear_space(
+      start,
+      stop,
+      num_samples,
+      closed: "both",
+      eager: false
+    )
+      start_rbexpr = Utils.parse_into_expression(start)
+      end_rbexpr = Utils.parse_into_expression(stop)
+      num_samples_rbexpr = Utils.parse_into_expression(num_samples)
+      result = Utils.wrap_expr(
+        Plr.linear_space(start_rbexpr, end_rbexpr, num_samples_rbexpr, closed)
+      )
+
+      if eager
+        return F.select(result).to_series
+      end
+
+      result
+    end
+
     # Generate a sequence of evenly-spaced values for each row between `start` and `end`.
     #
     # The number of values in each sequence is determined by `num_samples`.

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/database.rb

@@ -75,9 +75,8 @@ module Polars
             String
           when :time
             Time
-          # TODO fix issue with null
-          # when :json, :jsonb
-          #   Struct
+          when :json, :jsonb
+            Struct
           end
 
         schema_overrides[k] ||= polars_type if polars_type

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,
@@ -267,7 +267,7 @@ module Polars
             try_parse_hive_dates: try_parse_hive_dates,
             rechunk: rechunk,
             cache: cache,
-            storage_options: !storage_options.nil? ? storage_options.to_a : nil,
+            storage_options: storage_options,
             credential_provider: credential_provider_builder
           )
         )

data/lib/polars/io/lines.rb

@@ -0,0 +1,172 @@
+module Polars
+  module IO
+    # Read lines into a string column from a file.
+    #
+    # @note
+    #   This functionality is considered **unstable**. It may be changed
+    #   at any point without it being considered a breaking change.
+    #
+    # @param source [Object]
+    #   Path(s) to a file or directory
+    #   When needing to authenticate for scanning cloud locations, see the
+    #   `storage_options` parameter.
+    # @param name [String]
+    #   Name to use for the output column.
+    # @param n_rows [Integer]
+    #   Stop reading from parquet file after reading `n_rows`.
+    # @param row_index_name [String]
+    #   If not nil, this will insert a row index column with the given name into the
+    #   DataFrame
+    # @param row_index_offset [Integer]
+    #   Offset to start the row index column (only used if the name is set)
+    # @param glob [Boolean]
+    #   Expand path given via globbing rules.
+    # @param storage_options [Hash]
+    #   Options that indicate how to connect to a cloud provider.
+    #
+    #   The cloud providers currently supported are AWS, GCP, and Azure.
+    #   See supported keys here:
+    #
+    #   * [aws](https://docs.rs/object_store/latest/object_store/aws/enum.AmazonS3ConfigKey.html)
+    #   * [gcp](https://docs.rs/object_store/latest/object_store/gcp/enum.GoogleConfigKey.html)
+    #   * [azure](https://docs.rs/object_store/latest/object_store/azure/enum.AzureConfigKey.html)
+    #   * Hugging Face (`hf://`): Accepts an API key under the `token` parameter: \
+    #     `{'token': '...'}`, or by setting the `HF_TOKEN` environment variable.
+    #
+    #   If `storage_options` is not provided, Polars will try to infer the information
+    #   from environment variables.
+    # @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 include_file_paths [String]
+    #   Include the path of the source file(s) as a column with this name.
+    #
+    # @return [DataFrame]
+    #
+    # @example
+    #   Polars.read_lines(StringIO.new("Hello\nworld"))
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌───────┐
+    #   # │ line  │
+    #   # │ ---   │
+    #   # │ str   │
+    #   # ╞═══════╡
+    #   # │ Hello │
+    #   # │ world │
+    #   # └───────┘
+    def read_lines(
+      source,
+      name: "line",
+      n_rows: nil,
+      row_index_name: nil,
+      row_index_offset: 0,
+      glob: true,
+      storage_options: nil,
+      credential_provider: "auto",
+      include_file_paths: nil
+    )
+      scan_lines(
+        source,
+        name: name,
+        n_rows: n_rows,
+        row_index_name: row_index_name,
+        row_index_offset: row_index_offset,
+        glob: glob,
+        storage_options: storage_options,
+        credential_provider: credential_provider,
+        include_file_paths: include_file_paths
+      ).collect
+    end
+
+    # Construct a LazyFrame which scans lines into a string column from a file.
+    #
+    # @note
+    #   This functionality is considered **unstable**. It may be changed
+    #   at any point without it being considered a breaking change.
+    #
+    # @param source [Object]
+    #   Path(s) to a file or directory
+    #   When needing to authenticate for scanning cloud locations, see the
+    #   `storage_options` parameter.
+    # @param name [String]
+    #   Name to use for the output column.
+    # @param n_rows [Integer]
+    #   Stop reading from parquet file after reading `n_rows`.
+    # @param row_index_name [String]
+    #   If not nil, this will insert a row index column with the given name into the
+    #   DataFrame
+    # @param row_index_offset
+    #   Offset to start the row index column (only used if the name is set)
+    # @param glob [Boolean]
+    #   Expand path given via globbing rules.
+    # @param storage_options [Hash]
+    #   Options that indicate how to connect to a cloud provider.
+    #
+    #   The cloud providers currently supported are AWS, GCP, and Azure.
+    #   See supported keys here:
+    #
+    #   * [aws](https://docs.rs/object_store/latest/object_store/aws/enum.AmazonS3ConfigKey.html)
+    #   * [gcp](https://docs.rs/object_store/latest/object_store/gcp/enum.GoogleConfigKey.html)
+    #   * [azure](https://docs.rs/object_store/latest/object_store/azure/enum.AzureConfigKey.html)
+    #   * Hugging Face (`hf://`): Accepts an API key under the `token` parameter: \
+    #     `{'token': '...'}`, or by setting the `HF_TOKEN` environment variable.
+    #
+    #   If `storage_options` is not provided, Polars will try to infer the information
+    #   from environment variables.
+    # @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 include_file_paths [String]
+    #   Include the path of the source file(s) as a column with this name.
+    #
+    # @return [LazyFrame]
+    #
+    # @example
+    #   Polars.scan_lines(StringIO.new("Hello\nworld")).collect
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌───────┐
+    #   # │ line  │
+    #   # │ ---   │
+    #   # │ str   │
+    #   # ╞═══════╡
+    #   # │ Hello │
+    #   # │ world │
+    #   # └───────┘
+    def scan_lines(
+      source,
+      name: "line",
+      n_rows: nil,
+      row_index_name: nil,
+      row_index_offset: 0,
+      glob: true,
+      storage_options: nil,
+      credential_provider: "auto",
+      include_file_paths: nil
+    )
+      sources = get_sources(source)
+
+      credential_provider_builder = _init_credential_provider_builder(
+        credential_provider, sources, storage_options, "scan_lines"
+      )
+
+      rblf = RbLazyFrame.new_from_scan_lines(
+        sources,
+        ScanOptions.new(
+          row_index: !row_index_name.nil? ? [row_index_name, row_index_offset] : nil,
+          pre_slice: !n_rows.nil? ? [0, n_rows] : nil,
+          include_file_paths: include_file_paths,
+          glob: glob,
+          storage_options: storage_options,
+          credential_provider: credential_provider_builder
+        ),
+        name
+      )
+
+      Utils.wrap_ldf(rblf)
+    end
+  end
+end

data/lib/polars/io/parquet.rb

@@ -350,7 +350,7 @@ module Polars
             try_parse_hive_dates: try_parse_hive_dates,
             rechunk: rechunk,
             cache: cache,
-            storage_options: storage_options ? storage_options.map { |k, v| [k.to_s, v.to_s] } : nil,
+            storage_options: storage_options,
             credential_provider: credential_provider_builder,
             column_mapping: _column_mapping,
             default_values: _default_values,

data/lib/polars/io/sink_options.rb

@@ -2,23 +2,26 @@ 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
 
+  # @private
   SinkOptions = IO::SinkOptions
 end