Measure PolarsCursor on Polars 1.x and 2.x in the benchmark harness (GH-1292)

.gitignore

@@ -27,6 +27,7 @@ cover/
 .mypy_cache
 .tmp
 .venv
+.venv-*
 docs/_build
 .mcp.json
 .serena

benchmarks/README.md

@@ -41,6 +41,10 @@ The benchmark is a non-packaged member of the repository's uv workspace and runs
 It shares the root `uv.lock` with PyAthena, which it uses from the workspace checkout.
 Locally and in CI, its virtual environment is `benchmarks/.venv`, separate from the root `.venv`.
 On the EC2 instance, the dedicated checkout's root `/opt/pyathena/.venv` holds the benchmark environment.
+The `polars1` and `polars2` dependency groups select Polars 1.x or 2.x, which PolarsCursor reads with different code paths.
+The environments above have `polars1`; a second environment beside each, `benchmarks/.venv-polars2` or `/opt/pyathena/.venv-polars2`, has `polars2`.
+The root `pyproject.toml` declares both groups and the root `dev` group as conflicting, so `uv.lock` resolves each of them separately.
+The benchmark's default dependency groups include `polars1`; a sync with `--no-default-groups` needs `--group polars1` or `--group polars2`, or it installs whichever Polars version the lock resolves without them.
 The root `uv build -v` still builds only PyAthena, including a wheel built from its sdist.
 Benchmark dependencies do not become PyAthena runtime dependencies, and `uv sync --group dev` at the root does not install them.
 
@@ -51,7 +55,7 @@ just benchmark lint
 just benchmark test
 ```
 
-The recipes select `benchmarks/.venv` and Python 3.12, and sync that environment from the locked workspace.
+The recipes select `benchmarks/.venv` and Python 3.12, and sync that environment from the locked workspace; `just benchmark test` also runs the tests in `benchmarks/.venv-polars2`.
 The tests use local data, fake AWS clients, and child processes.
 They do not execute Athena queries or create AWS resources.
 `just benchmark format` formats the benchmark separately from the parent project.
@@ -67,6 +71,13 @@ export UV_PROJECT_ENVIRONMENT="$PWD/.venv" UV_PYTHON=3.12
 uv run --locked python -m pyathena_bench plan --suite single --scale small
 ```
 
