cudf-polars-cu12 25.2.2__py3-none-any.whl → 25.6.0__py3-none-any.whl

cudf_polars/dsl/translate.py

@@ -22,8 +22,14 @@ import pylibcudf as plc
 
 from cudf_polars.dsl import expr, ir
 from cudf_polars.dsl.to_ast import insert_colrefs
-from cudf_polars.typing import NodeTraverser
-from cudf_polars.utils import dtypes, sorting
+from cudf_polars.dsl.utils.aggregations import decompose_single_agg
+from cudf_polars.dsl.utils.groupby import rewrite_groupby
+from cudf_polars.dsl.utils.naming import unique_names
+from cudf_polars.dsl.utils.replace import replace
+from cudf_polars.dsl.utils.rolling import rewrite_rolling
+from cudf_polars.dsl.utils.windows import offsets_to_windows
+from cudf_polars.typing import Schema
+from cudf_polars.utils import config, dtypes, sorting
 
 if TYPE_CHECKING:
     from polars import GPUEngine
@@ -41,13 +47,13 @@ class Translator:
     ----------
     visitor
         Polars NodeTraverser object
-    config
+    engine
         GPU engine configuration.
     """
 
-    def __init__(self, visitor: NodeTraverser, config: GPUEngine):
+    def __init__(self, visitor: NodeTraverser, engine: GPUEngine):
         self.visitor = visitor
-        self.config = config
+        self.config_options = config.ConfigOptions.from_polars_engine(engine)
         self.errors: list[Exception] = []
 
     def translate_ir(self, *, n: int | None = None) -> ir.IR:
@@ -84,7 +90,7 @@ class Translator:
         # IR is versioned with major.minor, minor is bumped for backwards
         # compatible changes (e.g. adding new nodes), major is bumped for
         # incompatible changes (e.g. renaming nodes).
-        if (version := self.visitor.version()) >= (5, 1):
+        if (version := self.visitor.version()) >= (7, 1):
             e = NotImplementedError(
                 f"No support for polars IR {version=}"
             )  # pragma: no cover; no such version for now.
@@ -120,7 +126,7 @@ class Translator:
 
             return result
 
-    def translate_expr(self, *, n: int) -> expr.Expr:
+    def translate_expr(self, *, n: int, schema: Schema) -> expr.Expr:
         """
         Translate a polars-internal expression IR into our representation.
 
@@ -128,6 +134,8 @@ class Translator:
         ----------
         n
             Node to translate, an integer referencing a polars internal node.
+        schema
+            Schema of the IR node this expression uses as evaluation context.
 
         Returns
         -------
@@ -143,7 +151,7 @@ class Translator:
         node = self.visitor.view_expression(n)
         dtype = dtypes.from_polars(self.visitor.get_dtype(n))
         try:
-            return _translate_expr(node, self, dtype)
+            return _translate_expr(node, self, dtype, schema)
         except Exception as e:
             self.errors.append(e)
             return expr.ErrorExpr(dtype, str(e))
@@ -168,6 +176,7 @@ class set_node(AbstractContextManager[None]):
 
     __slots__ = ("n", "visitor")
     visitor: NodeTraverser
+
     n: int
 
     def __init__(self, visitor: NodeTraverser, n: int) -> None:
@@ -187,30 +196,26 @@ noop_context: nullcontext[None] = nullcontext()
 
 
 @singledispatch
-def _translate_ir(
-    node: Any, translator: Translator, schema: dict[str, plc.DataType]
-) -> ir.IR:
+def _translate_ir(node: Any, translator: Translator, schema: Schema) -> ir.IR:
     raise NotImplementedError(
         f"Translation for {type(node).__name__}"
     )  # pragma: no cover
 
 
 @_translate_ir.register
-def _(
-    node: pl_ir.PythonScan, translator: Translator, schema: dict[str, plc.DataType]
-) -> ir.IR:
+def _(node: pl_ir.PythonScan, translator: Translator, schema: Schema) -> ir.IR:
     scan_fn, with_columns, source_type, predicate, nrows = node.options
     options = (scan_fn, with_columns, source_type, nrows)
     predicate = (
-        translate_named_expr(translator, n=predicate) if predicate is not None else None
+        translate_named_expr(translator, n=predicate, schema=schema)
+        if predicate is not None
+        else None
     )
     return ir.PythonScan(schema, options, predicate)
 
 
 @_translate_ir.register
