UNION ALLs in MSQ #14981

LakshSingla · 2023-09-13T22:18:55Z

Description

(Note: The description might use union and union all interchangeably unless specified, both of which mean union all in SQL)

This PR updates the following:

UnionDataSource can have data sources apart from the TableDataSource. This will be used for MSQ only, since MSQ, in theory, can plan arbitrary unions. Also, it is required to plan unions in MSQ in the DataSourcePlan.
Disallow top-level union alls in MSQ. This is because the SQL layer executes the top-level unions sequentially, which doesn't make sense for an async engine like MSQ. More info and examples for this added later.
Add the ability to plan a UnionDataSource in MSQ. This will provide MSQ feature parity with the native engine, as it will allow the
(Will be broken into a separate PR) Add the ability to plan arbitrary data sources as unions (this will be an engine feature), however, this requires alignment as to what leeways are we willing to support - Do the unions with different column names get planned, do the unions with different types get planned, etc.

1

MSQ needs to plan the individual data sources of the union and perform a replace operation so that each data source can be represented by the input specs that it requires. This warrants UnionDataSource to accept other data sources as its children. The current methods in the UnionDataSource have been refactored to perform the original checks only when the data source is used in certain contexts, like the native stack.

2

MSQ currently doesn't support UNION queries. However, in the query stack, there are two types of UNIONs:

UnionDataSource - Very limited. MSQ detects this and throws a QueryNotSupported fault, which is the expected behavior
Top-level union - Works around the shortcomings of 1.

However 2) is executed sequentially by the SQL layer and the results are appended sequentially. For a simple query like

SELECT * FROM foo
UNION ALL
SELECT * FROM foo2

SQL would execute SELECT * FROM foo and SELECT * FROM foo2 and concat the results together.
This works fine for working with engines producing results synchronously like sql-native where we return the results, however, for MSQ, which produces results asynchronously, the concatenation logic doesn't work as expected since don't wait for the query to finish, fetch the results and submit the second query.

To make matters worse, the SQL layer submits the first query, gets the query ID back as the result, and then executes the second query (that fails). Therefore we only submit the partial query successfully and we might even get the incorrect results back.

This PR introduces the engine feature ALLOW_TOP_LEVEL_UNION_ALL that dictates whether the planner can plan the query using top-level union alls. MSQ disallows this, so the queries are forced to plan using the union data source, which will return query not supported exception.

This flag will also be useful once we start supporting unions in MSQ, which we'd want to exclusively execute using UnionDataSource, and the flag would seamlessly tie in with the query paths we'd wanna take when planning unions then.

With the change, the following query:

native tasks plan query as before (top-level union all)

MSQ tasks plan can't plan query with top-level union all, therefore use the UnionDataSource to plan the query, which then ultimately fails with QueryNotSupported in MSQ

3

Check out the changes in the DataSourcePlan which allows the union to be planned in the SQL stack

4

TBD

Release note

MSQ can execute UNION ALL queries with UnionDataSource.

This PR has:

LakshSingla · 2023-09-14T17:27:07Z

DruidSortUnionRule has a defensive check, therefore, can't add more tests to satisfy code coverage. I think there's little value in trying to satisfy it.

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQFaultsTest.java

gianm · 2023-09-14T17:52:48Z

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQFaultsTest.java

+                + "SELECT * FROM foo\n")
+        .setExpectedRowSignature(rowSignature)
+        .setExpectedDataSource("foo1")
+        .setExpectedMSQFault(QueryNotSupportedFault.instance())


Hmm… I am confused about why this yields a QueryNotSupportedFault. Shouldn't it fail to plan, and generate a planner error instead of an MSQ fault?

It does plan using the UnionDataSource, which then goes into MSQ.

Added a comment on how it is getting planned.

I guess this test can be removed. It should be planned using a UnionDataSource.
Can we also add the NativeQuery for assertion ?

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQFaultsTest.java

gianm · 2023-09-14T17:58:43Z

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQFaultsTest.java

