polars-df 0.9.0 → 0.10.0

data/ext/polars/src/map/dataframe.rs

@@ -265,9 +265,9 @@ pub fn apply_lambda_with_rows_output<'a>(
                         // to the row. Before we mutate the row buf again, the reference is dropped.
                         // we only cannot prove it to the compiler.
                         // we still do this because it saves a Vec allocation in a hot loop.
-                        unsafe { &*ptr }
+                        Ok(unsafe { &*ptr })
                     }
-                    None => &null_row,
+                    None => Ok(&null_row),
                 }
             }
             Err(e) => panic!("ruby function failed {}", e),
@@ -277,22 +277,30 @@ pub fn apply_lambda_with_rows_output<'a>(
     // first rows for schema inference
     let mut buf = Vec::with_capacity(inference_size);
     buf.push(first_value);
-    buf.extend((&mut row_iter).take(inference_size).cloned());
-    let schema = rows_to_schema_first_non_null(&buf, Some(50));
+    for v in (&mut row_iter).take(inference_size) {
+        buf.push(v?.clone());
+    }
+
+    let schema = rows_to_schema_first_non_null(&buf, Some(50))?;
 
     if init_null_count > 0 {
         // Safety: we know the iterators size
         let iter = unsafe {
             (0..init_null_count)
-                .map(|_| &null_row)
-                .chain(buf.iter())
+                .map(|_| Ok(&null_row))
+                .chain(buf.iter().map(Ok))
                 .chain(row_iter)
                 .trust_my_length(df.height())
         };
-        DataFrame::from_rows_iter_and_schema(iter, &schema)
+        DataFrame::try_from_rows_iter_and_schema(iter, &schema)
     } else {
         // Safety: we know the iterators size
-        let iter = unsafe { buf.iter().chain(row_iter).trust_my_length(df.height()) };
-        DataFrame::from_rows_iter_and_schema(iter, &schema)
+        let iter = unsafe {
+            buf.iter()
+                .map(Ok)
+                .chain(row_iter)
+                .trust_my_length(df.height())
+        };
+        DataFrame::try_from_rows_iter_and_schema(iter, &schema)
     }
 }

data/ext/polars/src/series/mod.rs

@@ -233,8 +233,18 @@ impl RbSeries {
         }
     }
 
-    pub fn sort(&self, descending: bool, nulls_last: bool) -> Self {
-        (self.series.borrow_mut().sort(descending, nulls_last)).into()
+    pub fn sort(&self, descending: bool, nulls_last: bool, multithreaded: bool) -> RbResult<Self> {
+        Ok(self
+            .series
+            .borrow_mut()
+            .sort(
+                SortOptions::default()
+                    .with_order_descending(descending)
+                    .with_nulls_last(nulls_last)
+                    .with_multithreaded(multithreaded),
+            )
+            .map_err(RbPolarsErr::from)?
+            .into())
     }
 
     pub fn value_counts(&self, sorted: bool) -> RbResult<RbDataFrame> {

data/lib/polars/array_expr.rb

@@ -333,6 +333,10 @@ module Polars
     #
     # @param index [Integer]
     #   Index to return per sub-array
+    # @param null_on_oob [Boolean]
+    #   Behavior if an index is out of bounds:
+    #   true -> set as null
+    #   false -> raise an error
     #
     # @return [Expr]
     #
@@ -353,9 +357,9 @@ module Polars
     #   # │ [4, 5, 6]     ┆ -2  ┆ 5    │
     #   # │ [7, 8, 9]     ┆ 4   ┆ null │
     #   # └───────────────┴─────┴──────┘
-    def get(index)
+    def get(index, null_on_oob: true)
       index = Utils.parse_as_expression(index)
-      Utils.wrap_expr(_rbexpr.arr_get(index))
+      Utils.wrap_expr(_rbexpr.arr_get(index, null_on_oob))
     end
 
     # Get the first value of the sub-arrays.

data/lib/polars/batched_csv_reader.rb

@@ -27,7 +27,8 @@ module Polars
       row_count_offset: 0,
       sample_size: 1024,
       eol_char: "\n",
-      new_columns: nil
+      new_columns: nil,
+      truncate_ragged_lines: false
     )
       if Utils.pathlike?(file)
         path = Utils.normalise_filepath(file)
@@ -75,7 +76,8 @@ module Polars
         skip_rows_after_header,
         Utils._prepare_row_count_args(row_count_name, row_count_offset),
         sample_size,
-        eol_char
+        eol_char,
+        truncate_ragged_lines
       )
       self.new_columns = new_columns
     end

