polars-df 0.4.0-x86_64-linux → 0.6.0-x86_64-linux

data/lib/polars/series.rb

@@ -65,8 +65,14 @@ module Polars
           )
           .rename(name, in_place: true)
           ._s
-      elsif values.is_a?(Array)
+      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,7 +312,15 @@ 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)
+        if item < 0
+          item = len + item
+        end
+
         return _s.get_idx(item)
       end
 
@@ -297,7 +339,7 @@ module Polars
     #
     # @return [Object]
     def []=(key, value)
-      if value.is_a?(Array)
+      if value.is_a?(::Array)
         if is_numeric || is_datelike
           set_at_idx(key, value)
           return
@@ -315,7 +357,7 @@ module Polars
         else
           raise Todo
         end
-      elsif key.is_a?(Array)
+      elsif key.is_a?(::Array)
         s = Utils.wrap_s(sequence_to_rbseries("", key, dtype: UInt32))
         self[s] = value
       elsif key.is_a?(Range)
@@ -369,16 +411,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.
     #
@@ -667,8 +719,8 @@ module Polars
     #   # │ 0   ┆ 1   ┆ 0   │
     #   # │ 0   ┆ 0   ┆ 1   │
     #   # └─────┴─────┴─────┘
-    def to_dummies(separator: "_")
-      Utils.wrap_df(_s.to_dummies(separator))
+    def to_dummies(separator: "_", drop_first: false)
+      Utils.wrap_df(_s.to_dummies(separator, drop_first))
     end
 
     # Count the unique values in a Series.
@@ -1076,7 +1128,7 @@ module Polars
     #   #         3
     #   # ]
     def filter(predicate)
-      if predicate.is_a?(Array)
+      if predicate.is_a?(::Array)
         predicate = Series.new("", predicate)
       end
       Utils.wrap_s(_s.filter(predicate._s))
@@ -1314,6 +1366,7 @@ module Polars
     def unique(maintain_order: false)
       super
     end
+    alias_method :uniq, :unique
 
     # Take values by index.
     #
@@ -1535,6 +1588,7 @@ module Polars
     def is_in(other)
       super
     end
+    alias_method :in?, :is_in
 
     # Get index values where Boolean Series evaluate `true`.
     #
@@ -1660,6 +1714,7 @@ module Polars
     end
     alias_method :count, :len
     alias_method :length, :len
+    alias_method :size, :len
 
     # Cast between data types.
     #
@@ -1779,6 +1834,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 +2489,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.
     #
@@ -2759,7 +2817,8 @@ module Polars
       window_size,
       weights: nil,
       min_periods: nil,
-      center: false
+      center: false,
+      ddof: 1
     )
       to_frame
         .select(
@@ -2767,7 +2826,8 @@ module Polars
             window_size,
             weights: weights,
             min_periods: min_periods,
-            center: center
+            center: center,
+            ddof: ddof
           )
         )
         .to_series
@@ -2810,7 +2870,8 @@ module Polars
       window_size,
       weights: nil,
       min_periods: nil,
-      center: false
+      center: false,
+      ddof: 1
     )
       to_frame
         .select(
@@ -2818,7 +2879,8 @@ module Polars
             window_size,
             weights: weights,
             min_periods: min_periods,
-            center: center
+            center: center,
+            ddof: ddof
           )
         )
         .to_series
@@ -3483,7 +3545,7 @@ module Polars
     #   #         99
     #   # ]
     def extend_constant(value, n)
-      super
+      Utils.wrap_s(_s.extend_constant(value, n))
     end
 
     # Flags the Series as sorted.
@@ -3527,10 +3589,17 @@ module Polars
     # Create an object namespace of all list related methods.
     #
     # @return [ListNameSpace]
-    def arr
+    def list
       ListNameSpace.new(self)
     end
 
+    # Create an object namespace of all array related methods.
+    #
+    # @return [ArrayNameSpace]
+    def arr
+      ArrayNameSpace.new(self)
+    end
+
     # Create an object namespace of all binary related methods.
     #
     # @return [BinaryNameSpace]
