Add distributed CTAS output for Presto benchmarks

common/testing/normalize_ctas_results.py

@@ -0,0 +1,121 @@
+#!/usr/bin/env python3
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Convert temporary distributed CTAS tables into canonical qN.parquet files."""
+
+import argparse
+import os
+import re
+import shutil
+import sys
+from pathlib import Path
+
+import pandas as pd
+import pyarrow.parquet as pq
+import sqlglot
+from sqlglot import exp
+
+# Allow importing from the repo root when this file is executed directly.
+sys.path.insert(0, str(Path(__file__).parent.parent.parent))
+
+from common.testing.result_comparison import get_orderby_sort_spec, restore_orderby
+from common.testing.test_utils import get_queries
+
+
+def _query_directories(source_dir: Path) -> list[Path]:
+    return sorted(
+        (path for path in source_dir.iterdir() if path.is_dir() and re.fullmatch(r"q\d+", path.name)),
+        key=lambda path: int(path.name[1:]),
+    )
+
+
+def _source_output_names(query_sql: str, physical_names: list[str]) -> list[str]:
+    """Return names from an explicit SELECT projection for ORDER BY lookup."""
+    # Explicit CTAS projections are stored with unique positional names (c1,
+    # c2, ...), but ORDER BY still refers to names from the original SELECT.
+    # Recover those logical names so ORDER BY expressions can be mapped to the
+    # correct physical column positions without changing the stored schema.
+    parsed = sqlglot.parse_one(query_sql, read="presto")
+    select = next(parsed.find_all(exp.Select))
+    has_wildcard = any(
+        isinstance(projection := expression.this if isinstance(expression, exp.Alias) else expression, exp.Star)
+        or (isinstance(projection, exp.Column) and isinstance(projection.this, exp.Star))
+        for expression in select.expressions
+    )
+    if len(select.expressions) != len(physical_names) or has_wildcard:
+        return physical_names
+    return [expression.alias_or_name or physical_names[position] for position, expression in enumerate(select.expressions)]
+
+
+def _write_canonical_result(query_dir: Path, target: Path, query_sql: str | None) -> None:
+    parquet_parts = sorted(query_dir.rglob("*.parquet"))
+    sort_spec: tuple[list[int], list[bool], list[bool]] = ([], [], [])
+
+    if parquet_parts and query_sql is not None:
+        column_names = _source_output_names(query_sql, pq.read_schema(parquet_parts[0]).names)
+        sort_spec = get_orderby_sort_spec(query_sql, column_names)
+
+    temporary = target.with_name(f".{target.name}.tmp-{os.getpid()}")
+    try:
+        # A single unordered part is already a canonical Parquet file. Copying
+        # it avoids decoding and rewriting large results.
+        if len(parquet_parts) == 1 and not sort_spec[0]:
+            shutil.copy2(parquet_parts[0], temporary)
+        else:
+            frame = pd.read_parquet(query_dir)
+            if sort_spec[0]:
+                frame = restore_orderby(frame, *sort_spec)
+            frame.to_parquet(temporary, index=False)
+        os.replace(temporary, target)
+    finally:
+        temporary.unlink(missing_ok=True)
+
+
+def normalize_ctas_results(
+    source_dir: Path,
+    output_dir: Path,
+    queries: dict[str, str],
+) -> list[Path]:
+    """Convert each committed qN CTAS table into one canonical Parquet file."""
+    query_dirs = _query_directories(source_dir)
+    if not query_dirs:
+        raise ValueError(f"No committed CTAS result tables found in {source_dir}")
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    normalized: list[Path] = []
+    for query_dir in query_dirs:
+        target = output_dir / f"{query_dir.name}.parquet"
+        _write_canonical_result(query_dir, target, queries.get(query_dir.name.upper()))
+        shutil.rmtree(query_dir)
+        normalized.append(target)
+    return normalized
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--source-dir", required=True, type=Path, help="Directory containing temporary qN CTAS tables.")
+    parser.add_argument(
+        "--output-dir",
+        required=True,
+        type=Path,
+        help="Canonical query_results directory that will receive qN.parquet files.",
+    )
+    parser.add_argument("--benchmark-type", required=True, choices=["tpch", "tpcds"])
+    parser.add_argument("--queries-file", default=None, help="Optional custom query-definition JSON file.")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    try:
+        outputs = normalize_ctas_results(
+            args.source_dir,
+            args.output_dir,
+            get_queries(args.benchmark_type, args.queries_file),
+        )
+    except (OSError, ValueError) as error:
+        print(f"Error: Failed to normalize CTAS results: {error}", file=sys.stderr)
+        sys.exit(1)
+    for output in outputs:
+        print(f"Normalized CTAS result: {output}")

common/testing/normalize_ctas_results_test.py

@@ -0,0 +1,56 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from pathlib import Path
+
+import pandas as pd
+
+from common.testing.normalize_ctas_results import normalize_ctas_results
+
+
+def test_normalize_ctas_results_merges_parts_and_restores_order(tmp_path: Path):
+    scratch = tmp_path / "scratch"
+    query_dir = scratch / "q1"
+    query_dir.mkdir(parents=True)
+    # CTAS assigns positional storage names, while ORDER BY still refers to
+    # the original SELECT output name.
+    pd.DataFrame({"c1": [3, 1]}).to_parquet(query_dir / "part-00000.parquet", index=False)
+    pd.DataFrame({"c1": [4, 2]}).to_parquet(query_dir / "part-00001.parquet", index=False)
+
+    output = tmp_path / "output"
+    normalized = normalize_ctas_results(
+        scratch,
+        output,
+        {"Q1": "SELECT value FROM source ORDER BY value"},
+    )
+
+    assert normalized == [output / "q1.parquet"]
+    assert pd.read_parquet(normalized[0])["c1"].tolist() == [1, 2, 3, 4]
+    assert not query_dir.exists()
+
+
+def test_normalize_ctas_results_promotes_single_unordered_part(tmp_path: Path):
+    scratch = tmp_path / "scratch"
+    query_dir = scratch / "q6"
+    query_dir.mkdir(parents=True)
+    expected = pd.DataFrame({"revenue": [42.0]})
+    expected.to_parquet(query_dir / "part-00000.parquet", index=False)
+
+    output = tmp_path / "output"
+    normalize_ctas_results(scratch, output, {"Q6": "SELECT revenue FROM source"})
+
+    pd.testing.assert_frame_equal(pd.read_parquet(output / "q6.parquet"), expected)
+    assert not query_dir.exists()
+
+
+def test_normalize_ctas_results_writes_empty_result(tmp_path: Path):
+    scratch = tmp_path / "scratch"
+    query_dir = scratch / "q15"
+    query_dir.mkdir(parents=True)
+    (query_dir / ".prestoSchema").write_text("{}")
+
+    output = tmp_path / "output"
+    normalize_ctas_results(scratch, output, {"Q15": "SELECT supplier_no FROM source"})
+
+    assert pd.read_parquet(output / "q15.parquet").empty
+    assert not query_dir.exists()

common/testing/result_comparison.py

@@ -34,22 +34,19 @@
 # ---------------------------------------------------------------------------
 
 
-def get_orderby_col_indices(query_sql: str, expected_col_names: list[str]) -> tuple[list[int], list[bool]]:
-    """
-    Extract ORDER BY column positions and directions from SQL using sqlglot.
-
-    Returns (indices, ascending) where ascending[i] is True iff indices[i]
-    is sorted ASC. Returns ([], []) when there is no ORDER BY, or when any
-    ORDER BY expression is too complex to map to a result column (CASE,
-    aggregate, etc.).
-    """
-    expr = sqlglot.parse_one(query_sql)
+def get_orderby_sort_spec(
+    query_sql: str,
+    expected_col_names: list[str],
+) -> tuple[list[int], list[bool], list[bool]]:
+    """Resolve ORDER BY expressions to positional columns and sort options."""
+    expr = sqlglot.parse_one(query_sql, read="presto")
     order = next((e for e in expr.find_all(sqlglot.exp.Order)), None)
     if not order:
-        return [], []
+        return [], [], []
 
     sort_col_indices: list[int] = []
     ascending: list[bool] = []
+    nulls_first: list[bool] = []
 
     for ordered in order.expressions:
         key = ordered.this
@@ -76,10 +73,25 @@ def get_orderby_col_indices(query_sql: str, expected_col_names: list[str]) -> tu
                 f"ORDER BY expression {key.sql()!r} couldn't be mapped to a result column; "
                 "ORDER BY validation and tie-boundary handling will be skipped for this query."
             )
-            return [], []
+            return [], [], []
 
         sort_col_indices.append(resolved_idx)
         ascending.append(not is_desc)
+        nulls_first.append(bool(ordered.args.get("nulls_first")))
+
+    return sort_col_indices, ascending, nulls_first
+
+
+def get_orderby_col_indices(query_sql: str, expected_col_names: list[str]) -> tuple[list[int], list[bool]]:
+    """
+    Extract ORDER BY column positions and directions from SQL using sqlglot.
+
+    Returns (indices, ascending) where ascending[i] is True iff indices[i]
+    is sorted ASC. Returns ([], []) when there is no ORDER BY, or when any
+    ORDER BY expression is too complex to map to a result column (CASE,
+    aggregate, etc.).
+    """
+    sort_col_indices, ascending, _ = get_orderby_sort_spec(query_sql, expected_col_names)
 
     return sort_col_indices, ascending
 
@@ -223,11 +235,10 @@ def _validate_orderby(df: pd.DataFrame, sort_col_indices: list[int], ascending:
     itself.
 
     Null handling: SQL leaves the position of NULLs in ORDER BY engine-
-    defined (Presto/Velox default to NULLS LAST for ASC, NULLS FIRST for
-    DESC, others differ), so adjacent pairs where either side is NULL are
-    not counted as violations. Two NULLs are treated as tied (do not break
-    the outer tie group); a NULL adjacent to a non-NULL value is treated as
-    a change for tie-group tracking only.
+    defined unless NULLS FIRST/LAST is explicit, so adjacent pairs where
+    either side is NULL are not counted as violations. Two NULLs are treated
+    as tied (do not break the outer tie group); a NULL adjacent to a non-NULL
+    value is treated as a change for tie-group tracking only.
 
     Raises AssertionError pointing at the first offending row. A no-op when
     sort_col_indices is empty or the frame has 0-1 rows.
@@ -288,6 +299,32 @@ def _validate_orderby(df: pd.DataFrame, sort_col_indices: list[int], ascending:
         prior_changed = prior_changed | value_changed
 
 
+def restore_orderby(
+    df: pd.DataFrame,
+    sort_col_indices: list[int],
+    ascending: list[bool],
+    nulls_first: list[bool],
+) -> pd.DataFrame:
+    """Reconstruct SQL result order after unordered distributed storage."""
+    if not sort_col_indices or len(df) <= 1:
+        return df.reset_index(drop=True)
+
+    # Apply stable sorts from the least-significant key to the most-significant
+    # key. Access columns positionally because SQL results may contain duplicate
+    # column labels.
+    row_order = np.arange(len(df))
+    for col_idx, asc, null_first in reversed(list(zip(sort_col_indices, ascending, nulls_first))):
+        values = df.iloc[row_order, col_idx].reset_index(drop=True)
+        permutation = values.sort_values(
+            ascending=asc,
+            na_position="first" if null_first else "last",
+            kind="stable",
+        ).index.to_numpy()
+        row_order = row_order[permutation]
+
+    return df.iloc[row_order].reset_index(drop=True)
+
+
 def _canonical_sort(df: pd.DataFrame) -> pd.DataFrame:
     """
     Sort df by [all non-float columns, all float columns] for a deterministic,
