polars-df 0.21.0 → 0.21.1

data/lib/polars/data_type_expr.rb

@@ -0,0 +1,52 @@
+module Polars
+  # A lazily instantiated `DataType` that can be used in an `Expr`.
+  class DataTypeExpr
+    # @private
+    attr_accessor :_rbdatatype_expr
+
+    # @private
+    def self._from_rbdatatype_expr(rbdatatype_expr)
+      slf = new
+      slf._rbdatatype_expr = rbdatatype_expr
+      slf
+    end
+
+    # Materialize the `DataTypeExpr` in a specific context.
+    #
+    # This is a useful function when debugging datatype expressions.
+    #
+    # @return [DataType]
+    #
+    # @example
+    #   lf = Polars::LazyFrame.new(
+    #     {
+    #       "a" => [1, 2, 3]
+    #     }
+    #   )
+    #   Polars.dtype_of("a").collect_dtype(lf)
+    #   # => Polars::Int64
+    #
+    # @example
+    #   Polars.dtype_of("a").collect_dtype({"a" => Polars::String})
+    #   # => Polars::String
+    def collect_dtype(
+      context
+    )
+      schema = nil
+      if context.is_a?(Schema)
+         schema = context
+      elsif context.is_a?(Hash)
+        schema = Schema.new(context)
+      elsif context.is_a?(DataFrame)
+        schema = context.schema
+      elsif context.is_a?(LazyFrame)
+        schema = context.collect_schema
+      else
+        msg = "DataTypeExpr.collect_dtype did not expect #{context.inspect}"
+        raise TypeError, msg
+      end
+
+      _rbdatatype_expr.collect_dtype(schema.to_h)
+    end
+  end
+end

data/lib/polars/data_types.rb

@@ -99,7 +99,18 @@ module Polars
       self < NestedType
     end
 
-    [:numeric?, :decimal?, :integer?, :signed_integer?, :unsigned_integer?, :float?, :temporal?, :nested?].each do |v|
+    # Return a `DataTypeExpr` with a static `DataType`.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   Polars::Int16.new.to_dtype_expr.collect_dtype({})
+    #   # => Polars::Int16
+    def self.to_dtype_expr
+      DataTypeExpr._from_rbdatatype_expr(RbDataTypeExpr.from_dtype(self))
+    end
+
+    [:numeric?, :decimal?, :integer?, :signed_integer?, :unsigned_integer?, :float?, :temporal?, :nested?, :to_dtype_expr].each do |v|
       define_method(v) do
         self.class.public_send(v)
       end
@@ -317,6 +328,7 @@ module Polars
       raise Todo
     end
 
+    # @private
     def self._from_rb_categories(rb_categories)
       slf = new
       slf._categories = rb_categories
@@ -401,7 +413,7 @@ module Polars
   class Object < DataType
   end
 
-  # Type representing Null / None values.
+  # Type representing Null / nil values.
   class Null < DataType
   end
 

data/lib/polars/date_time_expr.rb

@@ -9,6 +9,57 @@ module Polars
       self._rbexpr = expr._rbexpr
     end
 
+    # Offset by `n` business days.
+    #
+    # @note
+    #   This functionality is considered **unstable**. It may be changed
+    #   at any point without it being considered a breaking change.
+    #
+    # @param n
+    #   Number of business days to offset by. Can be a single number of an
+    #   expression.
+    # @param week_mask
+    #   Which days of the week to count. The default is Monday to Friday.
+    #   If you wanted to count only Monday to Thursday, you would pass
+    #   `[true, true, true, true, false, false, false]`.
+    # @param roll
+    #   What to do when the start date lands on a non-business day. Options are:
+    #
+    #   - `'raise'`: raise an error
+    #   - `'forward'`: move to the next business day
+    #   - `'backward'`: move to the previous business day
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"start" => [Date.new(2020, 1, 1), Date.new(2020, 1, 2)]})
+    #   df.with_columns(result: Polars.col("start").dt.add_business_days(5))
+    #   # =>
+    #   # shape: (2, 2)
+    #   # ┌────────────┬────────────┐
+    #   # │ start      ┆ result     │
+    #   # │ ---        ┆ ---        │
+    #   # │ date       ┆ date       │
+    #   # ╞════════════╪════════════╡
+    #   # │ 2020-01-01 ┆ 2020-01-08 │
+    #   # │ 2020-01-02 ┆ 2020-01-09 │
+    #   # └────────────┴────────────┘
+    def add_business_days(
+      n,
+      week_mask: [true, true, true, true, true, false, false],
+      roll: "raise"
+    )
+      n_rbexpr = Utils.parse_into_expression(n)
+      Utils.wrap_expr(
+        _rbexpr.dt_add_business_days(
+          n_rbexpr,
+          week_mask,
+          [],
+          roll
+        )
+      )
+    end
+
     # Divide the date/datetime range into buckets.
     #
     # Each date/datetime is mapped to the start of its bucket.
