polars-df 0.21.1-x86_64-linux → 0.22.0-x86_64-linux

data/lib/polars/data_frame.rb

@@ -72,6 +72,43 @@ module Polars
       end
     end
 
+    # Read a serialized DataFrame from a file.
+    #
+    # @param source [Object]
+    #     Path to a file or a file-like object (by file-like object, we refer to
+    #     objects that have a `read` method, such as a file handler or `StringIO`).
+    #
+    # @return [DataFrame]
+    #
+    # @note
+    #   Serialization is not stable across Polars versions: a LazyFrame serialized
+    #   in one Polars version may not be deserializable in another Polars version.
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"a" => [1, 2, 3], "b" => [4.0, 5.0, 6.0]})
+    #   bytes = df.serialize
+    #   Polars::DataFrame.deserialize(StringIO.new(bytes))
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌─────┬─────┐
+    #   # │ a   ┆ b   │
+    #   # │ --- ┆ --- │
+    #   # │ i64 ┆ f64 │
+    #   # ╞═════╪═════╡
+    #   # │ 1   ┆ 4.0 │
+    #   # │ 2   ┆ 5.0 │
+    #   # │ 3   ┆ 6.0 │
+    #   # └─────┴─────┘
+    def self.deserialize(source)
+      if Utils.pathlike?(source)
+        source = Utils.normalize_filepath(source)
+      end
+
+      deserializer = RbDataFrame.method(:deserialize_binary)
+
+      _from_rbdf(deserializer.(source))
+    end
+
     # @private
     def self._from_rbdf(rb_df)
       df = DataFrame.allocate
@@ -562,8 +599,6 @@ module Polars
 
     # Convert every row to a hash.
     #
-    # Note that this is slow.
-    #
     # @return [Array]
     #
     # @example
@@ -572,12 +607,7 @@ module Polars
     #   # =>
     #   # [{"foo"=>1, "bar"=>4}, {"foo"=>2, "bar"=>5}, {"foo"=>3, "bar"=>6}]
     def to_hashes
-      rbdf = _df
-      names = columns
-
-      height.times.map do |i|
-        names.zip(rbdf.row_tuple(i)).to_h
-      end
+      rows(named: true)
     end
 
     # Convert DataFrame to a 2D Numo array.
@@ -634,6 +664,44 @@ module Polars
       Utils.wrap_s(_df.select_at_idx(index))
     end
 
+    # Serialize this DataFrame to a file or string.
+    #
+    # @param file [Object]
+    #   File path or writable file-like object to which the result will be written.
+    #   If set to `nil` (default), the output is returned as a string instead.
+    #
+    # @return [Object]
+    #
+    # @note
+    #   Serialization is not stable across Polars versions: a LazyFrame serialized
+    #   in one Polars version may not be deserializable in another Polars version.
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "foo" => [1, 2, 3],
+    #       "bar" => [6, 7, 8]
+    #     }
+    #   )
+    #   bytes = df.serialize
+    #   Polars::DataFrame.deserialize(StringIO.new(bytes))
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌─────┬─────┐
+    #   # │ foo ┆ bar │
+    #   # │ --- ┆ --- │
+    #   # │ i64 ┆ i64 │
+    #   # ╞═════╪═════╡
+    #   # │ 1   ┆ 6   │
+    #   # │ 2   ┆ 7   │
+    #   # │ 3   ┆ 8   │
+    #   # └─────┴─────┘
+    def serialize(file = nil)
+      serializer = _df.method(:serialize_binary)
+
+      Utils.serialize_polars_object(serializer, file)
+    end
+
     # Serialize to JSON representation.
     #
     # @param file [String]
@@ -1148,6 +1216,40 @@ module Polars
       end
     end
 
