polars-df 0.4.0-x86_64-linux → 0.5.0-x86_64-linux

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.
@@ -3637,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)
@@ -3655,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 = {
@@ -3695,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?
@@ -3738,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),
@@ -3834,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)
     }

data/lib/polars/string_expr.rb

@@ -11,8 +11,8 @@ module Polars
 
     # Parse a Utf8 expression to a Date/Datetime/Time type.
     #
-    # @param datatype [Symbol]
-    #   `:date`, `:dateime`, or `:time`.
+    # @param dtype [Object]
+    #   The data type to convert into. Can be either Date, Datetime, or Time.
     # @param fmt [String]
     #   Format to use, refer to the
     #   [chrono strftime documentation](https://docs.rs/chrono/latest/chrono/format/strftime/index.html)
@@ -33,57 +33,56 @@ module Polars
     #   the format string, if given, eg: "%F %T%.3f" => Datetime("ms"). If
     #   no fractional second component is found then the default is "us".
     #
-    # @example
+    # @example Dealing with a consistent format:
+    #   s = Polars::Series.new(["2020-01-01 01:00Z", "2020-01-01 02:00Z"])
+    #   s.str.strptime(Polars::Datetime, "%Y-%m-%d %H:%M%#z")
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: '' [datetime[μs, +00:00]]
+    #   # [
+    #   #         2020-01-01 01:00:00 +00:00
+    #   #         2020-01-01 02:00:00 +00:00
+    #   # ]
+    #
+    # @example Dealing with different formats.
     #   s = Polars::Series.new(
     #     "date",
     #     [
     #       "2021-04-22",
     #       "2022-01-04 00:00:00",
     #       "01/31/22",
-    #       "Sun Jul  8 00:34:60 2001"
+    #       "Sun Jul  8 00:34:60 2001",
     #     ]
     #   )
-    #   s.to_frame.with_column(
-    #     Polars.col("date")
-    #       .str.strptime(:date, "%F", strict: false)
-    #       .fill_null(
-    #         Polars.col("date").str.strptime(:date, "%F %T", strict: false)
-    #       )
-    #       .fill_null(Polars.col("date").str.strptime(:date, "%D", strict: false))
-    #       .fill_null(Polars.col("date").str.strptime(:date, "%c", strict: false))
-    #   )
+    #   s.to_frame.select(
+    #     Polars.coalesce(
+    #       Polars.col("date").str.strptime(Polars::Date, "%F", strict: false),
+    #       Polars.col("date").str.strptime(Polars::Date, "%F %T", strict: false),
+    #       Polars.col("date").str.strptime(Polars::Date, "%D", strict: false),
+    #       Polars.col("date").str.strptime(Polars::Date, "%c", strict: false)
+    #     )
+    #   ).to_series
     #   # =>
-    #   # shape: (4, 1)
-    #   # ┌────────────┐
-    #   # │ date       │
-    #   # │ ---        │
-    #   # │ date       │
-    #   # ╞════════════╡
-    #   # │ 2021-04-22 │
-    #   # │ 2022-01-04 │
-    #   # │ 2022-01-31 │
-    #   # │ 2001-07-08 │
-    #   # └────────────┘
-    def strptime(datatype, fmt = nil, strict: true, exact: true, cache: true, tz_aware: false, utc: false)
-      if !Utils.is_polars_dtype(datatype)
-        raise ArgumentError, "expected: {DataType} got: #{datatype}"
-      end
-
-      if datatype == :date
+    #   # shape: (4,)
+    #   # Series: 'date' [date]
+    #   # [
+    #   #         2021-04-22
+    #   #         2022-01-04
+    #   #         2022-01-31
+    #   #         2001-07-08
+    #   # ]
+    def strptime(dtype, fmt = nil, strict: true, exact: true, cache: true, tz_aware: false, utc: false)
+      if dtype == Date
         Utils.wrap_expr(_rbexpr.str_parse_date(fmt, strict, exact, cache))
-      elsif datatype == :datetime
-        # TODO fix
-        tu = nil # datatype.tu
-        dtcol = Utils.wrap_expr(_rbexpr.str_parse_datetime(fmt, strict, exact, cache, tz_aware, utc))
-        if tu.nil?
-          dtcol
-        else
-          dtcol.dt.cast_time_unit(tu)
-        end
-      elsif datatype == :time
+      elsif dtype == Datetime || dtype.is_a?(Datetime)
+        dtype = Datetime.new if dtype == Datetime
+        time_unit = dtype.time_unit
+        time_zone = dtype.time_zone
+        Utils.wrap_expr(_rbexpr.str_parse_datetime(fmt, time_unit, time_zone, strict, exact, cache, tz_aware, utc))
+      elsif dtype == Time
         Utils.wrap_expr(_rbexpr.str_parse_time(fmt, strict, exact, cache))
       else
-        raise ArgumentError, "dtype should be of type :date, :datetime, or :time"
+        raise ArgumentError, "dtype should be of type {Date, Datetime, Time}"
       end
     end
 
@@ -521,6 +520,40 @@ module Polars
       Utils.wrap_expr(_rbexpr.str_starts_with(sub))
     end
 
