polars-df 0.23.0 → 0.24.0

data/lib/polars/io/database.rb

@@ -85,6 +85,5 @@ module Polars
 
       DataFrame.new(data, schema_overrides: schema_overrides)
     end
-    alias_method :read_sql, :read_database
   end
 end

data/lib/polars/io/delta.rb

@@ -21,19 +21,23 @@ module Polars
       source,
       version: nil,
       columns: nil,
-      rechunk: false,
+      rechunk: nil,
       storage_options: nil,
       delta_table_options: nil
     )
-      dl_tbl =
-        _get_delta_lake_table(
+      df =
+        scan_delta(
           source,
           version: version,
           storage_options: storage_options,
-          delta_table_options: delta_table_options
+          delta_table_options: delta_table_options,
+          rechunk: rechunk
         )
 
-      dl_tbl.to_polars(columns: columns, rechunk: rechunk)
+      if !columns.nil?
+        df = df.select(columns)
+      end
+      df.collect
     end
 
     # Lazily read from a Delta lake table.
@@ -46,13 +50,17 @@ module Polars
     #   Extra options for the storage backends supported by `deltalake-rb`.
     # @param delta_table_options [Hash]
     #   Additional keyword arguments while reading a Delta lake Table.
+    # @param rechunk [Boolean]
+    #   Make sure that all columns are contiguous in memory by
+    #   aggregating the chunks into a single array.
     #
     # @return [LazyFrame]
     def scan_delta(
       source,
       version: nil,
       storage_options: nil,
-      delta_table_options: nil
+      delta_table_options: nil,
+      rechunk: nil
     )
       dl_tbl =
         _get_delta_lake_table(
@@ -62,7 +70,7 @@ module Polars
           delta_table_options: delta_table_options
         )
 
-      dl_tbl.to_polars(eager: false)
+      dl_tbl.to_polars(eager: false, rechunk: rechunk || false)
     end
 
     private

data/lib/polars/io/ipc.rb

@@ -15,10 +15,10 @@ module Polars
     #   Only uncompressed IPC files can be memory mapped.
     # @param storage_options [Hash]
     #   Extra options that make sense for a particular storage connection.
-    # @param row_count_name [String]
+    # @param row_index_name [String]
     #   If not nil, this will insert a row count column with give name into the
     #   DataFrame.
-    # @param row_count_offset [Integer]
+    # @param row_index_offset [Integer]
     #   Offset to start the row_count column (only use if the name is set).
     # @param rechunk [Boolean]
     #   Make sure that all data is contiguous.
@@ -30,8 +30,8 @@ module Polars
       n_rows: nil,
       memory_map: true,
       storage_options: nil,
-      row_count_name: nil,
-      row_count_offset: 0,
+      row_index_name: nil,
+      row_index_offset: 0,
       rechunk: true
     )
       storage_options ||= {}
@@ -40,8 +40,8 @@ module Polars
           data,
           columns: columns,
           n_rows: n_rows,
-          row_count_name: row_count_name,
-          row_count_offset: row_count_offset,
+          row_index_name: row_index_name,
+          row_index_offset: row_index_offset,
           rechunk: rechunk,
           memory_map: memory_map
         )
@@ -53,8 +53,8 @@ module Polars
       file,
       columns: nil,
       n_rows: nil,
-      row_count_name: nil,
-      row_count_offset: 0,
+      row_index_name: nil,
+      row_index_offset: 0,
       rechunk: true,
       memory_map: true
     )
@@ -76,7 +76,7 @@ module Polars
           columns,
           projection,
           n_rows,
-          Utils.parse_row_index_args(row_count_name, row_count_offset),
+          Utils.parse_row_index_args(row_index_name, row_index_offset),
           memory_map
         )
       Utils.wrap_df(rbdf)
@@ -182,15 +182,19 @@ module Polars
     #   Cache the result after reading.
     # @param rechunk [Boolean]
     #   Reallocate to contiguous memory when all chunks/ files are parsed.
-    # @param row_count_name [String]
+    # @param row_index_name [String]
     #   If not nil, this will insert a row count column with give name into the
     #   DataFrame.
-    # @param row_count_offset [Integer]
+    # @param row_index_offset [Integer]
     #   Offset to start the row_count column (only use if the name is set).
     # @param glob [Boolean]
     #   Expand path given via globbing rules.
     # @param storage_options [Hash]
     #   Extra options that make sense for a particular storage connection.
+    # @param credential_provider [Object]
+    #   Provide a function that can be called to provide cloud storage
+    #   credentials. The function is expected to return a hash of
+    #   credential keys along with an optional credential expiry time.
     # @param retries [Integer]
     #   Number of retries if accessing a cloud instance fails.
     # @param file_cache_ttl [Integer]
