Reuse aggregate-derived join domains as detail-side payloads

[pentschev]

This extends join_filter_pushdown to recognize when an aggregate-derived join domain can be reused as the payload for a later detail-side join input.

The motivating case is PDSH Q18. The query computes the set of l_orderkey values whose sum(l_quantity) > 300, and current planning can already use that result as a semi-join domain to filter orders. However, the final orders x lineitem join still reads the raw lineitem detail side again.

With this change, when the aggregate-derived domain is known to contain the join key and the payload columns needed by the later join, the planner can reuse that derived domain to feed the detail side as well. The original full join is still preserved, so the rewrite is a row-reduction optimization rather than a semantic shortcut.

Q18 plan before, relevant subtree

JOIN Inner ('o_orderkey',) ('l_orderkey',)
  JOIN Semi ('o_orderkey',) ('l_orderkey',)
    STREAMINGSCAN orders
    FILTER (sum_quantity > 300.0)
      GROUPBY ('l_orderkey',)
        STREAMINGSCAN lineitem ('l_orderkey', 'l_quantity')
  STREAMINGSCAN lineitem ('l_orderkey', 'l_quantity')

Q18 plan after, relevant subtree

JOIN Inner ('o_orderkey',) ('l_orderkey',)
  JOIN Semi ('o_orderkey',) ('l_orderkey',)
    STREAMINGSCAN orders
    FILTER (sum_quantity > 300.0)
      GROUPBY ('l_orderkey',)
        STREAMINGSCAN lineitem ('l_orderkey', 'l_quantity')
  FILTER BooleanFunction ('l_orderkey', 'l_quantity')
    SELECT ('l_orderkey', 'l_quantity')
      FILTER (sum_quantity > 300.0)
        GROUPBY ('l_orderkey',)
          STREAMINGSCAN lineitem ('l_orderkey', 'l_quantity')

Materially, this reduces Q18 memory pressure and improves runtime by avoiding the final unfiltered detail-side lineitem input. In the latest standalone validation, the full SF30K Q1-Q22 workflow completed successfully on 6xNVL4 nodes with validation passing in 146s lukewarm/128s hot.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[wence-]

I was trying to understand what this is doing, because the example rewrite in the PR description isn't valid on its own (the two query fragments are not equivalent).

I think I understand what is going on, but I am not fully convinced that this is a set of changes we want to incorporate into the streaming plan optimisation.

I think what is happening is the following.

Suppose we have:

q1 = l.group_by("k").agg(sum_v=sum("v")).filter("sum_v" > value)

# Which k's in o have sum("v") in l > value
q2 = o.join(q1, on="k", how="semi")

# Reconstruct the distinct v's.
q3 = q2.join(l, on="k")

# Aggregate the v's with sum
q4 = q3.group_by("...").agg(sum_v=sum("v"))

What we notice is that because of the algebraic structure of sum, we can push a filter onto the right hand side of the join in q3:

q3prime = q2.join(l.group_by("k").agg(v=sum("v")).filter("v" > value)

q4prime = q3prime.group_by("...").agg(sum_v=sum("v"))

Now, q3 and q3prime are not the same. But q4 and q4prime are

There are various conditions on v and the particular query structure that I think are not fully checked for in this rewrite.

In particular, in the example above, a requirement on the shape of the query is that in between the replaced right-side input to q3 and the final sum over v, no part of the query is allowed to depend on observing v.

Here's a test that demonstrates that we get the wrong answer with this optimisation as written:

def test_aggregated_domain_reuse_ok(
    tmp_path: pathlib.Path,
    engine: SPMDEngine,
) -> None:
    lineitem_path = tmp_path / "lineitem.parquet"
    orders_path = tmp_path / "orders.parquet"

    pl.DataFrame(
        {
            "k": [1, 1, 2, 2, 3, 3],
            "v": [1, 2, 0, 0, 0, 0],
        }
    ).write_parquet(lineitem_path)
    pl.DataFrame({"k": [1, 2, 3]}).write_parquet(orders_path)

    lineitem = pl.scan_parquet(lineitem_path)
    orders = pl.scan_parquet(orders_path)
    selected_keys = (
        lineitem.group_by("k")
        .agg(pl.col("v").sum().alias("sum_v"))
        .filter(pl.col("sum_v") > 0)
        .select("k")
    )

    q = (
        orders.join(selected_keys, on="k", how="semi")
        .join(lineitem, on="k")
        .select("k", "v", bucket="v")
        .group_by("bucket") # here we observe "v" with sum(v).
        .agg(pl.col("v").sum())
    )

    assert_gpu_result_equal(q, engine=engine, check_row_order=False)

It seems like a safer, but always correct, thing to do would be to deduce that the groupby-filter object should be used as a filtering domain for the lineitem join in Q18. That removes the unfiltered lineitem join, but doesn't rely on quite delicate algebraic conditions: we instead only need to find equivalence-sets of join keys and find a "filtering" join key.

I'm also surprised that the filter pushdown doesn't find the first orders.join(lineitem, how="semi") and use that to pre-filter the lineitem table. Even at 30K it only has about a million rows so seems like a good candidate.

I suppose in this case, because we will almost certainly do things as a broadcast join, this optimisation is not to reduce the communication volume but to avoid reading (parts of) lineitem twice (or buffering the whole lot behind a fanout node?).

Ah yeah, maybe it's this latter issue, the polars plan is (in part) for Q18:

          INNER JOIN:
          LEFT PLAN ON: [col("o_orderkey")]
            SEMI JOIN:
            LEFT PLAN ON: [col("o_orderkey")]
              Parquet SCAN [/home/coder/third-party/polars-benchmark/data/tables/scale-1.0/orders/part.0.parquet]
              PROJECT 4/9 COLUMNS
              ESTIMATED ROWS: 1500000
            RIGHT PLAN ON: [col("l_orderkey")]
              simple π 1/1 ["l_orderkey"]
                FILTER [(col("sum_quantity")) > (300.0)]
                FROM
                  AGGREGATE[maintain_order: false]
                    [col("l_quantity").sum().alias("sum_quantity")] BY [col("l_orderkey")]
                    FROM
                    CACHE[id: d82d0f89-be6d-4ee0-8c77-3c796cb2f175]
                      Parquet SCAN [/home/coder/third-party/polars-benchmark/data/tables/scale-1.0/lineitem/part.0.parquet]
                      PROJECT 2/16 COLUMNS
                      ESTIMATED ROWS: 6001215
            END SEMI JOIN
          RIGHT PLAN ON: [col("l_orderkey")]
            simple π 2/2 ["l_quantity", "l_orderkey"]
              CACHE[id: d82d0f89-be6d-4ee0-8c77-3c796cb2f175]
          END INNER JOIN

So we get cache reuse of the lineitem scan. But, because it's joined back, if we do this by broadcasting the orders side we must buffer the entire right hand "cached" table somewhere. Your groupby rewrite still requires us to cache something, but the thing we cache is the filtered grouped object, which is small.

[wence-]

I think that the rewrite optimize_with_stats removes Cache nodes, so we should never see Cache at this point.

But, if we need correct equality of Cache we should fix that on Cache, no?

[wence-]

Why can't we use the utilities in column_domain.py for this?

[wence-]

Why can't we use replace_at_path?

[wence-]

This thing only ever yields a single value (or no value at all). Does it need to be a generator at all?