polars-df 0.3.1-x86_64-linux → 0.5.0-x86_64-linux

data/lib/polars/list_expr.rb

@@ -27,7 +27,7 @@ module Polars
     #   # │ 1   │
     #   # └─────┘
     def lengths
-      Utils.wrap_expr(_rbexpr.arr_lengths)
+      Utils.wrap_expr(_rbexpr.list_lengths)
     end
 
     # Sum all the lists in the array.
@@ -48,7 +48,7 @@ module Polars
     #   # │ 5      │
     #   # └────────┘
     def sum
-      Utils.wrap_expr(_rbexpr.lst_sum)
+      Utils.wrap_expr(_rbexpr.list_sum)
     end
 
     # Compute the max value of the lists in the array.
@@ -69,7 +69,7 @@ module Polars
     #   # │ 3      │
     #   # └────────┘
     def max
-      Utils.wrap_expr(_rbexpr.lst_max)
+      Utils.wrap_expr(_rbexpr.list_max)
     end
 
     # Compute the min value of the lists in the array.
@@ -90,7 +90,7 @@ module Polars
     #   # │ 2      │
     #   # └────────┘
     def min
-      Utils.wrap_expr(_rbexpr.lst_min)
+      Utils.wrap_expr(_rbexpr.list_min)
     end
 
     # Compute the mean value of the lists in the array.
@@ -111,7 +111,7 @@ module Polars
     #   # │ 2.5    │
     #   # └────────┘
     def mean
-      Utils.wrap_expr(_rbexpr.lst_mean)
+      Utils.wrap_expr(_rbexpr.list_mean)
     end
 
     # Sort the arrays in the list.
@@ -136,7 +136,7 @@ module Polars
     #   # │ [1, 2, 9] │
     #   # └───────────┘
     def sort(reverse: false)
-      Utils.wrap_expr(_rbexpr.lst_sort(reverse))
+      Utils.wrap_expr(_rbexpr.list_sort(reverse))
     end
 
     # Reverse the arrays in the list.
@@ -161,7 +161,7 @@ module Polars
     #   # │ [2, 1, 9] │
     #   # └───────────┘
     def reverse
-      Utils.wrap_expr(_rbexpr.lst_reverse)
+      Utils.wrap_expr(_rbexpr.list_reverse)
     end
 
     # Get the unique/distinct values in the list.
@@ -184,8 +184,8 @@ module Polars
     #   # ╞═══════════╡
     #   # │ [1, 2]    │
     #   # └───────────┘
-    def unique
-      Utils.wrap_expr(_rbexpr.lst_unique)
+    def unique(maintain_order: false)
+      Utils.wrap_expr(_rbexpr.list_unique(maintain_order))
     end
 
     # Concat the arrays in a Series dtype List in linear time.
@@ -255,7 +255,7 @@ module Polars
     #   # └──────┘
     def get(index)
       index = Utils.expr_to_lit_or_expr(index, str_to_lit: false)._rbexpr
-      Utils.wrap_expr(_rbexpr.lst_get(index))
+      Utils.wrap_expr(_rbexpr.list_get(index))
     end
 
     # Get the value by index in the sublists.
@@ -265,6 +265,28 @@ module Polars
       get(item)
     end
 
+    # Take sublists by multiple indices.
+    #
+    # The indices may be defined in a single column, or by sublists in another
+    # column of dtype `List`.
+    #
+    # @param index [Object]
+    #   Indices 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
+    #   Note that defaulting to raising an error is much cheaper
+    #
+    # @return [Expr]
+    def take(index, null_on_oob: false)
+      if index.is_a?(Array)
+        index = Series.new(index)
+      end
+      index = Utils.expr_to_lit_or_expr(index, str_to_lit: false)._rbexpr
+      Utils.wrap_expr(_rbexpr.list_take(index, null_on_oob))
+    end
+
     # Get the first value of the sublists.
     #
     # @return [Expr]
@@ -331,7 +353,7 @@ module Polars
     #   # │ true  │
     #   # └───────┘
     def contains(item)