+    # Write DataFrame to an Iceberg table.
+    #
+    # @note
+    #   This functionality is currently considered **unstable**. It may be
+    #   changed at any point without it being considered a breaking change.
+    #
+    # @param target [Object]
+    #   Name of the table or the Table object representing an Iceberg table.
+    # @param mode ['append', 'overwrite']
+    #   How to handle existing data.
+    #
+    #   - If 'append', will add new data.
+    #   - If 'overwrite', will replace table with new data.
+    #
+    # @return [nil]
+    def write_iceberg(target, mode:)
+      require "iceberg"
+
+      table =
+        if target.is_a?(Iceberg::Table)
+          target
+        else
+          raise Todo
+        end
+
+      data = self
+
+      if mode == "append"
+        table.append(data)
+      else
+        raise Todo
+      end
+    end
+
     # Write DataFrame as delta table.
     #
     # @param target [Object]

data/lib/polars/data_types.rb

@@ -110,12 +110,23 @@ module Polars
       DataTypeExpr._from_rbdatatype_expr(RbDataTypeExpr.from_dtype(self))
     end
 
-    [:numeric?, :decimal?, :integer?, :signed_integer?, :unsigned_integer?, :float?, :temporal?, :nested?, :to_dtype_expr].each do |v|
+    [:numeric?, :decimal?, :integer?, :signed_integer?, :unsigned_integer?, :float?, :temporal?, :nested?].each do |v|
       define_method(v) do
         self.class.public_send(v)
       end
     end
 
+    # Return a `DataTypeExpr` with a static `DataType`.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   Polars::Int16.new.to_dtype_expr.collect_dtype({})
+    #   # => Polars::Int16
+    def to_dtype_expr
+      DataTypeExpr._from_rbdatatype_expr(RbDataTypeExpr.from_dtype(self))
+    end
+
     # Returns a string representing the data type.
     #
     # @return [String]
@@ -317,11 +328,9 @@ module Polars
   class Categories
     attr_accessor :_categories
 
-    def initialize
-      # TODO fix
-      name = nil
+    def initialize(name = nil)
       if name.nil? || name == ""
-        @_categories = RbCategories.global_categories
+        self._categories = RbCategories.global_categories
         return
       end
 

data/lib/polars/date_time_expr.rb

@@ -1188,7 +1188,7 @@ module Polars
       if Utils::DTYPE_TEMPORAL_UNITS.include?(time_unit)
         timestamp(time_unit)
       elsif time_unit == "s"
-        Utils.wrap_expr(_rbexpr.dt_epoch_seconds)
+        timestamp("ms").floordiv(F.lit(1000, dtype: Int64))
       elsif time_unit == "d"
         Utils.wrap_expr(_rbexpr).cast(:date).cast(:i32)
       else

data/lib/polars/expr.rb

@@ -146,6 +146,40 @@ module Polars
       wrap_expr(_rbexpr.neg)
     end
 
+    # Read a serialized expression from a file.
+    #
+    # @param source [Object]
+    #   Path to a file or a file-like object (by file-like object, we refer to
+    #   objects that have a `read` method, such as a file handler or `StringIO`).
+    #
+    # @return [Expr]
+    #
+    # @note
+    #   This function uses  marshaling if the logical plan contains Ruby UDFs,
+    #   and as such inherits the security implications. Deserializing can execute
+    #   arbitrary code, so it should only be attempted on trusted data.
+    #
+    # @note
+    #   Serialization is not stable across Polars versions: a LazyFrame serialized
+    #   in one Polars version may not be deserializable in another Polars version.
+    #
+    # @example
+    #   expr = Polars.col("foo").sum.over("bar")
+    #   bytes = expr.meta.serialize
+    #   Polars::Expr.deserialize(StringIO.new(bytes))
+    #   # => col("foo").sum().over([col("bar")])
+    def self.deserialize(source)
+      raise Todo unless RbExpr.respond_to?(:deserialize_binary)
+
+      if Utils.pathlike?(source)
+        source = Utils.normalize_filepath(source)
+      end
+
+      deserializer = RbExpr.method(:deserialize_binary)
+
+      _from_rbexpr(deserializer.(source))
+    end
+
     # Cast to physical representation of the logical dtype.
     #
     # - `:date` -> `:i32`
