diff --git a/polars/polars-lazy/polars-pipe/src/executors/operators/projection.rs b/polars/polars-lazy/polars-pipe/src/executors/operators/projection.rs index 2486e5a70605..7d2974bcc485 100644 --- a/polars/polars-lazy/polars-pipe/src/executors/operators/projection.rs +++ b/polars/polars-lazy/polars-pipe/src/executors/operators/projection.rs @@ -40,12 +40,29 @@ impl Operator for ProjectionOperator { context: &PExecutionContext, chunk: &DataChunk, ) -> PolarsResult { - let projected = self + let mut has_literals = false; + let mut projected = self .exprs .iter() - .map(|e| e.evaluate(chunk, context.execution_state.as_any())) + .map(|e| { + let s = e.evaluate(chunk, context.execution_state.as_any())?; + if s.len() == 1 { + has_literals = true; + } + Ok(s) + }) .collect::>>()?; + if has_literals { + let height = projected.iter().map(|s| s.len()).max().unwrap(); + for s in &mut projected { + let len = s.len(); + if len == 1 && len != height { + *s = s.new_from_index(0, height) + } + } + } + let chunk = chunk.with_data(DataFrame::new_no_checks(projected)); Ok(OperatorResult::Finished(chunk)) } diff --git a/py-polars/tests/unit/test_streaming.py b/py-polars/tests/unit/test_streaming.py index 69a4010d6fab..0338bf027ef5 100644 --- a/py-polars/tests/unit/test_streaming.py +++ b/py-polars/tests/unit/test_streaming.py @@ -243,3 +243,35 @@ def test_ooc_sort(monkeypatch: Any) -> None: ).to_series() assert_series_equal(out, s.sort(reverse=reverse)) + + +def test_streaming_literal_expansion() -> None: + df = pl.DataFrame( + { + "y": ["a", "b"], + "z": [1, 2], + } + ) + + q = df.lazy().select( + x=pl.lit("constant"), + y=pl.col("y"), + z=pl.col("z"), + ) + + assert q.collect(streaming=True).to_dict(False) == { + "x": ["constant", "constant"], + "y": ["a", "b"], + "z": [1, 2], + } + assert q.groupby(["x", "y"]).agg(pl.mean("z")).sort("y").collect( + streaming=True + ).to_dict(False) == { + "x": ["constant", "constant"], + "y": ["a", "b"], + "z": [1.0, 2.0], + } + assert q.groupby(["x"]).agg(pl.mean("z")).collect().to_dict(False) == { + "x": ["constant"], + "z": [1.5], + }