fix(python): add cluster_with_columns optimization toggle in python (p…

…ola-rs#16446)

py-polars/polars/functions/lazy.py

@@ -1747,6 +1747,7 @@ def collect_all(
  slice_pushdown: bool = True,
  comm_subplan_elim: bool = True,
  comm_subexpr_elim: bool = True,
+ cluster_with_columns: bool = True,
  streaming: bool = False,
 ) -> list[DataFrame]:
  """
@@ -1774,6 +1775,8 @@ def collect_all(
  Will try to cache branching subplans that occur on self-joins or unions.
  comm_subexpr_elim
  Common subexpressions will be cached and reused.
+ cluster_with_columns
+ Combine sequential independent calls to with_columns
  streaming
  Process the query in batches to handle larger-than-memory data.
  If set to `False` (default), the entire query is processed in a single
@@ -1798,6 +1801,7 @@ def collect_all(
  slice_pushdown = False
  comm_subplan_elim = False
  comm_subexpr_elim = False
+ cluster_with_columns = False
 
  if streaming:
  issue_unstable_warning("Streaming mode is considered unstable.")
@@ -1814,6 +1818,7 @@ def collect_all(
  slice_pushdown,
  comm_subplan_elim,
  comm_subexpr_elim,
+ cluster_with_columns,
  streaming,
  _eager=False,
  )
@@ -1840,6 +1845,7 @@ def collect_all_async(
  slice_pushdown: bool = True,
  comm_subplan_elim: bool = True,
  comm_subexpr_elim: bool = True,
+ cluster_with_columns: bool = True,
  streaming: bool = True,
 ) -> _GeventDataFrameResult[list[DataFrame]]: ...
 
@@ -1857,6 +1863,7 @@ def collect_all_async(
  slice_pushdown: bool = True,
  comm_subplan_elim: bool = True,
  comm_subexpr_elim: bool = True,
+ cluster_with_columns: bool = True,
  streaming: bool = False,
 ) -> Awaitable[list[DataFrame]]: ...
 
@@ -1874,6 +1881,7 @@ def collect_all_async(
  slice_pushdown: bool = True,
  comm_subplan_elim: bool = True,
  comm_subexpr_elim: bool = True,
+ cluster_with_columns: bool = True,
  streaming: bool = False,
 ) -> Awaitable[list[DataFrame]] | _GeventDataFrameResult[list[DataFrame]]:
  """
@@ -1912,6 +1920,8 @@ def collect_all_async(
  Will try to cache branching subplans that occur on self-joins or unions.
  comm_subexpr_elim
  Common subexpressions will be cached and reused.
+ cluster_with_columns
+ Combine sequential independent calls to with_columns
  streaming
  Process the query in batches to handle larger-than-memory data.
  If set to `False` (default), the entire query is processed in a single
@@ -1948,6 +1958,7 @@ def collect_all_async(
  slice_pushdown = False
  comm_subplan_elim = False
  comm_subexpr_elim = False
+ cluster_with_columns = False
 
  if streaming:
  issue_unstable_warning("Streaming mode is considered unstable.")
@@ -1964,6 +1975,7 @@ def collect_all_async(
  slice_pushdown,
  comm_subplan_elim,
  comm_subexpr_elim,
+ cluster_with_columns,
  streaming,
  _eager=False,
  )

py-polars/polars/lazyframe/frame.py

@@ -889,6 +889,7 @@ def explain(
  slice_pushdown: bool = True,
  comm_subplan_elim: bool = True,
  comm_subexpr_elim: bool = True,
+ cluster_with_columns: bool = True,
  streaming: bool = False,
  tree_format: bool = False,
  ) -> str:
@@ -918,6 +919,8 @@ def explain(
  Will try to cache branching subplans that occur on self-joins or unions.
  comm_subexpr_elim
  Common subexpressions will be cached and reused.
+ cluster_with_columns
+ Combine sequential independent calls to with_columns
  streaming
  Run parts of the query in a streaming fashion (this is in an alpha state)
  tree_format
@@ -945,6 +948,7 @@ def explain(
  slice_pushdown,
  comm_subplan_elim,
  comm_subexpr_elim,
+ cluster_with_columns,
  streaming,
  _eager=False,
  )
@@ -971,6 +975,7 @@ def show_graph(
  slice_pushdown: bool = True,
  comm_subplan_elim: bool = True,
  comm_subexpr_elim: bool = True,
+ cluster_with_columns: bool = True,
  streaming: bool = False,
  ) -> str | None:
  """
@@ -1005,6 +1010,8 @@ def show_graph(
  Will try to cache branching subplans that occur on self-joins or unions.
  comm_subexpr_elim
  Common subexpressions will be cached and reused.
+ cluster_with_columns
+ Combine sequential independent calls to with_columns
  streaming
  Run parts of the query in a streaming fashion (this is in an alpha state)
 
@@ -1029,6 +1036,7 @@ def show_graph(
  slice_pushdown,
  comm_subplan_elim,
  comm_subexpr_elim,
+ cluster_with_columns,
  streaming,
  _eager=False,
  )
@@ -1513,6 +1521,7 @@ def profile(
  slice_pushdown: bool = True,
  comm_subplan_elim: bool = True,
  comm_subexpr_elim: bool = True,
+ cluster_with_columns: bool = True,
  show_plot: bool = False,
  truncate_nodes: int = 0,
  figsize: tuple[int, int] = (18, 8),
@@ -1545,6 +1554,8 @@ def profile(
  Will try to cache branching subplans that occur on self-joins or unions.
  comm_subexpr_elim
  Common subexpressions will be cached and reused.
+ cluster_with_columns
+ Combine sequential independent calls to with_columns
  show_plot
  Show a gantt chart of the profiling result
  truncate_nodes
@@ -1593,6 +1604,7 @@ def profile(
  projection_pushdown = False
  comm_subplan_elim = False
  comm_subexpr_elim = False
+ cluster_with_columns = False
 
  ldf = self._ldf.optimization_toggle(
  type_coercion,
@@ -1602,6 +1614,7 @@ def profile(
  slice_pushdown,
  comm_subplan_elim,
  comm_subexpr_elim,
+ cluster_with_columns,
  streaming,
  _eager=False,
  )
@@ -1658,6 +1671,7 @@ def collect(
  slice_pushdown: bool = True,
  comm_subplan_elim: bool = True,
  comm_subexpr_elim: bool = True,
+ cluster_with_columns: bool = True,
  no_optimization: bool = False,
  streaming: bool = False,
  background: Literal[True],
@@ -1675,6 +1689,7 @@ def collect(
  slice_pushdown: bool = True,
  comm_subplan_elim: bool = True,
  comm_subexpr_elim: bool = True,
+ cluster_with_columns: bool = True,
  no_optimization: bool = False,
  streaming: bool = False,
  background: Literal[False] = False,
@@ -1691,6 +1706,7 @@ def collect(
  slice_pushdown: bool = True,
  comm_subplan_elim: bool = True,
  comm_subexpr_elim: bool = True,
+ cluster_with_columns: bool = True,
  no_optimization: bool = False,
  streaming: bool = False,
  background: bool = False,
@@ -1719,6 +1735,8 @@ def collect(
  Will try to cache branching subplans that occur on self-joins or unions.
  comm_subexpr_elim
  Common subexpressions will be cached and reused.
+ cluster_with_columns
+ Combine sequential independent calls to with_columns
  no_optimization
  Turn off (certain) optimizations.
  streaming
@@ -1792,6 +1810,7 @@ def collect(
  slice_pushdown = False
  comm_subplan_elim = False
  comm_subexpr_elim = False
+ cluster_with_columns = False
 
  if streaming:
  issue_unstable_warning("Streaming mode is considered unstable.")
@@ -1805,6 +1824,7 @@ def collect(
  slice_pushdown,
  comm_subplan_elim,
  comm_subexpr_elim,
+ cluster_with_columns,
  streaming,
  _eager,
  )
@@ -1829,6 +1849,7 @@ def collect_async(
  slice_pushdown: bool = True,
  comm_subplan_elim: bool = True,
  comm_subexpr_elim: bool = True,
+ cluster_with_columns: bool = True,
  streaming: bool = True,
  ) -> _GeventDataFrameResult[DataFrame]: ...
 
@@ -1845,6 +1866,7 @@ def collect_async(
  slice_pushdown: bool = True,
  comm_subplan_elim: bool = True,
  comm_subexpr_elim: bool = True,
+ cluster_with_columns: bool = True,
  streaming: bool = True,
  ) -> Awaitable[DataFrame]: ...
 
@@ -1860,6 +1882,7 @@ def collect_async(
  slice_pushdown: bool = True,
  comm_subplan_elim: bool = True,
  comm_subexpr_elim: bool = True,
+ cluster_with_columns: bool = True,
  streaming: bool = False,
  ) -> Awaitable[DataFrame] | _GeventDataFrameResult[DataFrame]:
  """
@@ -1896,6 +1919,8 @@ def collect_async(
  Will try to cache branching subplans that occur on self-joins or unions.
  comm_subexpr_elim
  Common subexpressions will be cached and reused.
+ cluster_with_columns
+ Combine sequential independent calls to with_columns
  streaming
  Process the query in batches to handle larger-than-memory data.
  If set to `False` (default), the entire query is processed in a single
@@ -1960,6 +1985,7 @@ def collect_async(
  slice_pushdown = False
  comm_subplan_elim = False
  comm_subexpr_elim = False
+ cluster_with_columns = False
 
  if streaming:
  issue_unstable_warning("Streaming mode is considered unstable.")
@@ -1973,6 +1999,7 @@ def collect_async(
  slice_pushdown,
  comm_subplan_elim,
  comm_subexpr_elim,
+ cluster_with_columns,
  streaming,
  _eager=False,
  )
@@ -2379,6 +2406,7 @@ def _set_sink_optimizations(
  slice_pushdown,
  comm_subplan_elim=False,
  comm_subexpr_elim=False,
+ cluster_with_columns=False,
  streaming=True,
  _eager=False,
  )
@@ -2395,6 +2423,7 @@ def fetch(
  slice_pushdown: bool = True,
  comm_subplan_elim: bool = True,
  comm_subexpr_elim: bool = True,
+ cluster_with_columns: bool = True,
  streaming: bool = False,
  ) -> DataFrame:
  """
@@ -2420,6 +2449,8 @@ def fetch(
  Will try to cache branching subplans that occur on self-joins or unions.
  comm_subexpr_elim
  Common subexpressions will be cached and reused.
+ cluster_with_columns
+ Combine sequential independent calls to with_columns
  streaming
  Run parts of the query in a streaming fashion (this is in an alpha state)
 
@@ -2467,6 +2498,7 @@ def fetch(
  slice_pushdown = False
  comm_subplan_elim = False
  comm_subexpr_elim = False
+ cluster_with_columns = False
 
  lf = self._ldf.optimization_toggle(
  type_coercion,
@@ -2476,6 +2508,7 @@ def fetch(
  slice_pushdown,
  comm_subplan_elim,
  comm_subexpr_elim,
+ cluster_with_columns,
  streaming,
  _eager=False,
  )

py-polars/src/lazyframe/mod.rs

@@ -457,6 +457,7 @@ impl PyLazyFrame {
  slice_pushdown: bool,
  comm_subplan_elim: bool,
  comm_subexpr_elim: bool,
+ cluster_with_columns: bool,
  streaming: bool,
  _eager: bool,
  ) -> Self {
@@ -466,6 +467,7 @@ impl PyLazyFrame {
  .with_predicate_pushdown(predicate_pushdown)
  .with_simplify_expr(simplify_expr)
  .with_slice_pushdown(slice_pushdown)
+ .with_cluster_with_columns(cluster_with_columns)
  .with_streaming(streaming)
  ._with_eager(_eager)
  .with_projection_pushdown(projection_pushdown);

py-polars/tests/unit/test_cwc.py

@@ -18,6 +18,21 @@ def test_basic_cwc() -> None:
  )
 
 
+def test_disable_cwc() -> None:
+ df = (
+ pl.LazyFrame({"a": [1, 2]})
+ .with_columns(pl.col("a").alias("b") * 2)
+ .with_columns(pl.col("a").alias("c") * 3)
+ .with_columns(pl.col("a").alias("d") * 4)
+ )
+
+ explain = df.explain(cluster_with_columns=False)
+
+ assert """[[(col("a")) * (2)].alias("b")]""" in explain
+ assert """[[(col("a")) * (3)].alias("c")]""" in explain
+ assert """[[(col("a")) * (4)].alias("d")]""" in explain
+
+
 def test_refuse_with_deps() -> None:
  df = (
  pl.LazyFrame({"a": [1, 2]})