+To run a command with Polars 2.x, select the second environment and its group:
+
+```bash
+UV_PROJECT_ENVIRONMENT="$PWD/.venv-polars2" uv run --locked --no-group polars1 --group polars2 \
+  python -m pyathena_bench plan --suite single --scale small --family polars
+```
+
 For local AWS access, configure the profile and region in the repository root's gitignored `.env` file:
 
 ```dotenv
@@ -117,7 +128,7 @@ As the results location, it removes leftovers by itself, but reports must be rec
 
 The deployment identity needs permission to create these resources and pass the instance role.
 Use an immutable, remotely accessible commit that contains this directory and the root `uv.lock`.
-Bootstrap checks out that commit, installs pinned uv and Python versions, and runs `uv sync --locked --no-dev` for the benchmark into `/opt/pyathena/.venv`.
+Bootstrap checks out that commit, installs pinned uv and Python versions, and syncs the benchmark with `uv sync --locked`: with the `polars1` group into `/opt/pyathena/.venv`, and with the `polars2` group into `/opt/pyathena/.venv-polars2`.
 It signals setup completion to CloudFormation but does not prepare data or start measurements.
 The EC2 commands use `--no-sync` to reuse this verified environment; they do not revalidate the lockfile on every invocation.
 
@@ -365,6 +376,9 @@ With `--split-pages`, a job whose single trial needs at least that many pages be
 `--chunksize` keeps only the given chunk sizes, with 0 for reads without chunks, in `plan`, `run`, and `jobs`; `jobs` passes it to each job and adds it to the job ID.
 `jobs` accepts `--warmups` and `--repetitions` too; with either option, every job that is not split carries both values as arguments, so one queue can combine jobs with different trial counts under one stored configuration.
 `queue --jobs` accepts several job files, rejects duplicate job IDs, and orders the combined jobs as `jobs` orders one selection.
+`jobs --environment polars2` runs the selected jobs with Polars 2.x: each job names the environment, and its ID ends with `-polars2`, so one queue can hold the same cases for both Polars versions.
+A worker runs such a job with the interpreter of `.venv-polars2` beside its own environment, and refuses to start if a job's environment is missing on its host; start the worker from the default environment, whose interpreter runs the other jobs.
+Select the families that use Polars, `polars` and `sqla_polars`, for this environment.
 
 Deploy a new stack with the fleet size under its own name; changing `FleetSize` on an existing stack does not wait for the added hosts to finish bootstrap.
 `QueueName` names the queue that every host consumes: after bootstrap, each host starts a worker with the stack's results location from its user data, also written to `results-location.txt` next to `stack-id.txt`, and the worker waits until a queue of that name is published; this includes the hosts that the group launches later to replace interrupted Spot Instances.

benchmarks/cloudformation/benchmark.yaml

@@ -361,7 +361,8 @@ Resources:
               git -C /opt/pyathena checkout --detach ${GitCommit}
               export UV_PYTHON_INSTALL_DIR=/opt/uv-python
               uv python install ${PythonVersion}
-              uv sync --directory /opt/pyathena/benchmarks --locked --no-dev --python ${PythonVersion}
+              uv sync --directory /opt/pyathena/benchmarks --locked --no-default-groups --group polars1 --python ${PythonVersion}
+              UV_PROJECT_ENVIRONMENT=/opt/pyathena/.venv-polars2 uv sync --directory /opt/pyathena/benchmarks --locked --no-default-groups --group polars2 --python ${PythonVersion}
               chmod -R a+rX /opt/uv-python
               chown -R ec2-user:ec2-user /opt/pyathena
               printf '%s\n' '${AWS::StackId}' > /opt/pyathena/benchmarks/stack-id.txt

benchmarks/pyathena_bench/__main__.py

@@ -33,6 +33,7 @@
 from pyathena_bench.config import Settings, read_json
 from pyathena_bench.fleet import (
     API_FAMILIES,
+    ENVIRONMENTS,
     Filters,
     Queue,
     SpotInterruptionError,
@@ -85,6 +86,11 @@ def parser() -> argparse.ArgumentParser:
                 help="Run each repetition as a separate job when one trial needs this many "
                 "GetQueryResults pages",
             )
+            command.add_argument(
+                "--environment",
+                choices=ENVIRONMENTS,
+                help="Virtual environment beside the default one that runs the jobs",
+            )
         else:
             command.add_argument("--suite", choices=SUITES, required=True)
             command.add_argument("--shape", choices=SHAPES, default="flat")
@@ -293,6 +299,7 @@ def main(argv: list[str] | None = None) -> int:
             filters,
             args.split_pages,
             explicit_trials=bool(overrides),
+            environment=args.environment,
         )
         sys.stdout.write(json.dumps(jobs, indent=1) + "\n")
     elif args.command in {"plan", "run"}:

benchmarks/pyathena_bench/fleet.py

@@ -57,6 +57,9 @@
 QUIESCE_RETRIES = 10
 # Times a worker quiesces its own failed job while Athena only throttles it.
 QUIESCE_ROUNDS = 3
+# Virtual environments that a job can name, each synced with the dependency group of
+# the same name beside the default environment, which has the polars1 group.
+ENVIRONMENTS = ("polars2",)
 
 
 @dataclass(frozen=True)
@@ -134,11 +137,14 @@ def merge_jobs(job_lists: Iterable[list[dict[str, Any]]]) -> list[dict[str, Any]
         All jobs, ordered as :func:`expand_jobs` orders one list.
 
     Raises:
-        ValueError: If job IDs are not unique.
+        ValueError: If job IDs are not unique, or a job names an unknown environment.
     """
     jobs = [job for jobs in job_lists for job in jobs]
     if len({j["id"] for j in jobs}) != len(jobs):
         raise ValueError("Job IDs must be unique")
+    unknown = [j["id"] for j in jobs if j.get("environment") not in (None, *ENVIRONMENTS)]
+    if unknown:
+        raise ValueError(f"Jobs name an unknown environment: {', '.join(unknown)}")
     return sorted(jobs, key=job_order)
 
 
@@ -150,6 +156,7 @@ def expand_jobs(
     filters: Filters,
     split_pages: int | None = None,
     explicit_trials: bool = False,
+    environment: str | None = None,
 ) -> list[dict[str, Any]]:
     """Split a selection into independent ``run`` invocations.
 
@@ -172,6 +179,8 @@ def expand_jobs(
         explicit_trials: Whether jobs that are not split carry the warmups and
             repetitions of ``settings`` as arguments, so a queue can mix trial
             counts under one stored configuration.
+        environment: One of ``ENVIRONMENTS`` that runs the jobs, which their IDs
+            name, or None for the default environment.
 
     Returns:
         Job dictionaries ordered with API-heavy jobs first and longer jobs earlier.
@@ -246,9 +255,29 @@ def expand_jobs(
                         )
     if not jobs:
         raise ValueError("No cases match the selected filters")
+    if environment:
+        for job in jobs:
+            job["id"] += f"-{environment}"
+            job["environment"] = environment
     return sorted(jobs, key=job_order)
 
 
+def interpreter(environment: str | None) -> Path:
+    """Return the Python interpreter that runs a job.
+
+    Args:
+        environment: The job's environment, a virtual environment named
+            ``.venv-<environment>`` beside the one of this process, or None for
+            this process's interpreter.
+
+    Returns:
+        Path of the interpreter.
+    """
+    if environment is None:
+        return Path(sys.executable)
+    return Path(sys.prefix).with_name(f".venv-{environment}") / "bin" / "python"
+
+
 def now() -> str:
     """Return the current UTC time in ISO 8601 format.
 
@@ -666,7 +695,8 @@ def run_job(
     ends, so no trial process outlives it.
 
     Args:
-        job: Job dictionary.
+        job: Job dictionary, whose environment selects the interpreter (see
+            :func:`interpreter`).
         workdir: Directory with the queue's ``config.toml`` and ``manifest.json``.
         output: New output directory for this attempt.
         resume: Output directories of earlier attempts whose completed trials
@@ -679,7 +709,8 @@ def run_job(
     """
     (workdir / "logs").mkdir(parents=True, exist_ok=True)
     command = [
-        sys.executable, "-m", "pyathena_bench", "--config", str(workdir / "config.toml"),
+        str(interpreter(job.get("environment"))), "-m", "pyathena_bench",
+        "--config", str(workdir / "config.toml"),
         "run", "--manifest", str(workdir / "manifest.json"), "--out", str(output),
         "--checkpoint", checkpoint, *job["args"],
     ]  # fmt: skip
@@ -1211,7 +1242,8 @@ def work(
 
     Raises:
         ValueError: If the queue has not been completely published and
-            ``wait`` is false, or a job needs more API slots than ``api_slots``.
+            ``wait`` is false, a job needs more API slots than ``api_slots``,
+            or a job's environment does not exist on this host.
         SpotInterruptionError: After an interruption notice, once the running attempt's
             records are uploaded and its claim is released as interrupted, or
             once the attempt of a failed job ends, if the notice arrived while
@@ -1240,6 +1272,9 @@ def notice() -> bool:
     too_wide = [j["id"] for j in jobs if j["api_heavy"] and j.get("api_weight", 1) > api_slots]
     if too_wide:
         raise ValueError(f"Jobs need more than {api_slots} API slots: {', '.join(too_wide)}")
+    missing = [j["id"] for j in jobs if not interpreter(j.get("environment")).exists()]
+    if missing:
+        raise ValueError(f"Jobs need an environment missing on this host: {', '.join(missing)}")
     if settle is None:
         settings = Settings.load(workdir / "config.toml")
 

benchmarks/pyproject.toml

@@ -23,12 +23,15 @@ dependencies = [
 
 [tool.uv]
 package = false
+default-groups = ["test", "polars1"]
 
 [tool.uv.sources]
 pyathena = { workspace = true }
 
 [dependency-groups]
-dev = ["pytest>=8", "pytest-asyncio>=1", "mypy", "types-psutil", "cfn-lint>=1"]
+test = ["pytest>=8", "pytest-asyncio>=1", "mypy", "types-psutil", "cfn-lint>=1"]
+polars1 = ["polars>=1.39.0,<2"]
+polars2 = ["polars>=2,<3"]
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]
@@ -37,7 +40,7 @@ asyncio_mode = "auto"
 
 [tool.ruff]
 extend = "../pyproject.toml"
-exclude = [".venv"]
+exclude = [".venv", ".venv-*"]
 target-version = "py312"
 
 [tool.mypy]

benchmarks/tests/test_cloudformation.py

@@ -187,4 +187,9 @@ def test_hosts_start_the_queue_worker_and_expose_their_tags():
     )
     assert variables["Results"] == template["Outputs"]["Results"]["Value"]
     assert "systemctl enable --now pyathena-bench-worker.service" in user_data
+    # Jobs of the polars2 environment run the interpreter beside the worker's environment.
+    assert (
+        "UV_PROJECT_ENVIRONMENT=/opt/pyathena/.venv-polars2 uv sync --directory "
+        "/opt/pyathena/benchmarks --locked --no-default-groups --group polars2"
+    ) in user_data
     assert template["Parameters"]["QueueName"]["Default"] == ""

benchmarks/tests/test_fleet.py

@@ -148,6 +148,22 @@ def test_row_cursor_jobs_are_split_by_arraysize_and_marked_api_heavy(self):
         assert jobs[0]["pages"] >= jobs[-1]["pages"]
         assert jobs[0]["args"][-2:] == ["--arraysize", "100"]
 
+    def test_jobs_of_another_environment_are_named_and_marked_for_it(self):
+        jobs = expand_jobs(
+            SETTINGS,
+            ["single"],
+            ["small"],
+            ["flat"],
+            Filters(family=["polars"], api=["sync"], transport=["unload"], output_kind="native"),
+            environment="polars2",
+        )
+        # The suffix keeps the claims, done markers, and results apart from the default's.
+        assert [(j["id"], j["environment"]) for j in jobs] == [
+            ("single-small-flat-polars-sync-unload-native-polars2", "polars2")
+        ]
+        default = expand_jobs(SETTINGS, ["single"], ["small"], ["flat"], Filters(family=["polars"]))
+        assert not any("environment" in j for j in default)
+
     def test_dataframe_chunk_variants_share_one_job(self):
         jobs = expand_jobs(
             SETTINGS,
@@ -398,6 +414,11 @@ def test_duplicate_ids_across_lists_are_rejected(self):
         with pytest.raises(ValueError, match="unique"):
             merge_jobs([[job("a")], [job("a")]])
 
+    def test_unknown_environments_are_rejected_before_publishing(self):
+        known = {**job("b"), "environment": "polars2"}
+        with pytest.raises(ValueError, match=r"unknown environment: c$"):
+            merge_jobs([[job("a"), known], [{**job("c"), "environment": "polars3"}]])
+
 
 METADATA = {
     "instance-id": "i-123",
@@ -491,6 +512,14 @@ def test_jobs_command_carries_trial_overrides(self, tmp_path, capsys):
         jobs = json.loads(capsys.readouterr().out)
         assert [j["args"][-4:] for j in jobs] == [["--warmups", "1", "--repetitions", "3"]]
 
+    def test_jobs_command_selects_the_environment(self, tmp_path, capsys):
+        config = tmp_path / "config.toml"
+        config.write_text("[scales]\nsmall = 10000\n")
+        argv = ["--config", str(config), "jobs", "--suite", "single", "--scale", "small"]
+        assert main([*argv, "--family", "polars", "--environment", "polars2"]) == 0
+        jobs = json.loads(capsys.readouterr().out)
+        assert {j["environment"] for j in jobs} == {"polars2"}
+
     def test_plan_honors_arraysize_filter_and_repetition_overrides(self, tmp_path, capsys):
         config = tmp_path / "config.toml"
         config.write_text("[scales]\nsmall = 10000\n")
@@ -610,6 +639,17 @@ def test_trial_evidence_reads_reported_queries_and_trial_prefixes(self, tmp_path
             assert trial_evidence(tmp_path) == ({"a"}, ("s3://b/runs/r/trials/t1/",))
         assert trial_evidence(tmp_path / "missing") == (set(), ())
 
+    def test_jobs_of_a_missing_environment_stop_the_worker_before_claiming(
+        self, tmp_path, monkeypatch
+    ):
+        monkeypatch.setattr("pyathena_bench.fleet.sys.prefix", str(tmp_path / ".venv"))
+        _, queue = queue_with(
+            [job("default"), {**job("other"), "environment": "polars2"}], tmp_path
+        )
+        with pytest.raises(ValueError, match=r"environment missing on this host: other$"):
+            work(queue, tmp_path / "w", api_slots=1, runner=lambda *a: 0, settle=no_queries)
+        assert not queue.names("claims")
+
     def test_jobs_wider_than_the_api_slots_stop_the_worker_before_claiming(self, tmp_path):
         _, queue = queue_with([job("wide", heavy=True, api_weight=11)], tmp_path)
         with pytest.raises(ValueError, match="more than 10 API slots"):
@@ -1154,15 +1194,22 @@ def wait(self, timeout=None):
 
 
 class TestRunJob:
-    def test_resumes_writes_checkpoints_and_interrupts_with_sigint(self, tmp_path, monkeypatch):
-        processes = []
+    @pytest.fixture
+    def processes(self, monkeypatch):
+        """Return the fake orchestrator processes that ``run_job()`` starts."""
+        started = []
 
         def popen(command, stdout, stderr, start_new_session):
-            processes.append(FakeProcess(command, stdout, stderr, start_new_session))
-            return processes[-1]
+            started.append(FakeProcess(command, stdout, stderr, start_new_session))
+            return started[-1]
 
-        killed = []
         monkeypatch.setattr("pyathena_bench.fleet.subprocess.Popen", popen)
+        return started
+
+    def test_resumes_writes_checkpoints_and_interrupts_with_sigint(
+        self, tmp_path, monkeypatch, processes
+    ):
+        killed = []
         monkeypatch.setattr(
             "pyathena_bench.fleet.os.killpg", lambda pid, signum: killed.append((pid, signum))
         )
@@ -1189,6 +1236,25 @@ def popen(command, stdout, stderr, start_new_session):
         assert killed == [(4321, signal.SIGKILL)]
         assert (tmp_path / "logs" / "a-attempt-2.log").exists()
 
+    def test_a_job_of_another_environment_runs_its_interpreter(
+        self, tmp_path, monkeypatch, processes
+    ):
+        monkeypatch.setattr("pyathena_bench.fleet.os.killpg", lambda pid, signum: None)
+        monkeypatch.setattr("pyathena_bench.fleet.sys.prefix", str(tmp_path / ".venv"))
+        # A set event stops the fake process at once.
+        stop = threading.Event()
+        stop.set()
+        run_job(
+            {"id": "a", "args": [], "environment": "polars2"},
+            tmp_path,
+            tmp_path / "results" / "a" / "attempt-1",
+            [],
+            stop,
+            "s3://bucket/fleet/q/results/a/attempt-1/",
+        )
+        [process] = processes
+        assert process.command[0] == str(tmp_path / ".venv-polars2" / "bin" / "python")
+
 
 class TestWork:
     def test_expired_claims_are_taken_over_and_live_ones_left(self, tmp_path):

justfile

@@ -85,8 +85,11 @@ _docs-lint:
 _docs-format:
     mise exec -- markdownlint-cli2 --fix
 
-# Benchmark workspace member with its own Python 3.12 environment
-BENCHMARK_UV := "UV_PROJECT_ENVIRONMENT='" + justfile_directory() + "/benchmarks/.venv' UV_PYTHON=3.12 uv run --directory benchmarks --locked"
+# Benchmark workspace member with its own Python 3.12 environments: benchmarks/.venv
+# with Polars 1.x, and benchmarks/.venv-polars2 with Polars 2.x
+BENCHMARK_RUN := "UV_PYTHON=3.12 uv run --directory benchmarks --locked"
+BENCHMARK_UV := "UV_PROJECT_ENVIRONMENT='" + justfile_directory() + "/benchmarks/.venv' " + BENCHMARK_RUN
+BENCHMARK_UV_POLARS2 := "UV_PROJECT_ENVIRONMENT='" + justfile_directory() + "/benchmarks/.venv-polars2' " + BENCHMARK_RUN + " --no-group polars1 --group polars2"
 
 # Standalone benchmark checks; these do not execute Athena queries
 benchmark target="lint":
@@ -104,6 +107,7 @@ _benchmark-lint:
 
 _benchmark-test: _benchmark-lint
     {{BENCHMARK_UV}} pytest
+    {{BENCHMARK_UV_POLARS2}} pytest
 
 # Install development tools
 tool: