polars-df 0.5.0-x86_64-linux → 0.7.0-x86_64-linux

data/lib/polars/data_types.rb

@@ -1,6 +1,25 @@
 module Polars
   # Base class for all Polars data types.
   class DataType
+    def self.base_type
+      self
+    end
+
+    def base_type
+      is_a?(DataType) ? self.class : self
+    end
+
+    def self.nested?
+      false
+    end
+
+    def nested?
+      self.class.nested?
+    end
+
+    def self.==(other)
+      eql?(other) || other.is_a?(self)
+    end
   end
 
   # Base class for numeric data types.
@@ -15,12 +34,19 @@ module Polars
   class FractionalType < NumericType
   end
 
+  # Base class for float data types.
+  class FloatType < FractionalType
+  end
+
   # Base class for temporal data types.
   class TemporalType < DataType
   end
 
   # Base class for nested data types.
   class NestedType < DataType
+    def self.nested?
+      true
+    end
   end
 
   # 8-bit signed integer type.
@@ -56,11 +82,37 @@ module Polars
   end
 
   # 32-bit floating point type.
-  class Float32 < FractionalType
+  class Float32 < FloatType
   end
 
   # 64-bit floating point type.
-  class Float64 < FractionalType
+  class Float64 < FloatType
+  end
+
+  # Decimal 128-bit type with an optional precision and non-negative scale.
+  #
+  # NOTE: this is an experimental work-in-progress feature and may not work as expected.
+  class Decimal < FractionalType
+    attr_reader :precision, :scale
+
+    def initialize(precision, scale)
+      @precision = precision
+      @scale = scale
+    end
+
+    def ==(other)
+      if other.eql?(Decimal)
+        true
+      elsif other.is_a?(Decimal)
+        precision == other.precision && scale == other.scale
+      else
+        false
+      end
+    end
+
+    def to_s
+      "#{self.class.name}(precision: #{precision.inspect}, scale: #{scale.inspect})"
+    end
   end
 
   # Boolean type.
@@ -71,17 +123,18 @@ module Polars
   class Utf8 < DataType
   end
 
-  # Nested list/array type.
-  class List < NestedType
-    def initialize(inner)
-      @inner = Utils.rb_type_to_dtype(inner)
-    end
+  # Binary type.
+  class Binary < DataType
   end
 
   # Calendar date type.
   class Date < TemporalType
   end
 
+  # Time of day type.
+  class Time < TemporalType
+  end
+
   # Calendar date and time type.
   class Datetime < TemporalType
     attr_reader :time_unit, :time_zone
@@ -91,6 +144,20 @@ module Polars
       @time_unit = time_unit || "us"
       @time_zone = time_zone
     end
+
+    def ==(other)
+      if other.eql?(Datetime)
+        true
+      elsif other.is_a?(Datetime)
+        time_unit == other.time_unit && time_zone == other.time_zone
+      else
+        false
+      end
+    end
+
+    def to_s
+      "#{self.class.name}(time_unit: #{time_unit.inspect}, time_zone: #{time_zone.inspect})"
+    end
   end
 
   # Time duration/delta type.
@@ -101,18 +168,85 @@ module Polars
     def initialize(time_unit = "us")
       @time_unit = time_unit
     end
+
+    def ==(other)
+      if other.eql?(Duration)
+        true
+      elsif other.is_a?(Duration)
+        time_unit == other.time_unit
+      else
+        false
+      end
+    end
+
+    def to_s
+      "#{self.class.name}(time_unit: #{time_unit.inspect})"
+    end
   end
 
-  # Time of day type.
-  class Time < TemporalType
+  # A categorical encoding of a set of strings.
+  class Categorical < DataType
   end
 
   # Type for wrapping arbitrary Ruby objects.
   class Object < DataType
   end
 
