diff --git a/src/lakebench/benchmarks/_load_and_query/_load_and_query.py b/src/lakebench/benchmarks/_load_and_query/_load_and_query.py index 39e1bb4..fd4cc35 100644 --- a/src/lakebench/benchmarks/_load_and_query/_load_and_query.py +++ b/src/lakebench/benchmarks/_load_and_query/_load_and_query.py @@ -230,7 +230,11 @@ def __init__( self.engine.extended_engine_metadata["analyze"] = analyze_mode self._power_test_uses_full_stream = query_list is None or query_list == ["*"] - if query_list is not None: + if self._power_test_uses_full_stream and self.QUERY_STREAMS: + query_plan = self._query_plan_for_stream(self.POWER_TEST_STREAM) + self.query_progress = [progress for progress, _ in query_plan] + self.query_list = [query_name for _, query_name in query_plan] + elif query_list is not None: expanded_query_list = [] for query in query_list: if query == "*": @@ -244,9 +248,10 @@ def __init__( f"Query list contains unsupported queries: {unsupported_queries}. Supported queries: {self.QUERY_REGISTRY}." ) self.query_list = expanded_query_list + self.query_progress = None else: self.query_list = self.QUERY_REGISTRY - self.query_progress = None + self.query_progress = None for base_engine, benchmark_impl in self.BENCHMARK_IMPL_REGISTRY.items(): if isinstance(engine, base_engine): diff --git a/src/lakebench/benchmarks/tpcds/tpcds.py b/src/lakebench/benchmarks/tpcds/tpcds.py index f4b9a07..cad6f2c 100644 --- a/src/lakebench/benchmarks/tpcds/tpcds.py +++ b/src/lakebench/benchmarks/tpcds/tpcds.py @@ -21,7 +21,7 @@ class TPCDS(_LoadAndQuery): scenario_name : str The name of the benchmark scenario. query_list : list of str, optional - List of queries to execute. Use '*' for all queries. If not specified, all queries will be run. + List of queries to execute. Use '*' or omit this parameter to run query stream 0. input_parquet_folder_uri : str, optional Path to the input parquet files. Must be the root directory containing a folder named after each table in TABLE_REGISTRY. @@ -37,9 +37,9 @@ class TPCDS(_LoadAndQuery): Runs the benchmark in the specified mode. Supported modes are: - 'load': Sequentially executes loading the 24 tables. - - 'query': Sequentially executes the 99 queries. + - 'query': Executes query stream 0 by default, or the explicitly requested query list. - 'power_test': Executes query stream 0 without loading data. - - 'load_and_query': Executes the load test followed by the query test. + - 'load_and_query': Executes the load test followed by the default stream or requested query list. _run_load_test() Loads the data for the benchmark. _run_query_test() diff --git a/src/lakebench/benchmarks/tpch/tpch.py b/src/lakebench/benchmarks/tpch/tpch.py index a906dd3..75a02c7 100644 --- a/src/lakebench/benchmarks/tpch/tpch.py +++ b/src/lakebench/benchmarks/tpch/tpch.py @@ -21,7 +21,7 @@ class TPCH(_LoadAndQuery): scenario_name : str The name of the benchmark scenario. query_list : list of str, optional - List of queries to execute. Use '*' for all queries. If not specified, all queries will be run. + List of queries to execute. Use '*' or omit this parameter to run query stream 0. input_parquet_folder_uri : str, optional Path to the input parquet files. Must be the root directory containing a folder named after each table in TABLE_REGISTRY. diff --git a/tests/test_query_streams.py b/tests/test_query_streams.py index e789c13..6e77c5b 100644 --- a/tests/test_query_streams.py +++ b/tests/test_query_streams.py @@ -7,6 +7,22 @@ from lakebench.benchmarks._load_and_query import _LoadAndQuery from lakebench.benchmarks.tpcds import TPCDS from lakebench.benchmarks.tpch import TPCH +from lakebench.engines.duckdb import DuckDB + + +def _benchmark_with_queries(benchmark_class, query_list): + engine = MagicMock(spec=DuckDB) + engine.SUPPORTS_MOUNT_PATH = True + engine.version = "test" + engine.extended_engine_metadata = {} + engine.get_total_cores.return_value = 1 + engine.get_compute_size.return_value = "test" + return benchmark_class( + engine=engine, + scenario_name="test", + query_list=query_list, + input_parquet_folder_uri="/tmp/data", + ) def test_power_test_runs_queries_without_loading(): @@ -50,6 +66,16 @@ def test_query_streams_are_complete_permutations(benchmark_class, stream_count, assert all(set(stream) == set(range(1, query_count + 1)) for stream in benchmark_class.QUERY_STREAMS) +@pytest.mark.parametrize("benchmark_class", [TPCDS, TPCH]) +@pytest.mark.parametrize("query_list", [None, ["*"]]) +def test_full_query_shortcuts_default_to_stream_zero(benchmark_class, query_list): + benchmark = _benchmark_with_queries(benchmark_class, query_list) + query_plan = benchmark_class._query_plan_for_stream(0) + + assert benchmark.query_list == [query_name for _, query_name in query_plan] + assert benchmark.query_progress == [progress for progress, _ in query_plan] + + def test_tpcds_power_test_uses_stream_zero_order(): benchmark = object.__new__(TPCDS) benchmark.query_list = ["q1"]