polars-df 0.15.0 → 0.26.0

data/lib/polars/array_name_space.rb

@@ -16,7 +16,7 @@ module Polars
     #
     # @example
     #   s = Polars::Series.new(
-    #     "a", [[1, 2], [4, 3]], dtype: Polars::Array.new(2, Polars::Int64)
+    #     "a", [[1, 2], [4, 3]], dtype: Polars::Array.new(Polars::Int64, 2)
     #   )
     #   s.arr.min
     #   # =>
@@ -36,7 +36,7 @@ module Polars
     #
     # @example
     #   s = Polars::Series.new(
-    #     "a", [[1, 2], [4, 3]], dtype: Polars::Array.new(2, Polars::Int64)
+    #     "a", [[1, 2], [4, 3]], dtype: Polars::Array.new(Polars::Int64, 2)
     #   )
     #   s.arr.max
     #   # =>
@@ -57,7 +57,7 @@ module Polars
     # @example
     #   df = Polars::DataFrame.new(
     #     {"a" => [[1, 2], [4, 3]]},
-    #     schema: {"a" => Polars::Array.new(2, Polars::Int64)}
+    #     schema: {"a" => Polars::Array.new(Polars::Int64, 2)}
     #   )
     #   df.select(Polars.col("a").arr.sum)
     #   # =>
@@ -74,6 +74,78 @@ module Polars
       super
     end
 
+    # Compute the mean of the values of the sub-arrays.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("a", [[1, 2], [4, 3]], dtype: Polars::Array.new(Polars::Int64, 2))
+    #   s.arr.mean
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: 'a' [f64]
+    #   # [
+    #   #         1.5
+    #   #         3.5
+    #   # ]
+    def mean
+      super
+    end
+
+    # Compute the std of the values of the sub-arrays.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("a", [[1, 2], [4, 3]], dtype: Polars::Array.new(Polars::Int64, 2))
+    #   s.arr.std
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: 'a' [f64]
+    #   # [
+    #   #         0.707107
+    #   #         0.707107
+    #   # ]
+    def std(ddof: 1)
+      super
+    end
+
+    # Compute the var of the values of the sub-arrays.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("a", [[1, 2], [4, 3]], dtype: Polars::Array.new(Polars::Int64, 2))
+    #   s.arr.var
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: 'a' [f64]
+    #   # [
+    #   #         0.5
+    #   #         0.5
+    #   # ]
+    def var(ddof: 1)
+      super
+    end
+
+    # Compute the median of the values of the sub-arrays.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("a", [[1, 2], [4, 3]], dtype: Polars::Array.new(Polars::Int64, 2))
+    #   s.arr.median
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: 'a' [f64]
+    #   # [
+    #   #         1.5
+    #   #         3.5
+    #   # ]
+    def median
+      super
+    end
+
     # Get the unique/distinct values in the array.
     #
     # @param maintain_order [Boolean]
@@ -102,6 +174,24 @@ module Polars
       super
     end
 
+    # Count the number of unique values in every sub-arrays.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("a", [[1, 2], [4, 4]], dtype: Polars::Array.new(Polars::Int64, 2))
+    #   s.arr.n_unique
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: 'a' [u32]
+    #   # [
+    #   #         2
+    #   #         1
+    #   # ]
+    def n_unique
+      super
+    end
+
     # Convert an Array column into a List column with the same inner data type.
     #
     # @return [Series]
@@ -122,6 +212,13 @@ module Polars
 
     # Evaluate whether any boolean value is true for every subarray.
     #
+    # @param ignore_nulls [Boolean]
+    #   * If set to `true` (default), null values are ignored. If there
+    #     are no non-null values, the output is `false`.
+    #   * If set to `false`, [Kleene logic](https://en.wikipedia.org/wiki/Three-valued_logic) is used to deal with nulls:
+    #     if the column contains any null values and no `true` values,
+    #     the output is null.
+    #
     # @return [Series]
     #
     # @example
@@ -140,12 +237,161 @@ module Polars
     #   #         false
     #   #         null
     #   # ]