-  # A categorical encoding of a set of strings.
-  class Categorical < DataType
+  # Type representing Null / None values.
+  class Null < DataType
+  end
+
+  # Type representing Datatype values that could not be determined statically.
+  class Unknown < DataType
+  end
+
+  # Nested list/array type.
+  class List < NestedType
+    attr_reader :inner
+
+    def initialize(inner)
+      @inner = Utils.rb_type_to_dtype(inner)
+    end
+
+    def ==(other)
+      if other.eql?(List)
+        true
+      elsif other.is_a?(List)
+        @inner.nil? || other.inner.nil? || @inner == other.inner
+      else
+        false
+      end
+    end
+
+    def to_s
+      "#{self.class.name}(#{inner})"
+    end
+  end
+
+  # Nested list/array type.
+  class Array < NestedType
+    attr_reader :width, :inner
+
+    def initialize(width, inner = nil)
+      @width = width
+      @inner = Utils.rb_type_to_dtype(inner) if inner
+    end
+
+    # TODO check width?
+    def ==(other)
+      if other.eql?(Array)
+        true
+      elsif other.is_a?(Array)
+        @inner.nil? || other.inner.nil? || @inner == other.inner
+      else
+        false
+      end
+    end
+
+    # TODO add width?
+    def to_s
+      "#{self.class.name}(#{inner})"
+    end
   end
 
   # Definition of a single field within a `Struct` DataType.
@@ -124,9 +258,12 @@ module Polars
       @dtype = Utils.rb_type_to_dtype(dtype)
     end
 
-    def inspect
-      class_name = self.class.name
-      "#{class_name}(#{@name}: #{@dtype})"
+    def ==(other)
+      name == other.name && dtype == other.dtype
+    end
+
+    def to_s
+      "#{self.class.name}(#{name.inspect}, #{dtype})"
     end
   end
 
@@ -142,25 +279,22 @@ module Polars
       end
     end
 
-    def inspect
-      class_name = self.class.name
-      "#{class_name}(#{@fields})"
+    def ==(other)
+      if other.eql?(Struct)
+        true
+      elsif other.is_a?(Struct)
+        fields == other.fields
+      else
+        false
+      end
+    end
+
+    def to_s
+      "#{self.class.name}([#{fields.map(&:to_s).join("\n")}])"
     end
 
     def to_schema
       @fields.to_h { |f| [f.name, f.dtype] }
     end
   end
-
-  # Binary type.
-  class Binary < DataType
-  end
-
-  # Type representing Null / None values.
-  class Null < DataType
-  end
-
-  # Type representing Datatype values that could not be determined statically.
-  class Unknown < DataType
-  end
 end

data/lib/polars/date_time_expr.rb

@@ -97,15 +97,20 @@ module Polars
     #   # │ 2001-01-01 00:50:00 ┆ 2001-01-01 00:30:00 │
     #   # │ 2001-01-01 01:00:00 ┆ 2001-01-01 01:00:00 │
     #   # └─────────────────────┴─────────────────────┘
-    def truncate(every, offset: nil)
+    def truncate(every, offset: nil, use_earliest: nil)
       if offset.nil?
         offset = "0ns"
       end
 
+      if !every.is_a?(Expr)
+        every = Utils._timedelta_to_pl_duration(every)
+      end
+      every = Utils.parse_as_expression(every, str_as_lit: true)
+
       Utils.wrap_expr(
         _rbexpr.dt_truncate(
-          Utils._timedelta_to_pl_duration(every),
-          Utils._timedelta_to_pl_duration(offset)
+          every,
+          Utils._timedelta_to_pl_duration(offset),
         )
       )
     end
@@ -1026,21 +1031,10 @@ module Polars
     #   Time zone for the `Datetime` Series.
     #
     # @return [Expr]
