polars-df 0.8.0 → 0.9.0

data/lib/polars/functions/len.rb

@@ -0,0 +1,49 @@
+module Polars
+  module Functions
+    # Return the number of rows in the context.
+    #
+    # This is similar to `COUNT(*)` in SQL.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "a" => [1, 2, nil],
+    #       "b" => [3, nil, nil],
+    #       "c" => ["foo", "bar", "foo"]
+    #     }
+    #   )
+    #   df.select(Polars.len)
+    #   # =>
+    #   # shape: (1, 1)
+    #   # ┌─────┐
+    #   # │ len │
+    #   # │ --- │
+    #   # │ u32 │
+    #   # ╞═════╡
+    #   # │ 3   │
+    #   # └─────┘
+    #
+    # @example Generate an index column by using `len` in conjunction with `int_range`.
+    #   df.select([
+    #     Polars.int_range(Polars.len, dtype: Polars::UInt32).alias("index"),
+    #     Polars.all
+    #   ])
+    #   # =>
+    #   # shape: (3, 4)
+    #   # ┌───────┬──────┬──────┬─────┐
+    #   # │ index ┆ a    ┆ b    ┆ c   │
+    #   # │ ---   ┆ ---  ┆ ---  ┆ --- │
+    #   # │ u32   ┆ i64  ┆ i64  ┆ str │
+    #   # ╞═══════╪══════╪══════╪═════╡
+    #   # │ 0     ┆ 1    ┆ 3    ┆ foo │
+    #   # │ 1     ┆ 2    ┆ null ┆ bar │
+    #   # │ 2     ┆ null ┆ null ┆ foo │
+    #   # └───────┴──────┴──────┴─────┘
+    def len
+      Utils.wrap_expr(Plr.len)
+    end
+    alias_method :length, :len
+  end
+end

data/lib/polars/functions/lit.rb

@@ -0,0 +1,35 @@
+module Polars
+  module Functions
+    # Return an expression representing a literal value.
+    #
+    # @return [Expr]
+    def lit(value, dtype: nil, allow_object: nil)
+      if value.is_a?(::Time) || value.is_a?(::DateTime)
+        time_unit = dtype&.time_unit || "ns"
+        time_zone = dtype.&time_zone
+        e = lit(Utils._datetime_to_pl_timestamp(value, time_unit)).cast(Datetime.new(time_unit))
+        if time_zone
+          return e.dt.replace_time_zone(time_zone.to_s)
+        else
+          return e
+        end
+      elsif value.is_a?(::Date)
+        return lit(::Time.utc(value.year, value.month, value.day)).cast(Date)
+      elsif value.is_a?(Polars::Series)
+        name = value.name
+        value = value._s
+        e = Utils.wrap_expr(Plr.lit(value, allow_object))
+        if name == ""
+          return e
+        end
+        return e.alias(name)
+      elsif (defined?(Numo::NArray) && value.is_a?(Numo::NArray)) || value.is_a?(::Array)
+        return lit(Series.new("", value))
+      elsif dtype
+        return Utils.wrap_expr(Plr.lit(value, allow_object)).cast(dtype)
+      end
+
+      Utils.wrap_expr(Plr.lit(value, allow_object))
+    end
+  end
+end

data/lib/polars/functions/random.rb

@@ -0,0 +1,16 @@
+module Polars
+  module Functions
+    # Set the global random seed for Polars.
+    #
+    # This random seed is used to determine things such as shuffle ordering.
+    #
+    # @param seed [Integer]
+    #   A non-negative integer < 2**64 used to seed the internal global
+    #   random number generator.
+    #
+    # @return [nil]
+    def set_random_seed(seed)
+      Plr.set_random_seed(seed)
+    end
+  end
+end

data/lib/polars/functions/range/date_range.rb

