polars-df 0.15.0 → 0.26.0

data/lib/polars/group_by.rb

@@ -2,10 +2,12 @@ module Polars
   # Starts a new GroupBy operation.
   class GroupBy
     # @private
-    def initialize(df, by, maintain_order: false)
+    def initialize(df, *by, maintain_order:, predicates:, **named_by)
       @df = df
       @by = by
+      @named_by = named_by
       @maintain_order = maintain_order
+      @predicates = predicates
     end
 
     # Allows iteration over the groups of the group by operation.
@@ -16,7 +18,7 @@ module Polars
     #   df = Polars::DataFrame.new({"foo" => ["a", "a", "b"], "bar" => [1, 2, 3]})
     #   df.group_by("foo", maintain_order: true).each.to_h
     #   # =>
-    #   # {"a"=>shape: (2, 2)
+    #   # {["a"]=>shape: (2, 2)
     #   # ┌─────┬─────┐
     #   # │ foo ┆ bar │
     #   # │ --- ┆ --- │
@@ -24,7 +26,7 @@ module Polars
     #   # ╞═════╪═════╡
     #   # │ a   ┆ 1   │
     #   # │ a   ┆ 2   │
-    #   # └─────┴─────┘, "b"=>shape: (1, 2)
+    #   # └─────┴─────┘, ["b"]=>shape: (1, 2)
     #   # ┌─────┬─────┐
     #   # │ foo ┆ bar │
     #   # │ --- ┆ --- │
@@ -39,9 +41,9 @@ module Polars
       groups_df =
         @df.lazy
           .with_row_index(name: temp_col)
-          .group_by(@by, maintain_order: @maintain_order)
+          .group_by(@by, **@named_by, maintain_order: @maintain_order)
           .agg(Polars.col(temp_col))
-          .collect(no_optimization: true)
+          .collect(optimizations: QueryOptFlags.none)
 
       group_names = groups_df.select(Polars.all.exclude(temp_col))
 
@@ -65,46 +67,46 @@ module Polars
       end
     end
 
-    # Apply a custom/user-defined function (UDF) over the groups as a sub-DataFrame.
+    # Filter groups with a list of predicates after aggregation.
     #
-    # Implementing logic using a Ruby function is almost always _significantly_
-    # slower and more memory intensive than implementing the same logic using
-    # the native expression API because:
-
-    # - The native expression engine runs in Rust; UDFs run in Ruby.
-    # - Use of Ruby UDFs forces the DataFrame to be materialized in memory.
-    # - Polars-native expressions can be parallelised (UDFs cannot).
-    # - Polars-native expressions can be logically optimised (UDFs cannot).
+    # Using this method is equivalent to adding the predicates to the aggregation and
+    # filtering afterwards.
     #
-    # Wherever possible you should strongly prefer the native expression API
-    # to achieve the best performance.
+    # This method can be chained and all conditions will be combined using `&`.
     #
-    # @return [DataFrame]
+    # @param predicates [Array]
+    #   Expressions that evaluate to a boolean value for each group. Typically, this
+    #   requires the use of an aggregation function. Multiple predicates are
+    #   combined using `&`.
+    #
+    # @return [GroupBy]
     #
     # @example
     #   df = Polars::DataFrame.new(
     #     {
-    #       "id" => [0, 1, 2, 3, 4],
-    #       "color" => ["red", "green", "green", "red", "red"],
-    #       "shape" => ["square", "triangle", "square", "triangle", "square"]
+    #       "a" => ["a", "b", "a", "b", "c"]
     #     }
     #   )
-    #   df.group_by("color").apply { |group_df| group_df.sample(2) }
+    #   df.group_by("a").having(Polars.len > 1).agg
     #   # =>