-    def replace_time_zone(tz, use_earliest: nil)
-      Utils.wrap_expr(_rbexpr.dt_replace_time_zone(tz, use_earliest))
-    end
-
-    # Localize tz-naive Datetime Series to tz-aware Datetime Series.
-    #
-    # This method takes a naive Datetime Series and makes this time zone aware.
-    # It does not move the time to another time zone.
-    #
-    # @param tz [String]
-    #   Time zone for the `Datetime` Series.
-    #
-    # @return [Expr]
-    def tz_localize(tz)
-      Utils.wrap_expr(_rbexpr.dt_tz_localize(tz))
+    def replace_time_zone(tz, use_earliest: nil, ambiguous: "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))
     end
 
     # Extract the days from a Duration type.
@@ -1348,6 +1342,7 @@ module Polars
     #   # │ 2006-01-01 00:00:00 ┆ 2003-11-01 00:00:00 │
     #   # └─────────────────────┴─────────────────────┘
     def offset_by(by)
+      by = Utils.parse_as_expression(by, str_as_lit: true)
       Utils.wrap_expr(_rbexpr.dt_offset_by(by))
     end
 

data/lib/polars/date_time_name_space.rb

@@ -23,18 +23,8 @@ module Polars
     # @return [Object]
     #
     # @example
-    #   date = Polars.date_range(DateTime.new(2001, 1, 1), DateTime.new(2001, 1, 3), "1d")
-    #   # =>
-    #   # shape: (3,)
-    #   # Series: '' [datetime[μs]]
-    #   # [
-    #   #         2001-01-01 00:00:00
-    #   #         2001-01-02 00:00:00
-    #   #         2001-01-03 00:00:00
-    #   # ]
-    #
-    # @example
-    #   date.dt.min
+    #   s = Polars.date_range(DateTime.new(2001, 1, 1), DateTime.new(2001, 1, 3), "1d")
+    #   s.dt.min
     #   # => 2001-01-01 00:00:00 UTC
     def min
       Utils.wrap_s(_s).min
@@ -45,18 +35,8 @@ module Polars
     # @return [Object]
     #
     # @example
-    #   date = Polars.date_range(DateTime.new(2001, 1, 1), DateTime.new(2001, 1, 3), "1d")
-    #   # =>
-    #   # shape: (3,)
-    #   # Series: '' [datetime[μs]]
-    #   # [
-    #   #         2001-01-01 00:00:00
-    #   #         2001-01-02 00:00:00
-    #   #         2001-01-03 00:00:00
-    #   # ]
-    #
-    # @example
-    #   date.dt.max
+    #   s = Polars.date_range(DateTime.new(2001, 1, 1), DateTime.new(2001, 1, 3), "1d")
+    #   s.dt.max
     #   # => 2001-01-03 00:00:00 UTC
     def max
       Utils.wrap_s(_s).max
@@ -82,8 +62,15 @@ module Polars
     #   # => 2001-01-02 00:00:00 UTC
     def median
       s = Utils.wrap_s(_s)
-      out = s.median.to_i
-      Utils._to_ruby_datetime(out, s.dtype, tu: s.time_unit)
+      out = s.median
+      if !out.nil?
+        if s.dtype == Date
+          return Utils._to_ruby_date(out.to_i)
+        else
+          return Utils._to_ruby_datetime(out.to_i, s.time_unit)
+        end
+      end
+      nil
     end
 
     # Return mean as Ruby object.
@@ -107,7 +94,14 @@ module Polars
     def mean
       s = Utils.wrap_s(_s)
       out = s.mean.to_i
-      Utils._to_ruby_datetime(out, s.dtype, tu: s.time_unit)
+      if !out.nil?
+        if s.dtype == Date
+          return Utils._to_ruby_date(out.to_i)
+        else
+          return Utils._to_ruby_datetime(out.to_i, s.time_unit)
+        end
+      end
+      nil
     end
 
     # Format Date/datetime with a formatting rule.
@@ -1386,7 +1380,7 @@ module Polars
     #   #         2001-01-01 00:30:00
     #   #         2001-01-01 01:00:00
     #   # ]
-    def truncate(every, offset: nil)
+    def truncate(every, offset: nil, use_earliest: nil)
       super
     end
 

data/lib/polars/dynamic_group_by.rb

@@ -2,7 +2,7 @@ module Polars
   # A dynamic grouper.
   #
   # This has an `.agg` method which allows you to run all polars expressions in a
-  # groupby context.
+  # group by context.
   class DynamicGroupBy
     def initialize(
       df,
@@ -34,7 +34,7 @@ module Polars
 
     def agg(aggs)
       @df.lazy
-        .groupby_dynamic(
+        .group_by_dynamic(
           @time_column,
           every: @every,
           period: @period,