+                DruidException.Persona.ADMIN,
+                DruidException.Category.INVALID_INPUT,
+                "general"
+            ).expectMessageIs("Query planning failed for unknown reason, our best guess is this "


Where is this error coming from? Looking at the code for the union rule, I would think it can't happen, because it's generated by isCompatible, which isn't called when the ALLOW_TOP_LEVEL_UNION_ALL feature is missing. The error should be something about UNION ALL being unsupported for this engine.

It tries to plan the query using the UnionDataSourceRule, goes into the isCompatible then, and then rewrites the already set planning error.
This gets executed using UnionDataSourceRule since the column names match, isCompatible returns and not the top-level union all.

I'll rename the two test cases that I added

Added comments and renamed the test cases. Hope they clarify the confusion

sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidUnionRule.java

sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidFreeUnionDataSourceRule.java

sql/src/main/java/org/apache/druid/sql/calcite/rel/DruidFreeUnionDataSourceRel.java

soumyava · 2023-10-02T05:16:58Z

sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidFreeUnionDataSourceRule.java

+      // No need to set the planning error here
+      return false;
+    }
+    if (!firstColumnNames.equals(secondColumnNames)) {


Due to casting Calcite might change the name to something like EXPR$0. In such a case this does not allow the onMatch to trigger. What's our plan of action for handling such cases ?

I am debating whether to act stringently and only allow the same names and types (no implicit casts) when using union all with this rule as well. Is there a way to realize that there is an implicit cast done, if so, then we can build something around it to remap it to the original variable name, otherwise it would be a hassle for the user since he won't be able to reference the original column.

Sidenote: I am debating whether to pull the planning changes out of the PR into their separate PR.

If the union is nested inside a subquery, it would be difficult for the upper callers to reference it once it has changed into the form EXPR$0, due to implicit casting. Therefore if we can identify that it has been cast implicitly, then we should remap it back to the original column name, else I am debating that we should include name and the type check,

sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidFreeUnionDataSourceRule.java

cryptoe

Went through the initial code. Left some initial review.

cryptoe · 2023-10-03T09:35:57Z

processing/src/main/java/org/apache/druid/query/UnionDataSource.java

@@ -36,13 +37,17 @@
 import java.util.function.Function;
 import java.util.stream.Collectors;

+/**
+ * TODO(laksh):


Lets add a java doc here.

cryptoe · 2023-10-03T09:40:10Z

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQFaultsTest.java

+                + "SELECT * FROM foo\n")
+        .setExpectedRowSignature(rowSignature)
+        .setExpectedDataSource("foo1")
+        .setExpectedMSQFault(QueryNotSupportedFault.instance())


I guess this test can be removed. It should be planned using a UnionDataSource.
Can we also add the NativeQuery for assertion ?

cryptoe · 2023-10-03T09:42:05Z

...core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteUnionQueryMSQTest.java

+   */
+  @Test
+  @Override
+  public void testUnionIsUnplannable()


Is this still required ?

cryptoe · 2023-10-03T09:43:01Z

...core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteUnionQueryMSQTest.java

+  }
+
+  @Test
+  public void testUnionOnSubqueries()


I guess this test can be marked with ignore till we have the new calcite rule in place.

cryptoe · 2023-10-03T09:48:40Z

...sions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/DataSourcePlan.java

@@ -170,6 +171,18 @@ public static DataSourcePlan forDataSource(
          minStageNumber,
          broadcast
      );