-      Utils.wrap_expr(_rbexpr.arr_contains(Utils.expr_to_lit_or_expr(item)._rbexpr))
+      Utils.wrap_expr(_rbexpr.list_contains(Utils.expr_to_lit_or_expr(item)._rbexpr))
     end
 
     # Join all string items in a sublist and place a separator between them.
@@ -357,7 +379,7 @@ module Polars
     #   # │ x y   │
     #   # └───────┘
     def join(separator)
-      Utils.wrap_expr(_rbexpr.lst_join(separator))
+      Utils.wrap_expr(_rbexpr.list_join(separator))
     end
 
     # Retrieve the index of the minimal value in every sublist.
@@ -382,7 +404,7 @@ module Polars
     #   # │ 1   │
     #   # └─────┘
     def arg_min
-      Utils.wrap_expr(_rbexpr.lst_arg_min)
+      Utils.wrap_expr(_rbexpr.list_arg_min)
     end
 
     # Retrieve the index of the maximum value in every sublist.
@@ -407,7 +429,7 @@ module Polars
     #   # │ 0   │
     #   # └─────┘
     def arg_max
-      Utils.wrap_expr(_rbexpr.lst_arg_max)
+      Utils.wrap_expr(_rbexpr.list_arg_max)
     end
 
     # Calculate the n-th discrete difference of every sublist.
@@ -426,11 +448,11 @@ module Polars
     #   # shape: (2,)
     #   # Series: 'a' [list[i64]]
     #   # [
-    #   #         [null, 1, ... 1]
+    #   #         [null, 1, … 1]
     #   #         [null, -8, -1]
     #   # ]
     def diff(n: 1, null_behavior: "ignore")
-      Utils.wrap_expr(_rbexpr.lst_diff(n, null_behavior))
+      Utils.wrap_expr(_rbexpr.list_diff(n, null_behavior))
     end
 
     # Shift values by the given period.
@@ -447,11 +469,11 @@ module Polars
     #   # shape: (2,)
     #   # Series: 'a' [list[i64]]
     #   # [
-    #   #         [null, 1, ... 3]
+    #   #         [null, 1, … 3]
     #   #         [null, 10, 2]
     #   # ]
     def shift(periods = 1)
-      Utils.wrap_expr(_rbexpr.lst_shift(periods))
+      Utils.wrap_expr(_rbexpr.list_shift(periods))
     end
 
     # Slice every sublist.
@@ -477,7 +499,7 @@ module Polars
     def slice(offset, length = nil)
       offset = Utils.expr_to_lit_or_expr(offset, str_to_lit: false)._rbexpr
       length = Utils.expr_to_lit_or_expr(length, str_to_lit: false)._rbexpr
-      Utils.wrap_expr(_rbexpr.lst_slice(offset, length))
+      Utils.wrap_expr(_rbexpr.list_slice(offset, length))
     end
 
     # Slice the first `n` values of every sublist.
@@ -523,6 +545,33 @@ module Polars
       slice(offset, n)
     end
 
+    # Count how often the value produced by ``element`` occurs.
+    #
+    # @param element [Expr]
+    #   An expression that produces a single value
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"listcol" => [[0], [1], [1, 2, 3, 2], [1, 2, 1], [4, 4]]})
+    #   df.select(Polars.col("listcol").arr.count_match(2).alias("number_of_twos"))
+    #   # =>
+    #   # shape: (5, 1)
+    #   # ┌────────────────┐
+    #   # │ number_of_twos │
+    #   # │ ---            │
+    #   # │ u32            │
+    #   # ╞════════════════╡
+    #   # │ 0              │
+    #   # │ 0              │
+    #   # │ 2              │
+    #   # │ 1              │
+    #   # │ 0              │
+    #   # └────────────────┘
+    def count_match(element)
+      Utils.wrap_expr(_rbexpr.list_count_match(Utils.expr_to_lit_or_expr(element)._rbexpr))
+    end
+
     # Convert the series of type `List` to a series of type `Struct`.
     #
     # @param n_field_strategy ["first_non_null", "max_width"]
@@ -548,7 +597,7 @@ module Polars
     #   # └────────────┘
     def to_struct(n_field_strategy: "first_non_null", name_generator: nil)
       raise Todo if name_generator
