polars-df 0.10.0-x86_64-linux → 0.11.0-x86_64-linux

data/lib/polars/io/database.rb

@@ -0,0 +1,73 @@
+module Polars
+  module IO
+    # Read a SQL query into a DataFrame.
+    #
+    # @param query [Object]
+    #   ActiveRecord::Relation or ActiveRecord::Result.
+    # @param schema_overrides [Hash]
+    #   A hash mapping column names to dtypes, used to override the schema
+    #   inferred from the query.
+    #
+    # @return [DataFrame]
+    def read_database(query, schema_overrides: nil)
+      if !defined?(ActiveRecord)
+        raise Error, "Active Record not available"
+      end
+
+      result =
+        if query.is_a?(ActiveRecord::Result)
+          query
+        elsif query.is_a?(ActiveRecord::Relation)
+          query.connection.select_all(query.to_sql)
+        elsif query.is_a?(::String)
+          ActiveRecord::Base.connection.select_all(query)
+        else
+          raise ArgumentError, "Expected ActiveRecord::Relation, ActiveRecord::Result, or String"
+        end
+
+      data = {}
+      schema_overrides = (schema_overrides || {}).transform_keys(&:to_s)
+
+      result.columns.each_with_index do |k, i|
+        column_type = result.column_types[i]
+
+        data[k] =
+          if column_type
+            result.rows.map { |r| column_type.deserialize(r[i]) }
+          else
+            result.rows.map { |r| r[i] }
+          end
+
+        polars_type =
+          case column_type&.type
+          when :binary
+            Binary
+          when :boolean
+            Boolean
+          when :date
+            Date
+          when :datetime, :timestamp
+            Datetime
+          when :decimal
+            Decimal
+          when :float
+            Float64
+          when :integer
+            Int64
+          when :string, :text
+            String
+          when :time
+            Time
+          # TODO fix issue with null
+          # when :json, :jsonb
+          #   Struct
+          end
+
+        schema_overrides[k] ||= polars_type if polars_type
+      end
+
+      DataFrame.new(data, schema_overrides: schema_overrides)
+    end
+    alias_method :read_sql, :read_database
+  end
+end

data/lib/polars/io/ipc.rb