@@ -377,8 +411,6 @@ module Polars
       wrap_expr(_rbexpr._alias(name))
     end
 
-    # TODO support symbols for exclude
-
     # Exclude certain columns from a wildcard/regex selection.
     #
     # You may also use regexes in the exclude list. They must start with `^` and end
@@ -1787,7 +1819,7 @@ module Polars
       wrap_expr(_rbexpr.arg_min)
     end
 
-    # Get the index of the first occurrence of a value, or ``None`` if it's not found.
+    # Get the index of the first occurrence of a value, or `nil` if it's not found.
     #
     # @param element [Object]
     #   Value to find.
@@ -7571,7 +7603,8 @@ module Polars
     #   # │ 1.584963 │
     #   # └──────────┘
     def log(base = Math::E)
-      wrap_expr(_rbexpr.log(base))
+      base_rbexpr = Utils.parse_into_expression(base)
+      wrap_expr(_rbexpr.log(base_rbexpr))
     end
 
     # Compute the natural logarithm of each element plus one.
@@ -7743,33 +7776,9 @@ module Polars
     # This can be used to reduce memory pressure.
     #
     # @return [Expr]
-    #
-    # @example
-    #   Polars::DataFrame.new(
-    #     {
-    #       "a" => [1, 2, 3],
-    #       "b" => [1, 2, 2 << 32],
-    #       "c" => [-1, 2, 1 << 30],
-    #       "d" => [-112, 2, 112],
-    #       "e" => [-112, 2, 129],
-    #       "f" => ["a", "b", "c"],
-    #       "g" => [0.1, 1.32, 0.12],
-    #       "h" => [true, nil, false]
-    #     }
-    #   ).select(Polars.all.shrink_dtype)
-    #   # =>
-    #   # shape: (3, 8)
-    #   # ┌─────┬────────────┬────────────┬──────┬──────┬─────┬──────┬───────┐
-    #   # │ a   ┆ b          ┆ c          ┆ d    ┆ e    ┆ f   ┆ g    ┆ h     │
-    #   # │ --- ┆ ---        ┆ ---        ┆ ---  ┆ ---  ┆ --- ┆ ---  ┆ ---   │
-    #   # │ i8  ┆ i64        ┆ i32        ┆ i8   ┆ i16  ┆ str ┆ f32  ┆ bool  │
-    #   # ╞═════╪════════════╪════════════╪══════╪══════╪═════╪══════╪═══════╡
-    #   # │ 1   ┆ 1          ┆ -1         ┆ -112 ┆ -112 ┆ a   ┆ 0.1  ┆ true  │
-    #   # │ 2   ┆ 2          ┆ 2          ┆ 2    ┆ 2    ┆ b   ┆ 1.32 ┆ null  │
-    #   # │ 3   ┆ 8589934592 ┆ 1073741824 ┆ 112  ┆ 129  ┆ c   ┆ 0.12 ┆ false │
-    #   # └─────┴────────────┴────────────┴──────┴──────┴─────┴──────┴───────┘
     def shrink_dtype
-      wrap_expr(_rbexpr.shrink_dtype)
+      warn "`Expr.shrink_dtype` is deprecated and is a no-op; use `Series.shrink_dtype` instead."
+      self
     end
 
     # Bin values into buckets and count their occurrences.

data/lib/polars/functions/business.rb

