fix: multi-column keys are broken in some scenarios when rearranged #7477

agavra · 2021-05-05T23:39:03Z

Description

There are a few situations where key columns are re-arranged (differ from the source/partition by). You can reproduce like this:

CREATE STREAM inputStream (id INT KEY, a INT, b INT, c INT) WITH (kafka_topic='input', partitions='4');

INSERT INTO input VALUES (1, 11, 21, 31);
INSERT INTO input VALUES (2, 12, 22, 32);
INSERT INTO input VALUES (1, 11, 21, 31);

CRATE STRAM repartitionedStream AS SELECT * FROM inputStream PARTITION BY c+5, b;
SELECT * FROM repartitionedStream EMIT CHANGES;

+-----------------+-----------------+-----------------+-----------------+-----------------+
|B                |KSQL_COL_0       |A                |C                |ID               |
+-----------------+-----------------+-----------------+-----------------+-----------------+
|36               |21               |11               |31               |1                |
|37               |22               |12               |32               |2                |
|36               |21               |11               |31               |1                |

This is caused because of a mismatch in the source schema and projection schema. This PR fixes this issue in two places:

in PARTITION BY, a SELECT * is resolved using the source schema and doesn't take into account the ordering of the keys selected in the PARTITION BY clause
in any SELECT (projection) clause where the keys are not ordered the same as the source, this PR just reorders it for the user so that it matches the source schema

Testing done

QTT tests

Reviewer checklist

Ensure docs are updated if necessary. (eg. if a user visible feature is being added or changed).
Ensure relevant issues are linked (description should include text like "Fixes #")

mjsax · 2021-05-05T23:50:56Z

ksqldb-engine/src/main/java/io/confluent/ksql/planner/plan/PlanNode.java

-    // Switch them around here:
-    final Stream<Column> keys = columns.stream()
-        .filter(c -> schema.isKeyColumn(c.name()));
+    // but are added at the back during processing for performance reasons. Furthermore,


For my own education. What does this comment mean: but are added at the back during processing for performance reasons -- what is the perf impact and why?

I think the context here is that it avoids shuffling data (adding at the end of an array is easier than shuffling the whole array)

ksqldb-engine/src/main/java/io/confluent/ksql/planner/plan/PlanNode.java

mjsax · 2021-05-05T23:56:42Z

ksqldb-engine/src/main/java/io/confluent/ksql/planner/plan/PlanNode.java

@@ -192,10 +195,13 @@ static void throwKeysNotIncludedError(
      final LogicalSchema schema


To better understand the PR. What is columns and what is schema exactly? Is columns the list of selected columns from the SELECT clause (in the corresponding order) ? Is schema the "physical" layout of the input (what is schema for a join query?)

honestly, this code is super weird to me; I was thinking the same thing, but erred against changing it. It's only ever called like this: orderColumns(getSchema().value(), getSchema()); so really, it could be done by just passing in one schema and then ensuring it's properly ordered.

mjsax · 2021-05-05T23:57:47Z

ksqldb-engine/src/main/java/io/confluent/ksql/planner/plan/PlanNode.java

+    // they should be selected in the order of the key, if that is not the same as the
+    // ordering within the value. Switch them around here:
+    final ImmutableMap<ColumnName, Column> columnsByName = Maps.uniqueIndex(columns, Column::name);
+    final Stream<Column> keys = schema.key().stream()


The main change (for the fix) seems to be to use schema instead of columns to find the keys?

spot on, that's the first fix (the second fix is in the SelectionUtil

mjsax · 2021-05-05T23:58:20Z

ksqldb-engine/src/main/java/io/confluent/ksql/planner/plan/PlanNode.java

+    final ImmutableMap<ColumnName, Column> columnsByName = Maps.uniqueIndex(columns, Column::name);
+    final Stream<Column> keys = schema.key().stream()
+        .map(key -> columnsByName.get(key.name()))
+        .filter(Objects::nonNull);


Wondering why we need this filter?

um, we might not need this one here after digging into it more - but from the API of this method, there's nothing prevent columns from not covering everything in schema, which would mean columnsByName.get(key.name()) would return empty. Without refactoring this method signature I think it's better to be defensive here.

Reading further below, I feel this logic need to be kept as-is even when we have #6374 in place with a clear separation (I'm assuming schema references the physical schema of the source, and columns references the logical schema mapped from the phyiscal schema by then).

I think it's better to be defensive here.

I guess that is the question? We know that it should never be null, so should not fail fast to expose a bug?

we don't "know" that - I believe methods should not make assumptions about their callers. If you look only at this method there's no guarantee that every field in the schema has a corresponding schema that's passed in. If we wanted to prevent bugs, that would belong in the caller (checking that every column that's passed in has a corresponding schema entry)