@@ -203,6 +254,93 @@ module Polars
       Utils.wrap_expr(_rbexpr.dt_round(every))
     end
 
+    # Replace time unit.
+    #
+    # @param year [Object]
+    #   Column or literal.
+    # @param month [Object]
+    #   Column or literal, ranging from 1-12.
+    # @param day [Object]
+    #   Column or literal, ranging from 1-31.
+    # @param hour [Object]
+    #   Column or literal, ranging from 0-23.
+    # @param minute [Object]
+    #   Column or literal, ranging from 0-59.
+    # @param second [Object]
+    #   Column or literal, ranging from 0-59.
+    # @param microsecond [Object]
+    #   Column or literal, ranging from 0-999999.
+    # @param ambiguous [String]
+    #   Determine how to deal with ambiguous datetimes:
+    #
+    #   - `'raise'` (default): raise
+    #   - `'earliest'`: use the earliest datetime
+    #   - `'latest'`: use the latest datetime
+    #   - `'null'`: set to null
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "date" => [Date.new(2024, 4, 1), Date.new(2025, 3, 16)],
+    #       "new_day" => [10, 15]
+    #     }
+    #   )
+    #   df.with_columns(Polars.col("date").dt.replace(day: "new_day").alias("replaced"))
+    #   # =>
+    #   # shape: (2, 3)
+    #   # ┌────────────┬─────────┬────────────┐
+    #   # │ date       ┆ new_day ┆ replaced   │
+    #   # │ ---        ┆ ---     ┆ ---        │
+    #   # │ date       ┆ i64     ┆ date       │
+    #   # ╞════════════╪═════════╪════════════╡
+    #   # │ 2024-04-01 ┆ 10      ┆ 2024-04-10 │
+    #   # │ 2025-03-16 ┆ 15      ┆ 2025-03-15 │
+    #   # └────────────┴─────────┴────────────┘
+    #
+    # @example
+    #   df.with_columns(Polars.col("date").dt.replace(year: 1800).alias("replaced"))
+    #   # =>
+    #   # shape: (2, 3)
+    #   # ┌────────────┬─────────┬────────────┐
+    #   # │ date       ┆ new_day ┆ replaced   │
+    #   # │ ---        ┆ ---     ┆ ---        │
+    #   # │ date       ┆ i64     ┆ date       │
+    #   # ╞════════════╪═════════╪════════════╡
+    #   # │ 2024-04-01 ┆ 10      ┆ 1800-04-01 │
+    #   # │ 2025-03-16 ┆ 15      ┆ 1800-03-16 │
+    #   # └────────────┴─────────┴────────────┘
+    def replace(
+      year: nil,
+      month: nil,
+      day: nil,
+      hour: nil,
+      minute: nil,
+      second: nil,
+      microsecond: nil,
+      ambiguous: "raise"
+    )
+      day, month, year, hour, minute, second, microsecond = (
+        Utils.parse_into_list_of_expressions(
+          day, month, year, hour, minute, second, microsecond
+        )
+      )
+      ambiguous_expr = Utils.parse_into_expression(ambiguous, str_as_lit: true)
+      Utils.wrap_expr(
+        _rbexpr.dt_replace(
+          year,
+          month,
+          day,
+          hour,
+          minute,
+          second,
+          microsecond,
+          ambiguous_expr
+        )
+      )
+    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,
@@ -317,6 +455,82 @@ module Polars
       Utils.wrap_expr(_rbexpr.strftime(fmt))
     end
 
