feat: express TPC-H plan generation with existing infra (#172)

This PR uses `substrait_producer_sql_test` from `common.py` to express the TPC-H test targets that regenerate the Substrait plans using Isthmus. This replaces test-specific code with generic code and is, thus, a steps towards reducing the complexity of the tests. As a consequence, the paths of some snapshots change; however, note in the corresponding git commit that all files are renamed without content changes. Signed-off-by: Ingo Müller <[email protected]>
substrait-io · Dec 17, 2024 · 9aa6cdf · 9aa6cdf
1 parent 6659d6f
commit 9aa6cdf
Show file tree

Hide file tree

Showing 47 changed files with 46 additions and 30 deletions.
diff --git a/substrait_consumer/functional/common.py b/substrait_consumer/functional/common.py
@@ -21,9 +21,11 @@
 RELATION_SNAPSHOT_DIR = (
     Path(__file__).parent.parent / "tests" / "functional" / "relations"
 )
+TPCH_SNAPSHOT_DIR = Path(__file__).parent.parent / "tests" / "functional" / "tpch"
 SNAPSHOT_DIR = {
     "function": FUNCTION_SNAPSHOT_DIR,
     "relation": RELATION_SNAPSHOT_DIR,
+    "tpch": TPCH_SNAPSHOT_DIR,
 }
 
 

diff --git a/...rait_consumer/tests/functional/tpch/q01_snapshots/IsthmusProducer/q01-isthmus_outcome.txt b/...rait_consumer/tests/functional/tpch/q01_snapshots/IsthmusProducer/q01-isthmus_outcome.txt
@@ -0,0 +1 @@
+True
diff --git a/...s/tpch_substrait_plans/query_01_plan.json → ...1_snapshots/IsthmusProducer/q01_plan.json b/...s/tpch_substrait_plans/query_01_plan.json → ...1_snapshots/IsthmusProducer/q01_plan.json
diff --git a/...rait_consumer/tests/functional/tpch/q02_snapshots/IsthmusProducer/q02-isthmus_outcome.txt b/...rait_consumer/tests/functional/tpch/q02_snapshots/IsthmusProducer/q02-isthmus_outcome.txt
@@ -0,0 +1 @@
+True
diff --git a/...s/tpch_substrait_plans/query_02_plan.json → ...2_snapshots/IsthmusProducer/q02_plan.json b/...s/tpch_substrait_plans/query_02_plan.json → ...2_snapshots/IsthmusProducer/q02_plan.json
diff --git a/...rait_consumer/tests/functional/tpch/q03_snapshots/IsthmusProducer/q03-isthmus_outcome.txt b/...rait_consumer/tests/functional/tpch/q03_snapshots/IsthmusProducer/q03-isthmus_outcome.txt
@@ -0,0 +1 @@
+True
diff --git a/...s/tpch_substrait_plans/query_03_plan.json → ...3_snapshots/IsthmusProducer/q03_plan.json b/...s/tpch_substrait_plans/query_03_plan.json → ...3_snapshots/IsthmusProducer/q03_plan.json
diff --git a/...rait_consumer/tests/functional/tpch/q04_snapshots/IsthmusProducer/q04-isthmus_outcome.txt b/...rait_consumer/tests/functional/tpch/q04_snapshots/IsthmusProducer/q04-isthmus_outcome.txt
@@ -0,0 +1 @@
+True
diff --git a/...s/tpch_substrait_plans/query_04_plan.json → ...4_snapshots/IsthmusProducer/q04_plan.json b/...s/tpch_substrait_plans/query_04_plan.json → ...4_snapshots/IsthmusProducer/q04_plan.json
diff --git a/...rait_consumer/tests/functional/tpch/q05_snapshots/IsthmusProducer/q05-isthmus_outcome.txt b/...rait_consumer/tests/functional/tpch/q05_snapshots/IsthmusProducer/q05-isthmus_outcome.txt
@@ -0,0 +1 @@
+True
diff --git a/...s/tpch_substrait_plans/query_05_plan.json → ...5_snapshots/IsthmusProducer/q05_plan.json b/...s/tpch_substrait_plans/query_05_plan.json → ...5_snapshots/IsthmusProducer/q05_plan.json
diff --git a/...rait_consumer/tests/functional/tpch/q06_snapshots/IsthmusProducer/q06-isthmus_outcome.txt b/...rait_consumer/tests/functional/tpch/q06_snapshots/IsthmusProducer/q06-isthmus_outcome.txt
@@ -0,0 +1 @@
+True
diff --git a/...s/tpch_substrait_plans/query_06_plan.json → ...6_snapshots/IsthmusProducer/q06_plan.json b/...s/tpch_substrait_plans/query_06_plan.json → ...6_snapshots/IsthmusProducer/q06_plan.json
diff --git a/...rait_consumer/tests/functional/tpch/q07_snapshots/IsthmusProducer/q07-isthmus_outcome.txt b/...rait_consumer/tests/functional/tpch/q07_snapshots/IsthmusProducer/q07-isthmus_outcome.txt
@@ -0,0 +1 @@
+True
diff --git a/...s/tpch_substrait_plans/query_07_plan.json → ...7_snapshots/IsthmusProducer/q07_plan.json b/...s/tpch_substrait_plans/query_07_plan.json → ...7_snapshots/IsthmusProducer/q07_plan.json
diff --git a/...rait_consumer/tests/functional/tpch/q08_snapshots/IsthmusProducer/q08-isthmus_outcome.txt b/...rait_consumer/tests/functional/tpch/q08_snapshots/IsthmusProducer/q08-isthmus_outcome.txt
@@ -0,0 +1 @@
+True
diff --git a/...s/tpch_substrait_plans/query_08_plan.json → ...8_snapshots/IsthmusProducer/q08_plan.json b/...s/tpch_substrait_plans/query_08_plan.json → ...8_snapshots/IsthmusProducer/q08_plan.json
diff --git a/...rait_consumer/tests/functional/tpch/q09_snapshots/IsthmusProducer/q09-isthmus_outcome.txt b/...rait_consumer/tests/functional/tpch/q09_snapshots/IsthmusProducer/q09-isthmus_outcome.txt
@@ -0,0 +1 @@
+True
diff --git a/...s/tpch_substrait_plans/query_09_plan.json → ...9_snapshots/IsthmusProducer/q09_plan.json b/...s/tpch_substrait_plans/query_09_plan.json → ...9_snapshots/IsthmusProducer/q09_plan.json
diff --git a/...rait_consumer/tests/functional/tpch/q10_snapshots/IsthmusProducer/q10-isthmus_outcome.txt b/...rait_consumer/tests/functional/tpch/q10_snapshots/IsthmusProducer/q10-isthmus_outcome.txt
@@ -0,0 +1 @@
+True
diff --git a/...s/tpch_substrait_plans/query_10_plan.json → ...0_snapshots/IsthmusProducer/q10_plan.json b/...s/tpch_substrait_plans/query_10_plan.json → ...0_snapshots/IsthmusProducer/q10_plan.json
diff --git a/...rait_consumer/tests/functional/tpch/q11_snapshots/IsthmusProducer/q11-isthmus_outcome.txt b/...rait_consumer/tests/functional/tpch/q11_snapshots/IsthmusProducer/q11-isthmus_outcome.txt
@@ -0,0 +1 @@
+True
diff --git a/...s/tpch_substrait_plans/query_11_plan.json → ...1_snapshots/IsthmusProducer/q11_plan.json b/...s/tpch_substrait_plans/query_11_plan.json → ...1_snapshots/IsthmusProducer/q11_plan.json
diff --git a/...rait_consumer/tests/functional/tpch/q12_snapshots/IsthmusProducer/q12-isthmus_outcome.txt b/...rait_consumer/tests/functional/tpch/q12_snapshots/IsthmusProducer/q12-isthmus_outcome.txt
@@ -0,0 +1 @@
+True
diff --git a/...s/tpch_substrait_plans/query_12_plan.json → ...2_snapshots/IsthmusProducer/q12_plan.json b/...s/tpch_substrait_plans/query_12_plan.json → ...2_snapshots/IsthmusProducer/q12_plan.json
diff --git a/...rait_consumer/tests/functional/tpch/q13_snapshots/IsthmusProducer/q13-isthmus_outcome.txt b/...rait_consumer/tests/functional/tpch/q13_snapshots/IsthmusProducer/q13-isthmus_outcome.txt
@@ -0,0 +1 @@
+True
diff --git a/...s/tpch_substrait_plans/query_13_plan.json → ...3_snapshots/IsthmusProducer/q13_plan.json b/...s/tpch_substrait_plans/query_13_plan.json → ...3_snapshots/IsthmusProducer/q13_plan.json
diff --git a/...rait_consumer/tests/functional/tpch/q14_snapshots/IsthmusProducer/q14-isthmus_outcome.txt b/...rait_consumer/tests/functional/tpch/q14_snapshots/IsthmusProducer/q14-isthmus_outcome.txt
@@ -0,0 +1 @@
+True
diff --git a/...s/tpch_substrait_plans/query_14_plan.json → ...4_snapshots/IsthmusProducer/q14_plan.json b/...s/tpch_substrait_plans/query_14_plan.json → ...4_snapshots/IsthmusProducer/q14_plan.json
diff --git a/...rait_consumer/tests/functional/tpch/q15_snapshots/IsthmusProducer/q15-isthmus_outcome.txt b/...rait_consumer/tests/functional/tpch/q15_snapshots/IsthmusProducer/q15-isthmus_outcome.txt
@@ -0,0 +1 @@
+<java class 'org.apache.calcite.runtime.CalciteContextException'>
diff --git a/...s/tpch_substrait_plans/query_15_plan.json → ...5_snapshots/IsthmusProducer/q15_plan.json b/...s/tpch_substrait_plans/query_15_plan.json → ...5_snapshots/IsthmusProducer/q15_plan.json
diff --git a/...rait_consumer/tests/functional/tpch/q16_snapshots/IsthmusProducer/q16-isthmus_outcome.txt b/...rait_consumer/tests/functional/tpch/q16_snapshots/IsthmusProducer/q16-isthmus_outcome.txt
@@ -0,0 +1 @@
+True
diff --git a/...s/tpch_substrait_plans/query_16_plan.json → ...6_snapshots/IsthmusProducer/q16_plan.json b/...s/tpch_substrait_plans/query_16_plan.json → ...6_snapshots/IsthmusProducer/q16_plan.json
diff --git a/...rait_consumer/tests/functional/tpch/q17_snapshots/IsthmusProducer/q17-isthmus_outcome.txt b/...rait_consumer/tests/functional/tpch/q17_snapshots/IsthmusProducer/q17-isthmus_outcome.txt
@@ -0,0 +1 @@
+True
diff --git a/...s/tpch_substrait_plans/query_17_plan.json → ...7_snapshots/IsthmusProducer/q17_plan.json b/...s/tpch_substrait_plans/query_17_plan.json → ...7_snapshots/IsthmusProducer/q17_plan.json
diff --git a/...rait_consumer/tests/functional/tpch/q18_snapshots/IsthmusProducer/q18-isthmus_outcome.txt b/...rait_consumer/tests/functional/tpch/q18_snapshots/IsthmusProducer/q18-isthmus_outcome.txt
@@ -0,0 +1 @@
+True
diff --git a/...s/tpch_substrait_plans/query_18_plan.json → ...8_snapshots/IsthmusProducer/q18_plan.json b/...s/tpch_substrait_plans/query_18_plan.json → ...8_snapshots/IsthmusProducer/q18_plan.json
diff --git a/...rait_consumer/tests/functional/tpch/q19_snapshots/IsthmusProducer/q19-isthmus_outcome.txt b/...rait_consumer/tests/functional/tpch/q19_snapshots/IsthmusProducer/q19-isthmus_outcome.txt
@@ -0,0 +1 @@
+True
diff --git a/...s/tpch_substrait_plans/query_19_plan.json → ...9_snapshots/IsthmusProducer/q19_plan.json b/...s/tpch_substrait_plans/query_19_plan.json → ...9_snapshots/IsthmusProducer/q19_plan.json
diff --git a/...rait_consumer/tests/functional/tpch/q20_snapshots/IsthmusProducer/q20-isthmus_outcome.txt b/...rait_consumer/tests/functional/tpch/q20_snapshots/IsthmusProducer/q20-isthmus_outcome.txt
@@ -0,0 +1 @@
+True
diff --git a/...s/tpch_substrait_plans/query_20_plan.json → ...0_snapshots/IsthmusProducer/q20_plan.json b/...s/tpch_substrait_plans/query_20_plan.json → ...0_snapshots/IsthmusProducer/q20_plan.json
diff --git a/...rait_consumer/tests/functional/tpch/q21_snapshots/IsthmusProducer/q21-isthmus_outcome.txt b/...rait_consumer/tests/functional/tpch/q21_snapshots/IsthmusProducer/q21-isthmus_outcome.txt
@@ -0,0 +1 @@
+True
diff --git a/...s/tpch_substrait_plans/query_21_plan.json → ...1_snapshots/IsthmusProducer/q21_plan.json b/...s/tpch_substrait_plans/query_21_plan.json → ...1_snapshots/IsthmusProducer/q21_plan.json
diff --git a/...rait_consumer/tests/functional/tpch/q22_snapshots/IsthmusProducer/q22-isthmus_outcome.txt b/...rait_consumer/tests/functional/tpch/q22_snapshots/IsthmusProducer/q22-isthmus_outcome.txt
@@ -0,0 +1 @@
+True
diff --git a/...s/tpch_substrait_plans/query_22_plan.json → ...2_snapshots/IsthmusProducer/q22_plan.json b/...s/tpch_substrait_plans/query_22_plan.json → ...2_snapshots/IsthmusProducer/q22_plan.json
diff --git a/substrait_consumer/tests/integration/test_isthmus_on_acero.py b/substrait_consumer/tests/integration/test_isthmus_on_acero.py
@@ -8,7 +8,7 @@
 from substrait_consumer.functional.utils import load_json
 from substrait_consumer.producers.duckdb_producer import DuckDBProducer
 
-PLAN_DIR = Path(__file__).parent / "queries" / "tpch_substrait_plans"
+PLAN_DIR = Path(__file__).parent.parent / "functional" / "tpch"
 
 CONFIG_DIR = Path(__file__).parent.parent / "integration"
 TPCH_CONFIG_DIR = CONFIG_DIR / "tpch"
@@ -68,7 +68,12 @@ def test_isthmus_substrait_plan(
     outcome_path = f"query_{tpch_num:02d}_outcome.txt"
 
     # Load Isthmus plan from file.
-    substrait_plan_path = PLAN_DIR / f"query_{tpch_num:02d}_plan.json"
+    substrait_plan_path = (
+        PLAN_DIR
+        / f"q{tpch_num:02d}_snapshots"
+        / "IsthmusProducer"
+        / f"q{tpch_num:02d}_plan.json"
+    )
     with open(substrait_plan_path, "r") as f:
         proto_bytes = f.read()
 

diff --git a/substrait_consumer/tests/integration/test_tpch_plans_valid.py b/substrait_consumer/tests/integration/test_tpch_plans_valid.py
@@ -4,15 +4,11 @@
 import pytest
 from pytest_snapshot.plugin import Snapshot
 
-from substrait_consumer.functional.common import check_match
+from substrait_consumer.functional.common import substrait_producer_sql_test
 from substrait_consumer.functional.utils import load_json
 from substrait_consumer.producers.duckdb_producer import DuckDBProducer
 from substrait_consumer.producers.isthmus_producer import IsthmusProducer
 
-PLAN_SNAPSHOT_DIR = (
-    Path(__file__).parent / "queries" / "tpch_substrait_plans"
-)
-
 CONFIG_DIR = Path(__file__).parent.parent / "integration"
 TPCH_CONFIG_DIR = CONFIG_DIR / "tpch"
 TEST_CASE_PATHS = list(
@@ -23,38 +19,29 @@
 
 @pytest.mark.parametrize(["path"], TEST_CASE_PATHS, ids=IDS)
 @pytest.mark.usefixtures("prepare_tpch_parquet_data")
+@pytest.mark.produce_substrait_snapshot
 def test_isthmus_substrait_plan_generation(
     path: Path,
     snapshot: Snapshot,
+    record_property,
     db_con: duckdb.DuckDBPyConnection,
-) -> None:
-    """
-    Generate the substrait plans using Isthmus.
-    """
+):
     test_case = load_json(CONFIG_DIR / path)
-    test_name = test_case["test_name"]
     local_files = test_case["local_files"]
     named_tables = test_case["named_tables"]
-    sql_query, supported_producers = test_case["sql_query"]
-
-    assert "isthmus" in supported_producers
-
-    tpch_num = test_name.split("_")[-1].zfill(2)
-    snapshot.snapshot_dir = PLAN_SNAPSHOT_DIR
-
+    sql_query = test_case["sql_query"]
     producer = IsthmusProducer()
-    producer.setup(db_con, local_files, named_tables)
-
-    try:
-        substrait_query = producer.produce_substrait(sql_query)
-    except BaseException as e:
-        snapshot.assert_match(str(type(e)), f"query_{tpch_num}_outcome.txt")
-        return
-
-    match_result = check_match(
-        snapshot, str(substrait_query), f"query_{tpch_num}_plan.json"
+    substrait_producer_sql_test(
+        path,
+        snapshot,
+        record_property,
+        db_con,
+        local_files,
+        named_tables,
+        sql_query,
+        producer,
+        validate=False,
     )
-    snapshot.assert_match(str(match_result), f"query_{tpch_num}_outcome.txt")
 
 
 @pytest.mark.parametrize(["path"], TEST_CASE_PATHS, ids=IDS)
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		<java class 'org.apache.calcite.runtime.CalciteContextException'>