Make run query generic

pola-rs · Apr 5, 2024 · 8188896 · 8188896
1 parent ca7703e
commit 8188896
Show file tree

Hide file tree

Showing 6 changed files with 74 additions and 110 deletions.
diff --git a/queries/common_utils.py b/queries/common_utils.py
@@ -2,6 +2,7 @@
 
 import re
 import sys
+from importlib.metadata import version
 from pathlib import Path
 from subprocess import run
 from typing import TYPE_CHECKING, Any
@@ -11,6 +12,8 @@
 from settings import Settings
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     import pandas as pd
     import polars as pl
 
@@ -93,6 +96,37 @@ def _get_query_numbers(library_name: str) -> list[int]:
     return sorted(query_numbers)
 
 
+def run_query_generic(
+    query: Callable[..., Any],
+    query_number: int,
+    library_name: str,
+    query_checker: Callable[..., None] | None = None,
+) -> None:
+    """Execute a query."""
+    with CodeTimer(name=f"Run {library_name} query {query_number}", unit="s") as timer:
+        result = query()
+
+    if settings.run.log_timings:
+        log_query_timing(
+            solution=library_name,
+            version=version(library_name),
+            query_number=query_number,
+            time=timer.took,
+        )
+
+    if settings.run.check_results:
+        if query_checker is None:
+            msg = "cannot check results if no query checking function is provided"
+            raise ValueError(msg)
+        if settings.scale_factor != 1:
+            msg = f"cannot check results when scale factor is not 1, got {settings.scale_factor}"
+            raise RuntimeError(msg)
+        query_checker(result, query_number)
+
+    if settings.run.show_results:
+        print(result)
+
+
 def check_query_result_pl(result: pl.DataFrame, query_number: int) -> None:
     """Assert that the Polars result of the query is correct."""
     from polars.testing import assert_frame_equal

diff --git a/queries/dask/utils.py b/queries/dask/utils.py
@@ -4,9 +4,12 @@
 
 import dask
 import dask.dataframe as dd
-from linetimer import CodeTimer
 
-from queries.common_utils import check_query_result_pd, log_query_timing, on_second_call
+from queries.common_utils import (
+    check_query_result_pd,
+    on_second_call,
+    run_query_generic,
+)
 from settings import Settings
 
 if TYPE_CHECKING:
@@ -77,22 +80,4 @@ def get_part_supp_ds() -> DataFrame:
 
 
 def run_query(query_number: int, query: Callable[..., Any]) -> None:
-    with CodeTimer(name=f"Run Dask query {query_number}", unit="s") as timer:
-        result = query()
-
-    if settings.run.log_timings:
-        log_query_timing(
-            solution="dask",
-            version=dask.__version__,
-            query_number=query_number,
-            time=timer.took,
-        )
-
-    if settings.run.check_results:
-        if settings.scale_factor != 1:
-            msg = f"cannot check results when scale factor is not 1, got {settings.scale_factor}"
-            raise RuntimeError(msg)
-        check_query_result_pd(result, query_number)
-
-    if settings.run.show_results:
-        print(result)
+    run_query_generic(query, query_number, "dask", query_checker=check_query_result_pd)
diff --git a/queries/duckdb/utils.py b/queries/duckdb/utils.py
@@ -1,11 +1,9 @@
-from importlib.metadata import version
 from pathlib import Path
 
 import duckdb
 from duckdb import DuckDBPyRelation
-from linetimer import CodeTimer
 
-from queries.common_utils import check_query_result_pl, log_query_timing
+from queries.common_utils import check_query_result_pl, run_query_generic
 from settings import Settings
 
 settings = Settings()
@@ -65,22 +63,7 @@ def get_part_supp_ds() -> str:
 
 
 def run_query(query_number: int, context: DuckDBPyRelation) -> None:
-    with CodeTimer(name=f"Run DuckDB query {query_number}", unit="s") as timer:
-        result = context.pl()  # Force DuckDB to materialize
-
-    if settings.run.log_timings:
-        log_query_timing(
-            solution="duckdb",
-            version=version("duckdb"),
-            query_number=query_number,
-            time=timer.took,
-        )
-
-    if settings.run.check_results:
-        if settings.scale_factor != 1:
-            msg = f"cannot check results when scale factor is not 1, got {settings.scale_factor}"
-            raise RuntimeError(msg)
-        check_query_result_pl(result, query_number)
-
-    if settings.run.show_results:
-        print(result)
+    query = context.pl
+    run_query_generic(
+        query, query_number, "duckdb", query_checker=check_query_result_pl
+    )
diff --git a/queries/pandas/utils.py b/queries/pandas/utils.py
@@ -3,9 +3,12 @@
 from typing import TYPE_CHECKING, Any
 
 import pandas as pd