+    # Extract the millennium from underlying representation.
+    #
+    # Applies to Date and Datetime columns.
+    #
+    # Returns the millennium number in the calendar date.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "date" => [
+    #         Date.new(999, 12, 31),
+    #         Date.new(1897, 5, 7),
+    #         Date.new(2000, 1, 1),
+    #         Date.new(2001, 7, 5),
+    #         Date.new(3002, 10, 20)
+    #       ]
+    #     }
+    #   )
+    #   df.with_columns(mlnm: Polars.col("date").dt.millennium)
+    #   # =>
+    #   # shape: (5, 2)
+    #   # ┌────────────┬──────┐
+    #   # │ date       ┆ mlnm │
+    #   # │ ---        ┆ ---  │
+    #   # │ date       ┆ i32  │
+    #   # ╞════════════╪══════╡
+    #   # │ 0999-12-31 ┆ 1    │
+    #   # │ 1897-05-07 ┆ 2    │
+    #   # │ 2000-01-01 ┆ 2    │
+    #   # │ 2001-07-05 ┆ 3    │
+    #   # │ 3002-10-20 ┆ 4    │
+    #   # └────────────┴──────┘
+    def millennium
+      Utils.wrap_expr(_rbexpr.dt_millennium)
+    end
+
+    # Extract the century from underlying representation.
+    #
+    # Applies to Date and Datetime columns.
+    #
+    # Returns the century number in the calendar date.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "date" => [
+    #         Date.new(999, 12, 31),
+    #         Date.new(1897, 5, 7),
+    #         Date.new(2000, 1, 1),
+    #         Date.new(2001, 7, 5),
+    #         Date.new(3002, 10, 20)
+    #       ]
+    #     }
+    #   )
+    #   df.with_columns(cent: Polars.col("date").dt.century)
+    #   # =>
+    #   # shape: (5, 2)
+    #   # ┌────────────┬──────┐
+    #   # │ date       ┆ cent │
+    #   # │ ---        ┆ ---  │
+    #   # │ date       ┆ i32  │
+    #   # ╞════════════╪══════╡
+    #   # │ 0999-12-31 ┆ 10   │
+    #   # │ 1897-05-07 ┆ 19   │
+    #   # │ 2000-01-01 ┆ 20   │
+    #   # │ 2001-07-05 ┆ 21   │
+    #   # │ 3002-10-20 ┆ 31   │
+    #   # └────────────┴──────┘
+    def century
+      Utils.wrap_expr(_rbexpr.dt_century)
+    end
+
     # Extract year from underlying Date representation.
     #
     # Applies to Date and Datetime columns.
@@ -348,6 +562,43 @@ module Polars
       Utils.wrap_expr(_rbexpr.dt_year)
     end
 
+    # Determine whether each day lands on a business day.
+    #
+    # @note
+    #   This functionality is considered **unstable**. It may be changed
+    #   at any point without it being considered a breaking change.
+    #
+    # @param week_mask [Array]
+    #   Which days of the week to count. The default is Monday to Friday.
+    #   If you wanted to count only Monday to Thursday, you would pass
+    #   `[true, true, true, true, false, false, false]`.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"start" => [Date.new(2020, 1, 3), Date.new(2020, 1, 5)]})
+    #   df.with_columns(is_business_day: Polars.col("start").dt.is_business_day)
+    #   # =>
+    #   # shape: (2, 2)
+    #   # ┌────────────┬─────────────────┐
+    #   # │ start      ┆ is_business_day │
+    #   # │ ---        ┆ ---             │
+    #   # │ date       ┆ bool            │
+    #   # ╞════════════╪═════════════════╡
+    #   # │ 2020-01-03 ┆ true            │
+    #   # │ 2020-01-05 ┆ false           │
+    #   # └────────────┴─────────────────┘
+    def is_business_day(
+      week_mask: [true, true, true, true, true, false, false]
+    )
+      Utils.wrap_expr(
+        _rbexpr.dt_is_business_day(
+          week_mask,
+          []
+        )
+      )
+    end
+
     # Determine whether the year of the underlying date is a leap year.
     #
     # Applies to Date and Datetime columns.

data/lib/polars/date_time_name_space.rb

@@ -18,6 +18,68 @@ module Polars
       s[item]
     end
 