data/lib/polars/data_frame.rb

@@ -91,7 +91,8 @@ module Polars
       row_count_name: nil,
       row_count_offset: 0,
       sample_size: 1024,
-      eol_char: "\n"
+      eol_char: "\n",
+      truncate_ragged_lines: false
     )
       if Utils.pathlike?(file)
         path = Utils.normalise_filepath(file)
@@ -147,7 +148,8 @@ module Polars
           skip_rows_after_header: skip_rows_after_header,
           row_count_name: row_count_name,
           row_count_offset: row_count_offset,
-          eol_char: eol_char
+          eol_char: eol_char,
+          truncate_ragged_lines: truncate_ragged_lines
         )
         if columns.nil?
           return _from_rbdf(scan.collect._df)
@@ -186,7 +188,8 @@ module Polars
           skip_rows_after_header,
           Utils._prepare_row_count_args(row_count_name, row_count_offset),
           sample_size,
-          eol_char
+          eol_char,
+          truncate_ragged_lines
         )
       )
     end
@@ -814,8 +817,6 @@ module Polars
 
     # Serialize to JSON representation.
     #
-    # @return [nil]
-    #
     # @param file [String]
     #   File path to which the result should be written.
     # @param pretty [Boolean]
@@ -823,17 +824,45 @@ module Polars
     # @param row_oriented [Boolean]
     #   Write to row oriented json. This is slower, but more common.
     #
-    # @see #write_ndjson
+    # @return [nil]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "foo" => [1, 2, 3],
+    #       "bar" => [6, 7, 8]
+    #     }
+    #   )
+    #   df.write_json
+    #   # => "{\"columns\":[{\"name\":\"foo\",\"datatype\":\"Int64\",\"bit_settings\":\"\",\"values\":[1,2,3]},{\"name\":\"bar\",\"datatype\":\"Int64\",\"bit_settings\":\"\",\"values\":[6,7,8]}]}"
+    #
+    # @example
+    #   df.write_json(row_oriented: true)
+    #   # => "[{\"foo\":1,\"bar\":6},{\"foo\":2,\"bar\":7},{\"foo\":3,\"bar\":8}]"
     def write_json(
-      file,
+      file = nil,
       pretty: false,
       row_oriented: false
     )
       if Utils.pathlike?(file)
         file = Utils.normalise_filepath(file)
       end
-
-      _df.write_json(file, pretty, row_oriented)
+      to_string_io = !file.nil? && file.is_a?(StringIO)
+      if file.nil? || to_string_io
+        buf = StringIO.new
+        buf.set_encoding(Encoding::BINARY)
+        _df.write_json(buf, pretty, row_oriented)
+        json_bytes = buf.string
+
+        json_str = json_bytes.force_encoding(Encoding::UTF_8)
+        if to_string_io
+          file.write(json_str)
+        else
+          return json_str
+        end
+      else
+        _df.write_json(file, pretty, row_oriented)
+      end
       nil
     end
 
@@ -843,12 +872,36 @@ module Polars
     #   File path to which the result should be written.
     #
     # @return [nil]
-    def write_ndjson(file)
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "foo" => [1, 2, 3],
+    #       "bar" => [6, 7, 8]
+    #     }
+    #   )
+    #   df.write_ndjson()
+    #   # => "{\"foo\":1,\"bar\":6}\n{\"foo\":2,\"bar\":7}\n{\"foo\":3,\"bar\":8}\n"
+    def write_ndjson(file = nil)
       if Utils.pathlike?(file)
         file = Utils.normalise_filepath(file)
       end
-
-      _df.write_ndjson(file)
+      to_string_io = !file.nil? && file.is_a?(StringIO)
+      if file.nil? || to_string_io
+        buf = StringIO.new
+        buf.set_encoding(Encoding::BINARY)
+        _df.write_ndjson(buf)
+        json_bytes = buf.string
+
+        json_str = json_bytes.force_encoding(Encoding::UTF_8)
+        if to_string_io
+          file.write(json_str)
+        else
+          return json_str
+        end
+      else
+        _df.write_ndjson(file)
+      end
       nil
     end
 
