fix: Handle lit broadcasting

Author: dangotbanned

narwhals/_plan/impl_arrow.py

@@ -38,12 +38,13 @@ def cols(node: expr.Columns, frame: NativeFrame) -> Evaluated:
 
 
 @evaluate.register(expr.Literal)
-def lit(node: expr.Literal, frame: NativeFrame) -> Evaluated:  # noqa: ARG001
+def lit(node: expr.Literal, frame: NativeFrame) -> Evaluated:
     import pyarrow as pa
 
     if is_scalar_literal(node.value):
-        scalar = node.value.unwrap()
-        return [pa.chunked_array([[scalar]])]
+        lit: t.Any = pa.scalar
+        array = pa.repeat(lit(node.value.unwrap()), len(frame))
+        return [pa.chunked_array([array])]
     elif is_series_literal(node.value):
         ca = node.value.unwrap().to_native()
         return [t.cast("NativeSeries", ca)]

tests/plan/to_compliant_test.py

@@ -46,9 +46,9 @@ def test_to_compliant(backend: BackendName, expr: DummyExpr) -> None:
     [
         (nwd.col("a"), ["A", "B", "A"]),
         (nwd.col("a", "b"), [["A", "B", "A"], [1, 2, 3]]),
-        (nwd.lit(1), [1]),
-        (nwd.lit(2.0), [2.0]),
-        (nwd.lit(None, nw.String()), [None]),
+        (nwd.lit(1), [1, 1, 1]),
+        (nwd.lit(2.0), [2.0, 2.0, 2.0]),
+        (nwd.lit(None, nw.String()), [None, None, None]),
     ],
     ids=_ids_ir,
 )