polars-df 0.9.0-x86_64-linux → 0.10.0-x86_64-linux

data/lib/polars/functions/whenthen.rb

@@ -4,7 +4,7 @@ module Polars
     #
     # @return [When]
     #
-    # @example
+    # @example Below we add a column with the value 1, where column "foo" > 2 and the value -1 where it isn't.
     #   df = Polars::DataFrame.new({"foo" => [1, 3, 4], "bar" => [3, 4, 0]})
     #   df.with_column(Polars.when(Polars.col("foo") > 2).then(Polars.lit(1)).otherwise(Polars.lit(-1)))
     #   # =>
@@ -18,10 +18,79 @@ module Polars
     #   # │ 3   ┆ 4   ┆ 1       │
     #   # │ 4   ┆ 0   ┆ 1       │
     #   # └─────┴─────┴─────────┘
-    def when(expr)
-      expr = Utils.expr_to_lit_or_expr(expr)
-      pw = Plr.when(expr._rbexpr)
-      When.new(pw)
+    #
+    # @example Or with multiple when-then operations chained:
+    #   df.with_columns(
+    #     Polars.when(Polars.col("foo") > 2)
+    #     .then(1)
+    #     .when(Polars.col("bar") > 2)
+    #     .then(4)
+    #     .otherwise(-1)
+    #     .alias("val")
+    #   )
+    #   # =>
+    #   # shape: (3, 3)
+    #   # ┌─────┬─────┬─────┐
+    #   # │ foo ┆ bar ┆ val │
+    #   # │ --- ┆ --- ┆ --- │
+    #   # │ i64 ┆ i64 ┆ i32 │
+    #   # ╞═════╪═════╪═════╡
+    #   # │ 1   ┆ 3   ┆ 4   │
+    #   # │ 3   ┆ 4   ┆ 1   │
+    #   # │ 4   ┆ 0   ┆ 1   │
+    #   # └─────┴─────┴─────┘
+    #
+    # @example The `otherwise` at the end is optional. If left out, any rows where none of the `when` expressions evaluate to True, are set to `null`:
+    #   df.with_columns(Polars.when(Polars.col("foo") > 2).then(1).alias("val"))
+    #   # =>
+    #   # shape: (3, 3)
+    #   # ┌─────┬─────┬──────┐
+    #   # │ foo ┆ bar ┆ val  │
+    #   # │ --- ┆ --- ┆ ---  │
+    #   # │ i64 ┆ i64 ┆ i32  │
+    #   # ╞═════╪═════╪══════╡
+    #   # │ 1   ┆ 3   ┆ null │
+    #   # │ 3   ┆ 4   ┆ 1    │
+    #   # │ 4   ┆ 0   ┆ 1    │
+    #   # └─────┴─────┴──────┘
+    #
+    # @example Pass multiple predicates, each of which must be met:
+    #   df.with_columns(
+    #     val: Polars.when(
+    #       Polars.col("bar") > 0,
+    #       Polars.col("foo") % 2 != 0
+    #     )
+    #     .then(99)
+    #     .otherwise(-1)
+    #   )
+    #   # =>
+    #   # shape: (3, 3)
+    #   # ┌─────┬─────┬─────┐
+    #   # │ foo ┆ bar ┆ val │
+    #   # │ --- ┆ --- ┆ --- │
+    #   # │ i64 ┆ i64 ┆ i32 │
+    #   # ╞═════╪═════╪═════╡
+    #   # │ 1   ┆ 3   ┆ 99  │
+    #   # │ 3   ┆ 4   ┆ 99  │
+    #   # │ 4   ┆ 0   ┆ -1  │
+    #   # └─────┴─────┴─────┘
+    #
+    # @example Pass conditions as keyword arguments:
+    #   df.with_columns(val: Polars.when(foo: 4, bar: 0).then(99).otherwise(-1))
+    #   # =>
+    #   # shape: (3, 3)
+    #   # ┌─────┬─────┬─────┐
+    #   # │ foo ┆ bar ┆ val │
+    #   # │ --- ┆ --- ┆ --- │
+    #   # │ i64 ┆ i64 ┆ i32 │
+    #   # ╞═════╪═════╪═════╡
+    #   # │ 1   ┆ 3   ┆ -1  │
+    #   # │ 3   ┆ 4   ┆ -1  │
+    #   # │ 4   ┆ 0   ┆ 99  │
+    #   # └─────┴─────┴─────┘
+    def when(*predicates, **constraints)
+      condition = Utils.parse_when_inputs(*predicates, **constraints)
+      When.new(Plr.when(condition))
     end
   end
 end