-    def any
+    def any(ignore_nulls: true)
+      super
+    end
+
+    # Return the number of elements in each array.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("a", [[1, 2], [4, 3]], dtype: Polars::Array.new(Polars::Int8, 2))
+    #   s.arr.len
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: 'a' [u32]
+    #   # [
+    #   #         2
+    #   #         2
+    #   # ]
+    def len
+      super
+    end
+
+    # Slice the sub-arrays.
+    #
+    # @param offset [Integer]
+    #   The starting index of the slice.
+    # @param length [Integer]
+    #   The length of the slice.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new(
+    #     [[1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12]],
+    #     dtype: Polars::Array.new(Polars::Int64, 6)
+    #   )
+    #   s.arr.slice(1)
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: '' [list[i64]]
+    #   # [
+    #   #         [2, 3, … 6]
+    #   #         [8, 9, … 12]
+    #   # ]
+    #
+    # @example
+    #   s.arr.slice(1, 3, as_array: true)
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: '' [array[i64, 3]]
+    #   # [
+    #   #         [2, 3, 4]
+    #   #         [8, 9, 10]
+    #   # ]
+    #
+    # @example
+    #   s.arr.slice(-2)
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: '' [list[i64]]
+    #   # [
+    #   #         [5, 6]
+    #   #         [11, 12]
+    #   # ]
+    def slice(
+      offset,
+      length = nil,
+      as_array: false
+    )
+      super
+    end
+
+    # Get the first `n` elements of the sub-arrays.
+    #
+    # @param n [Integer]
+    #   Number of values to return for each sublist.
+    # @param as_array [Boolean]
+    #   Return result as a fixed-length `Array`, otherwise as a `List`.
+    #   If true `n` must be a constant value.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new(
+    #     [[1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12]],
+    #     dtype: Polars::Array.new(Polars::Int64, 6)
+    #   )
+    #   s.arr.head
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: '' [list[i64]]
+    #   # [
+    #   #         [1, 2, … 5]
+    #   #         [7, 8, … 11]
+    #   # ]
+    #
+    # @example
+    #   s.arr.head(3, as_array: true)
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: '' [array[i64, 3]]
+    #   # [
+    #   #         [1, 2, 3]
+    #   #         [7, 8, 9]
+    #   # ]
+    def head(n = 5, as_array: false)
+      super
+    end
+
+    # Slice the last `n` values of every sublist.
+    #
+    # @param n [Integer]
+    #   Number of values to return for each sublist.
+    # @param as_array [Boolean]
+    #   Return result as a fixed-length `Array`, otherwise as a `List`.
+    #   If true `n` must be a constant value.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new(
+    #     [[1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12]],
+    #     dtype: Polars::Array.new(Polars::Int64, 6)
+    #   )
+    #   s.arr.tail
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: '' [list[i64]]
+    #   # [
+    #   #         [2, 3, … 6]
+    #   #         [8, 9, … 12]
+    #   # ]
+    #
+    # @example
+    #   s.arr.tail(3, as_array: true)
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: '' [array[i64, 3]]
+    #   # [
+    #   #         [4, 5, 6]
+    #   #         [10, 11, 12]
+    #   # ]
+    def tail(n = 5, as_array: false)
       super
     end
 
     # Evaluate whether all boolean values are true for every subarray.
     #
+    # @param ignore_nulls [Boolean]
+    #   * If set to `true` (default), null values are ignored. If there
+    #     are no non-null values, the output is `true`.
+    #   * If set to `false`, [Kleene logic](https://en.wikipedia.org/wiki/Three-valued_logic) is used to deal with nulls:
+    #     if the column contains any null values and no `false` values,
+    #     the output is null.
+    #
     # @return [Series]
     #
     # @example
@@ -164,7 +410,7 @@ module Polars
     #   #         true
     #   #         null
     #   # ]
-    def all
+    def all(ignore_nulls: true)
       super
     end
 
@@ -263,6 +509,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 [Series]
     #
@@ -270,7 +520,7 @@ module Polars
     #   s = Polars::Series.new(
     #     "a", [[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype: Polars::Array.new(Polars::Int32, 3)
     #   )
-    #   s.arr.get(Polars::Series.new([1, -2, 4]))
+    #   s.arr.get(Polars::Series.new([1, -2, 4]), null_on_oob: true)
     #   # =>
     #   # shape: (3,)
     #   # Series: 'a' [i32]
@@ -279,7 +529,7 @@ module Polars
     #   #         5
     #   #         null
     #   # ]
-    def get(index)
+    def get(index, null_on_oob: false)
       super
     end
 
@@ -334,7 +584,7 @@ module Polars
     # @param ignore_nulls [Boolean]
     #   Ignore null values (default).
     #
-    #   If set to `False`, null values will be propagated.
+    #   If set to `false`, null values will be propagated.
     #   If the sub-list contains any null values, the output is `nil`.
     #
     # @return [Series]
@@ -355,6 +605,11 @@ module Polars
 
     # Returns a column with a separate row for every array element.
     #