@@ -0,0 +1,247 @@
+module Polars
+  module IO
+    # Read into a DataFrame from Arrow IPC (Feather v2) file.
+    #
+    # @param source [Object]
+    #   Path to a file or a file-like object.
+    # @param columns [Object]
+    #   Columns to select. Accepts a list of column indices (starting at zero) or a list
+    #   of column names.
+    # @param n_rows [Integer]
+    #   Stop reading from IPC file after reading `n_rows`.
+    # @param memory_map [Boolean]
+    #   Try to memory map the file. This can greatly improve performance on repeated
+    #   queries as the OS may cache pages.
+    #   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]
+    #   If not nil, this will insert a row count column with give name into the
+    #   DataFrame.
+    # @param row_count_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.
+    #
+    # @return [DataFrame]
+    def read_ipc(
+      source,
+      columns: nil,
+      n_rows: nil,
+      memory_map: true,
+      storage_options: nil,
+      row_count_name: nil,
+      row_count_offset: 0,
+      rechunk: true
+    )
+      storage_options ||= {}
+      _prepare_file_arg(source, **storage_options) do |data|
+        _read_ipc_impl(
+          data,
+          columns: columns,
+          n_rows: n_rows,
+          row_count_name: row_count_name,
+          row_count_offset: row_count_offset,
+          rechunk: rechunk,
+          memory_map: memory_map
+        )
+      end
+    end
+
+    # @private
+    def _read_ipc_impl(
+      file,
+      columns: nil,
+      n_rows: nil,
+      row_count_name: nil,
+      row_count_offset: 0,
+      rechunk: true,
+      memory_map: true
+    )
+      if Utils.pathlike?(file)
+        file = Utils.normalize_filepath(file)
+      end
+      if columns.is_a?(::String)
+        columns = [columns]
+      end
+
+      if file.is_a?(::String) && file.include?("*")
+        raise Todo
+      end
+
+      projection, columns = Utils.handle_projection_columns(columns)
+      rbdf =
+        RbDataFrame.read_ipc(
+          file,
+          columns,
+          projection,
+          n_rows,
+          Utils._prepare_row_count_args(row_count_name, row_count_offset),
+          memory_map
+        )
+      Utils.wrap_df(rbdf)
+    end
+
+    # Read into a DataFrame from Arrow IPC record batch stream.
+    #
+    # See "Streaming format" on https://arrow.apache.org/docs/python/ipc.html.
+    #
+    # @param source [Object]
+    #   Path to a file or a file-like object.
+    # @param columns [Array]
+    #   Columns to select. Accepts a list of column indices (starting at zero) or a list
+    #   of column names.
+    # @param n_rows [Integer]
+    #   Stop reading from IPC stream after reading `n_rows`.
+    # @param storage_options [Hash]
+    #   Extra options that make sense for a particular storage connection.
+    # @param row_index_name [String]
+    #   Insert a row index column with the given name into the DataFrame as the first
+    #   column. If set to `nil` (default), no row index column is created.
+    # @param row_index_offset [Integer]
+    #   Start the row index at this offset. Cannot be negative.
+    #   Only used if `row_index_name` is set.
+    # @param rechunk [Boolean]
+    #   Make sure that all data is contiguous.
+    #
+    # @return [DataFrame]
+    def read_ipc_stream(
+      source,
+      columns: nil,
+      n_rows: nil,
+      storage_options: nil,
+      row_index_name: nil,
+      row_index_offset: 0,
+      rechunk: true
+    )
+      storage_options ||= {}
+      _prepare_file_arg(source, **storage_options) do |data|
+        _read_ipc_stream_impl(
+          data,
+          columns: columns,
+          n_rows: n_rows,
+          row_index_name: row_index_name,
+          row_index_offset: row_index_offset,
+          rechunk: rechunk
+        )
+      end
+    end
+
+    # @private
+    def _read_ipc_stream_impl(
+      source,
+      columns: nil,
+      n_rows: nil,
+      row_index_name: nil,
+      row_index_offset: 0,
+      rechunk: true
+    )
+      if Utils.pathlike?(source)
+        source = Utils.normalize_filepath(source)
+      end
+      if columns.is_a?(String)
+        columns = [columns]
+      end
+
+      projection, columns = Utils.handle_projection_columns(columns)
+      pydf = RbDataFrame.read_ipc_stream(
+        source,
+        columns,
+        projection,
+        n_rows,
+        Utils._prepare_row_count_args(row_index_name, row_index_offset),
+        rechunk
+      )
+      Utils.wrap_df(pydf)
+    end
+
+    # Get a schema of the IPC file without reading data.
+    #
+    # @param source [Object]
+    #   Path to a file or a file-like object.
+    #
+    # @return [Hash]
+    def read_ipc_schema(source)
+      if Utils.pathlike?(source)
+        source = Utils.normalize_filepath(source)
+      end
+
+      Plr.ipc_schema(source)
+    end
+
+    # Lazily read from an Arrow IPC (Feather v2) file or multiple files via glob patterns.
+    #
+    # This allows the query optimizer to push down predicates and projections to the scan
+    # level, thereby potentially reducing memory overhead.
+    #
+    # @param source [String]
+    #   Path to a IPC file.
+    # @param n_rows [Integer]
+    #   Stop reading from IPC file after reading `n_rows`.
+    # @param cache [Boolean]
+    #   Cache the result after reading.
+    # @param rechunk [Boolean]
+    #   Reallocate to contiguous memory when all chunks/ files are parsed.
+    # @param row_count_name [String]
+    #   If not nil, this will insert a row count column with give name into the
+    #   DataFrame.
+    # @param row_count_offset [Integer]
+    #   Offset to start the row_count column (only use if the name is set).
+    # @param storage_options [Hash]
+    #   Extra options that make sense for a particular storage connection.
+    # @param memory_map [Boolean]
+    #   Try to memory map the file. This can greatly improve performance on repeated
+    #   queries as the OS may cache pages.
+    #   Only uncompressed IPC files can be memory mapped.
+    #
+    # @return [LazyFrame]
+    def scan_ipc(
+      source,
+      n_rows: nil,
+      cache: true,
+      rechunk: true,
+      row_count_name: nil,
+      row_count_offset: 0,
+      storage_options: nil,
+      memory_map: true
+    )
+      _scan_ipc_impl(
+        source,
+        n_rows: n_rows,
+        cache: cache,
+        rechunk: rechunk,
+        row_count_name: row_count_name,
+        row_count_offset: row_count_offset,
+        storage_options: storage_options,
+        memory_map: memory_map
+      )
+    end
+
+    # @private
+    def _scan_ipc_impl(
+      file,
+      n_rows: nil,
+      cache: true,
+      rechunk: true,
+      row_count_name: nil,
+      row_count_offset: 0,
+      storage_options: nil,
+      memory_map: true
+    )
+      if Utils.pathlike?(file)
+        file = Utils.normalize_filepath(file)
+      end
+
+      rblf =
+        RbLazyFrame.new_from_ipc(
+          file,
+          n_rows,
+          cache,
+          rechunk,
+          Utils._prepare_row_count_args(row_count_name, row_count_offset),
+          memory_map
+        )
+      Utils.wrap_ldf(rblf)
+    end
+  end
+end