+    } else if (dataSource instanceof UnionDataSource) {


Please update the MSQ known issues and the docs where ever we are calling union all as unsupported in MSQ.

cryptoe

Changes lgtm!!

cryptoe · 2023-10-06T11:55:02Z

...core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteUnionQueryMSQTest.java

+
+  /**
+   * Doesn't pass through Druid however the planning error is different as it rewrites to a union datasource.
+   * This test is disabled because MSQ wants to support union datasources, and it makes little sense to add highly


This comment seems outdated.

cryptoe · 2023-10-06T11:56:54Z

processing/src/main/java/org/apache/druid/query/UnionDataSource.java

+                        if (!(input instanceof TableDataSource)) {
+                          throw DruidException.defensive("should be table");
+                        }
+                        return Iterables.getOnlyElement(input.getTableNames());


Nit:Lets avoid using Iterables.getOnlyElement(). Lets use CollectionUtils.getOnlyElement()

MSQ now supports UNION ALL with UnionDataSource

init

bc40808

github-actions bot added Area - Batch Ingestion Area - Querying MSQ labels Sep 13, 2023

LakshSingla added 4 commits September 14, 2023 10:49

fix test, checkstyle

a96dcb6

remove CoreRules.UNION_TO_DISTINCT

448b583

add CoreRules.UNION_TO_DISTINCT

2f7b160

Trigger Build

971589b

LakshSingla added Area - MSQ For multi stage queries - https://github.com/apache/druid/issues/12262 and removed MSQ labels Sep 14, 2023

gianm reviewed Sep 14, 2023

View reviewed changes

review comments

54a664a

github-actions bot added the MSQ label Sep 14, 2023

LakshSingla requested a review from gianm September 14, 2023 20:40

fix CalcitePlannerModuleTest

75652e6

LakshSingla removed the MSQ label Sep 15, 2023

LakshSingla added 4 commits September 26, 2023 10:19

Merge branch 'master' into msq-disallow-top-level-union-all

fd652a7

use real engine instead of mock

9b31be6

changes

7e60702

add new rules, fix some bugs

a6ee632

github-advanced-security bot found potential problems Sep 28, 2023

View reviewed changes

LakshSingla changed the title ~~Disallow top-level UNION ALLs in MSQ~~ Unions in MSQ Sep 29, 2023

LakshSingla changed the title ~~Unions in MSQ~~ UNION ALLs in MSQ Sep 29, 2023

soumyava reviewed Oct 2, 2023

View reviewed changes

changes to the new rules

9620c4e

github-advanced-security bot found potential problems Oct 3, 2023

View reviewed changes

sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidFreeUnionDataSourceRule.java Fixed Show fixed Hide fixed

cryptoe reviewed Oct 3, 2023

View reviewed changes

LakshSingla added this to the 28.0 milestone Oct 3, 2023

remove rule and rel

0ee8d50

LakshSingla added 2 commits October 5, 2023 18:49

review

92f35c0

Merge branch 'master' into msq-disallow-top-level-union-all

feb8875

cryptoe approved these changes Oct 6, 2023

View reviewed changes

LakshSingla added 3 commits October 9, 2023 11:14

test cases

368ae08

readd decoupled ignore

b09e300

tests

7c4e048

LakshSingla merged commit 549ef56 into apache:master Oct 9, 2023
81 checks passed

LakshSingla mentioned this pull request Oct 9, 2023

Fix compilation failure in master #15111

Merged

ektravel pushed a commit to ektravel/druid that referenced this pull request Oct 16, 2023

UNION ALLs in MSQ (apache#14981)

3a52f8b

MSQ now supports UNION ALL with UnionDataSource

LakshSingla mentioned this pull request Nov 4, 2023

[DRAFT] 28.0.0 release notes #15326

Closed

CaseyPan pushed a commit to CaseyPan/druid that referenced this pull request Nov 17, 2023

UNION ALLs in MSQ (apache#14981)

9f3c7d7

MSQ now supports UNION ALL with UnionDataSource

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

UNION ALLs in MSQ #14981

UNION ALLs in MSQ #14981

LakshSingla commented Sep 13, 2023 •

edited

Loading

LakshSingla commented Sep 14, 2023

gianm Sep 14, 2023

LakshSingla Sep 14, 2023

LakshSingla Sep 14, 2023

cryptoe Oct 3, 2023

gianm Sep 14, 2023

LakshSingla Sep 14, 2023 •

edited

Loading

LakshSingla Sep 14, 2023 •

edited

Loading

LakshSingla Sep 14, 2023

soumyava Oct 2, 2023

LakshSingla Oct 2, 2023

LakshSingla Oct 2, 2023

cryptoe left a comment

cryptoe Oct 3, 2023

cryptoe Oct 3, 2023

cryptoe Oct 3, 2023

cryptoe Oct 3, 2023

cryptoe Oct 3, 2023

cryptoe left a comment

cryptoe Oct 6, 2023

cryptoe Oct 6, 2023

UNION ALLs in MSQ #14981

UNION ALLs in MSQ #14981

Conversation

LakshSingla commented Sep 13, 2023 • edited Loading

Description

1

2

3

4

Release note

LakshSingla commented Sep 14, 2023

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

LakshSingla Sep 14, 2023 • edited Loading

Choose a reason for hiding this comment

LakshSingla Sep 14, 2023 • 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

cryptoe 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

cryptoe left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

LakshSingla commented Sep 13, 2023 •

edited

Loading

LakshSingla Sep 14, 2023 •

edited

Loading

LakshSingla Sep 14, 2023 •

edited

Loading