@@ -1010,7 +1063,7 @@ module Polars
 
     # Write to Apache Parquet file.
     #
-    # @param file [String]
+    # @param file [String, Pathname, StringIO]
     #   File path to which the file should be written.
     # @param compression ["lz4", "uncompressed", "snappy", "gzip", "lzo", "brotli", "zstd"]
     #   Choose "zstd" for good compression performance.
@@ -1027,10 +1080,9 @@ module Polars
     # @param statistics [Boolean]
     #   Write statistics to the parquet headers. This requires extra compute.
     # @param row_group_size [Integer, nil]
-    #   Size of the row groups in number of rows.
-    #   If `nil` (default), the chunks of the DataFrame are
-    #   used. Writing in smaller chunks may reduce memory pressure and improve
-    #   writing speeds.
+    #   Size of the row groups in number of rows. Defaults to 512^2 rows.
+    # @param data_page_size [Integer, nil]
+    #   Size of the data page in bytes. Defaults to 1024^2 bytes.
     #
     # @return [nil]
     def write_parquet(
@@ -1038,7 +1090,8 @@ module Polars
       compression: "zstd",
       compression_level: nil,
       statistics: false,
-      row_group_size: nil
+      row_group_size: nil,
+      data_page_size: nil
     )
       if compression.nil?
         compression = "uncompressed"
@@ -1048,7 +1101,7 @@ module Polars
       end
 
       _df.write_parquet(
-        file, compression, compression_level, statistics, row_group_size
+        file, compression, compression_level, statistics, row_group_size, data_page_size
       )
     end
 
@@ -1084,7 +1137,7 @@ module Polars
     #   df.estimated_size
     #   # => 25888898
     #   df.estimated_size("mb")
-    #   # => 26.702880859375
+    #   # => 17.0601749420166
     def estimated_size(unit = "b")
       sz = _df.estimated_size
       Utils.scale_bytes(sz, to: unit)
@@ -2161,12 +2214,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"]           │
@@ -2621,26 +2675,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 Replaced
     #   df.with_column(Polars.col("a") ** 2)
     #   # =>
     #   # 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)
       lazy
         .with_column(column)
@@ -2807,16 +2861,36 @@ module Polars
     #   # │ 2   ┆ 7.0 │
     #   # │ 3   ┆ 8.0 │
     #   # └─────┴─────┘
-    def drop(columns)
-      if columns.is_a?(::Array)
-        df = clone
-        columns.each do |n|
-          df._df.drop_in_place(n)
-        end
-        df
-      else
-        _from_rbdf(_df.drop(columns))
-      end
+    #
+    # @example Drop multiple columns by passing a list of column names.
+    #   df.drop(["bar", "ham"])
+    #   # =>
+    #   # shape: (3, 1)
+    #   # ┌─────┐
+    #   # │ foo │
+    #   # │ --- │
+    #   # │ i64 │
+    #   # ╞═════╡
+    #   # │ 1   │
+    #   # │ 2   │
+    #   # │ 3   │
+    #   # └─────┘
+    #
+    # @example Use positional arguments to drop multiple columns.
+    #   df.drop("foo", "ham")
+    #   # =>
+    #   # shape: (3, 1)
+    #   # ┌─────┐
+    #   # │ bar │
+    #   # │ --- │
+    #   # │ f64 │
+    #   # ╞═════╡
+    #   # │ 6.0 │
+    #   # │ 7.0 │
+    #   # │ 8.0 │
+    #   # └─────┘
+    def drop(*columns)
+      lazy.drop(*columns).collect(_eager: true)
     end
 
     # Drop in place.
@@ -3735,16 +3809,16 @@ module Polars
     #   df.with_columns((Polars.col("a") ** 2).alias("a^2"))
     #   # =>
     #   # shape: (4, 4)
