feat(frontend): prune order key using functional dependencies #16204

kwannoel · 2024-04-08T15:12:10Z

I hereby agree to the terms of the RisingWave Labs, Inc. Contributor License Agreement.

What's changed and what's your intention?

Closes #16148 .

The main idea in this PR is to prune order key in areas where it is consumed:

In derive_pk in the final part for create materialized view, create sink, create subscription.

Note that we did not do so in BatchSort, since functional dependency may not be supported there yet (?) I tried to add it, but seems it did not prune the order key, as it did in stream.

We specifically do it in these areas, rather than add a new optimizer rule + optimizer pass, since the call to FunctionalDependencySet::minimize_key can be quite expensive, and we only need it in a few areas.

We also don't call minimize_key for Index, since it reuses the order keys as distribution key, rather than as a key. Pruning it can lead to incorrect distribution, for instance when:

pk = []

Then,
order key = [0]

After pruning,
order key = [].

But we want index to be distributed on col 0. So it is incorrect.

risingwave/src/frontend/src/handler/create_index.rs

Lines 385 to 400 in a14ff37

    
           PlanRoot::new( 
        
               logical_project, 
        
               RequiredDist::PhysicalDist(Distribution::HashShard( 
        
                   (0..distributed_by_columns_len).collect(), 
        
               )), 
        
               Order::new( 
        
                   index_columns 
        
                       .iter() 
        
                       .enumerate() 
        
                       .map(|(i, (_, order))| ColumnOrder::new(i, *order)) 
        
                       .collect(), 
        
               ), 
        
               project_required_cols, 
        
               out_names, 
        
           ) 
        
           .gen_index_plan(index_name, definition, retention_seconds)

Finally, we use a new minimizing function, minimize_order_key, rather than minimize_key, since the properties of order key is different from that of minimize_key.

For instance, given the following order key:

fun dep: (0, 1) -> 2
order key: [2, 0, 1]

If we simply apply minimize_key to it, it will treat 2 is obsolete, since the remainder can still form a key due to the functional dependency of (0, 1) -> 2.

But this breaks the ordering properties, since we no longer order by 2 first, before (0, 1).

So we introduce a new minimizing function which will minimize by using the functional dependencies of the prefixes of the order key to prune the suffixes.

Checklist