@@ -0,0 +1,95 @@
+module Polars
+  module Functions
+    # Count the number of business days between `start` and `end` (not including `end`).
+    #
+    # @note
+    #   This functionality is considered **unstable**. It may be changed
+    #   at any point without it being considered a breaking change.
+    #
+    # @param start [Object]
+    #   Start dates.
+    # @param stop [Object]
+    #   End dates.
+    # @param week_mask [Array]
+    #   Which days of the week to count. The default is Monday to Friday.
+    #   If you wanted to count only Monday to Thursday, you would pass
+    #   `[true, true, true, true, false, false, false]`.
+    # @param holidays [Array]
+    #   Holidays to exclude from the count.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "start" => [Date.new(2020, 1, 1), Date.new(2020, 1, 2)],
+    #       "end" => [Date.new(2020, 1, 2), Date.new(2020, 1, 10)]
+    #     }
+    #   )
+    #   df.with_columns(
+    #     business_day_count: Polars.business_day_count("start", "end")
+    #   )
+    #   # =>
+    #   # shape: (2, 3)
+    #   # ┌────────────┬────────────┬────────────────────┐
+    #   # │ start      ┆ end        ┆ business_day_count │
+    #   # │ ---        ┆ ---        ┆ ---                │
+    #   # │ date       ┆ date       ┆ i32                │
+    #   # ╞════════════╪════════════╪════════════════════╡
+    #   # │ 2020-01-01 ┆ 2020-01-02 ┆ 1                  │
+    #   # │ 2020-01-02 ┆ 2020-01-10 ┆ 6                  │
+    #   # └────────────┴────────────┴────────────────────┘
+    #
+    # @example You can pass a custom weekend - for example, if you only take Sunday off:
+    #   week_mask = [true, true, true, true, true, true, false]
+    #   df.with_columns(
+    #     business_day_count: Polars.business_day_count(
+    #       "start", "end", week_mask: week_mask
+    #     )
+    #   )
+    #   # =>
+    #   # shape: (2, 3)
+    #   # ┌────────────┬────────────┬────────────────────┐
+    #   # │ start      ┆ end        ┆ business_day_count │
+    #   # │ ---        ┆ ---        ┆ ---                │
+    #   # │ date       ┆ date       ┆ i32                │
+    #   # ╞════════════╪════════════╪════════════════════╡
+    #   # │ 2020-01-01 ┆ 2020-01-02 ┆ 1                  │
+    #   # │ 2020-01-02 ┆ 2020-01-10 ┆ 7                  │
+    #   # └────────────┴────────────┴────────────────────┘
+    #
+    # @example You can also pass a list of holidays to exclude from the count:
+    #   holidays = [Date.new(2020, 1, 1), Date.new(2020, 1, 2)]
+    #   df.with_columns(
+    #     business_day_count: Polars.business_day_count("start", "end", holidays: holidays)
+    #   )
+    #   # =>
+    #   # shape: (2, 3)
+    #   # ┌────────────┬────────────┬────────────────────┐
+    #   # │ start      ┆ end        ┆ business_day_count │
+    #   # │ ---        ┆ ---        ┆ ---                │
+    #   # │ date       ┆ date       ┆ i32                │
+    #   # ╞════════════╪════════════╪════════════════════╡
+    #   # │ 2020-01-01 ┆ 2020-01-02 ┆ 0                  │
+    #   # │ 2020-01-02 ┆ 2020-01-10 ┆ 5                  │
+    #   # └────────────┴────────────┴────────────────────┘
+    def business_day_count(
+      start,
+      stop,
+      week_mask: [true, true, true, true, true, false, false],
+      holidays: []
+    )
+      start_rbexpr = Utils.parse_into_expression(start)
+      end_rbexpr = Utils.parse_into_expression(stop)
+      unix_epoch = ::Date.new(1970, 1, 1)
+      Utils.wrap_expr(
+        Plr.business_day_count(
+          start_rbexpr,
+          end_rbexpr,
+          week_mask,
+          holidays.map { |holiday| holiday - unix_epoch }
+        )
+      )
+    end
+  end
+end

data/lib/polars/functions/lazy.rb

@@ -823,7 +823,7 @@ module Polars
     #   # ┌─────┐
     #   # │ sum │
     #   # │ --- │