+    # Offset by `n` business days.
+    #
+    # @note
+    #   This functionality is considered **unstable**. It may be changed
+    #   at any point without it being considered a breaking change.
+    #
+    # @param n [Object]
+    #   Number of business days to offset by. Can be a single number of an
+    #   expression.
+    # @param week_mask [Array]
+    #   Which days of the week to count. The default is Monday to Friday.
+    #   If you wanted to count only Monday to Thursday, you would pass
+    #   `[true, true, true, true, false, false, false]`.
+    # roll
+    #   What to do when the start date lands on a non-business day. Options are:
+    #
+    #   - `'raise'`: raise an error
+    #   - `'forward'`: move to the next business day
+    #   - `'backward'`: move to the previous business day
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("start", [Date.new(2020, 1, 1), Date.new(2020, 1, 2)])
+    #   s.dt.add_business_days(5)
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: 'start' [date]
+    #   # [
+    #   #         2020-01-08
+    #   #         2020-01-09
+    #   # ]
+    #
+    # @example You can pass a custom weekend - for example, if you only take Sunday off:
+    #   week_mask = [true, true, true, true, true, true, false]
+    #   s.dt.add_business_days(5, week_mask: week_mask)
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: 'start' [date]
+    #   # [
+    #   #         2020-01-07
+    #   #         2020-01-08
+    #   # ]
+    #
+    # @example Roll all dates forwards to the next business day:
+    #   s = Polars::Series.new("start", [Date.new(2020, 1, 5), Date.new(2020, 1, 6)])
+    #   s.dt.add_business_days(0, roll: "forward")
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: 'start' [date]
+    #   # [
+    #   #         2020-01-06
+    #   #         2020-01-06
+    #   # ]
+    def add_business_days(
+      n,
+      week_mask: [true, true, true, true, true, false, false],
+      roll: "raise"
+    )
+      super
+    end
+
     # Return minimum as Ruby object.
     #
     # @return [Object]
@@ -139,6 +201,74 @@ module Polars
       super
     end
 
+    # Extract the millennium from underlying representation.
+    #
+    # Applies to Date and Datetime columns.
+    #
+    # Returns the millennium number in the calendar date.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new(
+    #     "dt",
+    #     [
+    #       Date.new(999, 12, 31),
+    #       Date.new(1897, 5, 7),
+    #       Date.new(2000, 1, 1),
+    #       Date.new(2001, 7, 5),
+    #       Date.new(3002, 10, 20)
+    #     ]
+    #   )
+    #   s.dt.millennium
+    #   # =>
+    #   # shape: (5,)
+    #   # Series: 'dt' [i32]
+    #   # [
+    #   #         1
+    #   #         2
+    #   #         2
+    #   #         3
+    #   #         4
+    #   # ]
+    def millennium
+      super
+    end
+
+    # Extract the century from underlying representation.
+    #
+    # Applies to Date and Datetime columns.
+    #
+    # Returns the century number in the calendar date.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new(
+    #     "dt",
+    #     [
+    #       Date.new(999, 12, 31),
+    #       Date.new(1897, 5, 7),
+    #       Date.new(2000, 1, 1),
+    #       Date.new(2001, 7, 5),
+    #       Date.new(3002, 10, 20)
+    #     ]
+    #   )
+    #   s.dt.century
+    #   # =>
+    #   # shape: (5,)
+    #   # Series: 'dt' [i32]
+    #   # [
+    #   #         10
+    #   #         19
+    #   #         20
+    #   #         21
+    #   #         31
+    #   # ]
+    def century
+      super
+    end
+
     # Extract the year from the underlying date representation.
     #
     # Applies to Date and Datetime columns.
@@ -161,6 +291,69 @@ module Polars
       super
     end
 