-    #   # shape: (4, 3)
-    #   # ┌─────┬───────┬──────────┐
-    #   # │ id  ┆ color ┆ shape    │
-    #   # │ --- ┆ ---   ┆ ---      │
-    #   # │ i64 ┆ str   ┆ str      │
-    #   # ╞═════╪═══════╪══════════╡
-    #   # │ 1   ┆ green ┆ triangle │
-    #   # │ 2   ┆ green ┆ square   │
-    #   # │ 4   ┆ red   ┆ square   │
-    #   # │ 3   ┆ red   ┆ triangle │
-    #   # └─────┴───────┴──────────┘
-    # def apply(&f)
-    #   _dataframe_class._from_rbdf(_df.group_by_apply(by, f))
-    # end
+    #   # shape: (2, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ str │
+    #   # ╞═════╡
+    #   # │ b   │
+    #   # │ a   │
+    #   # └─────┘
+    def having(*predicates)
+      GroupBy.new(
+        @df,
+        *@by,
+        maintain_order: @maintain_order,
+        predicates: Utils._chain_predicates(@predicates, predicates),
+        **@named_by
+      )
+    end
 
     # Compute aggregations for each group of a group by operation.
     #
@@ -201,10 +203,71 @@ module Polars
     #   # │ b   ┆ 5     ┆ 10.0           │
     #   # └─────┴───────┴────────────────┘
     def agg(*aggs, **named_aggs)
-      @df.lazy
-        .group_by(@by, maintain_order: @maintain_order)
+      _lgb
         .agg(*aggs, **named_aggs)
-        .collect(no_optimization: true)
+        .collect(optimizations: QueryOptFlags.none)
+    end
+
+    # Apply a custom/user-defined function (UDF) over the groups as a sub-DataFrame.
+    #
+    # @note
+    #   This method is much slower than the native expressions API.
+    #   Only use it if you cannot implement your logic otherwise.
+    #
+    # Implementing logic using a Ruby function is almost always *significantly*
+    # slower and more memory intensive than implementing the same logic using
+    # the native expression API because:
+    #
+    # - The native expression engine runs in Rust; UDFs run in Ruby.
+    # - Use of Ruby UDFs forces the DataFrame to be materialized in memory.
+    # - Polars-native expressions can be parallelised (UDFs cannot).
+    # - Polars-native expressions can be logically optimised (UDFs cannot).
+    #
+    # Wherever possible you should strongly prefer the native expression API
+    # to achieve the best performance.
+    #
+    # @return [DataFrame]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "id" => [0, 1, 2, 3, 4],
+    #       "color" => ["red", "green", "green", "red", "red"],
+    #       "shape" => ["square", "triangle", "square", "triangle", "square"]
+    #     }
+    #   )
+    #   df.group_by("color").map_groups { |group_df| group_df.sample(n: 2) }
+    #   # =>
+    #   # shape: (4, 3)
+    #   # ┌─────┬───────┬──────────┐
+    #   # │ id  ┆ color ┆ shape    │
+    #   # │ --- ┆ ---   ┆ ---      │
+    #   # │ i64 ┆ str   ┆ str      │
+    #   # ╞═════╪═══════╪══════════╡
+    #   # │ 1   ┆ green ┆ triangle │
+    #   # │ 2   ┆ green ┆ square   │
+    #   # │ 4   ┆ red   ┆ square   │
+    #   # │ 3   ┆ red   ┆ triangle │
+    #   # └─────┴───────┴──────────┘
+    def map_groups(&function)
+      if @predicates&.any?
+        msg = "cannot call `map_groups` when filtering groups with `having`"
+        raise TypeError, msg
+      end
+      if @named_by&.any?
+        msg = "cannot call `map_groups` when grouping by named expressions"
+        raise TypeError, msg
+      end
+      if !@by.all? { |c| Utils.strlike?(c) }
+        msg = "cannot call `map_groups` when grouping by an expression"
+        raise TypeError, msg
+      end
+
+      by_strs = @by.map(&:to_s)
+
+      @df.class._from_rbdf(
+        @df._df.group_by_map_groups(by_strs, function, @maintain_order)
+      )
     end
 
     # Get the first `n` rows of each group.
@@ -252,10 +315,7 @@ module Polars
     #   # │ c       ┆ 2   │
     #   # └─────────┴─────┘
     def head(n = 5)
-      @df.lazy
-        .group_by(@by, maintain_order: @maintain_order)
-        .head(n)
-        .collect(no_optimization: true)
+      _lgb.head(n).collect(optimizations: QueryOptFlags._eager)
     end
 
     # Get the last `n` rows of each group.
@@ -303,14 +363,78 @@ module Polars
     #   # │ c       ┆ 4   │
     #   # └─────────┴─────┘
     def tail(n = 5)
