Merge pull request #824 from pyathena-dev/fix/644-fleet-init-api-weight

Count initialization row validation as API-heavy fleet work

Author: laughingman7743

benchmarks/README.md

@@ -275,7 +275,7 @@ All hosts share the stack's scratch bucket and database: prepare once, and every
 Each host runs one worker, and each worker runs one orchestrator at a time, so trials on a host never overlap.
 
 Athena executes queries from different hosts independently, but API request rates are account-wide.
-Cursor and DictCursor row retrieval calls GetQueryResults for every page, about 5 to 10 pages per second per query in the recorded runs; check the account's GetQueryResults rate in Service Quotas (100 calls per second in the tested account).
+Cursor and DictCursor row retrieval calls GetQueryResults for every page, including the initialization suite's row-count validation, about 5 to 10 pages per second per query in the recorded runs; check the account's GetQueryResults rate in Service Quotas (100 calls per second in the tested account).
 `jobs` marks these jobs as API-heavy with an estimated page count and the number of simultaneous paging queries, which is the concurrency level for the concurrent suite.
 `worker --api-slots` bounds the simultaneous paging queries across the fleet (10 by default): a job takes one slot per paging query, and a worker refuses to start if any job needs more slots than that.
 While an API-heavy job waits for slots, workers do not start later API-heavy jobs; other jobs run on every free host.

benchmarks/pyathena_bench/fleet.py

@@ -68,13 +68,17 @@ def api_pages(case: Case, rows: int) -> int:
         case: Benchmark case.
         rows: Rows in the selected snapshot.
 
+    The initialization suite counts too: it validates the constructed result
+    set by reading every row.
+
     Returns:
         Page count summed over the trial's concurrent queries; 0 for cases that
         do not page through the API.
     """
-    if case.family not in API_FAMILIES or case.suite == "init":
+    if case.family not in API_FAMILIES:
         return 0
-    return -(-rows // case.arraysize) * case.concurrency
+    # The first page of a SELECT result also carries the column labels.
+    return -(-(rows + 1) // case.arraysize) * case.concurrency
 
 
 def expand_jobs(

benchmarks/tests/test_fleet.py

@@ -147,6 +147,25 @@ def test_row_output_jobs_keep_one_arraysize_for_every_family(self):
         assert jobs[0]["args"][-2:] == ["--arraysize", "100"]
         assert not jobs[0]["api_heavy"]
 
+    def test_initialization_validation_of_row_cursors_is_api_heavy(self):
+        jobs = expand_jobs(
+            SETTINGS, ["init"], ["small"], ["flat"], Filters(family=["cursor", "arrow"])
+        )
+        heavy = {j["id"]: (j["api_heavy"], j["api_weight"]) for j in jobs}
+        assert heavy["init-small-flat-cursor-direct-csv-rows-a1000"] == (True, 1)
+        assert heavy["init-small-flat-arrow-direct-csv-native"] == (False, 0)
+
+    def test_page_estimate_includes_the_column_label_row(self):
+        cursor = Filters(family=["cursor"], api=["direct"])
+        jobs = expand_jobs(SETTINGS, ["init"], ["small"], ["flat"], cursor, split_pages=11)
+        assert [j["id"] for j in jobs] == [
+            f"init-small-flat-cursor-direct-csv-rows-a1000-r{repetition}"
+            for repetition in range(1, 6)
+        ]
+        assert {j["pages"] for j in jobs} == {11}
+        unsplit = expand_jobs(SETTINGS, ["init"], ["small"], ["flat"], cursor, split_pages=12)
+        assert [j["id"] for j in unsplit] == ["init-small-flat-cursor-direct-csv-rows-a1000"]
+
     def test_concurrent_row_jobs_weigh_simultaneous_paging_queries(self):
         settings = Settings(scales={"small": 10000}, concurrency=[1, 10])
         jobs = expand_jobs(