data/lib/polars/io.rb

@@ -80,6 +80,8 @@ module Polars
     #   allocation needed.
     # @param eol_char [String]
     #   Single byte end of line character.
+    # @param truncate_ragged_lines [Boolean]
+    #   Truncate lines that are longer than the schema.
     #
     # @return [DataFrame]
     #
@@ -113,7 +115,8 @@ module Polars
       row_count_name: nil,
       row_count_offset: 0,
       sample_size: 1024,
-      eol_char: "\n"
+      eol_char: "\n",
+      truncate_ragged_lines: false
     )
       Utils._check_arg_is_1byte("sep", sep, false)
       Utils._check_arg_is_1byte("comment_char", comment_char, false)
@@ -161,7 +164,8 @@ module Polars
           row_count_name: row_count_name,
           row_count_offset: row_count_offset,
           sample_size: sample_size,
-          eol_char: eol_char
+          eol_char: eol_char,
+          truncate_ragged_lines: truncate_ragged_lines
         )
       end
 
@@ -239,6 +243,8 @@ module Polars
     #   the column remains of data type `:str`.
     # @param eol_char [String]
     #   Single byte end of line character.
+    # @param truncate_ragged_lines [Boolean]
+    #   Truncate lines that are longer than the schema.
     #
     # @return [LazyFrame]
     def scan_csv(
@@ -262,7 +268,8 @@ module Polars
       row_count_name: nil,
       row_count_offset: 0,
       parse_dates: false,
-      eol_char: "\n"
+      eol_char: "\n",
+      truncate_ragged_lines: false
     )
       Utils._check_arg_is_1byte("sep", sep, false)
       Utils._check_arg_is_1byte("comment_char", comment_char, false)
@@ -294,6 +301,7 @@ module Polars
         row_count_offset: row_count_offset,
         parse_dates: parse_dates,
         eol_char: eol_char,
+        truncate_ragged_lines: truncate_ragged_lines
       )
     end
 
@@ -520,7 +528,7 @@ module Polars
 
     # Read into a DataFrame from a parquet file.
     #
-    # @param source [Object]
+    # @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
@@ -755,6 +763,8 @@ module Polars
     #   allocation needed.
     # @param eol_char [String]
     #   Single byte end of line character.
+    # @param truncate_ragged_lines [Boolean]
+    #   Truncate lines that are longer than the schema.
     #
     # @return [BatchedCsvReader]
     #
@@ -787,7 +797,8 @@ module Polars
       row_count_name: nil,
       row_count_offset: 0,
       sample_size: 1024,
-      eol_char: "\n"
+      eol_char: "\n",
+      truncate_ragged_lines: false
     )
       projection, columns = Utils.handle_projection_columns(columns)
 
@@ -827,7 +838,8 @@ module Polars
         row_count_offset: row_count_offset,
         sample_size: sample_size,
         eol_char: eol_char,
-        new_columns: new_columns
+        new_columns: new_columns,
+        truncate_ragged_lines: truncate_ragged_lines
       )
     end
 

data/lib/polars/lazy_frame.rb

@@ -49,7 +49,8 @@ module Polars
       row_count_name: nil,
       row_count_offset: 0,
       parse_dates: false,
-      eol_char: "\n"
+      eol_char: "\n",
+      truncate_ragged_lines: true
     )
       dtype_list = nil
       if !dtypes.nil?
@@ -81,7 +82,8 @@ module Polars
           encoding,
           Utils._prepare_row_count_args(row_count_name, row_count_offset),
           parse_dates,
-          eol_char
+          eol_char,
+          truncate_ragged_lines
         )
       )
     end