-      @df.lazy
-        .group_by(@by, maintain_order: @maintain_order)
-        .tail(n)
-        .collect(no_optimization: true)
+      _lgb.tail(n).collect(optimizations: QueryOptFlags._eager)
+    end
+
+    # Aggregate the groups into Series.
+    #
+    # @return [DataFrame]
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"a" => ["one", "two", "one", "two"], "b" => [1, 2, 3, 4]})
+    #   df.group_by("a", maintain_order: true).all
+    #   # =>
+    #   # shape: (2, 2)
+    #   # ┌─────┬───────────┐
+    #   # │ a   ┆ b         │
+    #   # │ --- ┆ ---       │
+    #   # │ str ┆ list[i64] │
+    #   # ╞═════╪═══════════╡
+    #   # │ one ┆ [1, 3]    │
+    #   # │ two ┆ [2, 4]    │
+    #   # └─────┴───────────┘
+    def all
+      agg(F.all)
+    end
+
+    # Return the number of rows in each group.
+    #
+    # @param name [String]
+    #   Assign a name to the resulting column; if unset, defaults to "len".
+    #
+    # @return [DataFrame]
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"a" => ["Apple", "Apple", "Orange"], "b" => [1, nil, 2]})
+    #   df.group_by("a").len
+    #   # =>
+    #   # shape: (2, 2)
+    #   # ┌────────┬─────┐
+    #   # │ a      ┆ len │
+    #   # │ ---    ┆ --- │
+    #   # │ str    ┆ u32 │
+    #   # ╞════════╪═════╡
+    #   # │ Apple  ┆ 2   │
+    #   # │ Orange ┆ 1   │
+    #   # └────────┴─────┘
+    #
+    # @example
+    #   df.group_by("a").len(name: "n")
+    #   # =>
+    #   # shape: (2, 2)
+    #   # ┌────────┬─────┐
+    #   # │ a      ┆ n   │
+    #   # │ ---    ┆ --- │
+    #   # │ str    ┆ u32 │
+    #   # ╞════════╪═════╡
+    #   # │ Apple  ┆ 2   │
+    #   # │ Orange ┆ 1   │
+    #   # └────────┴─────┘
+    def len(name: nil)
+      len_expr = F.len
+      if !name.nil?
+        len_expr = len_expr.alias(name)
+      end
+      agg(len_expr)
     end
 
     # Aggregate the first values in the group.
     #
+    # @param ignore_nulls [Boolean]
+    #   Ignore null values (default `false`).
+    #   If set to `true`, the first non-null value for each aggregation is returned,
+    #   otherwise `nil` is returned if no non-null value exists.
+    #
     # @return [DataFrame]
     #
     # @example
@@ -334,12 +458,17 @@ module Polars
     #   # │ Orange ┆ 2   ┆ 0.5  ┆ true  │
     #   # │ Banana ┆ 4   ┆ 13.0 ┆ false │
     #   # └────────┴─────┴──────┴───────┘
-    def first
-      agg(Polars.all.first)
+    def first(ignore_nulls: false)
+      agg(F.all.first(ignore_nulls: ignore_nulls))
     end
 
     # Aggregate the last values in the group.
     #
+    # @param ignore_nulls [Boolean]
+    #   Ignore null values (default `false`).
+    #   If set to `true`, the last non-null value for each aggregation is returned,
+    #   otherwise `nil` is returned if no non-null value exists.
+    #
     # @return [DataFrame]
     #
     # @example
@@ -363,8 +492,8 @@ module Polars
     #   # │ Orange ┆ 2   ┆ 0.5  ┆ true  │
     #   # │ Banana ┆ 5   ┆ 14.0 ┆ true  │
     #   # └────────┴─────┴──────┴───────┘
-    def last
-      agg(Polars.all.last)
+    def last(ignore_nulls: false)
+      agg(F.all.last(ignore_nulls: ignore_nulls))
     end
 
     # Reduce the groups to the sum.
@@ -599,15 +728,16 @@ module Polars
       agg(Polars.all.median)
     end
 
-    # Plot data.
-    #
-    # @return [Vega::LiteChart]
-    def plot(*args, **options)
-      raise ArgumentError, "Multiple groups not supported" if @by.is_a?(::Array) && @by.size > 1
-      # same message as Ruby
-      raise ArgumentError, "unknown keyword: :group" if options.key?(:group)
+    private
 