-    #   # │ i64 │
+    #   # │ i32 │
     #   # ╞═════╡
     #   # │ 10  │
     #   # │ 13  │

data/lib/polars/io/iceberg.rb

@@ -0,0 +1,27 @@
+module Polars
+  module IO
+    # Lazily read from an Apache Iceberg table.
+    #
+    # @param source [Object]
+    #   A Iceberg Ruby table, or a direct path to the metadata.
+    # @param snapshot_id [Integer]
+    #   The snapshot ID to scan from.
+    # @param storage_options [Hash]
+    #   Extra options for the storage backends.
+    #
+    # @return [LazyFrame]
+    def scan_iceberg(
+      source,
+      snapshot_id: nil,
+      storage_options: nil
+    )
+      require "iceberg"
+
+      unless source.is_a?(Iceberg::Table)
+        raise Todo
+      end
+
+      source.to_polars(snapshot_id:, storage_options:)
+    end
+  end
+end

data/lib/polars/io/parquet.rb

@@ -117,14 +117,13 @@ module Polars
     # @param source [Object]
     #   Path to a file or a file-like object.
     #
-    # @return [Hash]
+    # @return [Schema]
     def read_parquet_schema(source)
       if Utils.pathlike?(source)
         source = Utils.normalize_filepath(source)
       end
 
-      # TODO return Schema
-      scan_parquet(source).collect_schema.to_h
+      scan_parquet(source).collect_schema
     end
 
     # Get file-level custom metadata of a Parquet file without reading data.
@@ -207,6 +206,9 @@ module Polars
     #   defined schema are encountered in the data:
     #     * `ignore`: Silently ignores.
     #     * `raise`: Raises an error.
+    # @param cast_options [Object]
+    #   Configuration for column type-casting during scans. Useful for datasets
+    #   containing files that have differing schemas.
     #
     # @return [LazyFrame]
     def scan_parquet(
@@ -230,6 +232,7 @@ module Polars
       include_file_paths: nil,
       allow_missing_columns: false,
       extra_columns: "raise",
+      cast_options: nil,
       _column_mapping: nil,
       _deletion_files: nil
     )