-      Utils.wrap_expr(_rbexpr.lst_to_struct(n_field_strategy, name_generator, 0))
+      Utils.wrap_expr(_rbexpr.list_to_struct(n_field_strategy, name_generator, 0))
     end
 
     # Run any polars expression against the lists' elements.
@@ -582,7 +631,7 @@ module Polars
     #   # │ 3   ┆ 2   ┆ [2.0, 1.0] │
     #   # └─────┴─────┴────────────┘
     def eval(expr, parallel: false)
-       Utils.wrap_expr(_rbexpr.lst_eval(expr._rbexpr, parallel))
+      Utils.wrap_expr(_rbexpr.list_eval(expr._rbexpr, parallel))
     end
   end
 end

data/lib/polars/list_name_space.rb

@@ -185,7 +185,7 @@ module Polars
     #   # shape: (2,)
     #   # Series: 'a' [list[i64]]
     #   # [
-    #   #         [null, 1, ... 1]
+    #   #         [null, 1, … 1]
     #   #         [null, -8, -1]
     #   # ]
     def diff(n: 1, null_behavior: "ignore")
@@ -206,7 +206,7 @@ module Polars
     #   # shape: (2,)
     #   # Series: 'a' [list[i64]]
     #   # [
-    #   #         [null, 1, ... 3]
+    #   #         [null, 1, … 3]
     #   #         [null, 10, 2]
     #   # ]
     def shift(periods = 1)

data/lib/polars/series.rb

@@ -67,6 +67,12 @@ module Polars
           ._s
       elsif values.is_a?(Array)
         self._s = sequence_to_rbseries(name, values, dtype: dtype, strict: strict, dtype_if_empty: dtype_if_empty)
+      elsif defined?(Numo::NArray) && values.is_a?(Numo::NArray)
+        self._s = numo_to_rbseries(name, values, strict: strict, nan_to_null: nan_to_null)
+
+        if !dtype.nil?
+          self._s = self.cast(dtype, strict: true)._s
+        end
       else
         raise ArgumentError, "Series constructor called with unsupported type; got #{values.class.name}"
       end
@@ -90,10 +96,14 @@ module Polars
     #
     # @return [Hash]
     def flags