-      @df.plot(*args, **options, group: @by)
+    def _lgb
+      group_by = @df.lazy.group_by(
+        *@by, **@named_by, maintain_order: @maintain_order
+      )
+      if @predicates&.any?
+        return group_by.having(@predicates)
+      end
+      group_by
     end
   end
 end

data/lib/polars/iceberg_dataset.rb

@@ -0,0 +1,108 @@
+module Polars
+  # @private
+  class IcebergDataset
+    def initialize(
+      source,
+      snapshot_id:,
+      storage_options:
+    )
+      @source = source
+      @snapshot_id = snapshot_id
+      @storage_options = storage_options
+    end
+
+    def to_lazyframe
+      scan = @source.scan(snapshot_id: @snapshot_id)
+      files = scan.plan_files
+
+      table = scan.table
+      snapshot = scan.snapshot
+      schema = snapshot ? table.schema_by_id(snapshot[:schema_id]) : table.current_schema
+
+      if files.empty?
+        # TODO improve
+        schema =
+          schema.fields.to_h do |field|
+            dtype =
+              case field[:type]
+              when "int"
+                Polars::Int32
+              when "long"
+                Polars::Int64
+              when "double"
+                Polars::Float64
+              when "string"
+                Polars::String
+              when "timestamp"
+                Polars::Datetime
+              else
+                raise Todo
+              end
+
+            [field[:name], dtype]
+          end
+
+        LazyFrame.new(schema: schema)
+      else
+        sources = files.map { |v| v[:data_file_path] }
+
+        column_mapping = [
+          "iceberg-column-mapping",
+          arrow_schema(schema)
+        ]
+
+        deletion_files = [
+          "iceberg-position-delete",
+          files.map.with_index
+            .select { |v, i| v[:deletes].any? }
+            .to_h { |v, i| [i, v[:deletes].map { |d| d[:file_path] }] }
+        ]
+
+        scan_options = {
+          storage_options: @storage_options,
+          cast_options: Polars::ScanCastOptions._default_iceberg,
+          missing_columns: "insert",
+          extra_columns: "ignore",
+          _column_mapping: column_mapping,
+          _deletion_files: deletion_files
+        }
+
+        Polars.scan_parquet(sources, **scan_options)
+      end
+    end
+
+    private
+
+    def arrow_schema(schema)
+      fields =
+        schema.fields.map do |field|
+          type =
+            case field[:type]
+            when "boolean"
+              "boolean"
+            when "int"
+              "int32"
+            when "long"
+              "int64"
+            when "float"
+              "float32"
+            when "double"
+              "float64"
+            else
+              raise Todo
+            end
+
+          {
+            name: field[:name],
+            type: type,
+            nullable: !field[:required],
+            metadata: {
+              "PARQUET:field_id" => field[:id].to_s
+            }
+          }
+        end
+
+      {fields: fields}
+    end
+  end
+end

data/lib/polars/in_process_query.rb

@@ -0,0 +1,37 @@
+module Polars
+  # A placeholder for an in process query.
+  #
+  # This can be used to do something else while a query is running.
+  # The queries can be cancelled. You can peek if the query is finished,
+  # or you can await the result.
+  class InProcessQuery
+    # @private
+    attr_accessor :_inner
+
+    def initialize(ipq)
+      self._inner = ipq
+    end
+
+    # Cancel the query at earliest convenience.
+    def cancel
+      _inner.cancel
+    end
+
+    # Fetch the result.
+    #
+    # If it is ready, a materialized DataFrame is returned.
+    # If it is not ready it will return `nil`.
+    def fetch
+      if !(out = _inner.fetch).nil?
+        Utils.wrap_df(out)
+      else
+        nil
+      end
+    end
+
+    # Await the result synchronously.
+    def fetch_blocking
+      Utils.wrap_df(_inner.fetch_blocking)
+    end
+  end
+end

data/lib/polars/io/cloud.rb

@@ -0,0 +1,18 @@
+module Polars
+  module IO
+    private
+
+    def _init_credential_provider_builder(
+      credential_provider,
+      source,
+      storage_options,
+      caller_name
+    )
+      if credential_provider && credential_provider != "auto"
+        raise Todo
+      end
+
+      nil
+    end
+  end
+end