-    #   # ┌─────┬──────┬───────┬──────┐
-    #   # │ a   ┆ b    ┆ c     ┆ a^2  │
-    #   # │ --- ┆ ---  ┆ ---   ┆ ---  │
-    #   # │ i64 ┆ f64  ┆ bool  ┆ f64  │
-    #   # ╞═════╪══════╪═══════╪══════╡
-    #   # │ 1   ┆ 0.5  ┆ true  ┆ 1.0  │
-    #   # │ 2   ┆ 4.0  ┆ true  ┆ 4.0  │
-    #   # │ 3   ┆ 10.0 ┆ false ┆ 9.0  │
-    #   # │ 4   ┆ 13.0 ┆ true  ┆ 16.0 │
-    #   # └─────┴──────┴───────┴──────┘
+    #   # ┌─────┬──────┬───────┬─────┐
+    #   # │ a   ┆ b    ┆ c     ┆ a^2 │
+    #   # │ --- ┆ ---  ┆ ---   ┆ --- │
+    #   # │ i64 ┆ f64  ┆ bool  ┆ i64 │
+    #   # ╞═════╪══════╪═══════╪═════╡
+    #   # │ 1   ┆ 0.5  ┆ true  ┆ 1   │
+    #   # │ 2   ┆ 4.0  ┆ true  ┆ 4   │
+    #   # │ 3   ┆ 10.0 ┆ false ┆ 9   │
+    #   # │ 4   ┆ 13.0 ┆ true  ┆ 16  │
+    #   # └─────┴──────┴───────┴─────┘
     #
     # @example Added columns will replace existing columns with the same name.
     #   df.with_columns(Polars.col("a").cast(Polars::Float64))
@@ -3771,16 +3845,16 @@ module Polars
     #   )
     #   # =>
     #   # 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 │
+    #   # └─────┴──────┴───────┴─────┴──────┴───────┘
     #
     # @example Multiple columns also can be added using positional arguments instead of a list.
     #   df.with_columns(
@@ -3790,16 +3864,16 @@ module Polars
     #   )
     #   # =>
     #   # 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 │
+    #   # └─────┴──────┴───────┴─────┴──────┴───────┘
     #
     # @example Use keyword arguments to easily name your expression inputs.
     #   df.with_columns(

data/lib/polars/date_time_expr.rb

@@ -1027,14 +1027,20 @@ module Polars
     # Different from `convert_time_zone`, this will also modify
     # the underlying timestamp,
     #
-    # @param tz [String]
-    #   Time zone for the `Datetime` Series.
+    # @param time_zone [String]
+    #     Time zone for the `Datetime` Series. Pass `nil` to unset time zone.
+    # @param use_earliest [Boolean]
+    #     Determine how to deal with ambiguous datetimes.
+    # @param ambiguous [String]
+    #     Determine how to deal with ambiguous datetimes.
+    # @param non_existent [String]
+    #     Determine how to deal with non-existent datetimes.
     #
     # @return [Expr]
-    def replace_time_zone(tz, use_earliest: nil, ambiguous: "raise")
+    def replace_time_zone(time_zone, use_earliest: nil, ambiguous: "raise", non_existent: "raise")
       ambiguous = Utils.rename_use_earliest_to_ambiguous(use_earliest, ambiguous)
       ambiguous = Polars.lit(ambiguous) unless ambiguous.is_a?(Expr)
-      Utils.wrap_expr(_rbexpr.dt_replace_time_zone(tz, ambiguous._rbexpr))
+      Utils.wrap_expr(_rbexpr.dt_replace_time_zone(time_zone, ambiguous._rbexpr, non_existent))
     end
 
     # Extract the days from a Duration type.

data/lib/polars/date_time_name_space.rb

@@ -910,8 +910,14 @@ module Polars
     # Different from `with_time_zone`, this will also modify
     # the underlying timestamp.
     #
-    # @param tz [String]
-    #   Time zone for the `Datetime` Series.
+    # @param time_zone [String]
+    #     Time zone for the `Datetime` Series. Pass `nil` to unset time zone.
+    # @param use_earliest [Boolean]
+    #     Determine how to deal with ambiguous datetimes.
+    # @param ambiguous [String]
+    #     Determine how to deal with ambiguous datetimes.
+    # @param non_existent [String]
+    #     Determine how to deal with non-existent datetimes.
     #
     # @return [Series]
     #
@@ -982,7 +988,7 @@ module Polars
     #   #         1585717200
     #   #         1588309200
     #   # ]
-    def replace_time_zone(tz)
+    def replace_time_zone(time_zone, use_earliest: nil, ambiguous: "raise", non_existent: "raise")
       super
     end