+    # @param empty_as_null [Boolean]
+    #   Explode an empty array into a `null`.
+    # @param keep_nulls [Boolean]
+    #   Explode a `null` array into a `null`.
+    #
     # @return [Series]
     #
     # @example
@@ -371,7 +626,7 @@ module Polars
     #   #         5
     #   #         6
     #   # ]
-    def explode
+    def explode(empty_as_null: true, keep_nulls: true)
       super
     end
 
@@ -379,6 +634,8 @@ module Polars
     #
     # @param item [Object]
     #   Item that will be checked for membership
+    # @param nulls_equal [Boolean]
+    #   If true, treat null as a distinct value. Null values will not propagate.
     #
     # @return [Series]
     #
@@ -395,7 +652,7 @@ module Polars
     #   #         true
     #   #         false
     #   # ]
-    def contains(item)
+    def contains(item, nulls_equal: true)
       super
     end
 
@@ -419,5 +676,122 @@ module Polars
     def count_matches(element)
       super
     end
+
+    # Convert the series of type `Array` to a series of type `Struct`.
+    #
+    # @param fields [Object]
+    #   If the name and number of the desired fields is known in advance
+    #   a list of field names can be given, which will be assigned by index.
+    #   Otherwise, to dynamically assign field names, a custom function can be
+    #   used; if neither are set, fields will be `field_0, field_1 .. field_n`.
+    #
+    # @return [Series]
+    #
+    # @example Convert array to struct with default field name assignment:
+    #   s1 = Polars::Series.new("n", [[0, 1, 2], [3, 4, 5]], dtype: Polars::Array.new(Polars::Int8, 3))
+    #   s2 = s1.arr.to_struct
+    #   s2.struct.fields
+    #   # => ["field_0", "field_1", "field_2"]
+    def to_struct(fields: nil)
+      s = Utils.wrap_s(_s)
+      s.to_frame.select(F.col(s.name).arr.to_struct(fields: fields)).to_series
+    end
+
+    # Shift array values by the given number of indices.
+    #
+    # @param n [Integer]
+    #   Number of indices to shift forward. If a negative value is passed, values
+    #   are shifted in the opposite direction instead.
+    #
+    # @return [Series]
+    #
+    # @note
+    #   This method is similar to the `LAG` operation in SQL when the value for `n`
+    #   is positive. With a negative value for `n`, it is similar to `LEAD`.
+    #
+    # @example By default, array values are shifted forward by one index.
+    #   s = Polars::Series.new([[1, 2, 3], [4, 5, 6]], dtype: Polars::Array.new(Polars::Int64, 3))
+    #   s.arr.shift
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: '' [array[i64, 3]]
+    #   # [
+    #   #         [null, 1, 2]
+    #   #         [null, 4, 5]
+    #   # ]
+    #
+    # @example Pass a negative value to shift in the opposite direction instead.
+    #   s.arr.shift(-2)
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: '' [array[i64, 3]]
+    #   # [
+    #   #         [3, null, null]
+    #   #         [6, null, null]
+    #   # ]
+    def shift(n = 1)
+      super
+    end
+
+    # Run any polars expression against the arrays' elements.
+    #
+    # @param expr [Expr]
+    #   Expression to run. Note that you can select an element with `pl.element()`
+    # @param as_list [Boolean]
+    #   Collect the resulting data as a list. This allows for expressions which
+    #   output a variable amount of data.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("a", [[1, 4], [8, 5], [3, 2]], dtype: Polars::Array.new(Polars::Int64, 2))
+    #   s.arr.eval(Polars.element.rank)
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'a' [array[f64, 2]]
+    #   # [
+    #   #         [1.0, 2.0]
+    #   #         [2.0, 1.0]
+    #   #         [2.0, 1.0]
+    #   # ]
+    def eval(expr, as_list: false)
+      s = Utils.wrap_s(_s)
+      s.to_frame.select(F.col(s.name).arr.eval(expr, as_list: as_list)).to_series
+    end
+
+    # Run any polars aggregation expression against the arrays' elements.
+    #
+    # @param expr [Expr]
+    #   Expression to run. Note that you can select an element with `Polars.element`.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new(
+    #     "a", [[1, nil], [42, 13], [nil, nil]], dtype: Polars::Array.new(Polars::Int64, 2)
+    #   )
+    #   s.arr.agg(Polars.element.null_count)
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'a' [u32]
+    #   # [
+    #   #         1
+    #   #         0
+    #   #         2
+    #   # ]
+    #
+    # @example
+    #   s.arr.agg(Polars.element.drop_nulls)
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'a' [list[i64]]
+    #   # [
+    #   #         [1]
+    #   #         [42, 13]
+    #   #         []
+    #   # ]
+    def agg(expr)
+      super
+    end
   end
 end