@@ -103,6 +105,7 @@ module Polars
       _from_rbldf(
         RbLazyFrame.new_from_parquet(
           file,
+          [],
           n_rows,
           cache,
           parallel,
@@ -110,7 +113,8 @@ module Polars
           Utils._prepare_row_count_args(row_count_name, row_count_offset),
           low_memory,
           use_statistics,
-          hive_partitioning
+          hive_partitioning,
+          nil
         )
       )
     end
@@ -400,16 +404,16 @@ module Polars
     #   # │ 2   ┆ 7.0 ┆ b   │
     #   # │ 1   ┆ 6.0 ┆ a   │
     #   # └─────┴─────┴─────┘
-    def sort(by, reverse: false, nulls_last: false, maintain_order: false)
+    def sort(by, reverse: false, nulls_last: false, maintain_order: false, multithreaded: true)
       if by.is_a?(::String)
-        return _from_rbldf(_ldf.sort(by, reverse, nulls_last, maintain_order))
+        return _from_rbldf(_ldf.sort(by, reverse, nulls_last, maintain_order, multithreaded))
       end
       if Utils.bool?(reverse)
         reverse = [reverse]
       end
 
       by = Utils.selection_to_rbexpr_list(by)
-      _from_rbldf(_ldf.sort_by_exprs(by, reverse, nulls_last, maintain_order))
+      _from_rbldf(_ldf.sort_by_exprs(by, reverse, nulls_last, maintain_order, multithreaded))
     end
 
     # def profile
@@ -1523,12 +1527,13 @@ module Polars
     #     closed: "right"
     #   ).agg(Polars.col("A").alias("A_agg_list"))
     #   # =>
-    #   # shape: (3, 4)
+    #   # shape: (4, 4)
     #   # ┌─────────────────┬─────────────────┬─────┬─────────────────┐
     #   # │ _lower_boundary ┆ _upper_boundary ┆ idx ┆ A_agg_list      │
     #   # │ ---             ┆ ---             ┆ --- ┆ ---             │
     #   # │ i64             ┆ i64             ┆ i64 ┆ list[str]       │
     #   # ╞═════════════════╪═════════════════╪═════╪═════════════════╡
+    #   # │ -2              ┆ 1               ┆ -2  ┆ ["A", "A"]      │
     #   # │ 0               ┆ 3               ┆ 0   ┆ ["A", "B", "B"] │
     #   # │ 2               ┆ 5               ┆ 2   ┆ ["B", "B", "C"] │
     #   # │ 4               ┆ 7               ┆ 4   ┆ ["C"]           │
@@ -1837,7 +1842,7 @@ module Polars
       if how == "cross"
         return _from_rbldf(
           _ldf.join(
-            other._ldf, [], [], allow_parallel, force_parallel, how, suffix
+            other._ldf, [], [], allow_parallel, join_nulls, force_parallel, how, suffix
           )
         )
       end
@@ -1891,16 +1896,16 @@ module Polars
     #   ).collect
     #   # =>
     #   # shape: (4, 6)
-    #   # ┌─────┬──────┬───────┬──────┬──────┬───────┐
-    #   # │ a   ┆ b    ┆ c     ┆ a^2  ┆ b/2  ┆ not c │
-    #   # │ --- ┆ ---  ┆ ---   ┆ ---  ┆ ---  ┆ ---   │
-    #   # │ i64 ┆ f64  ┆ bool  ┆ f64  ┆ f64  ┆ bool  │
-    #   # ╞═════╪══════╪═══════╪══════╪══════╪═══════╡
-    #   # │ 1   ┆ 0.5  ┆ true  ┆ 1.0  ┆ 0.25 ┆ false │
-    #   # │ 2   ┆ 4.0  ┆ true  ┆ 4.0  ┆ 2.0  ┆ false │
-    #   # │ 3   ┆ 10.0 ┆ false ┆ 9.0  ┆ 5.0  ┆ true  │
-    #   # │ 4   ┆ 13.0 ┆ true  ┆ 16.0 ┆ 6.5  ┆ false │
-    #   # └─────┴──────┴───────┴──────┴──────┴───────┘
+    #   # ┌─────┬──────┬───────┬─────┬──────┬───────┐
+    #   # │ a   ┆ b    ┆ c     ┆ a^2 ┆ b/2  ┆ not c │
+    #   # │ --- ┆ ---  ┆ ---   ┆ --- ┆ ---  ┆ ---   │
+    #   # │ i64 ┆ f64  ┆ bool  ┆ i64 ┆ f64  ┆ bool  │
+    #   # ╞═════╪══════╪═══════╪═════╪══════╪═══════╡
+    #   # │ 1   ┆ 0.5  ┆ true  ┆ 1   ┆ 0.25 ┆ false │
+    #   # │ 2   ┆ 4.0  ┆ true  ┆ 4   ┆ 2.0  ┆ false │
+    #   # │ 3   ┆ 10.0 ┆ false ┆ 9   ┆ 5.0  ┆ true  │
+    #   # │ 4   ┆ 13.0 ┆ true  ┆ 16  ┆ 6.5  ┆ false │
+    #   # └─────┴──────┴───────┴─────┴──────┴───────┘
     def with_columns(*exprs, **named_exprs)
       structify = ENV.fetch("POLARS_AUTO_STRUCTIFY", "0") != "0"
       rbexprs = Utils.parse_as_list_of_expressions(*exprs, **named_exprs, __structify: structify)
