test(substrait): update TPCH tests #12462

vbarua · 2024-09-13T23:23:33Z

Which issue does this PR close?

Part of fixing/implementing #12347

Rationale for this change

The work for #12347 caused Substrait TPCH test failures. Investigations showed that the vendored plans being used for them were quite old, and in fact wrong in some cases. As a result, the tests have been updated in the upstream in substrait-io/consumer-testing#105

I've chosen to update these tests as their own PR to make it easier to review the follow up change for #12347.

What changes are included in this PR?

Functional Changes

Allows for more flexible handling of IntervalDayToSeconds Substrait literals when the precision_mode is not set.

Test Changes

Brings in updated Substrait TPCH plans from the https://github.com/substrait-io/consumer-testing repo.
Updates TestSchemaCollector to traverse expressions in order to collect schemas that are found in subqueries.

Are these changes tested?

The changes are primarily to tests. Changes were made to allow the tests to continue running.

Are there any user-facing changes?

A specific IntervalDayToSeconds literal combinations that was rejected before is now allowed.

version: https://github.com/substrait-io/consumer-testing/tree/d9a0453

Substrait plans can contains relations within expressions, for example in Subqueries. As such, in order to fully collect schemas we must traverse through relation and expressions.

When subseconds is not set, allow for missing precision_mode

vbarua · 2024-09-13T23:27:30Z

datafusion/substrait/tests/cases/consumer_integration.rs

-            "tests/testdata/tpch/lineitem.csv",
-        )])
-        .await?;
-        let path = "tests/testdata/tpch_substrait_plans/query_1.json";


I noticed a lot of duplication so I encapsulated the test setup code in tpch_plan_to_string.

vbarua · 2024-09-13T23:29:14Z

datafusion/substrait/tests/cases/consumer_integration.rs

        let proto = serde_json::from_reader::<_, Plan>(BufReader::new(
            File::open(path).expect("file not found"),
        ))
        .expect("failed to parse json");

+        let ctx = add_plan_schemas_to_ctx(SessionContext::new(), &proto)?;


Instead of reading the CSV input files to generate schemas, we can generate them directly from the Substrait plans.

vbarua · 2024-09-13T23:32:57Z

datafusion/substrait/tests/utils.rs