data/lib/polars/io/json.rb

@@ -0,0 +1,18 @@
+module Polars
+  module IO
+    # Read into a DataFrame from a JSON file.
+    #
+    # @param source [Object]
+    #   Path to a file or a file-like object.
+    #
+    # @return [DataFrame]
+    def read_json(source)
+      if Utils.pathlike?(source)
+        source = Utils.normalize_filepath(source)
+      end
+
+      rbdf = RbDataFrame.read_json(source)
+      Utils.wrap_df(rbdf)
+    end
+  end
+end

data/lib/polars/io/ndjson.rb

@@ -0,0 +1,69 @@
+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.
+    #
+    # @return [DataFrame]
+    def read_ndjson(source)
+      if Utils.pathlike?(source)
+        source = Utils.normalize_filepath(source)
+      end
+
+      rbdf = RbDataFrame.read_ndjson(source)
+      Utils.wrap_df(rbdf)
+    end
+
+    # Lazily read from a newline delimited JSON file.
+    #
+    # This allows the query optimizer to push down predicates and projections to the scan
+    # level, thereby potentially reducing memory overhead.
+    #
+    # @param source [String]
+    #   Path to a file.
+    # @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_count_name [String]
+    #   If not nil, this will insert a row count column with give name into the
+    #   DataFrame.
+    # @param row_count_offset [Integer]
+    #   Offset to start the row_count column (only use if the name is set).
+    #
+    # @return [LazyFrame]
+    def scan_ndjson(
+      source,
+      infer_schema_length: 100,
+      batch_size: 1024,
+      n_rows: nil,
+      low_memory: false,
+      rechunk: true,
+      row_count_name: nil,
+      row_count_offset: 0
+    )
+      if Utils.pathlike?(source)
+        source = Utils.normalize_filepath(source)
+      end
+
+      rblf =
+        RbLazyFrame.new_from_ndjson(
+          source,
+          infer_schema_length,
+          batch_size,
+          n_rows,
+          low_memory,
+          rechunk,
+          Utils._prepare_row_count_args(row_count_name, row_count_offset)
+        )
+      Utils.wrap_ldf(rblf)
+    end
+  end
+end

data/lib/polars/io/parquet.rb