@@ -1965,26 +1970,26 @@ module Polars
     #   # ┌─────┬─────┬───────────┐
     #   # │ a   ┆ b   ┆ b_squared │
     #   # │ --- ┆ --- ┆ ---       │
-    #   # │ i64 ┆ i64 ┆ f64       │
+    #   # │ i64 ┆ i64 ┆ i64       │
     #   # ╞═════╪═════╪═══════════╡
-    #   # │ 1   ┆ 2   ┆ 4.0       │
-    #   # │ 3   ┆ 4   ┆ 16.0      │
-    #   # │ 5   ┆ 6   ┆ 36.0      │
+    #   # │ 1   ┆ 2   ┆ 4         │
+    #   # │ 3   ┆ 4   ┆ 16        │
+    #   # │ 5   ┆ 6   ┆ 36        │
     #   # └─────┴─────┴───────────┘
     #
     # @example
     #   df.with_column(Polars.col("a") ** 2).collect
     #   # =>
     #   # shape: (3, 2)
-    #   # ┌──────┬─────┐
-    #   # │ a    ┆ b   │
-    #   # │ ---  ┆ --- │
-    #   # │ f64  ┆ i64 │
-    #   # ╞══════╪═════╡
-    #   # │ 1.0  ┆ 2   │
-    #   # │ 9.0  ┆ 4   │
-    #   # │ 25.0 ┆ 6   │
-    #   # └──────┴─────┘
+    #   # ┌─────┬─────┐
+    #   # │ a   ┆ b   │
+    #   # │ --- ┆ --- │
+    #   # │ i64 ┆ i64 │
+    #   # ╞═════╪═════╡
+    #   # │ 1   ┆ 2   │
+    #   # │ 9   ┆ 4   │
+    #   # │ 25  ┆ 6   │
+    #   # └─────┴─────┘
     def with_column(column)
       with_columns([column])
     end
@@ -1996,11 +2001,9 @@ module Polars
     #   - List of column names.
     #
     # @return [LazyFrame]
-    def drop(columns)
-      if columns.is_a?(::String)
-        columns = [columns]
-      end
-      _from_rbldf(_ldf.drop(columns))
+    def drop(*columns)
+      drop_cols = Utils._expand_selectors(self, *columns)
+      _from_rbldf(_ldf.drop(drop_cols))
     end
 
     # Rename column names.

data/lib/polars/list_expr.rb

@@ -365,6 +365,10 @@ module Polars
     #
     # @param index [Integer]
     #   Index to return per sublist
+    # @param null_on_oob [Boolean]
+    #   Behavior if an index is out of bounds:
+    #   true -> set as null
+    #   false -> raise an error
     #
     # @return [Expr]
     #
@@ -382,9 +386,9 @@ module Polars
     #   # │ null │
     #   # │ 1    │
     #   # └──────┘
-    def get(index)
+    def get(index, null_on_oob: true)
       index = Utils.parse_as_expression(index)
-      Utils.wrap_expr(_rbexpr.list_get(index))
+      Utils.wrap_expr(_rbexpr.list_get(index, null_on_oob))
     end
 
     # Get the value by index in the sublists.

