feat: Add pl$repeat_()

R/000-wrappers.R

@@ -272,6 +272,14 @@ NULL
 }
 
 
+`repeat_` <- function(`value`, `n`, `dtype` = NULL) {
+  `value` <- .savvy_extract_ptr(`value`, "PlRExpr")
+  `n` <- .savvy_extract_ptr(`n`, "PlRExpr")
+  `dtype` <- .savvy_extract_ptr(`dtype`, "PlRDataType")
+  .savvy_wrap_PlRExpr(.Call(savvy_repeat___impl, `value`, `n`, `dtype`))
+}
+
+
 `sum_horizontal` <- function(`exprs`, `ignore_nulls`) {
   .savvy_wrap_PlRExpr(.Call(savvy_sum_horizontal__impl, `exprs`, `ignore_nulls`))
 }

R/expr-expr.R

@@ -4130,7 +4130,7 @@ expr__qcut <- function(
 #' df <- pl$DataFrame(a = c(1, 1, 2))
 #'
 #' # Create a Series with 3 nulls, append column a then rechunk
-#' df$select(pl$repeat(NA, 3)$append(pl$col("a"))$rechunk())
+#' df$select(pl$repeat_(NA, 3)$append(pl$col("a"))$rechunk())
 expr__rechunk <- function() {
   wrap({
     self$`_rexpr`$rechunk()

R/functions-repeat.R

@@ -0,0 +1,33 @@
+#' Construct a column of length n`` filled with the given value
+#'
+#' @inheritParams rlang::args_dots_empty
+#' @param value Value to repeat.
+#' @param n Length of the resulting column
+#' @param dtype Data type of the resulting column. If `NULL` (default), data
+#' type is inferred from the given value. Defaults to Int32 for integer values,
+#' unless Int64 is required to fit the given value. Defaults to Float64 for
+#' float values.
+#'
+#' @details
+#' If you want to construct a column in lazy mode and do not need a
+#' pre-determined length, use `pl$lit()` instead.
+#' @inherit as_polars_expr return
+#' @examples
+#' # Construct a column with a repeated value in a lazy context.
+#' pl$select(pl$repeat_("z", n = 3))
+#'
+#' # Specify an output dtype
+#' pl$select(pl$repeat_(3, n = 3, dtype = pl$Int8))
+pl__repeat_ <- function(value, n, ..., dtype = NULL) {
+  wrap({
+    check_polars_dtype(dtype, allow_null = TRUE)
+    if (is_integerish(n)) {
+      if (n < 0) {
+        abort("`n` must be greater than or equal to 0.")
+      }
+      n <- pl$lit(n)$cast(pl$Int64)
+    }
+    value <- as_polars_expr(value, as_lit = TRUE)
+    repeat_(value$`_rexpr`, n$`_rexpr`, dtype$`_dt`)
+  })
+}

src/init.c

@@ -234,6 +234,11 @@ SEXP savvy_min_horizontal__impl(SEXP c_arg__exprs) {
     return handle_result(res);
 }
 
+SEXP savvy_repeat___impl(SEXP c_arg__value, SEXP c_arg__n, SEXP c_arg__dtype) {
+    SEXP res = savvy_repeat___ffi(c_arg__value, c_arg__n, c_arg__dtype);
+    return handle_result(res);
+}
+
 SEXP savvy_sum_horizontal__impl(SEXP c_arg__exprs, SEXP c_arg__ignore_nulls) {
     SEXP res = savvy_sum_horizontal__ffi(c_arg__exprs, c_arg__ignore_nulls);
     return handle_result(res);
@@ -2796,6 +2801,7 @@ static const R_CallMethodDef CallEntries[] = {
     {"savvy_max_horizontal__impl", (DL_FUNC) &savvy_max_horizontal__impl, 1},
     {"savvy_mean_horizontal__impl", (DL_FUNC) &savvy_mean_horizontal__impl, 2},
     {"savvy_min_horizontal__impl", (DL_FUNC) &savvy_min_horizontal__impl, 1},
+    {"savvy_repeat___impl", (DL_FUNC) &savvy_repeat___impl, 3},
     {"savvy_sum_horizontal__impl", (DL_FUNC) &savvy_sum_horizontal__impl, 2},
     {"savvy_time_range__impl", (DL_FUNC) &savvy_time_range__impl, 4},
     {"savvy_time_ranges__impl", (DL_FUNC) &savvy_time_ranges__impl, 4},

src/rust/api.h

@@ -38,6 +38,7 @@ SEXP savvy_lit_null__ffi(void);
 SEXP savvy_max_horizontal__ffi(SEXP c_arg__exprs);
 SEXP savvy_mean_horizontal__ffi(SEXP c_arg__exprs, SEXP c_arg__ignore_nulls);
 SEXP savvy_min_horizontal__ffi(SEXP c_arg__exprs);
+SEXP savvy_repeat___ffi(SEXP c_arg__value, SEXP c_arg__n, SEXP c_arg__dtype);
 SEXP savvy_sum_horizontal__ffi(SEXP c_arg__exprs, SEXP c_arg__ignore_nulls);
 SEXP savvy_time_range__ffi(SEXP c_arg__start, SEXP c_arg__end, SEXP c_arg__every, SEXP c_arg__closed);
 SEXP savvy_time_ranges__ffi(SEXP c_arg__start, SEXP c_arg__end, SEXP c_arg__every, SEXP c_arg__closed);

src/rust/src/functions/lazy.rs

@@ -1,4 +1,4 @@
-use crate::{prelude::*, PlRDataFrame, PlRExpr, PlRLazyFrame, PlRSeries, RPolarsErr};
+use crate::{prelude::*, PlRDataFrame, PlRDataType, PlRExpr, PlRLazyFrame, PlRSeries, RPolarsErr};
 use polars::functions;
 use polars::lazy::dsl;
 use savvy::{savvy, ListSexp, NumericSexp, RawSexp, Result, StringSexp};
@@ -275,3 +275,14 @@ pub fn concat_str(s: ListSexp, separator: &str, ignore_nulls: bool) -> Result<Pl
 pub fn arg_where(condition: PlRExpr) -> Result<PlRExpr> {
     Ok(dsl::arg_where(condition.inner.clone()).into())
 }
+
+#[savvy]
+#[allow(non_snake_case)]
+pub fn repeat_(value: PlRExpr, n: PlRExpr, dtype: Option<&PlRDataType>) -> Result<PlRExpr> {
+    let mut value = value.inner;
+    let n = n.inner;
+    if let Some(dtype) = dtype {
+        value = value.cast(dtype.dt.clone());
+    }
+    Ok(dsl::repeat(value, n).into())
+}

tests/testthat/test-functions-repeat.R

@@ -0,0 +1,20 @@
+test_that("pl$repeat_() works", {
+  expect_equal(
+    pl$select(pl$repeat_("z", n = 3)),
+    pl$DataFrame(`repeat` = c("z", "z", "z"))
+  )
+  expect_equal(
+    pl$select(pl$repeat_(1, n = 3, dtype = pl$Int8)),
+    pl$DataFrame(`repeat` = c(1, 1, 1))$cast(pl$Int8)
+  )
+  # values can't be cast to dtype
+  expect_equal(
+    pl$select(pl$repeat_("a", n = 3, dtype = pl$Int8)),
+    pl$DataFrame(`repeat` = c(NA, NA, NA))$cast(pl$Int8)
+  )
+  # wrong `n`
+  expect_error(
+    pl$select(pl$repeat_("z", n = -1)),
+    "must be greater than or equal to 0"
+  )
+})