@@ -0,0 +1,103 @@
+module Polars
+  module Functions
+    # Create a range of type `Datetime` (or `Date`).
+    #
+    # @param start [Object]
+    #   Lower bound of the date range.
+    # @param stop [Object]
+    #   Upper bound of the date range.
+    # @param interval [Object]
+    #   Interval periods. It can be a polars duration string, such as `3d12h4m25s`
+    #   representing 3 days, 12 hours, 4 minutes, and 25 seconds.
+    # @param lazy [Boolean]
+    #   Return an expression.
+    # @param closed ["both", "left", "right", "none"]
+    #   Define whether the temporal window interval is closed or not.
+    # @param name [String]
+    #   Name of the output Series.
+    # @param time_unit [nil, "ns", "us", "ms"]
+    #   Set the time unit.
+    # @param time_zone [String]
+    #   Optional timezone
+    #
+    # @return [Object]
+    #
+    # @note
+    #   If both `low` and `high` are passed as date types (not datetime), and the
+    #   interval granularity is no finer than 1d, the returned range is also of
+    #   type date. All other permutations return a datetime Series.
+    #
+    # @example Using polars duration string to specify the interval
+    #   Polars.date_range(Date.new(2022, 1, 1), Date.new(2022, 3, 1), "1mo", name: "drange")
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'drange' [date]
+    #   # [
+    #   #         2022-01-01
+    #   #         2022-02-01
+    #   #         2022-03-01
+    #   # ]
+    #
+    # @example Using `timedelta` object to specify the interval:
+    #   Polars.date_range(
+    #       DateTime.new(1985, 1, 1),
+    #       DateTime.new(1985, 1, 10),
+    #       "1d12h",
+    #       time_unit: "ms"
+    #   )
+    #   # =>
+    #   # shape: (7,)
+    #   # Series: '' [datetime[ms]]
+    #   # [
+    #   #         1985-01-01 00:00:00
+    #   #         1985-01-02 12:00:00
+    #   #         1985-01-04 00:00:00
+    #   #         1985-01-05 12:00:00
+    #   #         1985-01-07 00:00:00
+    #   #         1985-01-08 12:00:00
+    #   #         1985-01-10 00:00:00
+    #   # ]
+    def date_range(
+      start,
+      stop,
+      interval,
+      lazy: false,
+      closed: "both",
+      name: nil,
+      time_unit: nil,
+      time_zone: nil
+    )
+      if defined?(ActiveSupport::Duration) && interval.is_a?(ActiveSupport::Duration)
+        raise Todo
+      else
+        interval = interval.to_s
+        if interval.include?(" ")
+          interval = interval.gsub(" ", "")
+        end
+      end
+
+      if time_unit.nil?
+        if interval.include?("ns")
+          time_unit = "ns"
+        else
+          time_unit = "us"
+        end
+      end
+
+      start_rbexpr = Utils.parse_as_expression(start)
+      stop_rbexpr = Utils.parse_as_expression(stop)
+
+      result = Utils.wrap_expr(
+        Plr.date_range(start_rbexpr, stop_rbexpr, interval, closed, time_unit, time_zone)
+      )
+
+      result = result.alias(name.to_s)
+
+      if !lazy
+        return select(result).to_series
+      end
+
+      result
+    end
+  end
+end

data/lib/polars/functions/range/int_range.rb

@@ -0,0 +1,51 @@
+module Polars
+  module Functions
+    # Create a range expression (or Series).
+    #
+    # This can be used in a `select`, `with_column`, etc. Be sure that the resulting
+    # range size is equal to the length of the DataFrame you are collecting.
+    #
+    # @param start [Integer, Expr, Series]
+    #   Lower bound of range.
+    # @param stop [Integer, Expr, Series]
+    #   Upper bound of range.
+    # @param step [Integer]
+    #   Step size of the range.
+    # @param eager [Boolean]
+    #   If eager evaluation is `True`, a Series is returned instead of an Expr.
+    # @param dtype [Symbol]
+    #   Apply an explicit integer dtype to the resulting expression (default is `Int64`).
+    #
+    # @return [Expr, Series]
+    #
+    # @example
+    #   Polars.arange(0, 3, eager: true)
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'arange' [i64]
+    #   # [
+    #   #         0
+    #   #         1
+    #   #         2
+    #   # ]
+    def int_range(start, stop = nil, step: 1, eager: false, dtype: nil)
+      if stop.nil?
+        stop = start
+        start = 0
+      end
+
+      start = Utils.parse_as_expression(start)
+      stop = Utils.parse_as_expression(stop)
+      dtype ||= Int64
+      dtype = dtype.to_s if dtype.is_a?(Symbol)
+      result = Utils.wrap_expr(Plr.int_range(start, stop, step, dtype)).alias("arange")
+
+      if eager
+        return select(result).to_series
+      end
+
+      result
+    end
+    alias_method :arange, :int_range
+  end
+end

data/lib/polars/functions/repeat.rb