@@ -0,0 +1,226 @@
+module Polars
+  module IO
+    # Read into a DataFrame from a parquet file.
+    #
+    # @param source [String, Pathname, StringIO]
+    #   Path to a file or a file-like object.
+    # @param columns [Object]
+    #   Columns to select. Accepts a list of column indices (starting at zero) or a list
+    #   of column names.
+    # @param n_rows [Integer]
+    #   Stop reading from parquet file after reading `n_rows`.
+    # @param storage_options [Hash]
+    #   Extra options that make sense for a particular storage connection.
+    # @param parallel ["auto", "columns", "row_groups", "none"]
+    #   This determines the direction of parallelism. 'auto' will try to determine the
+    #   optimal direction.
+    # @param row_count_name [String]
+    #   If not nil, this will insert a row count column with give name into the
+    #   DataFrame.
+    # @param row_count_offset [Integer]
+    #   Offset to start the row_count column (only use if the name is set).
+    # @param low_memory [Boolean]
+    #   Reduce memory pressure at the expense of performance.
+    # @param use_statistics [Boolean]
+    #   Use statistics in the parquet to determine if pages
+    #   can be skipped from reading.
+    # @param rechunk [Boolean]
+    #   Make sure that all columns are contiguous in memory by
+    #   aggregating the chunks into a single array.
+    #
+    # @return [DataFrame]
+    #
+    # @note
+    #   This operation defaults to a `rechunk` operation at the end, meaning that
+    #   all data will be stored continuously in memory.
+    #   Set `rechunk: false` if you are benchmarking the parquet-reader. A `rechunk` is
+    #   an expensive operation.
+    def read_parquet(
+      source,
+      columns: nil,
+      n_rows: nil,
+      storage_options: nil,
+      parallel: "auto",
+      row_count_name: nil,
+      row_count_offset: 0,
+      low_memory: false,
+      use_statistics: true,
+      rechunk: true
+    )
+      _prepare_file_arg(source) do |data|
+        _read_parquet_impl(
+          data,
+          columns: columns,
+          n_rows: n_rows,
+          parallel: parallel,
+          row_count_name: row_count_name,
+          row_count_offset: row_count_offset,
+          low_memory: low_memory,
+          use_statistics: use_statistics,
+          rechunk: rechunk
+        )
+      end
+    end
+
+    # @private
+    def _read_parquet_impl(
+      source,
+      columns: nil,
+      n_rows: nil,
+      parallel: "auto",
+      row_count_name: nil,
+      row_count_offset: 0,
+      low_memory: false,
+      use_statistics: true,
+      rechunk: true
+    )
+      if Utils.pathlike?(source)
+        source = Utils.normalize_filepath(source)
+      end
+      if columns.is_a?(::String)
+        columns = [columns]
+      end
+
+      if source.is_a?(::String) && source.include?("*") && Utils.local_file?(source)
+        scan =
+          scan_parquet(
+            source,
+            n_rows: n_rows,
+            rechunk: true,
+            parallel: parallel,
+            row_count_name: row_count_name,
+            row_count_offset: row_count_offset,
+            low_memory: low_memory
+          )
+
+        if columns.nil?
+          return scan.collect
+        elsif Utils.is_str_sequence(columns, allow_str: false)
+          return scan.select(columns).collect
+        else
+          raise ArgumentError, "cannot use glob patterns and integer based projection as `columns` argument; Use columns: Array[String]"
+        end
+      end
+
+      projection, columns = Utils.handle_projection_columns(columns)
+      rbdf =
+        RbDataFrame.read_parquet(
+          source,
+          columns,
+          projection,
+          n_rows,
+          parallel,
+          Utils._prepare_row_count_args(row_count_name, row_count_offset),
+          low_memory,
+          use_statistics,
+          rechunk
+        )
+      Utils.wrap_df(rbdf)
+    end
+
+    # Get a schema of the Parquet file without reading data.
+    #
+    # @param source [Object]
+    #   Path to a file or a file-like object.
+    #
+    # @return [Hash]
+    def read_parquet_schema(source)
+      if Utils.pathlike?(source)
+        source = Utils.normalize_filepath(source)
+      end
+
+      Plr.parquet_schema(source)
+    end
+
+    # Lazily read from a parquet file or multiple files via glob patterns.
+    #
+    # This allows the query optimizer to push down predicates and projections to the scan
+    # level, thereby potentially reducing memory overhead.
+    #
+    # @param source [String]
+    #   Path to a file.
+    # @param n_rows [Integer]
+    #   Stop reading from parquet file after reading `n_rows`.
+    # @param cache [Boolean]
+    #   Cache the result after reading.
+    # @param parallel ["auto", "columns", "row_groups", "none"]
+    #   This determines the direction of parallelism. 'auto' will try to determine the
+    #   optimal direction.
+    # @param rechunk [Boolean]
+    #   In case of reading multiple files via a glob pattern rechunk the final DataFrame
+    #   into contiguous memory chunks.
+    # @param row_count_name [String]
+    #   If not nil, this will insert a row count column with give name into the
+    #   DataFrame.
+    # @param row_count_offset [Integer]
+    #   Offset to start the row_count column (only use if the name is set).
+    # @param storage_options [Hash]
+    #   Extra options that make sense for a particular storage connection.
+    # @param low_memory [Boolean]
+    #   Reduce memory pressure at the expense of performance.
+    #
+    # @return [LazyFrame]
+    def scan_parquet(
+      source,
+      n_rows: nil,
+      cache: true,
+      parallel: "auto",
+      glob: true,
+      rechunk: true,
+      row_count_name: nil,
+      row_count_offset: 0,
+      storage_options: nil,
+      low_memory: false
+    )
+      if Utils.pathlike?(source)
+        source = Utils.normalize_filepath(source)
+      end
+
+      _scan_parquet_impl(
+        source,
+        n_rows:n_rows,
+        cache: cache,
+        parallel: parallel,
+        rechunk: rechunk,
+        row_count_name: row_count_name,
+        row_count_offset: row_count_offset,
+        storage_options: storage_options,
+        low_memory: low_memory,
+        glob: glob
+      )
+    end
+
+    # @private
+    def _scan_parquet_impl(
+      file,
+      n_rows: nil,
+      cache: true,
+      parallel: "auto",
+      rechunk: true,
+      row_count_name: nil,
+      row_count_offset: 0,
+      storage_options: nil,
+      low_memory: false,
+      use_statistics: true,
+      hive_partitioning: true,
+      glob: true
+    )
+      rblf =
+        RbLazyFrame.new_from_parquet(
+          file,
+          [],
+          n_rows,
+          cache,
+          parallel,
+          rechunk,
+          Utils._prepare_row_count_args(row_count_name, row_count_offset),
+          low_memory,
+          use_statistics,
+          hive_partitioning,
+          nil,
+          glob
+        )
+      Utils.wrap_ldf(rblf)
+    end
+  end
+end