+                if existing_table.schema() != schema {
+                    return substrait_err!(
+                        "Substrait plan contained the same table {} with different schemas.\nSchema 1: {}\nSchema 2: {}",
+                        table_reference, existing_table.schema(), schema);


I'm checking for the case of the same table having different schemas because I've been bitten by it before due to schema pruning shenanigans and it took a while to figure out what was hapenning.

vbarua · 2024-09-13T23:34:30Z

datafusion/substrait/tests/utils.rs

-        fn collect_schemas(plan: &Plan) -> Vec<(TableReference, Arc<dyn TableProvider>)> {
+        fn collect_schemas(
+            plan: &Plan,
+        ) -> Result<Vec<(TableReference, Arc<dyn TableProvider>)>> {


I switched this to be Result based because eventually I would like this to be part of the core library. Collecting schemas like this is super useful.

vbarua · 2024-09-13T23:40:12Z

datafusion/substrait/tests/utils.rs

-                RelType::Fetch(f) => self.apply(f.input.as_ref().map(|b| b.as_ref())),
-                RelType::Aggregate(a) => self.apply(a.input.as_ref().map(|b| b.as_ref())),
-                RelType::Sort(s) => self.apply(s.input.as_ref().map(|b| b.as_ref())),
+        fn collect_schemas_from_rel(&mut self, rel: &Rel) -> Result<()> {


meta: writing traversals like this is quite painful. It would be good go get some visitation machinery added to substrait-rust.

I wonder if https://docs.rs/datafusion/latest/datafusion/common/tree_node/trait.TreeNode.html could serve as inspiration. @peter-toth put a lot of effort into improving that API and I think i is pretty neat now

vbarua · 2024-09-13T23:45:48Z

datafusion/substrait/tests/utils.rs

+                        .subquery_type
+                        .as_ref()
+                        .ok_or(substrait_datafusion_err!("subquery_type must be set"))?;
+                    match subquery_type {


A number of the TPCH plans include Subquery expressions. Without this traversal, the schemas generated from the plans were missing tables.

vbarua · 2024-09-14T00:05:05Z

datafusion/substrait/tests/cases/consumer_integration.rs

-            \n            Aggregate: groupBy=[[]], aggr=[[min(FILENAME_PLACEHOLDER_5.ps_supplycost)]]\
-            \n              Projection: FILENAME_PLACEHOLDER_5.ps_supplycost\
-            \n                Filter: FILENAME_PLACEHOLDER_5.ps_partkey = FILENAME_PLACEHOLDER_5.ps_partkey AND FILENAME_PLACEHOLDER_6.s_suppkey = FILENAME_PLACEHOLDER_5.ps_suppkey AND FILENAME_PLACEHOLDER_6.s_nationkey = FILENAME_PLACEHOLDER_7.N_NATIONKEY AND FILENAME_PLACEHOLDER_7.N_REGIONKEY = FILENAME_PLACEHOLDER_8.R_REGIONKEY AND FILENAME_PLACEHOLDER_8.R_NAME = CAST(Utf8(\"EUROPE\") AS Utf8)\
-            \n                  Inner Join:  Filter: Boolean(true)\


Note that the Inner Join with an always true filter has been replaced by an explicit CrossJoin, which is equivalent.

vbarua · 2024-09-14T00:15:42Z

datafusion/substrait/tests/testdata/tpch_substrait_plans/query_1.json

-                "input": {
-                  "project": {
-                    "common": {
-


The newly copied plans set the remap field correctly. Their absence here didn't matter because DataFusion ignored them, but once I started adding remap support these plans broke because they were incorrect.

alamb

Thank you very much @vbarua -- this PR looks like a significant improvement of the substrait functionality to me. Also thank you for your attention to detail

cc @Blizzara @Lordworms and @waynexia who may also be interested in this PR

I suggest we file another ticket and list the TPCH queries that don't currently run to make the current state of substrait more visible

alamb · 2024-09-15T12:23:36Z

datafusion/substrait/tests/cases/consumer_integration.rs

-             \n      Projection: FILENAME_PLACEHOLDER_0.l_returnflag, FILENAME_PLACEHOLDER_0.l_linestatus, FILENAME_PLACEHOLDER_0.l_quantity, FILENAME_PLACEHOLDER_0.l_extendedprice, FILENAME_PLACEHOLDER_0.l_extendedprice * (CAST(Int32(1) AS Decimal128(19, 0)) - FILENAME_PLACEHOLDER_0.l_discount), FILENAME_PLACEHOLDER_0.l_extendedprice * (CAST(Int32(1) AS Decimal128(19, 0)) - FILENAME_PLACEHOLDER_0.l_discount) * (CAST(Int32(1) AS Decimal128(19, 0)) + FILENAME_PLACEHOLDER_0.l_tax), FILENAME_PLACEHOLDER_0.l_discount\
-             \n        Filter: FILENAME_PLACEHOLDER_0.l_shipdate <= Date32(\"1998-12-01\") - IntervalDayTime(\"IntervalDayTime { days: 120, milliseconds: 0 }\")\
-             \n          TableScan: FILENAME_PLACEHOLDER_0 projection=[l_orderkey, l_partkey, l_suppkey, l_linenumber, l_quantity, l_extendedprice, l_discount, l_tax, l_returnflag, l_linestatus, l_shipdate, l_commitdate, l_receiptdate, l_shipinstruct, l_shipmode, l_comment]"
+            "Projection: LINEITEM.L_RETURNFLAG, LINEITEM.L_LINESTATUS, sum(LINEITEM.L_QUANTITY) AS SUM_QTY, sum(LINEITEM.L_EXTENDEDPRICE) AS SUM_BASE_PRICE, sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT) AS SUM_DISC_PRICE, sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT * Int32(1) + LINEITEM.L_TAX) AS SUM_CHARGE, avg(LINEITEM.L_QUANTITY) AS AVG_QTY, avg(LINEITEM.L_EXTENDEDPRICE) AS AVG_PRICE, avg(LINEITEM.L_DISCOUNT) AS AVG_DISC, count(Int64(1)) AS COUNT_ORDER\


those certainly look nicer

alamb · 2024-09-15T12:24:31Z

datafusion/substrait/tests/cases/consumer_integration.rs

-        \n              TableScan: NATION projection=[N_NATIONKEY, N_NAME, N_REGIONKEY, N_COMMENT]\
-        \n            TableScan: REGION projection=[R_REGIONKEY, R_NAME, R_COMMENT]");
+    #[ignore]
+    #[tokio::test]


for other reviewers, these plans weren't included in the initial coverage either

alamb · 2024-09-15T12:25:26Z

datafusion/substrait/tests/cases/consumer_integration.rs

-        \n          TableScan: FILENAME_PLACEHOLDER_0 projection=[l_orderkey, l_partkey, l_suppkey, l_linenumber, l_quantity, l_extendedprice, l_discount, l_tax, l_returnflag, l_linestatus, l_shipdate, l_commitdate, l_receiptdate, l_shipinstruct, l_shipmode, l_comment]\
-        \n          TableScan: FILENAME_PLACEHOLDER_1 projection=[p_partkey, p_name, p_mfgr, p_brand, p_type, p_size, p_container, p_retailprice, p_comment]");
+    #[ignore]
+    #[tokio::test]


I wonder if we should be tracking in a ticket somewhere which TPCH plans aren't currently handled by the substrait implementation 🤔

alamb · 2024-09-15T12:27:28Z

datafusion/substrait/tests/utils.rs

-                RelType::Fetch(f) => self.apply(f.input.as_ref().map(|b| b.as_ref())),
-                RelType::Aggregate(a) => self.apply(a.input.as_ref().map(|b| b.as_ref())),
-                RelType::Sort(s) => self.apply(s.input.as_ref().map(|b| b.as_ref())),
+        fn collect_schemas_from_rel(&mut self, rel: &Rel) -> Result<()> {


I wonder if https://docs.rs/datafusion/latest/datafusion/common/tree_node/trait.TreeNode.html could serve as inspiration. @peter-toth put a lot of effort into improving that API and I think i is pretty neat now

Blizzara · 2024-09-16T10:03:17Z

LGTM, thanks!

datafusion/substrait/tests/cases/consumer_integration.rs

vbarua · 2024-09-16T16:57:04Z

I've created the following to track the issues with the TPCH queries:

alamb · 2024-09-16T17:02:35Z

Thanks again @vbarua and @Blizzara

vbarua added 7 commits September 13, 2024 15:45

test(substrait): rename vendored test files to match upstream

00226c0

test(substrait): update vendored TPCH plans from substrait-consumer

6bbb426

version: https://github.com/substrait-io/consumer-testing/tree/d9a0453

test(substrait): collect schemas through from expressions

364017d

Substrait plans can contains relations within expressions, for example in Subqueries. As such, in order to fully collect schemas we must traverse through relation and expressions.

feat(substrait): more flexible IntervalDayToSecond handling

508ed8d

When subseconds is not set, allow for missing precision_mode

test(substrait): update existing TPCH tests

99e1605

test(substrait): include new TPCH tests

a04ebd0

test(substrait): remove unused test csv files

00de470

github-actions bot added the substrait label Sep 13, 2024

test(substrait): minor TestSchemaCollector updates

73b6b90

vbarua marked this pull request as ready for review September 14, 2024 00:02

vbarua commented Sep 14, 2024

View reviewed changes

vbarua changed the title ~~test(substrait): update TPCH test~~ test(substrait): update TPCH tests Sep 14, 2024

vbarua commented Sep 14, 2024

View reviewed changes

clippy

9f00142

alamb approved these changes Sep 15, 2024

View reviewed changes

vbarua commented Sep 16, 2024

View reviewed changes

datafusion/substrait/tests/cases/consumer_integration.rs Show resolved Hide resolved

alamb merged commit 257e140 into apache:main Sep 16, 2024
24 checks passed

This pull request was closed.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

test(substrait): update TPCH tests #12462

test(substrait): update TPCH tests #12462

vbarua commented Sep 13, 2024 •

edited

Loading

vbarua Sep 13, 2024

vbarua Sep 13, 2024

vbarua Sep 13, 2024

vbarua Sep 13, 2024

vbarua Sep 13, 2024

alamb Sep 15, 2024

vbarua Sep 13, 2024

vbarua Sep 14, 2024

vbarua Sep 14, 2024 •

edited

Loading

alamb left a comment

alamb Sep 15, 2024

alamb Sep 15, 2024

alamb Sep 15, 2024

alamb Sep 15, 2024

Blizzara commented Sep 16, 2024

vbarua commented Sep 16, 2024

alamb commented Sep 16, 2024

test(substrait): update TPCH tests #12462

test(substrait): update TPCH tests #12462

Conversation

vbarua commented Sep 13, 2024 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Functional Changes

Test Changes

Are these changes tested?

Are there any user-facing changes?

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

vbarua Sep 14, 2024 • edited Loading

Choose a reason for hiding this comment

alamb 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

Blizzara commented Sep 16, 2024

vbarua commented Sep 16, 2024

alamb commented Sep 16, 2024

vbarua commented Sep 13, 2024 •

edited

Loading

vbarua Sep 14, 2024 •

edited

Loading