+    # Determine whether each day lands on a business day.
+    #
+    # @note
+    #   This functionality is considered **unstable**. It may be changed
+    #   at any point without it being considered a breaking change.
+    #
+    # @param week_mask [Array]
+    #   Which days of the week to count. The default is Monday to Friday.
+    #   If you wanted to count only Monday to Thursday, you would pass
+    #   `(True, True, True, True, False, False, False)`.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new([Date.new(2020, 1, 3), Date.new(2020, 1, 5)])
+    #   s.dt.is_business_day
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: '' [bool]
+    #   # [
+    #   #         true
+    #   #         false
+    #   # ]
+    #
+    # @example You can pass a custom weekend - for example, if you only take Sunday off:
+    #   week_mask = [true, true, true, true, true, true, false]
+    #   s.dt.is_business_day(week_mask: week_mask)
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: '' [bool]
+    #   # [
+    #   #         true
+    #   #         false
+    #   # ]
+    def is_business_day(
+      week_mask: [true, true, true, true, true, false, false]
+    )
+      super
+    end
+
+    # Determine whether the year of the underlying date representation is a leap year.
+    #
+    # Applies to Date and Datetime columns.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new(
+    #     "date", [Date.new(2000, 1, 1), Date.new(2001, 1, 1), Date.new(2002, 1, 1)]
+    #   )
+    #   s.dt.is_leap_year
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'date' [bool]
+    #   # [
+    #   #         true
+    #   #         false
+    #   #         false
+    #   # ]
+    def is_leap_year
+      super
+    end
+
     # Extract ISO year from underlying Date representation.
     #
     # Applies to Date and Datetime columns.
@@ -346,6 +539,48 @@ module Polars
       super
     end
 
+    # Extract (local) time.
+    #
+    # Applies to Date/Datetime/Time columns.
+    #
+    # @return
+    #
+    # @example
+    #   ser = Polars::Series.new([DateTime.new(2021, 1, 2, 5)]).dt.replace_time_zone(
+    #     "Asia/Kathmandu"
+    #   )
+    #   ser.dt.time
+    #   # =>
+    #   # shape: (1,)
+    #   # Series: '' [time]
+    #   # [
+    #   #         05:00:00
+    #   # ]
+    def time
+      super
+    end
+
+    # Extract (local) date.
+    #
+    # Applies to Date/Datetime columns.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   ser = Polars::Series.new([DateTime.new(2021, 1, 2, 5)]).dt.replace_time_zone(
+    #     "Asia/Kathmandu"
+    #   )
+    #   ser.dt.date
+    #   # =>
+    #   # shape: (1,)
+    #   # Series: '' [date]
+    #   # [
+    #   #         2021-01-02
+    #   # ]
+    def date
+      super
+    end
+
     # Extract the hour from the underlying DateTime representation.
     #
     # Applies to Datetime columns.
@@ -1276,6 +1511,21 @@ module Polars
       super
     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 Series of the same length as this Series.
+    # @param time_unit ['ns', 'us', 'ms']
+    #   Unit of time.
+    #
+    # @return [Series]
+    def combine(time, time_unit: "us")
+      super
+    end
+
     # Roll backward to the first day of the month.
     #
     # @return [Series]
@@ -1370,5 +1620,54 @@ module Polars
     def dst_offset
       super
     end
+
+    # Replace time unit.
+    #
+    # @param year [Object]
+    #   Literal or Series.
+    # @param month [Object]
+    #   Literal or Series, ranging from 1-12.
+    # @param day [Object]
+    #   Literal or Series, ranging from 1-31.
+    # @param hour [Object]
+    #   Literal or Series, ranging from 0-23.
+    # @param minute [Object]
+    #   Literal or Series, ranging from 0-59.
+    # @param second [Object]
+    #   Literal or Series, ranging from 0-59.
+    # @param microsecond [Object]
+    #   Literal or Series, ranging from 0-999999.
+    # @param ambiguous [String]
+    #   Determine how to deal with ambiguous datetimes:
+    #
+    #   - `'raise'` (default): raise
+    #   - `'earliest'`: use the earliest datetime
+    #   - `'latest'`: use the latest datetime
+    #   - `'null'`: set to null
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("date", [Date.new(2013, 1, 1), Date.new(2024, 1, 2)])
+    #   s.dt.replace(year: 1800)
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: 'date' [date]
+    #   # [
+    #   #         1800-01-01
+    #   #         1800-01-02
+    #   # ]
+    def replace(
+      year: nil,
+      month: nil,
+      day: nil,
+      hour: nil,
+      minute: nil,
+      second: nil,
+      microsecond: nil,
+      ambiguous: "raise"
+    )
+      super
+    end
   end
 end