Support pl.Expr.extend_constant in cudf_polars

Author: mroeschke

python/cudf_polars/cudf_polars/dsl/expressions/unary.py

@@ -105,6 +105,7 @@ class UnaryFunction(Expr):
         {
             "as_struct",
             "drop_nulls",
+            "extend_constant",
             "fill_null",
             "fill_null_with_strategy",
             "mask_nans",
@@ -538,6 +539,33 @@ def do_evaluate(
                 plc.copying.shift(column.obj, offset, fill_scalar, stream=df.stream),
                 dtype=self.dtype,
             )
+        elif self.name == "extend_constant":
+            column = self.children[0].evaluate(df, context=context)
+            value_expr = self.children[1]
+            n_expr = self.children[2]
+            if isinstance(n_expr, Literal):
+                count = n_expr.value
+            else:
+                count = (
+                    n_expr.evaluate(df, context=context)
+                    .obj_scalar(stream=df.stream)
+                    .to_py(stream=df.stream)
+                )
+            if count == 0:
+                return column
+            if isinstance(value_expr, Literal):
+                fill = plc.Scalar.from_py(
+                    value_expr.value, self.dtype.plc_type, stream=df.stream
+                )
+            else:
+                fill = value_expr.evaluate(df, context=context).obj_scalar(
+                    stream=df.stream
+                )
+            extension = plc.Column.from_scalar(fill, count, stream=df.stream)
+            return Column(
+                plc.concatenate.concatenate([column.obj, extension], stream=df.stream),
+                dtype=self.dtype,
+            )
         elif self.name in self._OP_MAPPING:
             column = self.children[0].evaluate(df, context=context)
             if column.dtype.plc_type.id() != self.dtype.id():

python/cudf_polars/tests/expressions/test_extend_constant.py

@@ -0,0 +1,40 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import pytest
+
+import polars as pl
+
+from cudf_polars.testing.asserts import assert_gpu_result_equal
+
+
+@pytest.mark.parametrize(
+    "series, value",
+    [
+        (pl.Series([1, 2, 3]), 0),
+        (pl.Series([1, 2, None]), None),
+        (pl.Series([1.5, float("nan"), -3.0]), 2.5),
+        (pl.Series([1.5, 2.5, -3.0]), None),
+        (pl.Series(["a", "bb", None]), "zz"),
+        (pl.Series(["a", "bb", "c"]), None),
+        (pl.Series([], dtype=pl.Int64), 0),
+        (pl.Series([None, None, None], dtype=pl.Int64), 7),
+        (pl.Series([None, None, None], dtype=pl.Int64), None),
+        (pl.Series([42], dtype=pl.Int64), 7),
+        (pl.Series([None], dtype=pl.Int64), 7),
+    ],
+)
+@pytest.mark.parametrize("n", [0, 1, 3])
+def test_extend_constant(
+    engine: pl.GPUEngine, series: pl.Series, value: int | float | str | None, n: int
+) -> None:
+    lf = pl.LazyFrame({"a": series})
+    q = lf.select(pl.col("a").extend_constant(value, n))
+    assert_gpu_result_equal(q, engine=engine)
+
+
+def test_extend_constant_non_literal_value_and_n(engine: pl.GPUEngine) -> None:
+    lf = pl.LazyFrame({"a": [1, 2, 3], "v": [9, 9, 9], "n": [2, 2, 2]})
+    q = lf.select(pl.col("a").extend_constant(pl.col("v").first(), pl.col("n").first()))
+    assert_gpu_result_equal(q, engine=engine)