Support querying partitioned dataset

Author: EnricoMi

queries/common_utils.py

@@ -23,7 +23,9 @@
 def get_table_path(table_name: str) -> Path:
     """Return the path to the given table."""
     ext = settings.run.io_type if settings.run.include_io else "parquet"
-    return settings.dataset_base_dir / f"{table_name}.{ext}"
+    if settings.num_batches is None:
+        return settings.dataset_base_dir / f"{table_name}.{ext}"
+    return settings.dataset_base_dir / str(settings.num_batches) / table_name / "*" / f"part.{ext}"
 
 
 def log_query_timing(

settings.py

@@ -77,6 +77,7 @@ class Plot(BaseSettings):
 
 class Settings(BaseSettings):
     scale_factor: float = 1.0
+    num_batches: int | None = 1
 
     paths: Paths = Paths()
     plot: Plot = Plot()