- Overview
- Setup
- How to Run Tests
- TPCH Tests
- Function Tests
- CLI Tool
- How to Add Producers
- How to Add Consumers
This testing repository provides instructions on how to add and run substrait integration tests as well as a CLI tool for generating substrait plans from adhoc SQL queries.
The tests are organized into two categories; tpch tests (which test common benchmark queries)
and substrait function tests (which test individual extension functions). Test data is created
using DuckDB at the start of the test class using the prepare_tpch_parquet_data
fixture,
which is located in substrait_consumer/conftest.py
.
It is important to understand that there are two levels of "testing" in this repository:
- We want to determine whether producers and consumers produce the plans, results, schemas, error messages, etc. given a certain input. This is arguably the main purpose of this repository.
- We want to test whether the testing infrastructure (which is non-trivial) is behaving as intended and no bugs have been introduced on that level (such as loading a test case that doesn't exist or an exception in the testing logic).
The rationale of this repository is to make pytest
and, hence, CI pass if and
only if the second class of tests passes, i.e., if the testing infrastructure
runs as expected and, thus, irrespective of whether or not consumers and
consumers adhere to the Substrait spec. The goal is, thus, to have CI green at
all times. To understand the level of conformance of producers and consumers,
the support matrix is used instead.
Create and activate your conda environment with python3.9:
conda create -y -n substrait_consumer_testing -c conda-forge python=3.9 openjdk
conda activate substrait_consumer_testing
Note: Java is used by Jpype to access the Isthmus producer.
JPype should work with all versions of Java but to see details on which versions are
officially supported see https://jpype.readthedocs.io/en/latest/install.html
Install requirements from the top level directory:
pip install -r requirements.txt
Install the project in the developer mode (recommended)
git submodule init
git submodule update --init
cd consumer-testing
./build-and-copy-isthmus-shadow-jar.sh
This shell script may not work on Windows environments.
python3 setup.py develop
TPCH tests are located in the substrait_consumer/tests/integration
folder and substrait function tests
are located in the substrait_consumer/tests/functional
folder.
Tests are run with pytest.
TPCH Tests:
cd substrait_consumer/tests/integration/
pytest test_acero_tpch.py
Function Tests:
cd substrait_consumer/tests/functional/
# Run all function tests:
pytest extension_functions
# Run a single function test:
pytest extension_functions/test_arithmetic_functions.py
TPCH test files are located in the substrait_consumer/tests/integration
folder.
Test case arguments are located in substrait_consumer/tests/integration/queries/tpch_test_cases.py
. They specify
the parquet files, the SQL query, and substrait query plan that will be used for the test cases.
query_1.py
TPCH_QUERY_TESTS = (
{
"test_name": "test_tpch_sql_1",
"local_files": {},
"named_tables": {"lineitem": "lineitem.parquet"},
"sql_query": get_sql("q1.sql"),
"substrait_query": get_substrait_plan("query_01_plan.json"),
},
{
"test_name": "test_tpch_sql_2",
"local_files": {},
"named_tables": {
"part": "part.parquet",
"supplier": "supplier.parquet",
"partsupp": "partsupp.parquet",
"nation": "nation.parquet",
"region": "region.parquet",
"partsupp": "partsupp.parquet",
"supplier": "supplier.parquet",
"nation": "nation.parquet",
"region": "region.parquet",
},
"sql_query": get_sql("q2.sql"),
"substrait_query": get_substrait_plan("query_02_plan.json"),
},
)
Substrait query plans are located in substrait_consumer/tests/integration/queries/tpch_substrait_plans
.
SQL queries are located in substrait_consumer/tests/integration/queries/tpch_sql
.
The SQL queries have named placeholders ('{customer}'
) where the table names or file paths will be inserted.
Table names are determined based on the "named_tables"
and "local_files"
in the test case args file.
The substrait function tests aim to test the functions available in Substrait. This is done by converting queries (SQL/Ibis expressions) into substrait plans via various producers and running the substrait plans on different consumers.
The producer tests take the SQL/Ibis query and generate a substrait plan in json format. These plans are saved as snapshots, using pytest-snapshot, to be used later on for verification as well as an input to the consumer tests.
The consumer tests read the saved substrait plan snapshots and generate results. These results are saved as a snapshot to be used for verification.
If there is a mismatch between results and a saved snapshot, the result will be considered incorrect.
Substrait function test files are located in the substrait_consumer/functional/extension_functions
folder.
Test case arguments located in substrait_consumer/functional/queries/{*_tests}.py
. They specify
the parquet files, an SQL query, and an ibis expression.
The tests also take in the consumer and producer as test input via the producer/consumer test fixtures,
which are defined in substrait_consumer/conftest.py
. The fixtures allow the tests to cycle through all combinations
of producers and consumers.
arithmetic_tests.py
SCALAR_FUNCTIONS = (
{
"test_name": "add",
"local_files": {},
"named_tables": {"partsupp": "partsupp.parquet"},
"sql_query": SQL_SCALAR["add"],
},
The SQL queries are located in substrait_consumer/functional/queries/sql
.
arithmetic_functions_sql.py
SQL_SCALAR = {
"add":
"""
SELECT PS_PARTKEY, PS_SUPPKEY, add(PS_PARTKEY, PS_SUPPKEY) AS ADD_KEY
FROM '{partsupp}';
""",
The Ibis expressions are located in substrait_consumer/functional/queries/ibis_expressions
.
arithmetic_functions_expr.py
def add_expr(partsupp, lineitem, t):
new_col = (partsupp.ps_partkey + partsupp.ps_suppkey).name("ADD_KEY")
return partsupp[partsupp.ps_partkey, partsupp.ps_suppkey, new_col]
IBIS_SCALAR = {
"add": add_expr,
}
For each test, there are typically two snapshots, one for each of the levels we are testing: (1) the snapshot of the behavior under test such as a Substrait plan, a query result, etc. and (2) the "outcome" of whether or not the system under test behaved as expected. A typical test, thus, consists of running the system under test, comparing the output with the first snapshot (where a diverging answer will not lead to immediate test failure), and then a comparison of that outcome with the previously registered outcome (where a diverging answer will lead to test failure). When updating snapshots we, thus, need to be conscious about which of the two snapshots we update.
Each producer has its own set of substrait plan snapshots that are stored in the *_snapshots
directory under substrait_consumer/tests/functional/extension_functions/
cd substrait_consumer/tests/functional/extension_functions/boolean_snapshots
ls
DuckDBProducer IbisProducer IsthmusProducer
cd DuckDBProducer
and-duckdb_outcome.txt bool_and_plan.json not-duckdb_outcome.txt or_plan.json
and_plan.json bool_or-duckdb_outcome.txt not_plan.json xor-duckdb_outcome.txt
bool_and-duckdb_outcome.txt bool_or_plan.json or-duckdb_outcome.txt xor_plan.json
Substrait plan snapshots are used to verify that producers are able to generate substrait plans.
These tests are marked with the produce_substrait_snapshot
pytest marker.
Snapshots can be updated with the following command:
cd substrait_consumer/tests/functional/extension_functions
pytest -m produce_substrait_snapshot --snapshot-update
You can update the snapshots from a single producer with the --producer
option as well as
for a single test by specifying the test file:
pytest -m produce_substrait_snapshot --producer isthmus --snapshot-update test_arithmetic_functions.py
Note that this updates both snapshots, the one with the expected behavior and the one recording whether or not the system behaved that way. After updating the snapshots, be sure that the new snapshots correspond to what you want and manually intervene if necessary.
Results Snapshots are generated by running the SQL query corresponding to the function under
test against DuckDB. Those tests are marked with the generate_function_snapshots
pytest
marker.
cd substrait_consumer/tests/functional/extension_functions
pytest -m generate_function_snapshots --snapshot-update test_arithmetic_functions.py
Results snapshots are saved in the function_test_results
under each
function grouping snapshots folder.
cd substrait_consumer/tests/functional/extension_functions/boolean_snapshots/function_test_results
ls
and-datafusion-acero_outcome.txt and-duckdb-acero_outcome.txt and-ibis-acero_outcome.txt and_outcome.txt
and-datafusion-datafusion_outcome.txt and-duckdb-duckdb_outcome.txt and-isthmus-acero_outcome.txt and_result.txt
...
The CLI tool for generating substrait plans from adhoc SQL queries and Ibis expression
is located in the substrait_consumer/tests/adhoc
directory. The SQL queries should be
written using the same TPCH data used in the integration tests. This tool will generate
the substrait plans for each supported producer and run that plan against all supported consumers.
If you are testing out an SQL query, copy your SQL query into substrait_consumer/tests/adhoc/query.sql
and run the following command (make sure to specify a producer that can convert SQL to Substrait):
cd substrait_consumer/tests/adhoc
pytest --adhoc_producer=isthmus test_adhoc_expression.py
If you are testing out an Ibis expression, copy your Ibis expression into
substrait_consumer/tests/adhoc/ibis_expr.py
and run the following command:
cd substrait_consumer/tests/adhoc
pytest --adhoc_producer=ibis test_adhoc_expression.py
*Note: If you're using the IbisProducer, make sure you do not edit the function name and arguments
already in line 2 of ibis_expr.py
. The test is expecting the specific name and arguments.
You can save the produced substrait plans with the --saveplan
option.
pytest --saveplan True --adhoc_producer=isthmus test_adhoc_expression.py
Plans will be saved as {producer_name}_substrait.json
ls *.json
IsthmusProducer_substrait.json
If you want to run the tests using specific producer/consumer pairs, you can use
the both the --adhoc_producer
and --consumer
options.
pytest --adhoc_producer=isthmus --consumer=acero test_adhoc_expression.py
Producers should be added to the substrait_consumer/producers
folder and provide
methods on how to produce the substrait query plan. Look at
substrait_consumer/producers/duckdb_producer.py
for an example implementation.
In order for the test to use the new producer, the producer class name should also be added
to the PRODUCERS list in substrait_consumer/conftest.py
.
Consumers should be added to the substrait_consumer/consumers
folder and provide
methods on how to run the substrait query plan against that consumer. Look at
substrait_consumer/producers/duckdb_consumer.py
for an example implementation.
In order for the test to use the new consumer, the consumer class name should also be added
to the CONSUMERS list in substrait_consumer/conftest.py
.
The versions of all dependencies in requirements.txt
are frozen in order to
make the tests reproducible. Similarly, the version of the Isthmus producer is
frozen via its version of the git submodule at substrait-java/
. Development
and CI should mainly happen with those versions in order to eliminate diverging
behavior due to version differences.
In order to update a specific producer or consumer or one or several of the dependencies, install the desired versions using:
pip install --upgrade package==1.23.4 # upgrade package to specific version
pip install --upgrade package # upgrade package to latest version
pip install --upgrade -r requirements-unlocked.txt # upgrade all dependencies to latest version
For dependencies in git submodules, update the corresponding submodules.
After the dependency was updated, run the tests and make sure that all changes in test outcomes are expected and fix potential problems and/or update the snapshots.