ksqldb-engine/src/main/java/io/confluent/ksql/planner/plan/SelectionUtil.java

mjsax · 2021-05-06T00:09:36Z

ksqldb-engine/src/main/java/io/confluent/ksql/planner/plan/SelectionUtil.java

+      for (final Pair<ColumnName, SqlType> aliasAndType : aliasAndTypes) {
+        if (aliasAndType != null) {
+          // can be null if the key was not selected - this will cause a failure
+          // down the line but here we just ignore it here


Why can it cause an error? And how would it surface to the user? -- Not sure what we don't want to catch this error right here?

you're required to select all keys, if you don't it'll complain :) I don't think we should add unnecessary redundant checks in places they don't belong. there's no reason why this specific code needs to know that you're required to select all keys. The code that requires that should enforce it.

ksqldb-engine/src/main/java/io/confluent/ksql/planner/plan/SelectionUtil.java

mjsax · 2021-05-06T00:14:48Z

ksqldb-functional-tests/src/test/resources/query-validation-tests/partition-by.json

+      "name": "multiple columns - select star - reorder columns",
+      "statements": [
+        "CREATE STREAM INPUT (NAME STRING, ID INT, AGE INT) with (kafka_topic='input', format='JSON');",
+        "CREATE STREAM OUTPUT AS select * from INPUT partition by AGE, ID;"


meta question: The lower/upper case style is confusing to me.

Should we not upper case all keywords, and lower case all identifiers? This is all a wild mix making it hard to read.

CREATE STREAM input (name STRING, id INT, age INT) with (kafka_topic='input', format='JSON'); CREATE STREAM output AS SELECT * FROM input PARTITION BY age, id;

Can we replicate this test (or can you point out existing tests) for:

input stream with existing multi-column key plus reordering: PARTITION BY key2,key1

PARTITION BY key, nonkey

PARTITION BY nonKey, key

all of the above replicated SELECT columnNames (instead of *) with all 4 reordering combination (either in SELECT or PARTITION BY); seems two are already covered

mixing key and values in select: SELECT key2, nonKey2, key1, nonKey1...

"name": "multiple key columns - reordered", already exists

"name": "multiple columns - some key some value" already exists

I'll add that

multiple key columns tests this

the one I mentioned above tests this as well

As far as the casing style I just copied existing tests - happy to change the new ones but I'm not going to go and change the old ones for this PR

'm not going to go and change the old ones for this PR

Sure :) (did not expect that)

vcrfxia

Thanks @agavra and @mjsax for fixing (and finding) this bug! Are we planning to get this into 0.18 as well?

ksqldb-engine/src/main/java/io/confluent/ksql/planner/plan/SelectionUtil.java

vcrfxia · 2021-05-06T15:25:36Z

ksqldb-engine/src/main/java/io/confluent/ksql/planner/plan/SelectionUtil.java

+   * The second pass goes through the list of projections again and builds the logical schema,
+   * but this time if we encounter a projection that references a key column, we instead take
+   * it from the list we built in the first pass (in order defined by the parent schema).
+   */


Thanks for the detailed comment! Super helpful.

vcrfxia · 2021-05-06T15:26:26Z

ksqldb-engine/src/main/java/io/confluent/ksql/planner/plan/SelectionUtil.java

+        final SelectExpression keyExp = keyExpressions.get(currKeyIdx).remove(0);
+        final SqlType type = expressionTypeManager.getExpressionSqlType(keyExp.getExpression());
+        builder.keyColumn(keyExp.getAlias(), type);
+        if (keyExpressions.get(currKeyIdx).isEmpty()) {


What happens if a key isn't selected? (Should this be a while-loop rather than a simple if-check? If so, we also need to update currKeyIdx = 0; on line 117.)

good call - made this a while loop and moved it up (before the remove) so no need to change line 117

vcrfxia · 2021-05-06T15:27:06Z

ksqldb-functional-tests/src/test/resources/query-validation-tests/partition-by.json

+          {"name": "OUTPUT", "type": "stream", "schema": "NAME STRING KEY, ID INT KEY, AGE INT"}
+        ]
+      }
+    },


Thanks for the test coverage! This is great. Do we also need to audit the test coverage for GROUP BY?

We probably should, though since GROUP BY doesn't use the same code we should do that separately (Matthias tested it locally and GROUP BYs have this "rearrange" behavior)

guozhangwang

The new algorithm looks good to me :) I think for #6374 itself we probably need to rely on this two-phase procedure anyways to separate logical columns from phyiscal schema, along with the maintained mapping between the two.

