From b0083cbf2aab93d400abae9441fa9c8f0b15786d Mon Sep 17 00:00:00 2001 From: Richard Tia Date: Thu, 12 Sep 2024 17:28:28 -0700 Subject: [PATCH] feat: add write relation tests (#106) --- .../queries/sql/relations/write_relations.py | 28 +++++ .../functional/write_relation_configs.py | 23 ++++ .../relations/test_write_relation.py | 103 ++++++++++++++++++ 3 files changed, 154 insertions(+) create mode 100644 substrait_consumer/functional/queries/sql/relations/write_relations.py create mode 100644 substrait_consumer/functional/write_relation_configs.py create mode 100644 substrait_consumer/tests/functional/relations/test_write_relation.py diff --git a/substrait_consumer/functional/queries/sql/relations/write_relations.py b/substrait_consumer/functional/queries/sql/relations/write_relations.py new file mode 100644 index 0000000..71af90c --- /dev/null +++ b/substrait_consumer/functional/queries/sql/relations/write_relations.py @@ -0,0 +1,28 @@ +from substrait_consumer.producers.duckdb_producer import DuckDBProducer +from substrait_consumer.producers.datafusion_producer import DataFusionProducer +from substrait_consumer.producers.isthmus_producer import IsthmusProducer + +WRITE_RELATIONS = { + "insert": ( + """ + INSERT INTO '{}' (r_regionkey, r_name, r_comment) + VALUES (99999, 'region_name', 'region comment'); + """, + [DuckDBProducer, DataFusionProducer, IsthmusProducer], + ), + "update": ( + """ + UPDATE '{}' + SET c_address = 'Substait Avenue', c_phone = '123-456-7890' + WHERE c_custkey = 1; + """, + [DuckDBProducer, DataFusionProducer, IsthmusProducer], + ), + "delete": ( + """ + DELETE FROM '{}' + WHERE c_custkey = 1; + """, + [DuckDBProducer, DataFusionProducer, IsthmusProducer], + ), +} diff --git a/substrait_consumer/functional/write_relation_configs.py b/substrait_consumer/functional/write_relation_configs.py new file mode 100644 index 0000000..c531b2f --- /dev/null +++ b/substrait_consumer/functional/write_relation_configs.py @@ -0,0 +1,23 @@ +from substrait_consumer.functional.queries.sql.relations.write_relations import ( + WRITE_RELATIONS) + +WRITE_RELATION_TESTS = ( + { + "test_name": "insert", + "file_names": ["region.parquet"], + "sql_query": WRITE_RELATIONS["insert"], + "ibis_expr": None + }, + { + "test_name": "update", + "file_names": ["customer.parquet"], + "sql_query": WRITE_RELATIONS["update"], + "ibis_expr": None + }, + { + "test_name": "delete", + "file_names": ["customer.parquet"], + "sql_query": WRITE_RELATIONS["delete"], + "ibis_expr": None + }, +) diff --git a/substrait_consumer/tests/functional/relations/test_write_relation.py b/substrait_consumer/tests/functional/relations/test_write_relation.py new file mode 100644 index 0000000..df499ec --- /dev/null +++ b/substrait_consumer/tests/functional/relations/test_write_relation.py @@ -0,0 +1,103 @@ +from typing import Callable, Iterable + +import duckdb +from ibis.expr.types.relations import Table +from ibis_substrait.tests.compiler.conftest import * + +from substrait_consumer.functional.write_relation_configs import ( + WRITE_RELATION_TESTS) +from substrait_consumer.functional.common import ( + generate_snapshot_results, + substrait_consumer_sql_test, substrait_producer_sql_test) +from substrait_consumer.parametrization import custom_parametrization + + +@pytest.fixture +def mark_producer_tests_as_xfail(request): + """Marks a subset of tests as expected to be fail.""" + test_case_name = request.node.callspec.id.split('-')[-1] + if test_case_name in ["insert", "update", "delete"]: + pytest.skip(reason='Creating substrait plans with write relations is not supported') + + +@pytest.fixture +def mark_consumer_tests_as_xfail(request): + """Marks a subset of tests as expected to be fail.""" + test_case_name = request.node.callspec.id.split('-')[-1] + if test_case_name in ["insert", "update", "delete"]: + pytest.skip(reason='Creating substrait plans with write relations is not supported') + + +@pytest.mark.usefixtures("prepare_tpch_parquet_data") +class TestWriteRelation: + """ + Test Class verifying different consumers are able to run substrait plans + that include substrait write relations. + """ + + @staticmethod + @pytest.fixture(scope="class", autouse=True) + def setup_teardown_class(request): + cls = request.cls + + cls.db_connection = duckdb.connect() + cls.db_connection.execute("install substrait") + cls.db_connection.execute("load substrait") + cls.created_tables = set() + + yield + + cls.db_connection.close() + + @custom_parametrization(WRITE_RELATION_TESTS) + @pytest.mark.produce_substrait_snapshot + @pytest.mark.usefixtures('mark_producer_tests_as_xfail') + def test_producer_write_relations( + self, + snapshot, + test_name: str, + file_names: Iterable[str], + sql_query: tuple, + ibis_expr: Callable[[Table], Table], + producer, + partsupp + ) -> None: + test_name = f"write_relation_snapshots:{test_name}" + substrait_producer_sql_test( + test_name, + snapshot, + self.db_connection, + self.created_tables, + file_names, + sql_query, + ibis_expr, + producer, + partsupp, + validate=True + ) + + @custom_parametrization(WRITE_RELATION_TESTS) + @pytest.mark.consume_substrait_snapshot + @pytest.mark.usefixtures('mark_consumer_tests_as_xfail') + def test_consumer_write_relations( + self, + snapshot, + test_name: str, + file_names: Iterable[str], + sql_query: tuple, + ibis_expr: Callable[[Table], Table], + producer, + consumer, + ) -> None: + test_name = f"write_relation_snapshots:{test_name}" + substrait_consumer_sql_test( + test_name, + snapshot, + self.db_connection, + self.created_tables, + file_names, + sql_query, + ibis_expr, + producer, + consumer, + )