-from linetimer import CodeTimer
 
-from queries.common_utils import check_query_result_pd, log_query_timing, on_second_call
+from queries.common_utils import (
+    check_query_result_pd,
+    on_second_call,
+    run_query_generic,
+)
 from settings import Settings
 
 if TYPE_CHECKING:
@@ -68,23 +71,7 @@ def get_part_supp_ds() -> pd.DataFrame:
     return _read_ds(settings.dataset_base_dir / "partsupp")
 
 
-def run_query(q_num: int, query: Callable[..., Any]) -> None:
-    with CodeTimer(name=f"Run pandas query {q_num}", unit="s") as timer:
-        result = query()
-
-    if settings.run.log_timings:
-        log_query_timing(
-            solution="pandas",
-            version=pd.__version__,
-            query_number=q_num,
-            time=timer.took,
-        )
-
-    if settings.run.check_results:
-        if settings.scale_factor != 1:
-            msg = f"cannot check results when scale factor is not 1, got {settings.scale_factor}"
-            raise RuntimeError(msg)
-        check_query_result_pd(result, q_num)
-
-    if settings.run.show_results:
-        print(result)
+def run_query(query_number: int, query: Callable[..., Any]) -> None:
+    run_query_generic(
+        query, query_number, "pandas", query_checker=check_query_result_pd
+    )
diff --git a/queries/polars/utils.py b/queries/polars/utils.py
@@ -1,9 +1,9 @@
+from functools import partial
 from pathlib import Path
 
 import polars as pl
-from linetimer import CodeTimer
 
-from queries.common_utils import check_query_result_pl, log_query_timing
+from queries.common_utils import check_query_result_pl, run_query_generic
 from settings import Settings
 
 settings = Settings()
@@ -56,25 +56,12 @@ def get_part_supp_ds() -> pl.LazyFrame:
 
 
 def run_query(query_number: int, lf: pl.LazyFrame) -> None:
+    streaming = settings.run.polars_streaming
+
     if settings.run.polars_show_plan:
-        print(lf.explain(streaming=settings.run.polars_streaming))
-
-    with CodeTimer(name=f"Run Polars query {query_number}", unit="s") as timer:
-        result = lf.collect(streaming=settings.run.polars_streaming)
-
-    if settings.run.log_timings:
-        log_query_timing(
-            solution="polars",
-            version=pl.__version__,
-            query_number=query_number,
-            time=timer.took,
-        )
-
-    if settings.run.check_results:
-        if settings.scale_factor != 1:
-            msg = f"cannot check results when scale factor is not 1, got {settings.scale_factor}"
-            raise RuntimeError(msg)
-        check_query_result_pl(result, query_number)
-
-    if settings.run.show_results:
-        print(result)
+        print(lf.explain(streaming=streaming))
+
+    query = partial(lf.collect, streaming=streaming)
+    run_query_generic(
+        query, query_number, "polars", query_checker=check_query_result_pl
+    )
diff --git a/queries/pyspark/utils.py b/queries/pyspark/utils.py
@@ -2,10 +2,13 @@
 
 from typing import TYPE_CHECKING
 
-from linetimer import CodeTimer
 from pyspark.sql import SparkSession
 
-from queries.common_utils import check_query_result_pd, log_query_timing, on_second_call
+from queries.common_utils import (
+    check_query_result_pd,
+    on_second_call,
+    run_query_generic,
+)
 from settings import Settings
 
 if TYPE_CHECKING:
@@ -83,23 +86,8 @@ def drop_temp_view() -> None:
     ]
 
 
-def run_query(query_number: int, query: SparkDF) -> None:
-    with CodeTimer(name=f"Run PySpark query {query_number}", unit="s") as timer:
-        result = query.toPandas()
-
-    if settings.run.log_timings:
-        log_query_timing(
-            solution="pyspark",
-            version=get_or_create_spark().version,
-            query_number=query_number,
-            time=timer.took,
-        )
-
-    if settings.run.check_results:
-        if settings.scale_factor != 1:
-            msg = f"cannot check results when scale factor is not 1, got {settings.scale_factor}"
-            raise RuntimeError(msg)
-        check_query_result_pd(result, query_number)
-
-    if settings.run.show_results:
-        print(result)
+def run_query(query_number: int, df: SparkDF) -> None:
+    query = df.toPandas
+    run_query_generic(
+        query, query_number, "pyspark", query_checker=check_query_result_pd
+    )