@@ -268,7 +271,7 @@ module Polars
           ScanOptions.new(
             row_index: !row_index_name.nil? ? [row_index_name, row_index_offset] : nil,
             pre_slice: !n_rows.nil? ? [0, n_rows] : nil,
-            # cast_options: cast_options,
+            cast_options: cast_options,
             extra_columns: extra_columns,
             missing_columns: missing_columns,
             include_file_paths: include_file_paths,

data/lib/polars/io/scan_options.rb

@@ -3,7 +3,8 @@ module Polars
     class ScanOptions
       attr_reader :row_index, :pre_slice, :cast_options, :extra_columns, :missing_columns,
         :include_file_paths, :glob, :hive_partitioning, :hive_schema, :try_parse_hive_dates,
-        :rechunk, :cache, :storage_options, :credential_provider, :retries, :column_mapping, :deletion_files
+        :rechunk, :cache, :storage_options, :credential_provider, :retries, :column_mapping,
+        :default_values, :deletion_files
 
       def initialize(
         row_index: nil,
@@ -22,6 +23,7 @@ module Polars
         credential_provider: nil,
         retries: 2,
         column_mapping: nil,
+        default_values: nil,
         deletion_files: nil
       )
         @row_index = row_index
@@ -40,6 +42,7 @@ module Polars
         @credential_provider = credential_provider
         @retries = retries
         @column_mapping = column_mapping
+        @default_values = default_values
         @deletion_files = deletion_files
       end
     end

data/lib/polars/lazy_frame.rb

@@ -27,9 +27,6 @@ module Polars
       ldf
     end
 
-    # def self.from_json
-    # end
-
     # Read a logical plan from a JSON file to construct a LazyFrame.
     #
     # @param file [String]
@@ -41,7 +38,49 @@ module Polars
         file = Utils.normalize_filepath(file)
       end
 
-      Utils.wrap_ldf(RbLazyFrame.read_json(file))
+      Utils.wrap_ldf(RbLazyFrame.deserialize_json(file))
+    end
+
+    # Read a logical plan from a file to construct a LazyFrame.
+    #
+    # @param source [Object]
+    #   Path to a file or a file-like object (by file-like object, we refer to
+    #   objects that have a `read` method, such as a file handler or `StringIO`).
+    #
+    # @return [LazyFrame]
+    #
+    # @note
+    #   This function uses marshaling if the logical plan contains Ruby UDFs,
+    #   and as such inherits the security implications. Deserializing can execute
+    #   arbitrary code, so it should only be attempted on trusted data.
+    #
+    # @note
+    #   Serialization is not stable across Polars versions: a LazyFrame serialized
+    #   in one Polars version may not be deserializable in another Polars version.
+    #
+    # @example
+    #   lf = Polars::LazyFrame.new({"a" => [1, 2, 3]}).sum
+    #   bytes = lf.serialize
+    #   Polars::LazyFrame.deserialize(StringIO.new(bytes)).collect
+    #   # =>
+    #   # shape: (1, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ i64 │
+    #   # ╞═════╡
+    #   # │ 6   │
+    #   # └─────┘
+    def self.deserialize(source)
+      raise Todo unless RbLazyFrame.respond_to?(:deserialize_binary)
+
+      if Utils.pathlike?(source)
+        source = Utils.normalize_filepath(source)
+      end
+
+      deserializer = RbLazyFrame.method(:deserialize_binary)
+
+      _from_rbldf(deserializer.(source))
     end
 
     # Get or set column names.
@@ -151,6 +190,38 @@ module Polars
       nil
     end
 
+    # Serialize the logical plan of this LazyFrame to a file or string.
+    #
+    # @param file [Object]
+    #   File path to which the result should be written. If set to `nil`
+    #   (default), the output is returned as a string instead.
+    #
+    # @return [Object]
+    #
+    # @note
+    #   Serialization is not stable across Polars versions: a LazyFrame serialized
+    #   in one Polars version may not be deserializable in another Polars version.
+    #
+    # @example Serialize the logical plan into a binary representation.
+    #   lf = Polars::LazyFrame.new({"a" => [1, 2, 3]}).sum
+    #   bytes = lf.serialize
+    #   Polars::LazyFrame.deserialize(StringIO.new(bytes)).collect
+    #   # =>
+    #   # shape: (1, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ i64 │
+    #   # ╞═════╡
+    #   # │ 6   │
+    #   # └─────┘
+    def serialize(file = nil)
+      raise Todo unless _ldf.respond_to?(:serialize_binary)
+
+      serializer = _ldf.method(:serialize_binary)
+      Utils.serialize_polars_object(serializer, file)
+    end
+
     # Offers a structured way to apply a sequence of user-defined functions (UDFs).
     #
     # @param func [Object]
@@ -774,6 +845,21 @@ module Polars
     # @param maintain_order [Boolean]
     #   Maintain the order in which data is processed.
     #   Setting this to `false` will  be slightly faster.
+    # @param storage_options [String]
+    #   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 retries [Integer]
+    #   Number of retries if accessing a cloud instance fails.
     # @param type_coercion [Boolean]
     #   Do type coercion optimization.
     # @param predicate_pushdown [Boolean]
@@ -806,6 +892,8 @@ module Polars
       path,
       compression: "zstd",
       maintain_order: true,
+      storage_options: nil,
+      retries: 2,
       type_coercion: true,
       predicate_pushdown: true,
       projection_pushdown: true,
@@ -816,10 +904,6 @@ module Polars
       mkdir: false,
       lazy: false
     )
-      # TODO support storage options in Rust
-      storage_options = nil
-      retries = 2
-
       lf = _set_sink_optimizations(
         type_coercion: type_coercion,
         predicate_pushdown: predicate_pushdown,