@@ -215,11 +219,12 @@ module Polars
       source,
       n_rows: nil,
       cache: true,
-      rechunk: true,
-      row_count_name: nil,
-      row_count_offset: 0,
+      rechunk: false,
+      row_index_name: nil,
+      row_index_offset: 0,
       glob: true,
       storage_options: nil,
+      credential_provider: "auto",
       retries: 2,
       file_cache_ttl: nil,
       hive_partitioning: nil,
@@ -227,11 +232,12 @@ module Polars
       try_parse_hive_dates: true,
       include_file_paths: nil
     )
-      row_index_name = row_count_name
-      row_index_offset = row_count_offset
-
       sources = get_sources(source)
 
+      credential_provider_builder = _init_credential_provider_builder(
+        credential_provider, sources, storage_options, "scan_parquet"
+      )
+
       rblf =
         RbLazyFrame.new_from_ipc(
           sources,
@@ -246,6 +252,7 @@ module Polars
             rechunk: rechunk,
             cache: cache,
             storage_options: !storage_options.nil? ? storage_options.to_a : nil,
+            credential_provider: credential_provider_builder,
             retries: retries
           ),
           file_cache_ttl

data/lib/polars/io/ndjson.rb

@@ -2,41 +2,106 @@ module Polars
   module IO
     # Read into a DataFrame from a newline delimited JSON file.
     #
-    # @param source [Object]
-    #   Path to a file or a file-like object.
+    # @param source [String]
+    #   Path to a file.
     # @param schema [Object]
     #   The DataFrame schema may be declared in several ways:
     #
-    #   * As a hash of \\\\{name:type} pairs; if type is nil, it will be auto-inferred.
-    #   * As an array of column names; in this case types are automatically inferred.
-    #   * As an array of [name,type] pairs; this is equivalent to the hash form.
+    #   * As a dict of \\\\{name:type} pairs; if type is nil, it will be auto-inferred.
+    #   * As a list of column names; in this case types are automatically inferred.
+    #   * As a list of (name,type) pairs; this is equivalent to the hash form.
     #
-    #   If you supply an array of column names that does not match the names in the
+    #   If you supply a list of column names that does not match the names in the
     #   underlying data, the names given here will overwrite them. The number
     #   of names given in the schema should match the underlying data dimensions.
     # @param schema_overrides [Hash]
     #   Support type specification or override of one or more columns; note that
     #   any dtypes inferred from the schema param will be overridden.
+    # @param infer_schema_length [Integer]
+    #   Infer the schema length from the first `infer_schema_length` rows.
+    # @param batch_size [Integer]
+    #   Number of rows to read in each batch.
+    # @param n_rows [Integer]
+    #   Stop reading from JSON file after reading `n_rows`.
+    # @param low_memory [Boolean]
+    #   Reduce memory pressure at the expense of performance.
+    # @param rechunk [Boolean]
+    #   Reallocate to contiguous memory when all chunks/ files are parsed.
+    # @param row_index_name [String]
+    #   If not nil, this will insert a row count column with give name into the
+    #   DataFrame.
+    # @param row_index_offset [Integer]
+    #   Offset to start the row_count column (only use if the name is set).
+    # @param ignore_errors [Boolean]
+    #   Return `Null` if parsing fails because of schema mismatches.
+    # @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 hash of
+    #   credential keys along with an optional credential expiry time.
+    # @param retries [Integer]
+    #   Number of retries if accessing a cloud instance fails.
+    # @param file_cache_ttl [Integer]
+    #   Amount of time to keep downloaded cloud files since their last access time,
+    #   in seconds. Uses the `POLARS_FILE_CACHE_TTL` environment variable
+    #   (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.
     #
     # @return [DataFrame]
     def read_ndjson(
       source,
       schema: nil,
       schema_overrides: nil,
-      ignore_errors: false
+      infer_schema_length: N_INFER_DEFAULT,
+      batch_size: 1024,
+      n_rows: nil,
+      low_memory: false,
+      rechunk: false,
+      row_index_name: nil,
+      row_index_offset: 0,
+      ignore_errors: false,
+      storage_options: nil,
+      credential_provider: "auto",
+      retries: 2,
+      file_cache_ttl: nil,
+      include_file_paths: nil
     )
-      if Utils.pathlike?(source)
-        source = Utils.normalize_filepath(source)
-      end
+      credential_provider_builder = _init_credential_provider_builder(
+        credential_provider, source, storage_options, "read_ndjson"
+      )
 