+    # Parse string values as JSON.
+    #
+    # Throw errors if encounter invalid JSON strings.
+    #
+    # @param dtype [Object]
+    #   The dtype to cast the extracted value to. If nil, the dtype will be
+    #   inferred from the JSON value.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"json" => ['{"a":1, "b": true}', nil, '{"a":2, "b": false}']}
+    #   )
+    #   dtype = Polars::Struct.new([Polars::Field.new("a", Polars::Int64), Polars::Field.new("b", Polars::Boolean)])
+    #   df.select(Polars.col("json").str.json_extract(dtype))
+    #   # =>
+    #   # shape: (3, 1)
+    #   # ┌─────────────┐
+    #   # │ json        │
+    #   # │ ---         │
+    #   # │ struct[2]   │
+    #   # ╞═════════════╡
+    #   # │ {1,true}    │
+    #   # │ {null,null} │
+    #   # │ {2,false}   │
+    #   # └─────────────┘
+    def json_extract(dtype = nil)
+      if !dtype.nil?
+        dtype = Utils.rb_type_to_dtype(dtype)
+      end
+      Utils.wrap_expr(_rbexpr.str_json_extract(dtype))
+    end
+
     # Extract the first match of json string with provided JSONPath expression.
     #
     # Throw errors if encounter invalid json strings.
@@ -846,10 +879,10 @@ module Polars
     #   # │ 1   ┆ 123ABC │
     #   # │ 2   ┆ abc456 │
     #   # └─────┴────────┘
-    def replace(pattern, value, literal: false)
+    def replace(pattern, value, literal: false, n: 1)
       pattern = Utils.expr_to_lit_or_expr(pattern, str_to_lit: true)
       value = Utils.expr_to_lit_or_expr(value, str_to_lit: true)
-      Utils.wrap_expr(_rbexpr.str_replace(pattern._rbexpr, value._rbexpr, literal))
+      Utils.wrap_expr(_rbexpr.str_replace_n(pattern._rbexpr, value._rbexpr, literal, n))
     end
 
     # Replace all matching regex/literal substrings with a new string value.
@@ -912,5 +945,78 @@ module Polars
     def slice(offset, length = nil)
       Utils.wrap_expr(_rbexpr.str_slice(offset, length))
     end
+
+    # Returns a column with a separate row for every string character.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"a": ["foo", "bar"]})
+    #   df.select(Polars.col("a").str.explode)
+    #   # =>
+    #   # shape: (6, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ str │
+    #   # ╞═════╡
+    #   # │ f   │
+    #   # │ o   │
+    #   # │ o   │
+    #   # │ b   │
+    #   # │ a   │
+    #   # │ r   │
+    #   # └─────┘
+    def explode
+      Utils.wrap_expr(_rbexpr.explode)
+    end
+
+    # Parse integers with base radix from strings.
+    #
+    # By default base 2. ParseError/Overflows become Nulls.
+    #
+    # @param radix [Integer]
+    #   Positive integer which is the base of the string we are parsing.
+    #   Default: 2.
+    # @param strict [Boolean]
+    #   Bool, Default=true will raise any ParseError or overflow as ComputeError.
+    #   False silently convert to Null.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"bin" => ["110", "101", "010", "invalid"]})
+    #   df.select(Polars.col("bin").str.parse_int(2, strict: false))
+    #   # =>
+    #   # shape: (4, 1)
+    #   # ┌──────┐
+    #   # │ bin  │
+    #   # │ ---  │
+    #   # │ i32  │
+    #   # ╞══════╡
+    #   # │ 6    │
+    #   # │ 5    │
+    #   # │ 2    │
+    #   # │ null │
+    #   # └──────┘
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"hex" => ["fa1e", "ff00", "cafe", nil]})
+    #   df.select(Polars.col("hex").str.parse_int(16, strict: true))
+    #   # =>
+    #   # shape: (4, 1)
+    #   # ┌───────┐
+    #   # │ hex   │
+    #   # │ ---   │
+    #   # │ i32   │
+    #   # ╞═══════╡
+    #   # │ 64030 │
+    #   # │ 65280 │
+    #   # │ 51966 │
+    #   # │ null  │
+    #   # └───────┘
+    def parse_int(radix = 2, strict: true)
+      Utils.wrap_expr(_rbexpr.str_parse_int(radix, strict))
+    end
   end
 end

data/lib/polars/string_name_space.rb

@@ -38,12 +38,12 @@ module Polars
     #   )
     #   s.to_frame.with_column(
     #     Polars.col("date")
-    #       .str.strptime(:date, "%F", strict: false)
+    #       .str.strptime(Polars::Date, "%F", strict: false)
     #       .fill_null(
-    #         Polars.col("date").str.strptime(:date, "%F %T", strict: false)
+    #         Polars.col("date").str.strptime(Polars::Date, "%F %T", strict: false)
     #       )
-    #       .fill_null(Polars.col("date").str.strptime(:date, "%D", strict: false))
-    #       .fill_null(Polars.col("date").str.strptime(:date, "%c", strict: false))
+    #       .fill_null(Polars.col("date").str.strptime(Polars::Date, "%D", strict: false))
+    #       .fill_null(Polars.col("date").str.strptime(Polars::Date, "%c", strict: false))
     #   )
     #   # =>
     #   # shape: (4, 1)