data/lib/polars/series.rb

@@ -1155,13 +1155,13 @@ module Polars
     #   s.cumulative_eval(Polars.element.first - Polars.element.last ** 2)
     #   # =>
     #   # shape: (5,)
-    #   # Series: 'values' [f64]
+    #   # Series: 'values' [i64]
     #   # [
-    #   #         0.0
-    #   #         -3.0
-    #   #         -8.0
-    #   #         -15.0
-    #   #         -24.0
+    #   #         0
+    #   #         -3
+    #   #         -8
+    #   #         -15
+    #   #         -24
     #   # ]
     def cumulative_eval(expr, min_periods: 1, parallel: false)
       super
@@ -1567,12 +1567,12 @@ module Polars
     #   #         2
     #   #         1
     #   # ]
-    def sort(reverse: false, nulls_last: false, in_place: false)
+    def sort(reverse: false, nulls_last: false, multithreaded: true, in_place: false)
       if in_place
-        self._s = _s.sort(reverse, nulls_last)
+        self._s = _s.sort(reverse, nulls_last, multithreaded)
         self
       else
-        Utils.wrap_s(_s.sort(reverse, nulls_last))
+        Utils.wrap_s(_s.sort(reverse, nulls_last, multithreaded))
       end
     end
 
@@ -3646,6 +3646,8 @@ module Polars
     #     on the order that the values occur in the Series.
     # @param reverse [Boolean]
     #   Reverse the operation.
+    # @param seed [Integer]
+    #   If `method: "random"`, use this as seed.
     #
     # @return [Series]
     #
@@ -3676,7 +3678,7 @@ module Polars
     #   #         2
     #   #         5
     #   # ]
-    def rank(method: "average", reverse: false)
+    def rank(method: "average", reverse: false, seed: nil)
       super
     end
 

data/lib/polars/string_expr.rb

@@ -1354,6 +1354,7 @@ module Polars
     #   # │ null ┆ null   │
     #   # └──────┴────────┘
     def to_integer(base: 10, strict: true)
+      base = Utils.parse_as_expression(base, str_as_lit: false)
       Utils.wrap_expr(_rbexpr.str_to_integer(base, strict))
     end
 

data/lib/polars/utils.rb

@@ -364,5 +364,59 @@ module Polars
         end
       end
     end
+
+    def self._expand_selectors(frame, *items)
+      items_iter = _parse_inputs_as_iterable(items)
+
+      expanded = []
+      items_iter.each do |item|
+        if is_selector(item)
+          selector_cols = expand_selector(frame, item)
+          expanded.concat(selector_cols)
+        else
+          expanded << item
+        end
+      end
+      expanded
+    end
+
+    # TODO
+    def self.is_selector(obj)
+      false
+    end
+
+    def self.parse_predicates_constraints_as_expression(*predicates, **constraints)
+      all_predicates = _parse_positional_inputs(predicates)
+
+      if constraints.any?
+        constraint_predicates = _parse_constraints(constraints)
+        all_predicates.concat(constraint_predicates)
+      end
+
+      _combine_predicates(all_predicates)
+    end
+
+    def self._parse_constraints(constraints)
+      constraints.map do |name, value|
+        Polars.col(name).eq(value)._rbexpr
+      end
+    end
+
+    def self._combine_predicates(predicates)
+      if !predicates.any?
+        msg = "at least one predicate or constraint must be provided"
+        raise TypeError, msg
+      end
+
+      if predicates.length == 1
+        return predicates[0]
+      end
+
+      Plr.all_horizontal(predicates)
+    end
+
+    def self.parse_when_inputs(*predicates, **constraints)
+      parse_predicates_constraints_as_expression(*predicates, **constraints)
+    end
   end
 end

data/lib/polars/version.rb

@@ -1,4 +1,4 @@
 module Polars
   # @private
-  VERSION = "0.9.0"
+  VERSION = "0.10.0"
 end

data/lib/polars/whenthen.rb