@@ -443,8 +480,8 @@ def compare_result_frames(
     if len(actual) != len(expected):
         raise AssertionError(f"Row count mismatch: {len(actual)} (actual) vs {len(expected)} (expected)")
 
-    # 4. Parse ORDER BY / LIMIT
-    sort_col_indices, ascending = get_orderby_col_indices(query_sql, expected_col_names)
+    # 4. Parse ORDER BY / LIMIT.
+    sort_col_indices, ascending, _ = get_orderby_sort_spec(query_sql, expected_col_names)
     limit = get_limit(query_sql)
 
     # 5. Per-frame ORDER BY validation — engine order is intact here.

common/testing/result_comparison_test.py

@@ -13,6 +13,8 @@
     _canonical_sort,
     _find_last_tie_start,
     _normalize_to_expected,
+    get_orderby_sort_spec,
+    restore_orderby,
     _validate_orderby,
 )
 
@@ -103,6 +105,38 @@ def test_validate_orderby_handles_duplicate_column_labels():
     _validate_orderby(df, sort_col_indices=[0], ascending=[True])
 
 
+def test_restore_orderby_uses_mixed_directions_and_explicit_null_placement():
+    df = pd.DataFrame(
+        {
+            "a": [2, 1, None, 1, 1],
+            "b": [4, 1, 9, None, 3],
+            "payload": ["a2", "a1-low", "null-a", "a1-null", "a1-high"],
+        }
+    )
+    indices, ascending, nulls_first = get_orderby_sort_spec(
+        "SELECT a, b, payload FROM source ORDER BY a ASC NULLS LAST, b DESC NULLS FIRST",
+        list(df.columns),
+    )
+
+    restored = restore_orderby(df, indices, ascending, nulls_first)
+
+    assert restored["payload"].tolist() == ["a1-null", "a1-high", "a1-low", "a2", "null-a"]
+
+
+def test_restore_orderby_defaults_to_presto_nulls_last_for_desc():
+    df = pd.DataFrame({"score": [None, 1, 2], "payload": ["null", "low", "high"]})
+    indices, ascending, nulls_first = get_orderby_sort_spec(
+        "SELECT score, payload FROM source ORDER BY score DESC",
+        list(df.columns),
+    )
+
+    restored = restore_orderby(df, indices, ascending, nulls_first)
+
+    assert ascending == [False]
+    assert nulls_first == [False]
+    assert restored["payload"].tolist() == ["high", "low", "null"]
+
+
 # ---------------------------------------------------------------------------
 # _validate_orderby null handling
 # ---------------------------------------------------------------------------
@@ -123,15 +157,15 @@ def test_validate_orderby_object_column_all_none_does_not_raise():
 
 
 def test_validate_orderby_object_column_null_then_value_does_not_raise():
-    # NULLS-FIRST style ordering (e.g. Presto DESC default): nulls then ASC
-    # of non-null values should pass.
+    # Explicit NULLS-FIRST style ordering: nulls then ASC of non-null values
+    # should pass.
     df = pd.DataFrame({0: [None, None, "a", "b", "c"]}, dtype=object)
     _validate_orderby(df, sort_col_indices=[0], ascending=[True])
 
 
 def test_validate_orderby_object_column_value_then_null_does_not_raise():
-    # NULLS-LAST style ordering (e.g. Presto ASC default): non-null ASC
-    # values then nulls at the tail should pass.
+    # Presto's default NULLS-LAST ordering: non-null values then nulls at the
+    # tail should pass regardless of sort direction.
     df = pd.DataFrame({0: ["a", "b", "c", None, None]}, dtype=object)
     _validate_orderby(df, sort_col_indices=[0], ascending=[True])