data/lib/polars/batched_csv_reader.rb

@@ -1,21 +1,22 @@
 module Polars
   # @private
   class BatchedCsvReader
-    attr_accessor :_reader, :new_columns
+    attr_accessor :_reader
 
     def initialize(
-      file,
+      source,
       has_header: true,
       columns: nil,
-      sep: ",",
-      comment_char: nil,
+      separator: ",",
+      comment_prefix: nil,
       quote_char: '"',
       skip_rows: 0,
-      dtypes: nil,
+      skip_lines: 0,
+      schema_overrides: nil,
       null_values: nil,
       missing_utf8_is_empty_string: false,
       ignore_errors: false,
-      parse_dates: false,
+      try_parse_dates: false,
       n_threads: nil,
       infer_schema_length: 100,
       batch_size: 50_000,
@@ -24,79 +25,60 @@ module Polars
       low_memory: false,
       rechunk: true,
       skip_rows_after_header: 0,
-      row_count_name: nil,
-      row_count_offset: 0,
+      row_index_name: nil,
+      row_index_offset: 0,
       eol_char: "\n",
       new_columns: nil,
       raise_if_empty: true,
       truncate_ragged_lines: false,
       decimal_comma: false
     )
-      if Utils.pathlike?(file)
-        path = Utils.normalize_filepath(file)
-      end
+      q = Polars.scan_csv(
+        source,
+        infer_schema_length: infer_schema_length,
+        has_header: has_header,
+        ignore_errors: ignore_errors,
+        n_rows: n_rows,
+        skip_rows: skip_rows,
+        skip_lines: skip_lines,
+        separator: separator,
+        rechunk: rechunk,
+        encoding: encoding,
+        schema_overrides: schema_overrides,
+        low_memory: low_memory,
+        comment_prefix: comment_prefix,
+        quote_char: quote_char,
+        null_values: null_values,
+        missing_utf8_is_empty_string: missing_utf8_is_empty_string,
+        try_parse_dates: try_parse_dates,
+        skip_rows_after_header: skip_rows_after_header,
+        row_index_name: row_index_name,
+        row_index_offset: row_index_offset,
+        eol_char: eol_char,
+        raise_if_empty: raise_if_empty,
+        truncate_ragged_lines: truncate_ragged_lines,
+        decimal_comma: decimal_comma,
+        new_columns: new_columns
+      )
 
-      dtype_list = nil
-      dtype_slice = nil
-      if !dtypes.nil?
-        if dtypes.is_a?(Hash)
-          dtype_list = []
-          dtypes.each do |k, v|
-            dtype_list << [k, Utils.rb_type_to_dtype(v)]
-          end
-        elsif dtypes.is_a?(::Array)
-          dtype_slice = dtypes
-        else
-          raise ArgumentError, "dtype arg should be list or dict"
-        end
+      if !columns.nil?
+        q = q.select(columns)
       end
 
-      processed_null_values = Utils._process_null_values(null_values)
-      projection, columns = Utils.handle_projection_columns(columns)
-
-      self._reader = RbBatchedCsv.new(
-        infer_schema_length,
-        batch_size,
-        has_header,
-        ignore_errors,
-        n_rows,
-        skip_rows,
-        projection,
-        sep,
-        rechunk,
-        columns,
-        encoding,
-        n_threads,
-        path,
-        dtype_list,
-        dtype_slice,
-        low_memory,
-        comment_char,
-        quote_char,
-        processed_null_values,
-        missing_utf8_is_empty_string,
-        parse_dates,
-        skip_rows_after_header,
-        Utils.parse_row_index_args(row_count_name, row_count_offset),
-        eol_char,
-        raise_if_empty,
-        truncate_ragged_lines,
-        decimal_comma
-      )
-      self.new_columns = new_columns
+      # Trigger empty data.
+      if raise_if_empty
+        q.collect_schema
+      end
+      self._reader = q.collect_batches(chunk_size: batch_size)
     end
 
     def next_batches(n)
-      batches = _reader.next_batches(n)
-      if !batches.nil?
-        if new_columns
-          batches.map { |df| Utils._update_columns(Utils.wrap_df(df), new_columns) }
-        else
-          batches.map { |df| Utils.wrap_df(df) }
-        end
-      else
-        nil
+      chunks = self._reader.take(n)
+
+      if chunks.length > 0
+        return chunks
       end
+      nil
     end
   end
 end