guozhangwang · 2021-05-06T17:48:58Z

ksqldb-engine/src/main/java/io/confluent/ksql/planner/plan/PlanNode.java

+    final ImmutableMap<ColumnName, Column> columnsByName = Maps.uniqueIndex(columns, Column::name);
+    final Stream<Column> keys = schema.key().stream()
+        .map(key -> columnsByName.get(key.name()))
+        .filter(Objects::nonNull);


Reading further below, I feel this logic need to be kept as-is even when we have #6374 in place with a clear separation (I'm assuming schema references the physical schema of the source, and columns references the logical schema mapped from the phyiscal schema by then).

mjsax · 2021-05-06T18:47:25Z

pom.xml

@@ -136,7 +136,7 @@
        <scala.version>2.13.2</scala.version>
        <apache.io.version>2.6</apache.io.version>
        <io.confluent.ksql.version>6.2.0-0</io.confluent.ksql.version>
-        <io.confluent.schema-registry.version>${confluent.version.range}</io.confluent.schema-registry.version>
+        <io.confluent.schema-registry.version>6.2.0-685</io.confluent.schema-registry.version>


Why this change?

oops! I didn't mean to check that in... but without this things take way too long locally...

agavra requested a review from a team as a code owner May 5, 2021 23:39

agavra requested review from mjsax and vcrfxia May 5, 2021 23:39

mjsax reviewed May 5, 2021

View reviewed changes

ksqldb-engine/src/main/java/io/confluent/ksql/planner/plan/PlanNode.java Outdated Show resolved Hide resolved

mjsax reviewed May 5, 2021

View reviewed changes

mjsax reviewed May 6, 2021

View reviewed changes

ksqldb-engine/src/main/java/io/confluent/ksql/planner/plan/SelectionUtil.java Outdated Show resolved Hide resolved

mjsax reviewed May 6, 2021

View reviewed changes

ksqldb-engine/src/main/java/io/confluent/ksql/planner/plan/SelectionUtil.java Outdated Show resolved Hide resolved

mjsax reviewed May 6, 2021

View reviewed changes

agavra force-pushed the 6.2.x branch from cff2299 to 014f763 Compare May 6, 2021 04:49

vcrfxia approved these changes May 6, 2021

View reviewed changes

guozhangwang approved these changes May 6, 2021

View reviewed changes

mjsax reviewed May 6, 2021

View reviewed changes

agavra force-pushed the 6.2.x branch from 16e0ea1 to f526192 Compare May 10, 2021 15:18

agavra added 4 commits May 11, 2021 08:18

fix: multi-column keys are broken in some scenarios when rearranged

3e5a03a

chore: historical tests

64e73db

chore: comments

b37740f

chore: change confluent version range

9052899

agavra force-pushed the 6.2.x branch from f526192 to 9052899 Compare May 11, 2021 15:19

agavra merged commit 453ca8b into confluentinc:6.2.x May 11, 2021

agavra deleted the 6.2.x branch May 11, 2021 17:01

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

fix: multi-column keys are broken in some scenarios when rearranged #7477

fix: multi-column keys are broken in some scenarios when rearranged #7477

agavra commented May 5, 2021

mjsax May 5, 2021

agavra May 6, 2021

mjsax May 5, 2021

agavra May 6, 2021

mjsax May 5, 2021

agavra May 6, 2021

mjsax May 5, 2021

agavra May 6, 2021

guozhangwang May 6, 2021

mjsax May 6, 2021

agavra May 11, 2021

mjsax May 6, 2021 •

edited

Loading

agavra May 6, 2021

mjsax May 6, 2021

mjsax May 6, 2021

agavra May 6, 2021

mjsax May 6, 2021

vcrfxia left a comment

vcrfxia May 6, 2021

vcrfxia May 6, 2021

agavra May 6, 2021

vcrfxia May 6, 2021

agavra May 6, 2021

guozhangwang left a comment

guozhangwang May 6, 2021

mjsax May 6, 2021

agavra May 6, 2021

		@@ -192,10 +195,13 @@ static void throwKeysNotIncludedError(
		final LogicalSchema schema

fix: multi-column keys are broken in some scenarios when rearranged #7477

fix: multi-column keys are broken in some scenarios when rearranged #7477

Conversation

agavra commented May 5, 2021

Description

Testing done

Reviewer checklist

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

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

mjsax May 6, 2021 • 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

vcrfxia left a 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

Choose a reason for hiding this comment

guozhangwang left a 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

mjsax May 6, 2021 •

edited

Loading