@@ -0,0 +1,144 @@
+module Polars
+  module Functions
+    # Repeat a single value n times.
+    #
+    # @param value [Object]
+    #   Value to repeat.
+    # @param n [Integer]
+    #   Repeat `n` times.
+    # @param eager [Boolean]
+    #   Run eagerly and collect into a `Series`.
+    # @param name [String]
+    #   Only used in `eager` mode. As expression, use `alias`.
+    #
+    # @return [Object]
+    #
+    # @example Construct a column with a repeated value in a lazy context.
+    #   Polars.select(Polars.repeat("z", 3)).to_series
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'repeat' [str]
+    #   # [
+    #   #         "z"
+    #   #         "z"
+    #   #         "z"
+    #   # ]
+    #
+    # @example Generate a Series directly by setting `eager: true`.
+    #   Polars.repeat(3, 3, dtype: Polars::Int8, eager: true)
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'repeat' [i8]
+    #   # [
+    #   #         3
+    #   #         3
+    #   #         3
+    #   # ]
+    def repeat(value, n, dtype: nil, eager: false, name: nil)
+      if !name.nil?
+        warn "the `name` argument is deprecated. Use the `alias` method instead."
+      end
+
+      if n.is_a?(Integer)
+        n = lit(n)
+      end
+
+      value = Utils.parse_as_expression(value, str_as_lit: true)
+      expr = Utils.wrap_expr(Plr.repeat(value, n._rbexpr, dtype))
+      if !name.nil?
+        expr = expr.alias(name)
+      end
+      if eager
+        return select(expr).to_series
+      end
+      expr
+    end
+
+    # Construct a column of length `n` filled with ones.
+    #
+    # This is syntactic sugar for the `repeat` function.
+    #
+    # @param n [Integer]
+    #   Length of the resulting column.
+    # @param dtype [Object]
+    #   Data type of the resulting column. Defaults to Float64.
+    # @param eager [Boolean]
+    #   Evaluate immediately and return a `Series`. If set to `false`,
+    #   return an expression instead.
+    #
+    # @return [Object]
+    #
+    # @example
+    #   Polars.ones(3, dtype: Polars::Int8, eager: true)
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'ones' [i8]
+    #   # [
+    #   #         1
+    #   #         1
+    #   #         1
+    #   # ]
+    def ones(n, dtype: nil, eager: true)
+      if (zero = _one_or_zero_by_dtype(1, dtype)).nil?
+        msg = "invalid dtype for `ones`; found #{dtype}"
+        raise TypeError, msg
+      end
+
+      repeat(zero, n, dtype: dtype, eager: eager).alias("ones")
+    end
+
+    # Construct a column of length `n` filled with zeros.
+    #
+    # This is syntactic sugar for the `repeat` function.
+    #
+    # @param n [Integer]
+    #   Length of the resulting column.
+    # @param dtype [Object]
+    #   Data type of the resulting column. Defaults to Float64.
+    # @param eager [Boolean]
+    #   Evaluate immediately and return a `Series`. If set to `false`,
+    #   return an expression instead.
+    #
+    # @return [Object]
+    #
+    # @example
+    #   Polars.zeros(3, dtype: Polars::Int8, eager: true)
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'zeros' [i8]
+    #   # [
+    #   #         0
+    #   #         0
+    #   #         0
+    #   # ]
+    def zeros(n, dtype: nil, eager: true)
+      if (zero = _one_or_zero_by_dtype(0, dtype)).nil?
+        msg = "invalid dtype for `zeros`; found #{dtype}"
+        raise TypeError, msg
+      end
+
+      repeat(zero, n, dtype: dtype, eager: eager).alias("zeros")
+    end
+
+    private
+
+    def _one_or_zero_by_dtype(value, dtype)
+      if dtype.integer?
+        value
+      elsif dtype.float?
+        value.to_f
+      elsif dtype == Boolean
+        value != 0
+      elsif dtype == Utf8
+        value.to_s
+      elsif dtype == Decimal
+        Decimal(value.to_s)
+      elsif [List, Array].include?(dtype)
+        arr_width = dtype.respond_to?(:width) ? dtype.width : 1
+        [_one_or_zero_by_dtype(value, dtype.inner)] * arr_width
+      else
+        nil
+      end
+    end
+  end
+end

data/lib/polars/functions/whenthen.rb

@@ -0,0 +1,27 @@
+module Polars
+  module Functions
+    # Start a "when, then, otherwise" expression.
+    #
+    # @return [When]
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"foo" => [1, 3, 4], "bar" => [3, 4, 0]})
+    #   df.with_column(Polars.when(Polars.col("foo") > 2).then(Polars.lit(1)).otherwise(Polars.lit(-1)))
+    #   # =>
+    #   # shape: (3, 3)
+    #   # ┌─────┬─────┬─────────┐
+    #   # │ foo ┆ bar ┆ literal │
+    #   # │ --- ┆ --- ┆ ---     │
+    #   # │ i64 ┆ i64 ┆ i32     │
+    #   # ╞═════╪═════╪═════════╡
+    #   # │ 1   ┆ 3   ┆ -1      │
+    #   # │ 3   ┆ 4   ┆ 1       │
+    #   # │ 4   ┆ 0   ┆ 1       │
+    #   # └─────┴─────┴─────────┘
+    def when(expr)
+      expr = Utils.expr_to_lit_or_expr(expr)
+      pw = Plr.when(expr._rbexpr)
+      When.new(pw)
+    end
+  end
+end