I have written necessary rustdoc comments
I have added necessary unit tests and integration tests
I have added test labels as necessary. See details.
I have added fuzzing tests or opened an issue to track them. (Optional, recommended for new SQL features Sqlsmith: Sql feature generation #7934).
My PR contains breaking changes. (If it deprecates some features, please create a tracking issue to remove them in the future).
All checks passed in ./risedev check (or alias, ./risedev c)
My PR changes performance-critical code. (Please run macro/micro-benchmarks and show the results.)

My PR contains critical fixes that are necessary to be merged into the latest release. (Please check out the details)

Documentation

My PR needs documentation updates. (Please use the Release note section below to summarize the impact on users)

Release note

If this PR includes changes that directly affect users or other significant modifications relevant to the community, kindly draft a release note to provide a concise summary of these changes. Please prioritize highlighting the impact these changes will have on users.

src/frontend/planner_test/tests/testdata/output/ch_benchmark.yaml

chenzl25 · 2024-04-09T08:27:55Z

src/frontend/src/optimizer/plan_node/stream_materialize.rs

+                &columns,
+                // For index, we can't prune the ORDER KEY,
+                // since it's also the distribution key.
+                table_type != TableType::Index,


Index actually is also a materialized view. This is a bit special if we can't unify them in this code path. Could you provide an example to illustrate the problem you met? Let's see whether we could resolve it together.

This is the failing query.

create table t1 (v1 int, v2 int); create materialized view v as select count(*) cnt from t1; explain (verbose) create index mv_idx on v(cnt);

This is its plan without pruning:

StreamMaterialize { columns: [cnt], stream_key: [], pk_columns: [cnt], pk_conflict: NoCheck } └─StreamExchange { dist: HashShard(v.cnt) } └─StreamTableScan { table: v, columns: [v.cnt], stream_scan_type: ArrangementBackfill, stream_key: [], pk: [], dist: Single }

This is its plan with pruning:

StreamMaterialize { columns: [cnt], stream_key: [], pk_columns: [cnt], pk_conflict: NoCheck } └─StreamExchange { dist: HashShard(v.cnt) } └─StreamTableScan { table: v, columns: [v.cnt], stream_scan_type: ArrangementBackfill, stream_key: [], pk: [], dist: Single }

It has the following functional dependency: [] -> all_columns, since it is a singleton.

As mentioned in the PR description, the following section of code:

risingwave/src/frontend/src/handler/create_index.rs

Lines 385 to 400 in a14ff37

PlanRoot::new(

logical_project,

RequiredDist::PhysicalDist(Distribution::HashShard(

(0..distributed_by_columns_len).collect(),

)),

Order::new(

index_columns

.iter()

.enumerate()

.map(|(i, (_, order))| ColumnOrder::new(i, *order))

.collect(),

),

project_required_cols,

out_names,

)

.gen_index_plan(index_name, definition, retention_seconds)

Will result in

order key: [cnt]

dist key: [0]

Now dist key refers to order key, you can see how it just takes the prefix:

RequiredDist::PhysicalDist(Distribution::HashShard( (0..distributed_by_columns_len).collect(), )),

This seems a little strange, because order key is not appended to pk at this point.
Subsequently, when we derive the pk, that is just stream key combined with order key.
Stream key is []. And after pruning, Order key is [] as well. So pk is just [].

Then distribution key is not a subset of pk at all.

There's 2 solutions I can think of here:

Just avoid index (stupid and simple but not elegant).

Separately append distribution key for index elsewhere. Instead of relying on order key to be the distribution key.

Maybe when we derive a pk, we need to consider the distribution key as well, because theoretically users can specify a distribution key for a materialized view as well, though we don't support it right now, but I remember this feature has been requested several times.

BTW, even for the normal materialized view, we need to ensure distribution key is a subset of pk.

Fixed it by appending the distribution key if it is not part of the pk.

kwannoel · 2024-04-09T13:16:12Z

src/frontend/planner_test/tests/testdata/output/limit.yaml

-    └─BatchSimpleAgg { aggs: [sum0(count)] }
-      └─BatchExchange { order: [], dist: Single }
-        └─BatchSimpleAgg { aggs: [count] }
-          └─BatchScan { table: t, columns: [], distribution: SomeShard }


Not sure why prune TopN columns can also rewrite the BatchTopN 🤔 . Regardless, here it could be rewritten to BatchLimit instead I think.

chenzl25

LGTM!

chenzl25 · 2024-04-10T10:10:42Z

src/frontend/src/optimizer/plan_node/derive.rs

+    // We need to ensure distribution key is part of pk.
+    // If it is not part of either stream_key or order_key,
+    // It must mean that it is only necessary for storage distribution.
+    // Such a case is rare, but it is possible,
+    // for example in the case of index created on singleton mv.
+    // In this case, we can simply append these columns to pk.
+    for &idx in input.distribution().dist_column_indices() {
+        if in_order.contains(idx) {
+            continue;
+        }
+        pk.push(ColumnOrder::new(idx, OrderType::ascending()));
+        in_order.insert(idx);
+    }


BTW, this behavior could break the assumption (not sure whether still holds ATM) we made, i.e. the distribution key should be the prefix of the pk. I remember we had some discussions on this topic @st1page @fuyufjh

My recollection is a bit hazy, but it seems to be for the performance consideration of batch scanning, the distribution key was restricted to the prefix of the primary key (pk), so that queries with a prefix can scan fewer partitions.

ok, now I prefer to not append the distribution key to the end when distribution key is eliminated by functional dependency and we can use the original one instead. @kwannoel

ok, now I prefer to not append the distribution key to the end when distribution key is eliminated by functional dependency and we can use the original one instead. @kwannoel

Hmm so just don't prune order key for indexes, since they are used as distribution key?

I think the tricky thing is that distribution key is always a prefix of PK, and so is order key in certain cases.

There isn't a clear way to determine how we should interleave this distribution key's missing columns into the current PK, preserving the above properties, unless we know the plans which generated the distribution key and order key.

In the case of derive pk, we don't have this info. So seems hard to make a generalised decision. Think we can just disable the order key pruning for index, seems to be the only one so far which has the overlap.

We can revisit if more edge cases show up.

I am afraid of this PR eliminating the distribution key and then to ensure the distribution key is part of pk, we generate a new pk with the distribution key in the suffix. For this case, we can just give up the optimization and use the original derived pk.

Fixed it by making distribution a constraint of minimizing order key. If pruning order key column also results in dist key column being pruned, we will not prune it.

kwannoel · 2024-04-10T14:59:32Z

src/frontend/src/planner/query.rs

-        // Optimize order key before using it for TopN.
-        let func_dep = plan.functional_dependency();
-        let order = func_dep.minimize_order_key(order);
-


Removed. Not clear what's the dist_key_indices to pass in here... Only impacts batch top n, we can leave it out for now.

We could pass in [] as dist_key_indices since top n is singleton dist.

Done in 4ee529c

chenzl25

LGTM!

kwannoel added 3 commits April 8, 2024 20:42

use minimize_order_key

cdc7fe8

handle case where input is not a key

ac95ac4

dapt

986df89

github-actions bot added the type/feature label Apr 8, 2024

dont prune index

c7112c9

kwannoel force-pushed the kwannoel/pk-prefix branch from e49f6a8 to c7112c9 Compare April 9, 2024 02:21

kwannoel requested review from chenzl25, hzxa21 and xxhZs April 9, 2024 03:38

chenzl25 reviewed Apr 9, 2024

View reviewed changes

src/frontend/planner_test/tests/testdata/output/ch_benchmark.yaml Outdated Show resolved Hide resolved

kwannoel added 5 commits April 9, 2024 15:44

use a different algorithm for order key

b297bd9

dapt

e7fa42d

add more planner test

9c2ea29

improve code

bfb385b

doc

be95574

kwannoel requested a review from chenzl25 April 9, 2024 08:02

add note on time complexity

0127d54

chenzl25 reviewed Apr 9, 2024

View reviewed changes

kwannoel added 3 commits April 9, 2024 17:22

revert index workaround

ac853a4

handle [] as a fun dep

5a46f37

handle dist key

aa11d51

kwannoel commented Apr 9, 2024

View reviewed changes

fix top n

73221fe

chenzl25 approved these changes Apr 10, 2024

View reviewed changes

chenzl25 reviewed Apr 10, 2024

View reviewed changes

minimize order key while making sure dist key is not pruned

1ab4f59

kwannoel force-pushed the kwannoel/pk-prefix branch from e113408 to 1ab4f59 Compare April 10, 2024 14:52

kwannoel requested review from chenzl25 and st1page April 10, 2024 14:58

kwannoel commented Apr 10, 2024

View reviewed changes

minimize order key for topn

4ee529c

chenzl25 approved these changes Apr 11, 2024

View reviewed changes

kwannoel added this pull request to the merge queue Apr 11, 2024

Merged via the queue into main with commit 132b4c9 Apr 11, 2024
27 of 28 checks passed

kwannoel deleted the kwannoel/pk-prefix branch April 11, 2024 04:19

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat(frontend): prune order key using functional dependencies #16204

feat(frontend): prune order key using functional dependencies #16204

kwannoel commented Apr 8, 2024 •

edited

Loading

chenzl25 Apr 9, 2024

kwannoel Apr 9, 2024 •

edited

Loading

chenzl25 Apr 9, 2024

chenzl25 Apr 9, 2024

kwannoel Apr 9, 2024

kwannoel Apr 9, 2024

kwannoel Apr 9, 2024

chenzl25 left a comment

chenzl25 Apr 10, 2024 •

edited

Loading

st1page Apr 10, 2024

chenzl25 Apr 10, 2024 •

edited

Loading

kwannoel Apr 10, 2024

kwannoel Apr 10, 2024 •

edited

Loading

chenzl25 Apr 10, 2024

kwannoel Apr 10, 2024

kwannoel Apr 10, 2024

kwannoel Apr 10, 2024

kwannoel Apr 10, 2024

chenzl25 left a comment

	PlanRoot::new(
	logical_project,
	RequiredDist::PhysicalDist(Distribution::HashShard(
	(0..distributed_by_columns_len).collect(),
	)),
	Order::new(
	index_columns
	.iter()
	.enumerate()
	.map(\|(i, (_, order))\| ColumnOrder::new(i, *order))
	.collect(),
	),
	project_required_cols,
	out_names,
	)
	.gen_index_plan(index_name, definition, retention_seconds)

feat(frontend): prune order key using functional dependencies #16204

feat(frontend): prune order key using functional dependencies #16204

Conversation

kwannoel commented Apr 8, 2024 • edited Loading

What's changed and what's your intention?

Checklist

Documentation

Release note

Choose a reason for hiding this comment

kwannoel Apr 9, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chenzl25 left a comment

Choose a reason for hiding this comment

chenzl25 Apr 10, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chenzl25 Apr 10, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kwannoel Apr 10, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chenzl25 left a comment

Choose a reason for hiding this comment

kwannoel commented Apr 8, 2024 •

edited

Loading

kwannoel Apr 9, 2024 •

edited

Loading

chenzl25 Apr 10, 2024 •

edited

Loading

chenzl25 Apr 10, 2024 •

edited

Loading

kwannoel Apr 10, 2024 •

edited

Loading