-      {
+      out = {
         "SORTED_ASC" => _s.is_sorted_flag,
         "SORTED_DESC" => _s.is_sorted_reverse_flag
       }
+      if dtype.is_a?(List)
+        out["FAST_EXPLODE"] = _s.can_fast_explode_flag
+      end
+      out
     end
 
     # Get the inner dtype in of a List typed Series.
@@ -222,14 +232,28 @@ module Polars
     #
     # @return [Series]
     def *(other)
-      _arithmetic(other, :mul)
+      if is_temporal
+        raise ArgumentError, "first cast to integer before multiplying datelike dtypes"
+      elsif other.is_a?(DataFrame)
+        other * self
+      else
+        _arithmetic(other, :mul)
+      end
     end
 
     # Performs division.
     #
     # @return [Series]
     def /(other)
-      _arithmetic(other, :div)
+      if is_temporal
+        raise ArgumentError, "first cast to integer before dividing datelike dtypes"
+      end
+
+      if is_float
+        return _arithmetic(other, :div)
+      end
+
+      cast(Float64) / other
     end
 
     # Returns the modulo.
@@ -252,6 +276,16 @@ module Polars
       to_frame.select(Polars.col(name).pow(power)).to_series
     end
 
+    # Performs boolean not.
+    #
+    # @return [Series]
+    def !
+      if dtype == Boolean
+        return Utils.wrap_s(_s.not)
+      end
+      raise NotImplementedError
+    end
+
     # Performs negation.
     #
     # @return [Series]
@@ -278,6 +312,10 @@ module Polars
         return Utils.wrap_s(_s.take_with_series(_pos_idxs(item)._s))
       end
 
+      if item.is_a?(Series) && item.bool?
+        return filter(item)
+      end
+
       if item.is_a?(Integer)
         return _s.get_idx(item)
       end
@@ -369,16 +407,26 @@ module Polars
     # Check if any boolean value in the column is `true`.
     #
     # @return [Boolean]
-    def any
-      to_frame.select(Polars.col(name).any).to_series[0]
+    def any?(&block)
+      if block_given?
+        apply(&block).any?
+      else
+        to_frame.select(Polars.col(name).any).to_series[0]
+      end
     end
+    alias_method :any, :any?
 
     # Check if all boolean values in the column are `true`.
     #
     # @return [Boolean]
-    def all
-      to_frame.select(Polars.col(name).all).to_series[0]
+    def all?(&block)
+      if block_given?
+        apply(&block).all?
+      else
+        to_frame.select(Polars.col(name).all).to_series[0]
+      end
     end
+    alias_method :all, :all?
 
     # Compute the logarithm to a given base.
     #
@@ -1314,6 +1362,7 @@ module Polars
     def unique(maintain_order: false)
       super
     end
+    alias_method :uniq, :unique
 
     # Take values by index.
     #
@@ -1535,6 +1584,7 @@ module Polars
     def is_in(other)
       super
     end
+    alias_method :in?, :is_in
 
     # Get index values where Boolean Series evaluate `true`.
     #
@@ -1660,6 +1710,7 @@ module Polars
     end
     alias_method :count, :len
     alias_method :length, :len
+    alias_method :size, :len
 
     # Cast between data types.
     #
@@ -1779,6 +1830,8 @@ module Polars
       [Date, Time].include?(dtype) || dtype.is_a?(Datetime) || dtype.is_a?(Duration)
     end
     alias_method :datelike?, :is_datelike
+    alias_method :is_temporal, :is_datelike
+    alias_method :temporal?, :is_datelike
 
     # Check if this Series has floating point numbers.
     #
@@ -2432,6 +2485,7 @@ module Polars
       end
       Utils.wrap_s(_s.apply_lambda(func, pl_return_dtype, skip_nulls))
     end
+    alias_method :map, :apply
 
     # Shift the values by a given period.
     #
@@ -3483,7 +3537,7 @@ module Polars
     #   #         99
     #   # ]
     def extend_constant(value, n)
-      super
+      Utils.wrap_s(_s.extend_constant(value, n))
     end
 
     # Flags the Series as sorted.
@@ -3531,6 +3585,13 @@ module Polars
       ListNameSpace.new(self)
     end
 
+    # Create an object namespace of all binary related methods.
+    #
+    # @return [BinaryNameSpace]
+    def bin
+      BinaryNameSpace.new(self)
+    end
+
     # Create an object namespace of all categorical related methods.
     #
     # @return [CatNameSpace]
@@ -3630,14 +3691,39 @@ module Polars
     end
 
     def _comp(other, op)
+      if dtype == Boolean && Utils.bool?(other) && [:eq, :neq].include?(op)
+        if (other == true && op == :eq) || (other == false && op == :neq)
+          return clone
+        elsif (other == false && op == :eq) || (other == true && op == :neq)
+          return !self
+        end
+      end
+
+      if other.is_a?(::Time) && dtype.is_a?(Datetime)
+        ts = Utils._datetime_to_pl_timestamp(other, time_unit)
+        f = ffi_func("#{op}_<>", Int64, _s)
+        fail if f.nil?
+        return Utils.wrap_s(f.call(ts))
+      elsif other.is_a?(::Date) && dtype == Date
+        d = Utils._date_to_pl_date(other)
+        f = ffi_func("#{op}_<>", Int32, _s)
+        fail if f.nil?
+        return Utils.wrap_s(f.call(d))
+      end
+
       if other.is_a?(Series)
         return Utils.wrap_s(_s.send(op, other._s))
       end
 
-      if dtype == Utf8
-        raise Todo
+      f = ffi_func("#{op}_<>", dtype, _s)
+      if f.nil?
+        raise NotImplementedError
       end
-      Utils.wrap_s(_s.send("#{op}_#{DTYPE_TO_FFINAME.fetch(dtype)}", other))
+      Utils.wrap_s(f.call(other))
+    end
+
+    def ffi_func(name, dtype, _s)
+      _s.method(name.sub("<>", DTYPE_TO_FFINAME.fetch(dtype))) if DTYPE_TO_FFINAME.key?(dtype)
     end
 
     def _arithmetic(other, op)
@@ -3648,14 +3734,16 @@ module Polars
         return Utils.wrap_s(_s.send(op, other._s))
       end
 
-      if other.is_a?(::Date) || other.is_a?(::DateTime) || other.is_a?(::Time) || other.is_a?(String)
-        raise Todo
-      end
-      if other.is_a?(Float) && !is_float
-        raise Todo
+      if (other.is_a?(Float) || other.is_a?(::Date) || other.is_a?(::DateTime) || other.is_a?(::Time) || other.is_a?(String)) && !is_float
+        _s2 = sequence_to_rbseries(name, [other])
+        return Utils.wrap_s(_s.send(op, _s2))
       end
 
-      Utils.wrap_s(_s.send("#{op}_#{DTYPE_TO_FFINAME.fetch(dtype)}", other))
+      f = ffi_func("#{op}_<>", dtype, _s)
+      if f.nil?
+        raise ArgumentError, "cannot do arithmetic with series of dtype: #{dtype} and argument of type: #{other.class.name}"
+      end
+      Utils.wrap_s(f.call(other))
     end
 
     DTYPE_TO_FFINAME = {
@@ -3688,25 +3776,57 @@ module Polars
       values._s
     end
 
+    def numo_to_rbseries(name, values, strict: true, nan_to_null: false)
+      # not needed yet
+      # if !values.contiguous?
+      # end
+
+      if values.shape.length == 1
+        values, dtype = numo_values_and_dtype(values)
+        strict = nan_to_null if [Numo::SFloat, Numo::DFloat].include?(dtype)
+        if dtype == Numo::RObject
+          sequence_to_rbseries(name, values.to_a, strict: strict)
+        else
+          constructor = numo_type_to_constructor(dtype)
+          # TODO improve performance
+          constructor.call(name, values.to_a, strict)
+        end
+      elsif values.shape.length == 2
+        raise Todo
+      else
+        raise Todo
+      end
+    end
+
+    def numo_values_and_dtype(values)
+      [values, values.class]
+    end
+
+    def numo_type_to_constructor(dtype)
+      {
+        Numo::Float32 => RbSeries.method(:new_opt_f32),
+        Numo::Float64 => RbSeries.method(:new_opt_f64),
+        Numo::Int8 => RbSeries.method(:new_opt_i8),
+        Numo::Int16 => RbSeries.method(:new_opt_i16),
+        Numo::Int32 => RbSeries.method(:new_opt_i32),
+        Numo::Int64 => RbSeries.method(:new_opt_i64),
+        Numo::UInt8 => RbSeries.method(:new_opt_u8),
+        Numo::UInt16 => RbSeries.method(:new_opt_u16),
+        Numo::UInt32 => RbSeries.method(:new_opt_u32),
+        Numo::UInt64 => RbSeries.method(:new_opt_u64)
+      }.fetch(dtype)
+    rescue KeyError
+      RbSeries.method(:new_object)
+    end
+
     def sequence_to_rbseries(name, values, dtype: nil, strict: true, dtype_if_empty: nil)
       ruby_dtype = nil
-      nested_dtype = nil
 
       if (values.nil? || values.empty?) && dtype.nil?
-        if dtype_if_empty
-          # if dtype for empty sequence could be guessed
-          # (e.g comparisons between self and other)
-          dtype = dtype_if_empty
-        else
-          # default to Float32 type
-          dtype = :f32
-        end
+        dtype = dtype_if_empty || Float32
       end
 
-      rb_temporal_types = []
-      rb_temporal_types << ::Date if defined?(::Date)
-      rb_temporal_types << ::DateTime if defined?(::DateTime)
-      rb_temporal_types << ::Time if defined?(::Time)
+      rb_temporal_types = [::Date, ::DateTime, ::Time]
 
       value = _get_first_non_none(values)
       if !value.nil?
@@ -3731,58 +3851,54 @@ module Polars
 
         # temporal branch
         if rb_temporal_types.include?(ruby_dtype)
-          # if dtype.nil?
-          #   dtype = rb_type_to_dtype(ruby_dtype)
-          # elsif rb_temporal_types.include?(dtype)
-          #   dtype = rb_type_to_dtype(dtype)
-          # end
-
-          if ruby_dtype == ::Date
-            RbSeries.new_opt_date(name, values, strict)
-          elsif ruby_dtype == ::Time
-            RbSeries.new_opt_datetime(name, values, strict)
-          elsif ruby_dtype == ::DateTime
-            RbSeries.new_opt_datetime(name, values.map(&:to_time), strict)
-          else
-            raise Todo
-          end
-        elsif ruby_dtype == Array
-          if nested_dtype.nil?
-            nested_value = _get_first_non_none(value)
-            nested_dtype = nested_value.nil? ? Float : nested_value.class
+          if dtype.nil?
+            dtype = Utils.rb_type_to_dtype(ruby_dtype)
+          elsif rb_temporal_types.include?(dtype)
+            dtype = Utils.rb_type_to_dtype(dtype)
           end
+          # TODO
+          time_unit = nil
 
-          if nested_dtype == Array
-            raise Todo
+          rb_series = RbSeries.new_from_anyvalues(name, values, strict)
+          if time_unit.nil?
+            s = Utils.wrap_s(rb_series)
+          else
+            s = Utils.wrap_s(rb_series).dt.cast_time_unit(time_unit)
           end
-
-          if value.is_a?(Array)
-            count = 0
-            equal_to_inner = true
-            values.each do |lst|
-              lst.each do |vl|
-                equal_to_inner = vl.class == nested_dtype
-                if !equal_to_inner || count > 50
-                  break
-                end
-                count += 1
+          return s._s
+        elsif defined?(Numo::NArray) && value.is_a?(Numo::NArray) && value.shape.length == 1
+          raise Todo
+        elsif ruby_dtype == Array
+          return sequence_from_anyvalue_or_object(name, values)
+        elsif ruby_dtype == Series
+          return RbSeries.new_series_list(name, values.map(&:_s), strict)
+        elsif ruby_dtype == RbSeries
+          return RbSeries.new_series_list(name, values, strict)
+        else
+          constructor =
+            if value.is_a?(String)
+              if value.encoding == Encoding::UTF_8
+                RbSeries.method(:new_str)
+              else
+                RbSeries.method(:new_binary)
               end
+            elsif value.is_a?(Integer) && values.any? { |v| v.is_a?(Float) }
+              # TODO improve performance
+              RbSeries.method(:new_opt_f64)
+            else
+              rb_type_to_constructor(value.class)
             end
-            if equal_to_inner
-              dtype = Utils.rb_type_to_dtype(nested_dtype)
-              # TODO rescue and fallback to new_object
-              return RbSeries.new_list(name, values, dtype)
-            end
-          end
-
-          RbSeries.new_object(name, values, strict)
-        else
-          constructor = rb_type_to_constructor(value.class)
           constructor.call(name, values, strict)
         end
       end
     end
 
+    def sequence_from_anyvalue_or_object(name, values)
+      RbSeries.new_from_anyvalues(name, values, true)
+    rescue
+      RbSeries.new_object(name, values, false)
+    end
+
     POLARS_TYPE_TO_CONSTRUCTOR = {
       Float32 => RbSeries.method(:new_opt_f32),
       Float64 => RbSeries.method(:new_opt_f64),
@@ -3795,7 +3911,8 @@ module Polars
       UInt32 => RbSeries.method(:new_opt_u32),
       UInt64 => RbSeries.method(:new_opt_u64),
       Boolean => RbSeries.method(:new_opt_bool),
-      Utf8 => RbSeries.method(:new_str)
+      Utf8 => RbSeries.method(:new_str),
+      Binary => RbSeries.method(:new_binary)
     }
 
     SYM_TYPE_TO_CONSTRUCTOR = {
@@ -3826,7 +3943,6 @@ module Polars
     RB_TYPE_TO_CONSTRUCTOR = {
       Float => RbSeries.method(:new_opt_f64),
       Integer => RbSeries.method(:new_opt_i64),
-      String => RbSeries.method(:new_str),
       TrueClass => RbSeries.method(:new_opt_bool),
       FalseClass => RbSeries.method(:new_opt_bool)
     }