polars-df 0.4.0-x86_64-linux → 0.6.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,46 +123,130 @@ 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 :tu
+    attr_reader :time_unit, :time_zone
+    alias_method :tu, :time_unit
 
     def initialize(time_unit = "us", time_zone = nil)
-      @tu = time_unit || "us"
+      @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.
   class Duration < TemporalType
-    attr_reader :tu
+    attr_reader :time_unit
+    alias_method :tu, :time_unit
 
     def initialize(time_unit = "us")
-      @tu = time_unit
+      @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.
@@ -122,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
 
@@ -140,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

@@ -218,6 +218,25 @@ module Polars
       )
     end
 
+    # Create a naive Datetime from an existing Date/Datetime expression and a Time.
+    #
+    # If the underlying expression is a Datetime then its time component is replaced,
+    # and if it is a Date then a new Datetime is created by combining the two values.
+    #
+    # @param time [Object]
+    #   A Ruby time literal or Polars expression/column that resolves to a time.
+    # @param time_unit ["ns", "us", "ms"]
+    #   Unit of time.
+    #
+    # @return [Expr]
+    def combine(time, time_unit: "us")
+      unless time.is_a?(Time) || time.is_a?(Expr)
+        raise TypeError, "expected 'time' to be a Ruby time or Polars expression, found #{time}"
+      end
+      time = Utils.expr_to_lit_or_expr(time)
+      Utils.wrap_expr(_rbexpr.dt_combine(time._rbexpr, time_unit))
+    end
+
     # Format Date/datetime with a formatting rule.
     #
     # See [chrono strftime/strptime](https://docs.rs/chrono/latest/chrono/format/strftime/index.html).
@@ -270,6 +289,34 @@ module Polars
       Utils.wrap_expr(_rbexpr.year)
     end
 
+    # Determine whether the year of the underlying date is a leap year.
+    #
+    # Applies to Date and Datetime columns.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   start = DateTime.new(2000, 1, 1)
+    #   stop = DateTime.new(2002, 1, 1)
+    #   df = Polars::DataFrame.new(
+    #     {"date" => Polars.date_range(start, stop, "1y")}
+    #   )
+    #   df.select(Polars.col("date").dt.is_leap_year)
+    #   # =>
+    #   # shape: (3, 1)
+    #   # ┌───────┐
+    #   # │ date  │
+    #   # │ ---   │
+    #   # │ bool  │
+    #   # ╞═══════╡
+    #   # │ true  │
+    #   # │ false │
+    #   # │ false │
+    #   # └───────┘
+    def is_leap_year
+      Utils.wrap_expr(_rbexpr.dt_is_leap_year)
+    end
+
     # Extract ISO year from underlying Date representation.
     #
     # Applies to Date and Datetime columns.
@@ -550,6 +597,27 @@ module Polars
       Utils.wrap_expr(_rbexpr.ordinal_day)
     end
 
+    # Time
+    #
+    # @return [Expr]
+    def time
+      Utils.wrap_expr(_rbexpr.dt_time)
+    end
+
+    # Date
+    #
+    # @return [Expr]
+    def date
+      Utils.wrap_expr(_rbexpr.dt_date)
+    end
+
+    # Datetime
+    #
+    # @return [Expr]
+    def datetime
+      Utils.wrap_expr(_rbexpr.dt_datetime)
+    end
+
     # Extract hour from underlying DateTime representation.
     #
     # Applies to Datetime columns.
@@ -958,8 +1026,8 @@ module Polars
     #   Time zone for the `Datetime` Series.
     #
     # @return [Expr]
-    def replace_time_zone(tz)
-      Utils.wrap_expr(_rbexpr.dt_replace_time_zone(tz))
+    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.
@@ -1282,5 +1350,77 @@ module Polars
     def offset_by(by)
       Utils.wrap_expr(_rbexpr.dt_offset_by(by))
     end
+
+    # Roll backward to the first day of the month.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "dates" => Polars.date_range(
+    #         DateTime.new(2000, 1, 15, 2),
+    #         DateTime.new(2000, 12, 15, 2),
+    #         "1mo"
+    #       )
+    #     }
+    #   )
+    #   df.select(Polars.col("dates").dt.month_start)
+    #   # =>
+    #   # shape: (12, 1)
+    #   # ┌─────────────────────┐
+    #   # │ dates               │
+    #   # │ ---                 │
+    #   # │ datetime[μs]        │
+    #   # ╞═════════════════════╡
+    #   # │ 2000-01-01 02:00:00 │
+    #   # │ 2000-02-01 02:00:00 │
+    #   # │ 2000-03-01 02:00:00 │
+    #   # │ 2000-04-01 02:00:00 │
+    #   # │ …                   │
+    #   # │ 2000-09-01 02:00:00 │
+    #   # │ 2000-10-01 02:00:00 │
+    #   # │ 2000-11-01 02:00:00 │
+    #   # │ 2000-12-01 02:00:00 │
+    #   # └─────────────────────┘
+    def month_start
+      Utils.wrap_expr(_rbexpr.dt_month_start)
+    end
+
+    # Roll forward to the last day of the month.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "dates" => Polars.date_range(
+    #         DateTime.new(2000, 1, 15, 2),
+    #         DateTime.new(2000, 12, 15, 2),
+    #         "1mo"
+    #       )
+    #     }
+    #   )
+    #   df.select(Polars.col("dates").dt.month_end)
+    #   # =>
+    #   # shape: (12, 1)
+    #   # ┌─────────────────────┐
+    #   # │ dates               │
+    #   # │ ---                 │
+    #   # │ datetime[μs]        │
+    #   # ╞═════════════════════╡
+    #   # │ 2000-01-31 02:00:00 │
+    #   # │ 2000-02-29 02:00:00 │
+    #   # │ 2000-03-31 02:00:00 │
+    #   # │ 2000-04-30 02:00:00 │
+    #   # │ …                   │
+    #   # │ 2000-09-30 02:00:00 │
+    #   # │ 2000-10-31 02:00:00 │
+    #   # │ 2000-11-30 02:00:00 │
+    #   # │ 2000-12-31 02:00:00 │
+    #   # └─────────────────────┘
+    def month_end
+      Utils.wrap_expr(_rbexpr.dt_month_end)
+    end
   end
 end

data/lib/polars/date_time_name_space.rb

@@ -82,8 +82,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 +114,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.