-def _(
-    node: pl_ir.Scan, translator: Translator, schema: dict[str, plc.DataType]
-) -> ir.IR:
+def _(node: pl_ir.Scan, translator: Translator, schema: Schema) -> ir.IR:
     typ, *options = node.scan_type
     if typ == "ndjson":
         (reader_options,) = map(json.loads, options)
@@ -219,118 +224,102 @@ def _(
         reader_options, cloud_options = map(json.loads, options)
     file_options = node.file_options
     with_columns = file_options.with_columns
-    n_rows = file_options.n_rows
-    if n_rows is None:
-        n_rows = -1  # All rows
-        skip_rows = 0  # Don't skip
+    row_index = file_options.row_index
+    include_file_paths = file_options.include_file_paths
+
+    pre_slice = file_options.n_rows
+    if pre_slice is None:
+        n_rows = -1
+        skip_rows = 0
     else:
-        # TODO: with versioning, rename on the rust side
-        skip_rows, n_rows = n_rows
+        skip_rows, n_rows = pre_slice
 
-    row_index = file_options.row_index
     return ir.Scan(
         schema,
         typ,
         reader_options,
         cloud_options,
-        translator.config.config.copy(),
+        translator.config_options,
         node.paths,
         with_columns,
         skip_rows,
         n_rows,
         row_index,
-        translate_named_expr(translator, n=node.predicate)
+        include_file_paths,
+        translate_named_expr(translator, n=node.predicate, schema=schema)
         if node.predicate is not None
         else None,
     )
 
 
 @_translate_ir.register
-def _(
-    node: pl_ir.Cache, translator: Translator, schema: dict[str, plc.DataType]
-) -> ir.IR:
-    return ir.Cache(schema, node.id_, translator.translate_ir(n=node.input))
+def _(node: pl_ir.Cache, translator: Translator, schema: Schema) -> ir.IR:
+    return ir.Cache(
+        schema, node.id_, node.cache_hits, translator.translate_ir(n=node.input)
+    )
 
 
 @_translate_ir.register
-def _(
-    node: pl_ir.DataFrameScan, translator: Translator, schema: dict[str, plc.DataType]
-) -> ir.IR:
+def _(node: pl_ir.DataFrameScan, translator: Translator, schema: Schema) -> ir.IR:
     return ir.DataFrameScan(
         schema,
         node.df,
         node.projection,
-        translator.config.config.copy(),
+        translator.config_options,
     )
 
 
 @_translate_ir.register
-def _(
-    node: pl_ir.Select, translator: Translator, schema: dict[str, plc.DataType]
-) -> ir.IR:
+def _(node: pl_ir.Select, translator: Translator, schema: Schema) -> ir.IR:
     with set_node(translator.visitor, node.input):
         inp = translator.translate_ir(n=None)
-        exprs = [translate_named_expr(translator, n=e) for e in node.expr]
+        exprs = [
+            translate_named_expr(translator, n=e, schema=inp.schema) for e in node.expr
+        ]
     return ir.Select(schema, exprs, node.should_broadcast, inp)
 
 
 @_translate_ir.register
-def _(
-    node: pl_ir.GroupBy, translator: Translator, schema: dict[str, plc.DataType]
-) -> ir.IR:
+def _(node: pl_ir.GroupBy, translator: Translator, schema: Schema) -> ir.IR:
     with set_node(translator.visitor, node.input):
         inp = translator.translate_ir(n=None)
-        aggs = [translate_named_expr(translator, n=e) for e in node.aggs]
-        keys = [translate_named_expr(translator, n=e) for e in node.keys]
-    return ir.GroupBy(
-        schema,
-        keys,
-        aggs,
-        node.maintain_order,
-        node.options,
-        inp,
-    )
+        keys = [
+            translate_named_expr(translator, n=e, schema=inp.schema) for e in node.keys
+        ]
+        original_aggs = [
+            translate_named_expr(translator, n=e, schema=inp.schema) for e in node.aggs
+        ]
+    is_rolling = node.options.rolling is not None
+    is_dynamic = node.options.dynamic is not None
+    if is_dynamic:
+        raise NotImplementedError("group_by_dynamic")
+    elif is_rolling:
+        return rewrite_rolling(
+            node.options, schema, keys, original_aggs, translator.config_options, inp
+        )
+    else:
+        return rewrite_groupby(
+            node, schema, keys, original_aggs, translator.config_options, inp
+        )
 
 
 @_translate_ir.register
-def _(
-    node: pl_ir.Join, translator: Translator, schema: dict[str, plc.DataType]
-) -> ir.IR:
+def _(node: pl_ir.Join, translator: Translator, schema: Schema) -> ir.IR:
     # Join key dtypes are dependent on the schema of the left and
     # right inputs, so these must be translated with the relevant
     # input active.
-    def adjust_literal_dtype(literal: expr.Literal) -> expr.Literal:
-        if literal.dtype.id() == plc.types.TypeId.INT32:
-            plc_int64 = plc.types.DataType(plc.types.TypeId.INT64)
-            return expr.Literal(
-                plc_int64,
-                pa.scalar(literal.value.as_py(), type=plc.interop.to_arrow(plc_int64)),
-            )
-        return literal
-
-    def maybe_adjust_binop(e) -> expr.Expr:
-        if isinstance(e.value, expr.BinOp):
-            left, right = e.value.children
-            if isinstance(left, expr.Col) and isinstance(right, expr.Literal):
-                e.value.children = (left, adjust_literal_dtype(right))
-            elif isinstance(left, expr.Literal) and isinstance(right, expr.Col):
-                e.value.children = (adjust_literal_dtype(left), right)
-        return e
-
-    def translate_expr_and_maybe_fix_binop_args(translator, exprs):
-        return [
-            maybe_adjust_binop(translate_named_expr(translator, n=e)) for e in exprs
-        ]
-
     with set_node(translator.visitor, node.input_left):
         inp_left = translator.translate_ir(n=None)
-        # TODO: There's bug in the polars type coercion phase. Use
-        # translate_named_expr directly once it is resolved.
-        # Tracking issue: https://github.com/pola-rs/polars/issues/20935
-        left_on = translate_expr_and_maybe_fix_binop_args(translator, node.left_on)
+        left_on = [
+            translate_named_expr(translator, n=e, schema=inp_left.schema)
+            for e in node.left_on
+        ]
     with set_node(translator.visitor, node.input_right):
         inp_right = translator.translate_ir(n=None)
-        right_on = translate_expr_and_maybe_fix_binop_args(translator, node.right_on)
+        right_on = [
+            translate_named_expr(translator, n=e, schema=inp_right.schema)
+            for e in node.right_on
+        ]
 
     if (how := node.options[0]) in {
         "Inner",
@@ -341,7 +330,15 @@ def _(
         "Semi",
         "Anti",
     }:
-        return ir.Join(schema, left_on, right_on, node.options, inp_left, inp_right)
+        return ir.Join(
+            schema,
+            left_on,
+            right_on,
+            node.options,
+            translator.config_options,
+            inp_left,
+            inp_right,
+        )
     else:
         how, op1, op2 = node.options[0]
         if how != "IEJoin":
@@ -385,29 +382,29 @@ def _(
 
 
 @_translate_ir.register
-def _(
-    node: pl_ir.HStack, translator: Translator, schema: dict[str, plc.DataType]
-) -> ir.IR:
+def _(node: pl_ir.HStack, translator: Translator, schema: Schema) -> ir.IR:
     with set_node(translator.visitor, node.input):
         inp = translator.translate_ir(n=None)
-        exprs = [translate_named_expr(translator, n=e) for e in node.exprs]
+        exprs = [
+            translate_named_expr(translator, n=e, schema=inp.schema) for e in node.exprs
+        ]
     return ir.HStack(schema, exprs, node.should_broadcast, inp)
 
 
 @_translate_ir.register
 def _(
-    node: pl_ir.Reduce, translator: Translator, schema: dict[str, plc.DataType]
+    node: pl_ir.Reduce, translator: Translator, schema: Schema
 ) -> ir.IR:  # pragma: no cover; polars doesn't emit this node yet
     with set_node(translator.visitor, node.input):
         inp = translator.translate_ir(n=None)
-        exprs = [translate_named_expr(translator, n=e) for e in node.expr]
+        exprs = [
+            translate_named_expr(translator, n=e, schema=inp.schema) for e in node.expr
+        ]
     return ir.Reduce(schema, exprs, inp)
 
 
 @_translate_ir.register
-def _(
-    node: pl_ir.Distinct, translator: Translator, schema: dict[str, plc.DataType]
-) -> ir.IR:
+def _(node: pl_ir.Distinct, translator: Translator, schema: Schema) -> ir.IR:
     (keep, subset, maintain_order, zlice) = node.options
     keep = ir.Distinct._KEEP_MAP[keep]
     subset = frozenset(subset) if subset is not None else None
@@ -422,12 +419,13 @@ def _(
 
 
 @_translate_ir.register
-def _(
-    node: pl_ir.Sort, translator: Translator, schema: dict[str, plc.DataType]
-) -> ir.IR:
+def _(node: pl_ir.Sort, translator: Translator, schema: Schema) -> ir.IR:
     with set_node(translator.visitor, node.input):
         inp = translator.translate_ir(n=None)
-        by = [translate_named_expr(translator, n=e) for e in node.by_column]
+        by = [
+            translate_named_expr(translator, n=e, schema=inp.schema)
+            for e in node.by_column
+        ]
     stable, nulls_last, descending = node.sort_options
     order, null_order = sorting.sort_order(
         descending, nulls_last=nulls_last, num_keys=len(by)
@@ -436,65 +434,101 @@ def _(
 
 
 @_translate_ir.register
-def _(
-    node: pl_ir.Slice, translator: Translator, schema: dict[str, plc.DataType]
-) -> ir.IR:
+def _(node: pl_ir.Slice, translator: Translator, schema: Schema) -> ir.IR:
     return ir.Slice(
         schema, node.offset, node.len, translator.translate_ir(n=node.input)
     )
 
 
 @_translate_ir.register
-def _(
-    node: pl_ir.Filter, translator: Translator, schema: dict[str, plc.DataType]
-) -> ir.IR:
+def _(node: pl_ir.Filter, translator: Translator, schema: Schema) -> ir.IR:
     with set_node(translator.visitor, node.input):
         inp = translator.translate_ir(n=None)
-        mask = translate_named_expr(translator, n=node.predicate)
+        mask = translate_named_expr(translator, n=node.predicate, schema=inp.schema)
     return ir.Filter(schema, mask, inp)
 
 
 @_translate_ir.register
-def _(
-    node: pl_ir.SimpleProjection,
-    translator: Translator,
-    schema: dict[str, plc.DataType],
-) -> ir.IR:
+def _(node: pl_ir.SimpleProjection, translator: Translator, schema: Schema) -> ir.IR:
     return ir.Projection(schema, translator.translate_ir(n=node.input))
 
 
 @_translate_ir.register
-def _(
-    node: pl_ir.MapFunction, translator: Translator, schema: dict[str, plc.DataType]
-) -> ir.IR:
+def _(node: pl_ir.MergeSorted, translator: Translator, schema: Schema) -> ir.IR:
+    key = node.key
+    inp_left = translator.translate_ir(n=node.input_left)
+    inp_right = translator.translate_ir(n=node.input_right)
+    return ir.MergeSorted(
+        schema,
+        key,
+        inp_left,
+        inp_right,
+    )
+
+
+@_translate_ir.register
+def _(node: pl_ir.MapFunction, translator: Translator, schema: Schema) -> ir.IR:
     name, *options = node.function
     return ir.MapFunction(
         schema,
         name,
         options,
-        # TODO: merge_sorted breaks this pattern
         translator.translate_ir(n=node.input),
     )
 
 
 @_translate_ir.register
-def _(
-    node: pl_ir.Union, translator: Translator, schema: dict[str, plc.DataType]
-) -> ir.IR:
+def _(node: pl_ir.Union, translator: Translator, schema: Schema) -> ir.IR:
     return ir.Union(
         schema, node.options, *(translator.translate_ir(n=n) for n in node.inputs)
     )
 
 
+@_translate_ir.register
+def _(node: pl_ir.HConcat, translator: Translator, schema: Schema) -> ir.IR:
+    return ir.HConcat(
+        schema,
+        False,  # noqa: FBT003
+        *(translator.translate_ir(n=n) for n in node.inputs),
+    )
+
+
 @_translate_ir.register
 def _(
-    node: pl_ir.HConcat, translator: Translator, schema: dict[str, plc.DataType]
+    node: pl_ir.Sink, translator: Translator, schema: dict[str, plc.DataType]
 ) -> ir.IR:
-    return ir.HConcat(schema, *(translator.translate_ir(n=n) for n in node.inputs))
+    payload = json.loads(node.payload)
+    try:
+        file = payload["File"]
+        sink_kind_options = file["file_type"]
+    except KeyError as err:  # pragma: no cover
+        raise NotImplementedError("Unsupported payload structure") from err
+    if isinstance(sink_kind_options, dict):
+        if len(sink_kind_options) != 1:  # pragma: no cover; not sure if this can happen
+            raise NotImplementedError("Sink options dict with more than one entry.")
+        sink_kind, options = next(iter(sink_kind_options.items()))
+    else:
+        raise NotImplementedError(
+            "Unsupported sink options structure"
+        )  # pragma: no cover
+
+    sink_options = file.get("sink_options", {})
+    cloud_options = file.get("cloud_options")
+
+    options.update(sink_options)
+
+    return ir.Sink(
+        schema=schema,
+        kind=sink_kind,
+        path=file["target"],
+        options=options,
+        cloud_options=cloud_options,
+        df=translator.translate_ir(n=node.input),
+    )
 
 
 def translate_named_expr(
-    translator: Translator, *, n: pl_expr.PyExprIR
+    translator: Translator, *, n: pl_expr.PyExprIR, schema: Schema
 ) -> expr.NamedExpr:
     """
     Translate a polars-internal named expression IR object into our representation.
@@ -505,6 +539,8 @@ def translate_named_expr(
         Translator object
     n
         Node to translate, a named expression node.
+    schema
+        Schema of the IR node this expression uses as evaluation context.
 
     Returns
     -------
@@ -522,12 +558,14 @@ def translate_named_expr(
     NotImplementedError
         If any translation fails due to unsupported functionality.
     """
-    return expr.NamedExpr(n.output_name, translator.translate_expr(n=n.node))
+    return expr.NamedExpr(
+        n.output_name, translator.translate_expr(n=n.node, schema=schema)
+    )
 
 
 @singledispatch
 def _translate_expr(
-    node: Any, translator: Translator, dtype: plc.DataType
+    node: Any, translator: Translator, dtype: plc.DataType, schema: Schema
 ) -> expr.Expr:
     raise NotImplementedError(
         f"Translation for {type(node).__name__}"
@@ -535,7 +573,9 @@ def _translate_expr(
 
 
 @_translate_expr.register
-def _(node: pl_expr.Function, translator: Translator, dtype: plc.DataType) -> expr.Expr:
+def _(
+    node: pl_expr.Function, translator: Translator, dtype: plc.DataType, schema: Schema
+) -> expr.Expr:
     name, *options = node.function_data
     options = tuple(options)
     if isinstance(name, pl_expr.StringFunction):
@@ -544,18 +584,20 @@ def _(node: pl_expr.Function, translator: Translator, dtype: plc.DataType) -> ex
             pl_expr.StringFunction.StripCharsStart,
             pl_expr.StringFunction.StripCharsEnd,
         }:
-            column, chars = (translator.translate_expr(n=n) for n in node.input)
+            column, chars = (
+                translator.translate_expr(n=n, schema=schema) for n in node.input
+            )
             if isinstance(chars, expr.Literal):
-                if chars.value == pa.scalar(""):
+                # We check for null first because we want to use the
+                # chars pyarrow type, but it is invalid to try and
+                # produce a string scalar with a null dtype.
+                if chars.value is None:
+                    # Polars uses None to mean "strip all whitespace"
+                    chars = expr.Literal(column.dtype, "")
+                elif chars.value == "":
                     # No-op in polars, but libcudf uses empty string
                     # as signifier to remove whitespace.
                     return column
-                elif chars.value == pa.scalar(None):
-                    # Polars uses None to mean "strip all whitespace"
-                    chars = expr.Literal(
-                        column.dtype,
-                        pa.scalar("", type=plc.interop.to_arrow(column.dtype)),
-                    )
             return expr.StringFunction(
                 dtype,
                 expr.StringFunction.Name.from_polars(name),
@@ -567,11 +609,13 @@ def _(node: pl_expr.Function, translator: Translator, dtype: plc.DataType) -> ex
             dtype,
             expr.StringFunction.Name.from_polars(name),
             options,
-            *(translator.translate_expr(n=n) for n in node.input),
+            *(translator.translate_expr(n=n, schema=schema) for n in node.input),
         )
     elif isinstance(name, pl_expr.BooleanFunction):
         if name == pl_expr.BooleanFunction.IsBetween:
-            column, lo, hi = (translator.translate_expr(n=n) for n in node.input)
+            column, lo, hi = (
+                translator.translate_expr(n=n, schema=schema) for n in node.input
+            )
             (closed,) = options
             lop, rop = expr.BooleanFunction._BETWEEN_OPS[closed]
             return expr.BinOp(
@@ -584,7 +628,7 @@ def _(node: pl_expr.Function, translator: Translator, dtype: plc.DataType) -> ex
             dtype,
             expr.BooleanFunction.Name.from_polars(name),
             options,
-            *(translator.translate_expr(n=n) for n in node.input),
+            *(translator.translate_expr(n=n, schema=schema) for n in node.input),
         )
     elif isinstance(name, pl_expr.TemporalFunction):
         # functions for which evaluation of the expression may not return
@@ -604,14 +648,14 @@ def _(node: pl_expr.Function, translator: Translator, dtype: plc.DataType) -> ex
             dtype,
             expr.TemporalFunction.Name.from_polars(name),
             options,
-            *(translator.translate_expr(n=n) for n in node.input),
+            *(translator.translate_expr(n=n, schema=schema) for n in node.input),
         )
         if name in needs_cast:
             return expr.Cast(dtype, result_expr)
         return result_expr
 
     elif isinstance(name, str):
-        children = (translator.translate_expr(n=n) for n in node.input)
+        children = (translator.translate_expr(n=n, schema=schema) for n in node.input)
         if name == "log":
             (base,) = options
             (child,) = children
@@ -619,10 +663,21 @@ def _(node: pl_expr.Function, translator: Translator, dtype: plc.DataType) -> ex
                 dtype,
                 plc.binaryop.BinaryOperator.LOG_BASE,
                 child,
-                expr.Literal(dtype, pa.scalar(base, type=plc.interop.to_arrow(dtype))),
+                expr.Literal(dtype, base),
             )
         elif name == "pow":
             return expr.BinOp(dtype, plc.binaryop.BinaryOperator.POW, *children)
+        elif name in "top_k":
+            (col, k) = children
+            assert isinstance(k, expr.Literal)
+            (descending,) = options
+            return expr.Slice(
+                dtype,
+                0,
+                k.value,
+                expr.Sort(dtype, (False, True, not descending), col),
+            )
+
         return expr.UnaryFunction(dtype, name, options, *children)
     raise NotImplementedError(
         f"No handler for Expr function node with {name=}"
@@ -630,73 +685,155 @@ def _(node: pl_expr.Function, translator: Translator, dtype: plc.DataType) -> ex
 
 
 @_translate_expr.register
-def _(node: pl_expr.Window, translator: Translator, dtype: plc.DataType) -> expr.Expr:
-    # TODO: raise in groupby?
+def _(
+    node: pl_expr.Window, translator: Translator, dtype: plc.DataType, schema: Schema
+) -> expr.Expr:
     if isinstance(node.options, pl_expr.RollingGroupOptions):
         # pl.col("a").rolling(...)
-        return expr.RollingWindow(
-            dtype, node.options, translator.translate_expr(n=node.function)
+        agg = translator.translate_expr(n=node.function, schema=schema)
+        name_generator = unique_names(schema)
+        aggs, named_post_agg = decompose_single_agg(
+            expr.NamedExpr(next(name_generator), agg), name_generator, is_top=True
         )
+        named_aggs = [agg for agg, _ in aggs]
+        orderby = node.options.index_column
+        orderby_dtype = schema[orderby]
+        if plc.traits.is_integral(orderby_dtype):
+            # Integer orderby column is cast in implementation to int64 in polars
+            orderby_dtype = plc.DataType(plc.TypeId.INT64)
+        preceding, following = offsets_to_windows(
+            orderby_dtype,
+            node.options.offset,
+            node.options.period,
+        )
+        closed_window = node.options.closed_window
+        if isinstance(named_post_agg.value, expr.Col):
+            (named_agg,) = named_aggs
+            return expr.RollingWindow(
+                named_agg.value.dtype,
+                preceding,
+                following,
+                closed_window,
+                orderby,
+                named_agg.value,
+            )
+        replacements: dict[expr.Expr, expr.Expr] = {
+            expr.Col(agg.value.dtype, agg.name): expr.RollingWindow(
+                agg.value.dtype,
+                preceding,
+                following,
+                closed_window,
+                orderby,
+                agg.value,
+            )
+            for agg in named_aggs
+        }
+        return replace([named_post_agg.value], replacements)[0]
     elif isinstance(node.options, pl_expr.WindowMapping):
         # pl.col("a").over(...)
         return expr.GroupedRollingWindow(
             dtype,
             node.options,
-            translator.translate_expr(n=node.function),
-            *(translator.translate_expr(n=n) for n in node.partition_by),
+            translator.translate_expr(n=node.function, schema=schema),
+            *(translator.translate_expr(n=n, schema=schema) for n in node.partition_by),
         )
     assert_never(node.options)
 
 
 @_translate_expr.register
-def _(node: pl_expr.Literal, translator: Translator, dtype: plc.DataType) -> expr.Expr:
+def _(
+    node: pl_expr.Literal, translator: Translator, dtype: plc.DataType, schema: Schema
+) -> expr.Expr:
     if isinstance(node.value, plrs.PySeries):
-        return expr.LiteralColumn(dtype, pl.Series._from_pyseries(node.value))
-    value = pa.scalar(node.value, type=plc.interop.to_arrow(dtype))
-    return expr.Literal(dtype, value)
+        data = pl.Series._from_pyseries(node.value).to_arrow(
+            compat_level=dtypes.TO_ARROW_COMPAT_LEVEL
+        )
+        return expr.LiteralColumn(
+            dtype, data.cast(dtypes.downcast_arrow_lists(data.type))
+        )
+    if dtype.id() == plc.TypeId.LIST:  # pragma: no cover
+        # TODO: Find an alternative to pa.infer_type
+        data = pa.array(node.value, type=pa.infer_type(node.value))
+        return expr.LiteralColumn(dtype, data)
+    return expr.Literal(dtype, node.value)
 
 
 @_translate_expr.register
-def _(node: pl_expr.Sort, translator: Translator, dtype: plc.DataType) -> expr.Expr:
+def _(
+    node: pl_expr.Sort, translator: Translator, dtype: plc.DataType, schema: Schema
+) -> expr.Expr:
     # TODO: raise in groupby
-    return expr.Sort(dtype, node.options, translator.translate_expr(n=node.expr))
+    return expr.Sort(
+        dtype, node.options, translator.translate_expr(n=node.expr, schema=schema)
+    )
 
 
 @_translate_expr.register
-def _(node: pl_expr.SortBy, translator: Translator, dtype: plc.DataType) -> expr.Expr:
+def _(
+    node: pl_expr.SortBy, translator: Translator, dtype: plc.DataType, schema: Schema
+) -> expr.Expr:
     options = node.sort_options
     return expr.SortBy(
         dtype,
         (options[0], tuple(options[1]), tuple(options[2])),
-        translator.translate_expr(n=node.expr),
-        *(translator.translate_expr(n=n) for n in node.by),
+        translator.translate_expr(n=node.expr, schema=schema),
+        *(translator.translate_expr(n=n, schema=schema) for n in node.by),
     )
 
 
 @_translate_expr.register
-def _(node: pl_expr.Gather, translator: Translator, dtype: plc.DataType) -> expr.Expr:
+def _(
+    node: pl_expr.Slice, translator: Translator, dtype: plc.DataType, schema: Schema
+) -> expr.Expr:
+    offset = translator.translate_expr(n=node.offset, schema=schema)
+    length = translator.translate_expr(n=node.length, schema=schema)
+    assert isinstance(offset, expr.Literal)
+    assert isinstance(length, expr.Literal)
+    return expr.Slice(
+        dtype,
+        offset.value,
+        length.value,
+        translator.translate_expr(n=node.input, schema=schema),
+    )
+
+
+@_translate_expr.register
+def _(
+    node: pl_expr.Gather, translator: Translator, dtype: plc.DataType, schema: Schema
+) -> expr.Expr:
     return expr.Gather(
         dtype,
-        translator.translate_expr(n=node.expr),
-        translator.translate_expr(n=node.idx),
+        translator.translate_expr(n=node.expr, schema=schema),
+        translator.translate_expr(n=node.idx, schema=schema),
     )
 
 
 @_translate_expr.register
-def _(node: pl_expr.Filter, translator: Translator, dtype: plc.DataType) -> expr.Expr:
+def _(
+    node: pl_expr.Filter, translator: Translator, dtype: plc.DataType, schema: Schema
+) -> expr.Expr:
     return expr.Filter(
         dtype,
-        translator.translate_expr(n=node.input),
-        translator.translate_expr(n=node.by),
+        translator.translate_expr(n=node.input, schema=schema),
+        translator.translate_expr(n=node.by, schema=schema),
     )
 
 
 @_translate_expr.register
-def _(node: pl_expr.Cast, translator: Translator, dtype: plc.DataType) -> expr.Expr:
-    inner = translator.translate_expr(n=node.expr)
+def _(
+    node: pl_expr.Cast, translator: Translator, dtype: plc.DataType, schema: Schema
+) -> expr.Expr:
+    inner = translator.translate_expr(n=node.expr, schema=schema)
     # Push casts into literals so we can handle Cast(Literal(Null))
     if isinstance(inner, expr.Literal):
-        return expr.Literal(dtype, inner.value.cast(plc.interop.to_arrow(dtype)))
+        plc_column = plc.Column.from_scalar(
+            plc.Scalar.from_py(inner.value, inner.dtype), 1
+        )
+        casted_column = plc.unary.cast(plc_column, dtype)
+        casted_py_scalar = plc.interop.to_arrow(
+            plc.copying.get_element(casted_column, 0)
+        ).as_py()
+        return expr.Literal(dtype, casted_py_scalar)
     elif isinstance(inner, expr.Cast):
         # Translation of Len/Count-agg put in a cast, remove double
         # casts if we have one.
@@ -705,17 +842,21 @@ def _(node: pl_expr.Cast, translator: Translator, dtype: plc.DataType) -> expr.E
 
 
 @_translate_expr.register
-def _(node: pl_expr.Column, translator: Translator, dtype: plc.DataType) -> expr.Expr:
+def _(
+    node: pl_expr.Column, translator: Translator, dtype: plc.DataType, schema: Schema
+) -> expr.Expr:
     return expr.Col(dtype, node.name)
 
 
 @_translate_expr.register
-def _(node: pl_expr.Agg, translator: Translator, dtype: plc.DataType) -> expr.Expr:
+def _(
+    node: pl_expr.Agg, translator: Translator, dtype: plc.DataType, schema: Schema
+) -> expr.Expr:
     value = expr.Agg(
         dtype,
         node.name,
         node.options,
-        *(translator.translate_expr(n=n) for n in node.arguments),
+        *(translator.translate_expr(n=n, schema=schema) for n in node.arguments),
     )
     if value.name == "count" and value.dtype.id() != plc.TypeId.INT32:
         return expr.Cast(value.dtype, value)
@@ -723,29 +864,36 @@ def _(node: pl_expr.Agg, translator: Translator, dtype: plc.DataType) -> expr.Ex
 
 
 @_translate_expr.register
-def _(node: pl_expr.Ternary, translator: Translator, dtype: plc.DataType) -> expr.Expr:
+def _(
+    node: pl_expr.Ternary, translator: Translator, dtype: plc.DataType, schema: Schema
+) -> expr.Expr:
     return expr.Ternary(
         dtype,
-        translator.translate_expr(n=node.predicate),
-        translator.translate_expr(n=node.truthy),
-        translator.translate_expr(n=node.falsy),
+        translator.translate_expr(n=node.predicate, schema=schema),
+        translator.translate_expr(n=node.truthy, schema=schema),
+        translator.translate_expr(n=node.falsy, schema=schema),
     )
 
 
 @_translate_expr.register
 def _(
-    node: pl_expr.BinaryExpr, translator: Translator, dtype: plc.DataType
+    node: pl_expr.BinaryExpr,
+    translator: Translator,
+    dtype: plc.DataType,
+    schema: Schema,
 ) -> expr.Expr:
     return expr.BinOp(
         dtype,
         expr.BinOp._MAPPING[node.op],
-        translator.translate_expr(n=node.left),
-        translator.translate_expr(n=node.right),
+        translator.translate_expr(n=node.left, schema=schema),
+        translator.translate_expr(n=node.right, schema=schema),
     )
 
 
 @_translate_expr.register
-def _(node: pl_expr.Len, translator: Translator, dtype: plc.DataType) -> expr.Expr:
+def _(
+    node: pl_expr.Len, translator: Translator, dtype: plc.DataType, schema: Schema
+) -> expr.Expr:
     value = expr.Len(dtype)
     if dtype.id() != plc.TypeId.INT32:
         return expr.Cast(dtype, value)