@@ -0,0 +1,83 @@
+module Polars
+  # @private
+  class When
+    attr_accessor :_when
+
+    def initialize(rbwhen)
+      self._when = rbwhen
+    end
+
+    def then(statement)
+      statement_rbexpr = Utils.parse_as_expression(statement)
+      Then.new(_when.then(statement_rbexpr))
+    end
+  end
+
+  # @private
+  class Then < Expr
+    attr_accessor :_then
+
+    def initialize(rbthen)
+      self._then = rbthen
+    end
+
+    def self._from_rbexpr(rbexpr)
+      Utils.wrap_expr(rbexpr)
+    end
+
+    def _rbexpr
+      _then.otherwise(Polars.lit(nil)._rbexpr)
+    end
+
+    def when(*predicates, **constraints)
+      condition_rbexpr = Utils.parse_when_inputs(*predicates, **constraints)
+      ChainedWhen.new(_then.when(condition_rbexpr))
+    end
+
+    def otherwise(statement)
+      statement_rbexpr = Utils.parse_as_expression(statement)
+      Utils.wrap_expr(_then.otherwise(statement_rbexpr))
+    end
+  end
+
+  # @private
+  class ChainedWhen
+    attr_accessor :_chained_when
+
+    def initialize(chained_when)
+      self._chained_when = chained_when
+    end
+
+    def then(statement)
+      statement_rbexpr = Utils.parse_as_expression(statement)
+      ChainedThen.new(_chained_when.then(statement_rbexpr))
+    end
+  end
+
+  # @private
+  class ChainedThen < Expr
+    attr_accessor :_chained_then
+
+    def initialize(chained_then)
+      self._chained_then = chained_then
+    end
+
+    def self._from_rbexpr(rbexpr)
+      Utils.wrap_expr(rbexpr)
+    end
+
+    def _rbexpr
+      _chained_then.otherwise(Polars.lit(nil)._rbexpr)
+    end
+
+    def when(*predicates, **constraints)
+      condition_rbexpr = Utils.parse_when_inputs(*predicates, **constraints)
+      ChainedWhen.new(_chained_then.when(condition_rbexpr))
+    end
+
+    def otherwise(statement)
+      statement_rbexpr = Utils.parse_as_expression(statement)
+      Utils.wrap_expr(_chained_then.otherwise(statement_rbexpr))
+    end
+  end
+end

data/lib/polars.rb

@@ -63,8 +63,7 @@ require_relative "polars/struct_name_space"
 require_relative "polars/testing"
 require_relative "polars/utils"
 require_relative "polars/version"
-require_relative "polars/when"
-require_relative "polars/when_then"
+require_relative "polars/whenthen"
 
 module Polars
   extend Convert

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: polars-df
 version: !ruby/object:Gem::Version
-  version: 0.9.0
+  version: 0.10.0
 platform: x86_64-linux
 authors:
 - Andrew Kane
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2024-03-03 00:00:00.000000000 Z
+date: 2024-05-03 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: bigdecimal
@@ -95,8 +95,7 @@ files:
 - lib/polars/testing.rb
 - lib/polars/utils.rb
 - lib/polars/version.rb
-- lib/polars/when.rb
-- lib/polars/when_then.rb
+- lib/polars/whenthen.rb
 homepage: https://github.com/ankane/polars-ruby
 licenses:
 - MIT

data/lib/polars/when.rb

@@ -1,16 +0,0 @@
-module Polars
-  # @private
-  class When
-    attr_accessor :_rbwhen
-
-    def initialize(rbwhen)
-      self._rbwhen = rbwhen
-    end
-
-    def then(expr)
-      expr = Utils.expr_to_lit_or_expr(expr)
-      rbwhenthen = _rbwhen._then(expr._rbexpr)
-      WhenThen.new(rbwhenthen)
-    end
-  end
-end

data/lib/polars/when_then.rb

@@ -1,19 +0,0 @@
-module Polars
-  # @private
-  class WhenThen
-    attr_accessor :_rbwhenthen
-
-    def initialize(rbwhenthen)
-      self._rbwhenthen = rbwhenthen
-    end
-
-    def when(predicate)
-      WhenThenThen.new(_rbwhenthen.when(predicate._rbexpr))
-    end
-
-    def otherwise(expr)
-      expr = Utils.expr_to_lit_or_expr(expr)
-      Utils.wrap_expr(_rbwhenthen.otherwise(expr._rbexpr))
-    end
-  end
-end