-      rbdf =
-        RbDataFrame.read_ndjson(
-          source,
-          ignore_errors,
-          schema,
-          schema_overrides
-        )
-      Utils.wrap_df(rbdf)
+      scan_ndjson(
+        source,
+        schema: schema,
+        schema_overrides: schema_overrides,
+        infer_schema_length: infer_schema_length,
+        batch_size: batch_size,
+        n_rows: n_rows,
+        low_memory: low_memory,
+        rechunk: rechunk,
+        row_index_name: row_index_name,
+        row_index_offset: row_index_offset,
+        ignore_errors: ignore_errors,
+        include_file_paths: include_file_paths,
+        retries: retries,
+        storage_options: storage_options,
+        credential_provider: credential_provider_builder,
+        file_cache_ttl: file_cache_ttl,
+      ).collect
     end
 
     # Lazily read from a newline delimited JSON file.
@@ -46,6 +111,19 @@ module Polars
     #
     # @param source [String]
     #   Path to a file.
+    # @param schema [Object]
+    #   The DataFrame schema may be declared in several ways:
+    #
+    #   * As a dict of \\\\{name:type} pairs; if type is nil, it will be auto-inferred.
+    #   * As a list of column names; in this case types are automatically inferred.
+    #   * As a list of (name,type) pairs; this is equivalent to the hash form.
+    #
+    #   If you supply a list of column names that does not match the names in the
+    #   underlying data, the names given here will overwrite them. The number
+    #   of names given in the schema should match the underlying data dimensions.
+    # @param schema_overrides [Hash]
+    #   Support type specification or override of one or more columns; note that
+    #   any dtypes inferred from the schema param will be overridden.
     # @param infer_schema_length [Integer]
     #   Infer the schema length from the first `infer_schema_length` rows.
     # @param batch_size [Integer]
@@ -56,22 +134,58 @@ module Polars
     #   Reduce memory pressure at the expense of performance.
     # @param rechunk [Boolean]
     #   Reallocate to contiguous memory when all chunks/ files are parsed.
-    # @param row_count_name [String]
+    # @param row_index_name [String]
     #   If not nil, this will insert a row count column with give name into the
     #   DataFrame.
-    # @param row_count_offset [Integer]
+    # @param row_index_offset [Integer]
     #   Offset to start the row_count column (only use if the name is set).
+    # @param ignore_errors [Boolean]
+    #   Return `Null` if parsing fails because of schema mismatches.
+    # @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 hash of
+    #   credential keys along with an optional credential expiry time.
+    # @param retries [Integer]
+    #   Number of retries if accessing a cloud instance fails.
+    # @param file_cache_ttl [Integer]
+    #   Amount of time to keep downloaded cloud files since their last access time,
+    #   in seconds. Uses the `POLARS_FILE_CACHE_TTL` environment variable
+    #   (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.
     #
     # @return [LazyFrame]
     def scan_ndjson(
       source,
+      schema: nil,
+      schema_overrides: nil,
       infer_schema_length: N_INFER_DEFAULT,
       batch_size: 1024,
       n_rows: nil,
       low_memory: false,
-      rechunk: true,
-      row_count_name: nil,
-      row_count_offset: 0
+      rechunk: false,
+      row_index_name: nil,
+      row_index_offset: 0,
+      ignore_errors: false,
+      storage_options: nil,
+      credential_provider: "auto",
+      retries: 2,
+      file_cache_ttl: nil,
+      include_file_paths: nil
     )
       sources = []
       if Utils.pathlike?(source)
@@ -86,16 +200,39 @@ module Polars
         source = nil
       end
 
+      if infer_schema_length == 0
+        msg = "'infer_schema_length' should be positive"
+        raise ArgumentError, msg
+      end
+
+      credential_provider_builder = _init_credential_provider_builder(
+        credential_provider, source, storage_options, "scan_ndjson"
+      )
+
+      if storage_options&.any?
+        storage_options = storage_options.map { |k, v| [k.to_s, v.to_s] }
+      else
+        storage_options = nil
+      end
+
       rblf =
         RbLazyFrame.new_from_ndjson(
           source,
           sources,
           infer_schema_length,
+          schema,
+          schema_overrides,
           batch_size,
           n_rows,
           low_memory,
           rechunk,
-          Utils.parse_row_index_args(row_count_name, row_count_offset)
+          Utils.parse_row_index_args(row_index_name, row_index_offset),
+          ignore_errors,
+          include_file_paths,
+          storage_options,
+          credential_provider_builder,
+          retries,
+          file_cache_ttl
         )
       Utils.wrap_ldf(rblf)
     end