@@ -3637,14 +3706,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)
@@ -3655,14 +3749,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 = {
@@ -3695,25 +3791,60 @@ 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
+      elsif dtype == List
+        ruby_dtype = ::Array
       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]
+      rb_temporal_types << ActiveSupport::TimeWithZone if defined?(ActiveSupport::TimeWithZone)
 
       value = _get_first_non_none(values)
       if !value.nil?
@@ -3722,9 +3853,20 @@ module Polars
         end
       end
 
-      if !dtype.nil? && Utils.is_polars_dtype(dtype) && ruby_dtype.nil?
+      if !dtype.nil? && ![List, Unknown].include?(dtype) && Utils.is_polars_dtype(dtype) && ruby_dtype.nil?
+        if dtype == Array && !dtype.is_a?(Array) && value.is_a?(::Array)
+          dtype = Array.new(value.size)
+        end
+
         constructor = polars_type_to_constructor(dtype)
         rbseries = constructor.call(name, values, strict)
+
+        base_type = dtype.is_a?(DataType) ? dtype.class : dtype
+        if [Date, Datetime, Duration, Time, Categorical, Boolean].include?(base_type)
+          if rbseries.dtype != dtype
+            rbseries = rbseries.cast(dtype, true)
+          end
+        end
         return rbseries
       else
         if ruby_dtype.nil?
@@ -3738,58 +3880,64 @@ 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
-              end
-            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)
+          return s._s
+        elsif defined?(Numo::NArray) && value.is_a?(Numo::NArray) && value.shape.length == 1
+          raise Todo
+        elsif ruby_dtype == ::Array
+          if dtype.is_a?(Object)
+            return RbSeries.new_object(name, values, strict)
+          end
+          if dtype
+            srs = sequence_from_anyvalue_or_object(name, values)
+            if dtype != srs.dtype
+              srs = srs.cast(dtype, strict: false)
             end
+            return srs
           end
-
-          RbSeries.new_object(name, values, strict)
+          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 = rb_type_to_constructor(value.class)
+          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
           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),
@@ -3801,9 +3949,17 @@ module Polars
       UInt16 => RbSeries.method(:new_opt_u16),
       UInt32 => RbSeries.method(:new_opt_u32),
       UInt64 => RbSeries.method(:new_opt_u64),
+      Decimal => RbSeries.method(:new_decimal),
+      Date => RbSeries.method(:new_from_anyvalues),
+      Datetime => RbSeries.method(:new_from_anyvalues),
+      Duration => RbSeries.method(:new_from_anyvalues),
+      Time => RbSeries.method(:new_from_anyvalues),
       Boolean => RbSeries.method(:new_opt_bool),
       Utf8 => RbSeries.method(:new_str),
-      Binary => RbSeries.method(:new_binary)
+      Object => RbSeries.method(:new_object),
+      Categorical => RbSeries.method(:new_str),
+      Binary => RbSeries.method(:new_binary),
+      Null => RbSeries.method(:new_null)
     }
 
     SYM_TYPE_TO_CONSTRUCTOR = {
@@ -3822,8 +3978,14 @@ module Polars
     }
 
     def polars_type_to_constructor(dtype)
-      if dtype.is_a?(Class) && dtype < DataType
+      if dtype.is_a?(Array)
+        lambda do |name, values, strict|
+          RbSeries.new_array(dtype.width, dtype.inner, name, values, strict)
+        end
+      elsif dtype.is_a?(Class) && dtype < DataType
         POLARS_TYPE_TO_CONSTRUCTOR.fetch(dtype)
+      elsif dtype.is_a?(DataType)
+        POLARS_TYPE_TO_CONSTRUCTOR.fetch(dtype.class)
       else
         SYM_TYPE_TO_CONSTRUCTOR.fetch(dtype.to_sym)
       end
@@ -3834,9 +3996,9 @@ 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)
+      FalseClass => RbSeries.method(:new_opt_bool),
+      BigDecimal => RbSeries.method(:new_decimal)
     }
 
     def rb_type_to_constructor(dtype)