From c4fa7818b7ec0928ebfb96847083a7ae276f7e39 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 19 Oct 2026 09:43:50 +0100 Subject: [PATCH 1/7] Add planner-owned selection and bounded ListLayout scans Rebase draft PR #10537 onto develop 92a7db079. Apply selections through a composable planner, resolve list offsets before selected elements, and split substantial fixed-width child projections at stored segment boundaries. The projection planner owns subdivision, byte-budget admission, child task execution and ordered collection. Unsupported splits retain the serial path. Keep Planner and Morsel trait methods unchanged. Replace the experimental completion callbacks, WorkFuture and ChildExecutor traits with one PlannerOutput::Await variant; morsels only publish the assembled result. Preserve develop's global local-file coalescing configuration. The 8 KiB experiment caused tiny-lookup regressions and is not part of this version. Preserve masked lazy lowering, bounded dense-list grouping, shared value caches, validated chunk offsets and the fieldless nullable-struct fix. Retain the diagnostic examples and regression benchmark workflow. Backtracking: previous PR head 8442d4695; experimental CI head 60fda25d2; local pre-review checkpoint a253e5727 and reviewed checkpoint e8fd21a4a. The byte budget covers in-flight decoded work, not retained output or encoded/coalesced buffers. Full performance revalidation remains necessary. Validation: 692 layout and 79 scan library tests passed, followed by all 7 await protocol tests including an additional await/IO regression. Scoped all-target/all-feature clippy passed with -D warnings; formatted with nightly-2026-09-10. Focused E2E measurements follow this checkpoint. Signed-off-by: Joe Isaacs --- .github/workflows/scan-v2-regression.yml | 110 ++++ Cargo.lock | 2 + benchmarks/random-access-bench/README.md | 11 + benchmarks/random-access-bench/src/lib.rs | 24 +- benchmarks/random-access-bench/src/main.rs | 4 + scripts/scan-v2-regression.py | 312 ++++++++++++ vortex-bench/src/random_access/take.rs | 34 +- vortex-layout/benches/pipeline_scan.rs | 1 + vortex-layout/src/layouts/chunked/mod.rs | 5 + vortex-layout/src/plan/physical/lower.rs | 12 +- vortex-layout/src/plan/physical/operator.rs | 22 + .../src/plan/physical/operators/list.rs | 393 +++++++++++++++ .../src/plan/physical/operators/list/work.rs | 411 +++++++++++++++ .../src/plan/physical/operators/mod.rs | 2 + .../src/plan/physical/operators/pack.rs | 13 +- .../src/plan/physical/pipeline/children.rs | 29 ++ .../src/plan/physical/pipeline/compile.rs | 71 ++- .../src/plan/physical/pipeline/mod.rs | 2 + .../src/plan/physical/pipeline/scan.rs | 172 ++++++- .../plan/physical/pipeline/scan/validate.rs | 6 + .../physical/pipeline/scheduling_tests.rs | 4 + .../src/plan/physical/pipeline/tests.rs | 264 ++++++++++ .../src/plan/physical/plans/concat.rs | 3 + .../src/plan/physical/plans/list_pack.rs | 76 +++ vortex-layout/src/plan/physical/tests.rs | 58 +++ vortex-layout/src/scan/execution/adaptor.rs | 18 +- vortex-layout/src/scan/execution/children.rs | 230 +++++++++ .../src/scan/execution/children/splits.rs | 152 ++++++ .../scan/execution/children/splits/tests.rs | 139 ++++++ .../src/scan/execution/children/tests.rs | 357 +++++++++++++ vortex-layout/src/scan/execution/mod.rs | 37 +- .../src/scan/execution/projection.rs | 146 +++--- vortex-layout/src/scan/v2/io.rs | 108 +++- vortex-layout/src/scan/v2/mod.rs | 98 ++-- vortex-layout/src/scan/v2/tests.rs | 237 ++++++++- vortex-layout/src/scan/v2/tests/child_work.rs | 101 ++++ vortex-scan/Cargo.toml | 1 + vortex-scan/src/planning/driver/mod.rs | 99 +++- vortex-scan/src/planning/driver/tests.rs | 30 ++ .../src/planning/driver/tests/tasks.rs | 175 +++++++ vortex-scan/src/planning/mod.rs | 1 + vortex-scan/src/planning/planner.rs | 6 + vortex-scan/src/planning/selection.rs | 210 ++++++++ vortex/Cargo.toml | 1 + vortex/examples/generate_list_fixture.rs | 183 +++++++ vortex/examples/rewrite_list_layout.rs | 151 ++++++ vortex/examples/selection_scan.rs | 468 ++++++++++++++++++ 47 files changed, 4775 insertions(+), 214 deletions(-) create mode 100644 .github/workflows/scan-v2-regression.yml create mode 100644 scripts/scan-v2-regression.py create mode 100644 vortex-layout/src/plan/physical/operators/list.rs create mode 100644 vortex-layout/src/plan/physical/operators/list/work.rs create mode 100644 vortex-layout/src/plan/physical/pipeline/children.rs create mode 100644 vortex-layout/src/scan/execution/children.rs create mode 100644 vortex-layout/src/scan/execution/children/splits.rs create mode 100644 vortex-layout/src/scan/execution/children/splits/tests.rs create mode 100644 vortex-layout/src/scan/execution/children/tests.rs create mode 100644 vortex-layout/src/scan/v2/tests/child_work.rs create mode 100644 vortex-scan/src/planning/driver/tests/tasks.rs create mode 100644 vortex-scan/src/planning/selection.rs create mode 100644 vortex/examples/generate_list_fixture.rs create mode 100644 vortex/examples/rewrite_list_layout.rs create mode 100644 vortex/examples/selection_scan.rs diff --git a/.github/workflows/scan-v2-regression.yml b/.github/workflows/scan-v2-regression.yml new file mode 100644 index 00000000000..5610959b28d --- /dev/null +++ b/.github/workflows/scan-v2-regression.yml @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +name: Scan V2 Regression Comparison +run-name: Scan V2 regression comparison on ${{ github.ref_name }} + +on: + push: + branches: ["ji/scan-v2-pr10540"] + workflow_dispatch: { } + +permissions: + contents: read + +concurrency: + group: scan-v2-regression-${{ github.ref }} + cancel-in-progress: true + +jobs: + compare: + if: github.repository == 'vortex-data/vortex' + timeout-minutes: 180 + runs-on: >- + runs-on=${{ github.run_id }}/runner=bench-dedicated/family=c8gd.8xlarge/image=ubuntu24-full-arm64-pre-v2/extras=s3-cache/tag=scan-v2-regression + env: + # Previous random-access branch, before rebasing onto PR #10540. + BASELINE_COMMIT: "ad83fcbfc05d1885c08b7ae5cc57302e412638f3" + CARGO_BUILD_JOBS: "8" + CARGO_TARGET_DIR: ${{ github.workspace }}/target + RUSTFLAGS: "-C target-cpu=native -C force-frame-pointers=yes" + RUST_LOG: warn + steps: + - uses: runs-on/action@v2 + with: + sccache: s3 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7 + with: + fetch-depth: 0 + - uses: ./.github/actions/setup-prebuild + with: + repo-token: ${{ secrets.GITHUB_TOKEN }} + enable-sccache: "true" + - uses: ./.github/actions/system-info + + - name: Build current scan and data preparation binaries + shell: bash + run: | + set -euo pipefail + mkdir -p scan-v2-binaries scan-v2-results + env -u VORTEX_SELECTION_METRICS cargo build --locked --profile release_debug \ + -p vortex --example selection_scan --example rewrite_list_layout + cp target/release_debug/examples/selection_scan scan-v2-binaries/current + cp target/release_debug/examples/rewrite_list_layout scan-v2-binaries/rewrite + env -u VORTEX_SELECTION_METRICS cargo build --locked --profile release_debug \ + -p random-access-bench + cp target/release_debug/random-access-bench scan-v2-binaries/standard + + - name: Build previous V2 from the pinned checkpoint + shell: bash + run: | + set -euo pipefail + git worktree add --detach "$RUNNER_TEMP/scan-v2-baseline" "$BASELINE_COMMIT" + ( + cd "$RUNNER_TEMP/scan-v2-baseline" + env -u VORTEX_SELECTION_METRICS cargo build --locked --profile release_debug \ + -p vortex --example selection_scan + ) + cp target/release_debug/examples/selection_scan scan-v2-binaries/baseline + sha256sum scan-v2-binaries/* > scan-v2-results/binaries.sha256 + + - name: Prepare shared datasets and verify ListLayout rewrites + shell: bash + env: + VORTEX_SCAN_V2: "0" + run: | + set -euo pipefail + scan-v2-binaries/standard --prepare-data --formats vortex + mkdir -p scan-v2-data + source_file=vortex-bench/data/random_access/nested_lists/nested_lists.vortex + scan-v2-binaries/rewrite "$source_file" scan-v2-data/flat-control.vortex flat \ + > scan-v2-results/layout-flat.json + scan-v2-binaries/rewrite "$source_file" scan-v2-data/elements-default.vortex list-elements 1048576 \ + > scan-v2-results/layout-1MiB.json + scan-v2-binaries/rewrite "$source_file" scan-v2-data/elements-64k.vortex list-elements 65536 \ + > scan-v2-results/layout-64KiB.json + + - name: Isolate benchmark CPUs + run: sudo bash scripts/setup-benchmark.sh + + - name: Compare Arrow and standard suites and repeat borderline regressions + shell: bash + env: + CANDIDATE_COMMIT: ${{ github.sha }} + run: | + set -euo pipefail + bash scripts/bench-taskset.sh python3 scripts/scan-v2-regression.py \ + --baseline scan-v2-binaries/baseline \ + --candidate scan-v2-binaries/current \ + --standard-benchmark scan-v2-binaries/standard \ + --baseline-commit "$BASELINE_COMMIT" \ + --candidate-commit "$CANDIDATE_COMMIT" \ + | tee scan-v2-results/run.log + + - name: Upload results, raw samples, and layout metadata + if: always() + uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7 + with: + name: scan-v2-regression-${{ github.sha }} + path: scan-v2-results/ + if-no-files-found: warn diff --git a/Cargo.lock b/Cargo.lock index 22aace294d8..6f3d8741fbf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11323,6 +11323,7 @@ dependencies = [ "arrow-schema 59.3.0", "codspeed-divan-compat", "fastlanes 0.7.2", + "futures", "mimalloc", "parquet 59.3.0", "rand 0.10.3", @@ -12396,6 +12397,7 @@ dependencies = [ "parking_lot", "roaring", "rstest", + "smallvec", "tracing", "vortex-array", "vortex-buffer", diff --git a/benchmarks/random-access-bench/README.md b/benchmarks/random-access-bench/README.md index 73f52f5b219..cddea03dd97 100644 --- a/benchmarks/random-access-bench/README.md +++ b/benchmarks/random-access-bench/README.md @@ -28,6 +28,17 @@ work in each timed iteration. CI drives the full matrix via cargo run -p random-access-bench --profile release_debug --features lance ``` +Pass `--scan-v2` to use Vortex's planner-driven V2 scan API. The default uses V1. Run each +version separately with the same dataset, pattern, and open-mode options; save their outputs +to separate files because measurement names stay the same: + +```bash +cargo run -p random-access-bench --profile release_debug -- \ + --formats vortex --display-format gh-json -o v1.json +cargo run -p random-access-bench --profile release_debug -- \ + --formats vortex --scan-v2 --display-format gh-json -o v2.json +``` + ## Running against S3 The same benchmark can read its data from an object store instead of local disk. The remote diff --git a/benchmarks/random-access-bench/src/lib.rs b/benchmarks/random-access-bench/src/lib.rs index 357e688ed53..466e7c980ea 100644 --- a/benchmarks/random-access-bench/src/lib.rs +++ b/benchmarks/random-access-bench/src/lib.rs @@ -153,10 +153,11 @@ async fn benchmark_random_access( storage: &str, reopen: bool, remote: Option<&RemoteDataDir>, + scan_v2: bool, ) -> Result { let time_limit = Duration::from_secs(time_limit_secs); let mut runs = Vec::new(); - let prepared_accessor = open_accessor(dataset, format, remote).await?; + let prepared_accessor = open_accessor(dataset, format, remote, scan_v2).await?; let cached_accessor = (!reopen).then_some(prepared_accessor); if let Some(accessor) = &cached_accessor { @@ -176,7 +177,7 @@ async fn benchmark_random_access( let arr = if let Some(accessor) = &cached_accessor { accessor.take(indices).await? } else { - open_accessor(dataset, format, remote) + open_accessor(dataset, format, remote, scan_v2) .await? .take(indices) .await? @@ -306,6 +307,7 @@ async fn open_accessor( dataset: &dyn BenchDataset, format: Format, remote: Option<&RemoteDataDir>, + scan_v2: bool, ) -> Result> { let name = format!( "random-access/{}/{}-{}", @@ -323,12 +325,13 @@ async fn open_accessor( } Format::OnDiskVortex | Format::VortexCompact => { let path = dataset_path(dataset, format).await?; - Ok(match remote { - Some(remote) => Box::new( - VortexRandomAccessor::open_object_store(remote, &path, name, format).await?, - ), - None => Box::new(VortexRandomAccessor::open(path, name, format).await?), - }) + let accessor = match remote { + Some(remote) => { + VortexRandomAccessor::open_object_store(remote, &path, name, format).await? + } + None => VortexRandomAccessor::open(path, name, format).await?, + }; + Ok(Box::new(accessor.with_scan_v2(scan_v2))) } Format::Parquet => { let path = dataset_path(dataset, format).await?; @@ -375,6 +378,8 @@ pub struct RunConfig { pub time_limit: u64, /// Whether each benchmark run reuses or reopens its file handle. pub open_mode: OpenMode, + /// Uses the explicit V2 scan API for Vortex formats when enabled. + pub scan_v2: bool, /// Format used to emit measurements. pub display_format: DisplayFormat, /// Optional path for the primary result output. @@ -393,6 +398,7 @@ pub async fn run(config: RunConfig) -> Result<()> { patterns, time_limit, open_mode, + scan_v2, display_format, output_path, ingest_output, @@ -446,6 +452,7 @@ pub async fn run(config: RunConfig) -> Result<()> { storage, reopen, remote, + scan_v2, ) .await?; @@ -475,6 +482,7 @@ pub async fn run(config: RunConfig) -> Result<()> { storage, reopen, remote, + scan_v2, ) .await?; diff --git a/benchmarks/random-access-bench/src/main.rs b/benchmarks/random-access-bench/src/main.rs index 3adb7e22beb..da0483a8bc6 100644 --- a/benchmarks/random-access-bench/src/main.rs +++ b/benchmarks/random-access-bench/src/main.rs @@ -88,6 +88,9 @@ struct Args { /// Whether to reopen the file on each iteration, use a cached handle, or run both. #[arg(long, value_enum, default_value_t = OpenMode::Both)] open_mode: OpenMode, + /// Use the explicit V2 scan API for Vortex formats. + #[arg(long)] + scan_v2: bool, /// Read the data files from this remote directory (e.g. `s3://bucket/prefix/`) instead of /// local disk. The directory must mirror the layout of the local benchmark data directory, /// as produced by `--prepare-data`. @@ -124,6 +127,7 @@ async fn main() -> Result<()> { patterns: args.patterns, time_limit: args.time_limit, open_mode: args.open_mode, + scan_v2: args.scan_v2, display_format: args.display_format, output_path: args.output_path, ingest_output: args.ingest_output, diff --git a/scripts/scan-v2-regression.py b/scripts/scan-v2-regression.py new file mode 100644 index 00000000000..17c637c27b0 --- /dev/null +++ b/scripts/scan-v2-regression.py @@ -0,0 +1,312 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +"""Compare V1, previous V2, and current V2 with full Arrow materialization.""" + +import argparse +import csv +import hashlib +import itertools +import json +import os +import platform +import statistics +import subprocess +from pathlib import Path + + +def digest(path: Path) -> str: + checksum = hashlib.sha256() + with path.open("rb") as source: + for chunk in iter(lambda: source.read(1024 * 1024), b""): + checksum.update(chunk) + return checksum.hexdigest() + + +def change(current: float, baseline: float) -> float: + return 100 * (current / baseline - 1) + + +def measure( + case: dict, + versions: dict[str, tuple[Path, str]], + files: dict[str, Path], + cpus: list[int], + output: Path, + iterations: int, + rounds: int, +) -> dict: + env = {k: v for k, v in os.environ.items() if not k.startswith(("VORTEX_", "TOKIO_"))} + env["RUST_LOG"] = "warn" + seeds = [] + for seed in [17, 42, 73]: + medians = {v: [] for v in versions} + expected_indices = None + for repeat in range(rounds): + order = list(versions)[:: 1 if repeat % 2 == 0 else -1] + for version in order: + binary, flag = versions[version] + stem = "-".join(str(v) for v in (*case.values(), seed, repeat, version)) + command = [ + "taskset", + "-c", + ",".join(map(str, cpus[: case["cores"]])), + str(binary), + str(files[case["dataset"]]), + str(iterations), + "100", + case["pattern"], + case["mode"], + "disk", + "decode", + str(seed), + ] + with (output / f"{stem}.json").open("w") as target, (output / f"{stem}.log").open("w") as log: + subprocess.run( + command, env=dict(env, VORTEX_SCAN_V2=flag), stdout=target, stderr=log, check=True, timeout=180 + ) + result = json.loads((output / f"{stem}.json").read_text()) + if result["v2"] != flag or len(result["runs"]) != iterations: + raise RuntimeError(f"Wrong scan mode or sample count: {stem}") + if any(run["instrumented"] for run in result["runs"]): + raise RuntimeError(f"Instrumented binary used for timing: {stem}") + if expected_indices is None: + expected_indices = result["indices"] + if result["indices"] != expected_indices: + raise RuntimeError(f"Versions selected different rows: {stem}") + samples = result["runs"][iterations // 8 :] + medians[version].append(statistics.median(x["elapsed_ns"] for x in samples) / 1e6) + times = {v: statistics.median(values) for v, values in medians.items()} + seeds.append(dict(seed=seed, **times, rounds=medians)) + row = dict(case, **{v + "_ms": statistics.median(s[v] for s in seeds) for v in versions}, seeds=seeds) + row["vs_old_v2_pct"] = change(row["new_v2_ms"], row["old_v2_ms"]) + if "v1" in versions: + row["vs_v1_pct"] = change(row["new_v2_ms"], row["v1_ms"]) + print(json.dumps({k: v for k, v in row.items() if k != "seeds"}), flush=True) + return row + + +def save(output: Path, metadata: dict, summary: list[dict], confirmations: list[dict]): + (output / "results.json").write_text( + json.dumps(dict(metadata, summary=summary, confirmations=confirmations), indent=2) + "\n" + ) + with (output / "results.csv").open("w") as target: + writer = csv.DictWriter(target, fieldnames=[k for k in summary[0] if k != "seeds"]) + writer.writeheader() + writer.writerows({k: v for k, v in row.items() if k != "seeds"} for row in summary) + + +def report(output: Path, metadata: dict, summary: list[dict], confirmations: list[dict]): + lines = [ + "# Scan V2 regression comparison", + "", + f"Candidate: `{metadata['candidate_commit']}`; previous V2: `{metadata['baseline_commit']}`.", + "", + "100 selected rows; full Arrow decode/drop; warm local disk. " + "Cached reuses file preparation; reopen includes opening/metadata.", + "Three seeds, two alternating rounds of 64 queries, first 8 discarded. Positive percentages mean slower.", + "", + f"Cases above +3% vs previous V2: {sum(r['vs_old_v2_pct'] > 3 for r in summary)}/{len(summary)}; " + f"above +5% vs V1: {sum(r['vs_v1_pct'] > 5 for r in summary)}/{len(summary)}.", + "", + "| Dataset | Cores | Selection | Open | V1 ms | Previous V2 ms | Current V2 ms | vs V1 | vs previous V2 |", + "|---|---:|---|---|---:|---:|---:|---:|---:|", + ] + for r in summary: + lines.append( + f"| {r['dataset']} | {r['cores']} | {r['pattern']} | {r['mode']} | {r['v1_ms']:.4f} | " + f"{r['old_v2_ms']:.4f} | {r['new_v2_ms']:.4f} | {r['vs_v1_pct']:+.1f}% | {r['vs_old_v2_pct']:+.1f}% |" + ) + if confirmations: + lines += [ + "", + "## Longer reruns", + "", + "Rerun when the first-pass aggregate exceeds +2%, or any seed exceeds +3%, versus previous V2.", + "Four alternating rounds of 128 queries, first 16 discarded, same three seeds.", + "", + "| Dataset | Cores | Selection | Open | Previous V2 ms | Current V2 ms | Change |", + "|---|---:|---|---|---:|---:|---:|", + ] + for r in confirmations: + lines.append( + f"| {r['dataset']} | {r['cores']} | {r['pattern']} | {r['mode']} | " + f"{r['old_v2_ms']:.4f} | {r['new_v2_ms']:.4f} | {r['vs_old_v2_pct']:+.1f}% |" + ) + text = "\n".join(lines) + "\n" + (output / "summary.md").write_text(text) + if step_summary := os.environ.get("GITHUB_STEP_SUMMARY"): + with Path(step_summary).open("a") as target: + target.write(text) + + +def compare_standard(binary: Path, cpus: list[int], output: Path): + """Compare the standard suite, including its legacy six-row taxi lookup.""" + output.mkdir(exist_ok=True) + binary = binary.resolve() + env = {k: v for k, v in os.environ.items() if not k.startswith(("VORTEX_", "TOKIO_"))} + env["RUST_LOG"] = "warn" + metadata = dict( + binary=str(binary), + sha256=digest(binary), + workdir=str(Path.cwd()), + cpus=cpus[:30], + rounds=2, + seconds_per_target=5, + seed=42, + ) + metadata["output"] = "standard canonical-array output; one-second cached warmup" + commands, summary = [], [] + for cores in [1, 30]: + samples = {} + for repeat in range(2): + for version in ["v1", "v2"] if repeat == 0 else ["v2", "v1"]: + stem = f"{cores}-{repeat}-{version}" + path = output / f"{stem}.jsonl" + command = [ + "taskset", + "-c", + ",".join(map(str, cpus[:cores])), + str(binary), + "--formats", + "vortex", + "--open-mode", + "both", + "--time-limit", + "5", + "--display-format", + "gh-json", + "-o", + str(path), + ] + if version == "v2": + command.append("--scan-v2") + commands.append(command) + print(f"Standard suite: {stem}", flush=True) + with (output / f"{stem}.log").open("w") as log: + subprocess.run(command, env=env, stdout=log, stderr=log, check=True, timeout=600) + rows = [json.loads(line) for line in path.read_text().splitlines() if line.strip()] + if len(rows) != 18 or len({row["name"] for row in rows}) != 18: + raise RuntimeError(f"Expected all 18 standard targets: {stem}") + for row in rows: + if row["unit"] != "ns": + raise RuntimeError(f"Unexpected standard benchmark unit: {row}") + samples.setdefault(row["name"], {"v1": [], "v2": []})[version].append(row["value"] / 1e6) + for name, rounds in samples.items(): + if any(len(values) != 2 for values in rounds.values()): + raise RuntimeError(f"Unpaired standard benchmark target: {name}") + times = {v: statistics.median(values) for v, values in rounds.items()} + row = dict( + cores=cores, + name=name, + v1_ms=times["v1"], + v2_ms=times["v2"], + vs_v1_pct=change(times["v2"], times["v1"]), + rounds=rounds, + ) + summary.append(row) + print(json.dumps(row), flush=True) + (output / "results.json").write_text( + json.dumps(dict(metadata, commands=commands, summary=summary), indent=2) + "\n" + ) + lines = [ + "## Standard random-access suite", + "", + "V1 versus V2 from the same binary; canonical-array output, seed 42, warm local disk.", + "Two alternating rounds, five seconds per target; positive percentages mean slower.", + "", + f"Cases above +3% vs V1: {sum(r['vs_v1_pct'] > 3 for r in summary)}/{len(summary)}.", + "", + "| Cores | Target | V1 ms | V2 ms | Change |", + "|---:|---|---:|---:|---:|", + ] + lines.extend( + f"| {r['cores']} | {r['name']} | {r['v1_ms']:.4f} | {r['v2_ms']:.4f} | {r['vs_v1_pct']:+.1f}% |" + for r in summary + ) + text = "\n".join(lines) + "\n" + (output / "summary.md").write_text(text) + if step_summary := os.environ.get("GITHUB_STEP_SUMMARY"): + with Path(step_summary).open("a") as target: + target.write(text) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--baseline", type=Path, required=True) + parser.add_argument("--candidate", type=Path, required=True) + parser.add_argument("--baseline-commit", required=True) + parser.add_argument("--candidate-commit", required=True) + parser.add_argument( + "--standard-benchmark", type=Path, help="Also compare V1/V2 with the standard random-access binary" + ) + parser.add_argument("--data-dir", type=Path, default=Path("vortex-bench/data")) + parser.add_argument("--list-dir", type=Path, default=Path("scan-v2-data")) + parser.add_argument("--output", type=Path, default=Path("scan-v2-results")) + args = parser.parse_args() + args.output.mkdir(parents=True, exist_ok=True) + raw = args.output / "raw" + raw.mkdir(exist_ok=True) + cpus = sorted(os.sched_getaffinity(0)) + if len(cpus) < 30: + raise RuntimeError(f"This comparison requires 30 available CPUs; found {len(cpus)}") + files = { + "taxi": args.data_dir / "taxi/taxi.vortex", + "feature-vectors": args.data_dir / "random_access/feature_vectors/feature_vectors.vortex", + "nested-structs": args.data_dir / "random_access/nested_structs/nested_structs.vortex", + "flat-lists": args.list_dir / "flat-control.vortex", + "list-1MiB": args.list_dir / "elements-default.vortex", + "list-64KiB": args.list_dir / "elements-64k.vortex", + } + versions = { + "v1": (args.candidate.resolve(), "0"), + "old_v2": (args.baseline.resolve(), "1"), + "new_v2": (args.candidate.resolve(), "1"), + } + metadata = dict( + baseline_commit=args.baseline_commit, + candidate_commit=args.candidate_commit, + platform=platform.platform(), + cpus=cpus[:30], + method=dict( + rows=100, + seeds=[17, 42, 73], + iterations=64, + warmups=8, + rounds=2, + output="full Arrow decode and drop", + storage="warm local disk", + aggregation="median of round medians, then median across seeds", + confirmation_iterations=128, + confirmation_warmups=16, + confirmation_rounds=4, + ), + files={k: dict(path=str(p), sha256=digest(p), bytes=p.stat().st_size) for k, p in files.items()}, + versions={k: dict(binary=str(p), sha256=digest(p), scan_v2=f) for k, (p, f) in versions.items()}, + ) + (args.output / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n") + summary = [] + confirmations = [] + for cores, dataset, pattern, mode in itertools.product( + [1, 30], files, ["clustered", "uniform"], ["cached", "reopen"] + ): + case = dict(dataset=dataset, cores=cores, pattern=pattern, mode=mode) + summary.append(measure(case, versions, files, cpus, raw, 64, 2)) + save(args.output, metadata, summary, confirmations) + reruns = args.output / "reruns" + reruns.mkdir(exist_ok=True) + for row in summary: + if row["vs_old_v2_pct"] <= 2 and all(change(s["new_v2"], s["old_v2"]) <= 3 for s in row["seeds"]): + continue + case = {k: row[k] for k in ["dataset", "cores", "pattern", "mode"]} + confirmations.append( + measure(case, {k: v for k, v in versions.items() if k != "v1"}, files, cpus, reruns, 128, 4) + ) + save(args.output, metadata, summary, confirmations) + report(args.output, metadata, summary, confirmations) + if args.standard_benchmark: + compare_standard(args.standard_benchmark, cpus, args.output / "standard") + + +if __name__ == "__main__": + main() diff --git a/vortex-bench/src/random_access/take.rs b/vortex-bench/src/random_access/take.rs index 93a66ccd49e..c11285f26f9 100644 --- a/vortex-bench/src/random_access/take.rs +++ b/vortex-bench/src/random_access/take.rs @@ -29,10 +29,12 @@ use tokio::fs::File; use vortex::array::Canonical; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; +use vortex::array::stream::ArrayStreamAdapter; use vortex::array::stream::ArrayStreamExt; use vortex::buffer::Buffer; use vortex::file::OpenOptionsSessionExt; use vortex::file::VortexFile; +use vortex::layout::scan::v2::ScanBuilder as V2ScanBuilder; use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; use vortex::utils::aliases::hash_map::HashMap; @@ -133,6 +135,7 @@ pub struct VortexRandomAccessor { name: String, format: Format, file: VortexFile, + scan_v2: bool, } impl VortexRandomAccessor { @@ -151,6 +154,7 @@ impl VortexRandomAccessor { name: name.into(), format, file, + scan_v2: false, }) } @@ -170,8 +174,15 @@ impl VortexRandomAccessor { name: name.into(), format, file, + scan_v2: false, }) } + + /// Selects the explicit planner-driven V2 scan API for this benchmark accessor. + pub fn with_scan_v2(mut self, enabled: bool) -> Self { + self.scan_v2 = enabled; + self + } } #[async_trait] @@ -186,13 +197,22 @@ impl RandomAccessor for VortexRandomAccessor { async fn take(&self, indices: &[u64]) -> anyhow::Result { let indices_buf: Buffer = Buffer::from(indices.to_vec()); - let array = self - .file - .scan()? - .with_row_indices(StrictSortedBuffer::try_new(indices_buf)?) - .into_array_stream()? - .read_all() - .await?; + let indices = StrictSortedBuffer::try_new(indices_buf)?; + let array = if self.scan_v2 { + let stream = V2ScanBuilder::new(SESSION.clone(), self.file.scan_file()) + .with_row_indices(indices) + .into_stream()?; + ArrayStreamAdapter::new(self.file.dtype().clone(), stream) + .read_all() + .await? + } else { + self.file + .scan()? + .with_row_indices(indices) + .into_array_stream()? + .read_all() + .await? + }; // We canonicalize / decompress for equivalence to Arrow's `RecordBatch`es. let mut ctx = SESSION.create_execution_ctx(); diff --git a/vortex-layout/benches/pipeline_scan.rs b/vortex-layout/benches/pipeline_scan.rs index db4beeaa043..e9980aa6f66 100644 --- a/vortex-layout/benches/pipeline_scan.rs +++ b/vortex-layout/benches/pipeline_scan.rs @@ -124,6 +124,7 @@ fn scan(bencher: Bencher, columns: u32) { Turn::Output(_, array) => { black_box(array); } + Turn::Child(_) => unreachable!("benchmark has no asynchronous child work"), Turn::Waiting => unreachable!("every read is delivered immediately"), Turn::Done => break, } diff --git a/vortex-layout/src/layouts/chunked/mod.rs b/vortex-layout/src/layouts/chunked/mod.rs index 459429046c2..55c268b295e 100644 --- a/vortex-layout/src/layouts/chunked/mod.rs +++ b/vortex-layout/src/layouts/chunked/mod.rs @@ -125,6 +125,11 @@ impl Layout { OwnedLayoutChildren::layout_children(children), ) } + + /// Validated chunk boundaries, including the terminal row count. + pub(crate) fn chunk_offsets(&self) -> &[u64] { + &self.chunk_offsets + } } fn chunk_offsets(children: &dyn LayoutChildren) -> VortexResult> { diff --git a/vortex-layout/src/plan/physical/lower.rs b/vortex-layout/src/plan/physical/lower.rs index 1ea284cc1ca..9f965206544 100644 --- a/vortex-layout/src/plan/physical/lower.rs +++ b/vortex-layout/src/plan/physical/lower.rs @@ -83,21 +83,13 @@ fn lower_flat(layout: &FlatLayout) -> SegmentScanPlan { } fn lower_chunked(layout: &ChunkedLayout) -> VortexResult { - let mut row_offsets = Vec::with_capacity(layout.nchildren()); - let mut row_count = 0u64; - for index in 0..layout.nchildren() { - row_offsets.push(row_count); - row_count = row_count - .checked_add(layout.child_row_count(index)) - .ok_or_else(|| vortex_err!("Chunked row count overflow"))?; - } // SAFETY: Chunked layout construction validates that every child has the parent dtype and - // that their row counts sum to the parent; offsets were computed with checked addition above. + // that these checked cumulative offsets end at the parent's row count. Ok(unsafe { ConcatPlan::from_children_unchecked( layout.dtype().clone(), layout.row_count(), - row_offsets.into(), + layout.chunk_offsets()[..layout.nchildren()].into(), lazy_children(layout.to_layout(), (0..layout.nchildren()).collect()), ) }) diff --git a/vortex-layout/src/plan/physical/operator.rs b/vortex-layout/src/plan/physical/operator.rs index 96a8d0cfbd7..de1b1e35560 100644 --- a/vortex-layout/src/plan/physical/operator.rs +++ b/vortex-layout/src/plan/physical/operator.rs @@ -15,6 +15,7 @@ use vortex_session::VortexSession; use super::pipeline::Arena; use super::pipeline::Inlet; use super::pipeline::PortId; +use crate::plan::physical::pipeline::ChildProjection; use crate::segments::SegmentId; /// The output of one non-blocking compute call. @@ -43,6 +44,8 @@ pub enum SourceReady { pub enum SourceBlocked { /// Wait for bytes from newly queued or outstanding reads. Io, + /// Wait for a child projection, possibly alongside pending reads. + Child, /// The inlet needs more batches or its sink to close it. It must be open and have room /// for its sink to push; debug builds check that the sink can still wake this source. Inlet(usize), @@ -157,9 +160,28 @@ pub struct Cx<'a> { pub(super) exec: &'a mut ExecutionCtx, pub(super) touched: &'a mut InletSet, pub(super) changed: &'a mut InletSet, + pub(super) child_requests: &'a mut Vec, + pub(super) child_inbox: &'a mut Option>, +} + +#[derive(Default)] +pub(super) struct ChildInbox { + pub(super) pending: usize, + pub(super) completed: VecDeque<(u64, Option>)>, } impl Cx<'_> { + /// Describes child work for the enclosing planner. Return [`SourceBlocked::Child`] while awaiting its completion. + /// Identifiers need only be unique among this source's outstanding child projections. + pub(crate) fn request_child(&mut self, child: ChildProjection) { + self.child_requests.push(child); + } + + /// Takes a completed child projection, oldest completion first. + pub(crate) fn take_child(&mut self) -> Option<(u64, Option>)> { + self.child_inbox.as_mut()?.completed.pop_front() + } + /// Takes an inlet index whose sink appended data or closed since the last readiness /// check. Each index is reported once, and all inlets are reported on the first check. /// Remaining notifications are cleared after [`Source::ready`] returns. A source must diff --git a/vortex-layout/src/plan/physical/operators/list.rs b/vortex-layout/src/plan/physical/operators/list.rs new file mode 100644 index 00000000000..7ccc49dabc8 --- /dev/null +++ b/vortex-layout/src/plan/physical/operators/list.rs @@ -0,0 +1,393 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Resolves selected list offsets before scanning their element ranges. + +use std::collections::VecDeque; +use std::ops::Range; + +use rustc_hash::FxHashMap; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::expr::stats::Precision; +use vortex_array::expr::stats::Stat; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_mask::AllOr; +use vortex_mask::Mask; +use vortex_utils::parallelism::get_available_parallelism; + +use super::join; +use super::join_inlet; +use crate::plan::physical::ListPackPlan; +use crate::plan::physical::operator::Cx; +use crate::plan::physical::operator::Input; +use crate::plan::physical::operator::Operator; +use crate::plan::physical::operator::Source; +use crate::plan::physical::operator::SourceBlocked; +use crate::plan::physical::operator::SourceReady; +use crate::plan::physical::operator::Step; +use crate::plan::physical::pipeline::ChildProjection; +use crate::plan::physical::pipeline::ReadId; +use crate::plan::physical::pipeline::Scan; +use crate::plan::physical::pipeline::Split; +use crate::plan::physical::pipeline::Turn; +use crate::plan::physical::pipeline::ValueCache; +use crate::segments::SegmentId; + +mod work; + +/// A list's offsets and validity arrive before its selected element ranges can be planned. +pub(crate) struct ListSource { + plan: ListPackPlan, + mask: Mask, + values: ValueCache, + plan_children: bool, + awaiting_child: bool, + serial_splits: Option<(Vec, bool)>, + initialized: bool, + offsets: Option, + validity: Option, + elements: Option, + reads: FxHashMap>, + pending_children: usize, + chunks: Vec<(usize, ArrayRef)>, + output: Option, + done: bool, +} + +impl ListSource { + pub(crate) fn new( + plan: ListPackPlan, + mask: Mask, + values: ValueCache, + plan_children: bool, + ) -> Self { + Self { + plan, + mask, + values, + plan_children, + awaiting_child: false, + serial_splits: None, + initialized: false, + offsets: None, + validity: None, + elements: None, + reads: FxHashMap::default(), + pending_children: 0, + chunks: Vec::new(), + output: None, + done: false, + } + } + + fn initialize(&mut self, cx: &mut Cx<'_>) -> VortexResult<()> { + let offsets = join_inlet(cx, 0, self.plan.offsets()?.dtype())? + .cast(DType::Primitive(PType::U64, false.into()))? + .execute::(cx.exec())?; + let validity = if self.plan.dtype().is_nullable() { + Validity::Array(join_inlet(cx, 1, &DType::Bool(false.into()))?) + } else { + Validity::NonNullable + }; + let valid = validity.execute_mask(self.mask.true_count(), cx.exec())?; + let elements = self.plan.elements()?; + let offsets = offsets.as_slice::(); + let (mut selected_offsets, mut ranges) = match self.mask.indices() { + AllOr::All => { + selected_ranges(0..self.mask.len(), offsets, &valid, elements.row_count())? + } + AllOr::Some(indices) => selected_ranges( + indices.iter().copied(), + offsets, + &valid, + elements.row_count(), + )?, + AllOr::None => (Buffer::from(vec![0_u64]), Vec::new()), + }; + // An all-selected scope can retain payload under null rows. For dense, bounded spans, + // this avoids filtering every element just to remove values hidden by list validity. + if self.mask.all_true() + && !valid.all_true() + && let Some(offsets) = work::preserve_null_span( + elements.dtype(), + offsets, + selected_offsets.last().copied().unwrap_or(0), + &mut ranges, + ) + { + selected_offsets = offsets; + } + let offsets = selected_offsets; + let last_offset = offsets.last().copied().unwrap_or(0); + let offsets = offsets.into_array(); + // These offsets start at zero and are checked sums of nonnegative lengths. Retain + // those facts so ListArray validation need not compute aggregates over them again. + let stats = offsets.statistics(); + stats.set(Stat::IsSorted, Precision::Exact(true.into())); + stats.set(Stat::Min, Precision::Exact(0_u64.into())); + stats.set(Stat::Max, Precision::Exact(last_offset.into())); + self.offsets = Some(offsets); + self.validity = Some(validity); + let overlap = work::can_overlap( + &elements, + &ranges, + self.mask.true_count(), + get_available_parallelism().unwrap_or(1), + )?; + let splits = work::element_splits(elements.dtype(), ranges)?; + // Avoid handing tiny or variable-width work to the outer planner. The planner + // decides whether the substantial fixed-width ranges can be split safely. + let width = match elements.dtype() { + DType::Primitive(ptype, _) => ptype.byte_width() as u64, + DType::Bool(_) => 1, + _ => 0, + }; + let bytes = splits.iter().fold(0_u64, |bytes, split| { + bytes.saturating_add((split.mask.true_count() as u64).saturating_mul(width)) + }); + if self.plan_children && bytes >= 512 * 1024 { + self.serial_splits = Some((splits.clone(), overlap)); + self.awaiting_child = true; + cx.request_child(ChildProjection { + id: 0, + plan: elements, + splits, + values: self.values.clone(), + }); + self.initialized = true; + return Ok(()); + } + self.start_elements(cx, splits, overlap)?; + self.initialized = true; + Ok(()) + } + + fn start_elements( + &mut self, + cx: &Cx<'_>, + splits: Vec, + overlap: bool, + ) -> VortexResult<()> { + let scan = Scan::try_new(cx.session().clone(), self.plan.elements()?, splits)? + .with_value_cache(self.values.clone()) + .with_child_planning(self.plan_children); + self.elements = Some(if overlap { + scan.with_max_active(2) + } else { + scan + }); + Ok(()) + } + + fn finish(&mut self) -> VortexResult<()> { + // Independent element ranges may finish out of order. Preserve batch order within + // each range while putting the ranges back in list order. + self.chunks.sort_by_key(|(split, _)| *split); + let elements = join( + self.plan.elements()?.dtype(), + std::mem::take(&mut self.chunks) + .into_iter() + .map(|(_, array)| array) + .collect(), + )?; + let offsets = self + .offsets + .take() + .ok_or_else(|| vortex_err!("Missing list offsets"))?; + let validity = self + .validity + .take() + .ok_or_else(|| vortex_err!("Missing list validity"))?; + self.output = Some(ListArray::try_new(elements, offsets, validity)?.into_array()); + self.elements = None; + Ok(()) + } +} + +impl Source for ListSource { + fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult { + if self.done || self.output.is_some() { + return Ok(SourceReady::Ready); + } + if !self.initialized { + for index in 0..1 + usize::from(self.plan.dtype().is_nullable()) { + if !cx.inlet(index).closed() { + return Ok(SourceReady::Blocked(SourceBlocked::Inlet(index))); + } + } + self.initialize(cx)?; + } + if self.awaiting_child { + let Some((_, arrays)) = cx.take_child() else { + return Ok(SourceReady::Blocked(SourceBlocked::Child)); + }; + self.awaiting_child = false; + let (splits, overlap) = self + .serial_splits + .take() + .ok_or_else(|| vortex_err!("Missing deferred list ranges"))?; + if let Some(arrays) = arrays { + self.chunks.extend(arrays.into_iter().enumerate()); + self.finish()?; + return Ok(SourceReady::Ready); + } + self.start_elements(cx, splits, overlap)?; + } + let scan = self + .elements + .as_mut() + .ok_or_else(|| vortex_err!("Missing list element scan"))?; + while let Some((segment, bytes)) = cx.take_segment() { + let request = self + .reads + .get_mut(&segment) + .and_then(VecDeque::pop_front) + .ok_or_else(|| vortex_err!("Unexpected list element segment {segment}"))?; + scan.deliver(request, bytes)?; + } + while let Some((id, arrays)) = cx.take_child() { + scan.complete_child(id, arrays)?; + self.pending_children -= 1; + } + loop { + match scan.step()? { + Turn::Read(read) => { + self.reads + .entry(read.segment_id) + .or_default() + .push_back(read.id); + cx.request_read(read.segment_id); + } + Turn::Output(split, array) => self.chunks.push((split, array)), + Turn::Child(child) => { + self.pending_children += 1; + cx.request_child(child); + } + Turn::Waiting => { + return Ok(SourceReady::Blocked(if self.pending_children > 0 { + SourceBlocked::Child + } else { + SourceBlocked::Io + })); + } + Turn::Done => { + self.finish()?; + return Ok(SourceReady::Ready); + } + } + } + } +} + +impl Operator for ListSource { + fn compute(&mut self, _input: Input, _cx: &mut Cx<'_>) -> VortexResult { + if self.done { + return Ok(Step::End); + } + self.done = true; + self.output + .take() + .map(Step::Final) + .ok_or_else(|| vortex_err!("List computed before its elements were ready")) + } +} + +/// Compacts selected offsets and merges adjacent element ranges. No mask spans the gaps between +/// distant lists, so planning memory scales with selected rows rather than their element span. +fn selected_ranges( + selected: impl Iterator, + offsets: &[u64], + valid: &Mask, + element_count: u64, +) -> VortexResult<(Buffer, Vec>)> { + let mut output_offsets = Vec::with_capacity(valid.len() + 1); + output_offsets.push(0_u64); + let mut ranges: Vec> = Vec::new(); + let mut previous = None; + let mut cursor = 0; + let mut total = 0_u64; + let mut last_end = 0; + for (index, is_valid) in selected.zip(valid.iter()) { + // Consecutive selected lists share one boundary in the offsets selection. + if previous == Some(index) { + cursor -= 1; + } + let pair = offsets + .get(cursor..cursor + 2) + .ok_or_else(|| vortex_err!("Missing selected list boundaries"))?; + let start = pair[0]; + let end = pair[1]; + vortex_ensure!( + last_end <= start && start <= end && end <= element_count, + "Invalid list element range {start}..{end}" + ); + last_end = end; + if is_valid { + total = total + .checked_add(end - start) + .ok_or_else(|| vortex_err!("Selected list elements overflow"))?; + if start < end { + match ranges.last_mut() { + Some(last) if last.end == start => last.end = end, + _ => ranges.push(start..end), + } + } + } + output_offsets.push(total); + previous = Some(index + 1); + cursor += 2; + } + vortex_ensure!( + cursor == offsets.len(), + "Unexpected selected list boundaries" + ); + Ok((Buffer::from(output_offsets), ranges)) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::adjacent(vec![0, 1], vec![0, 2, 5], vec![true, true], vec![0, 2, 5], vec![0..5])] + #[case::gaps(vec![0, 3], vec![2, 4, 100, 103], vec![true, true], vec![0, 2, 5], vec![2..4, 100..103])] + #[case::null(vec![0, 1, 4], vec![0, 3, 8, 100, 105], vec![true, false, true], vec![0, 3, 3, 8], vec![0..3, 100..105])] + #[case::empty_lists(vec![0, 1], vec![2, 2, 2], vec![true, true], vec![0, 0, 0], vec![])] + fn compacts_selected_elements( + #[case] selected: Vec, + #[case] boundaries: Vec, + #[case] valid: Vec, + #[case] offsets: Vec, + #[case] ranges: Vec>, + ) -> VortexResult<()> { + let (actual_offsets, actual_ranges) = selected_ranges( + selected.into_iter(), + &boundaries, + &Mask::from_iter(valid), + 1000, + )?; + assert_eq!(actual_offsets.as_slice(), offsets); + assert_eq!(actual_ranges, ranges); + Ok(()) + } + + #[rstest] + #[case::missing(vec![0, 1])] + #[case::negative_length(vec![0, 10, 9])] + #[case::out_of_bounds(vec![0, 9, 11])] + fn rejects_invalid_boundaries(#[case] boundaries: Vec) { + assert!(selected_ranges(0..2, &boundaries, &Mask::new_true(2), 10).is_err()); + } +} diff --git a/vortex-layout/src/plan/physical/operators/list/work.rs b/vortex-layout/src/plan/physical/operators/list/work.rs new file mode 100644 index 00000000000..e93595683d6 --- /dev/null +++ b/vortex-layout/src/plan/physical/operators/list/work.rs @@ -0,0 +1,411 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Bounded grouping and admission of list element work. + +use std::ops::Range; + +use vortex_array::dtype::DType; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use crate::plan::physical::Concat; +use crate::plan::physical::Filter; +use crate::plan::physical::LegacyStats; +use crate::plan::physical::PlanRef; +use crate::plan::physical::SegmentScan; +use crate::plan::physical::Zoned; +use crate::plan::physical::pipeline::Split; +use crate::plan::physical::pipeline::overlapping; + +const MAX_ELEMENT_BYTES: u64 = 8 * 1024 * 1024; +const MAX_MASK_FRAGMENT_BYTES: u64 = 32 * 1024; +/// Retain a dense, fixed-width element span across null lists. The caller has validated every +/// offset and selected every row in this scope; only null rows may occupy gaps in `ranges`. +pub(super) fn preserve_null_span( + dtype: &DType, + offsets: &[u64], + selected: u64, + ranges: &mut Vec>, +) -> Option> { + if ranges.len() < 2 { + return None; + } + let start = ranges.first()?.start; + let end = ranges.last()?.end; + let span = end - start; + let width = match dtype { + DType::Primitive(ptype, _) => ptype.byte_width() as u64, + DType::Bool(_) => 1, + _ => return None, + }; + let validity = if dtype.is_nullable() { + span.div_ceil(8) + } else { + 0 + }; + if selected < span.div_ceil(2) + || span.saturating_mul(width).saturating_add(validity) > MAX_ELEMENT_BYTES + { + return None; + } + // Trim payload under leading/trailing nulls while retaining internal gaps. + let offsets = offsets + .iter() + .map(|offset| (*offset).clamp(start, end) - start) + .collect(); + ranges.clear(); + ranges.push(start..end); + Some(offsets) +} + +/// Combine nearby ranges into a mask, keeping holes out of the output without a scan per hole. +/// Only dense fixed-width spans of at most 8 MiB are grouped. Large contiguous fragments keep +/// their slice path: applying a mask costs more than the scan setup it saves for those ranges. +pub(super) fn element_splits(dtype: &DType, ranges: Vec>) -> VortexResult> { + if ranges.len() < 2 { + return Ok(ranges.into_iter().map(Split::all).collect()); + } + let width = match dtype { + DType::Primitive(ptype, _) => ptype.byte_width() as u64, + DType::Bool(_) => 1, + _ => return Ok(ranges.into_iter().map(Split::all).collect()), + }; + let mut splits = Vec::new(); + let mut first = 0; + while first < ranges.len() { + let start = ranges[first].start; + let mut selected = ranges[first].end - start; + let mut end = first + 1; + while let Some(next) = ranges.get(end) { + let span = next.end - start; + let validity = if dtype.is_nullable() { + span.div_ceil(8) + } else { + 0 + }; + let bytes = span.saturating_mul(width).saturating_add(validity); + let total = selected + (next.end - next.start); + if bytes > MAX_ELEMENT_BYTES || total < span.div_ceil(2) { + break; + } + selected = total; + end += 1; + } + if end == first + 1 + || selected + .saturating_mul(width) + .div_ceil((end - first) as u64) + > MAX_MASK_FRAGMENT_BYTES + { + splits.extend(ranges[first..end].iter().cloned().map(Split::all)); + } else { + let rows = start..ranges[end - 1].end; + let slices = ranges[first..end] + .iter() + .map(|range| { + Ok(( + usize::try_from(range.start - start)?, + usize::try_from(range.end - start)?, + )) + }) + .collect::>>()?; + let mask = Mask::from_slices(usize::try_from(rows.end - start)?, slices); + splits.push(Split { rows, mask }); + } + first = end; + } + Ok(splits) +} + +/// Overlap grouped rows only when all touched chunks fit the decoded-element estimate. +/// This limits additional range concurrency, not total heap use or coalesced IO buffers. +pub(super) fn can_overlap( + elements: &PlanRef, + ranges: &[Range], + selected: usize, + parallelism: usize, +) -> VortexResult { + if parallelism <= 1 || ranges.len() <= 1 || selected / ranges.len() < 2 { + return Ok(false); + } + let mut bytes = 0_u64; + for range in ranges { + let Some(estimate) = range_bytes(elements, range)? else { + return Ok(false); + }; + // Charging a shared chunk again is conservative and avoids a per-scan set. + bytes = bytes.saturating_add(estimate); + if bytes > MAX_ELEMENT_BYTES { + return Ok(false); + } + } + Ok(true) +} + +fn range_bytes(plan: &PlanRef, rows: &Range) -> VortexResult> { + if let Some(filter) = plan.as_opt::() { + return range_bytes(&filter.child_plan()?, rows); + } + if let Some(zoned) = plan.as_opt::() { + return match zoned.data_plan()? { + Some(data) => range_bytes(&data, rows), + None => Ok(None), + }; + } + if let Some(stats) = plan.as_opt::() { + return range_bytes(&stats.data_plan()?, rows); + } + if let Some(concat) = plan.as_opt::() { + let mut bytes = 0_u64; + for (index, start, local) in overlapping(concat, rows) { + let Some(estimate) = range_bytes( + &concat.child_required(index)?, + &(local.start - start..local.end - start), + )? + else { + return Ok(None); + }; + bytes = bytes.saturating_add(estimate); + if bytes > MAX_ELEMENT_BYTES { + break; + } + } + return Ok(Some(bytes)); + } + if !plan.is::() { + return Ok(None); + } + let width = match plan.dtype() { + DType::Primitive(ptype, _) => ptype.byte_width() as u64, + DType::Bool(_) => 1, + _ => return Ok(None), + }; + let validity = if plan.dtype().is_nullable() { + plan.row_count().div_ceil(8) + } else { + 0 + }; + Ok(Some( + plan.row_count() + .saturating_mul(width) + .saturating_add(validity), + )) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::dtype::PType; + use vortex_mask::AllOr; + use vortex_session::registry::ReadContext; + + use super::*; + use crate::plan::physical::ConcatPlan; + use crate::plan::physical::LegacyStatsPlan; + use crate::plan::physical::SegmentScanPlan; + use crate::segments::SegmentId; + + fn segment(rows: u64, dtype: DType) -> PlanRef { + SegmentScanPlan::new(dtype, rows, SegmentId::from(0), ReadContext::new([]), None) + .into_plan() + } + + #[test] + fn retained_span_trims_null_edges_and_rebases_internal_offsets() { + let mut ranges = vec![10..18, 22..30]; + let offsets = preserve_null_span( + &DType::Primitive(PType::I64, true.into()), + &[0, 10, 18, 22, 30, 50], + 16, + &mut ranges, + ) + .unwrap(); + assert_eq!(offsets.as_slice(), &[0, 0, 8, 12, 20, 20]); + assert_eq!(ranges, vec![10..30]); + } + + #[rstest] + #[case::at_limit(MAX_ELEMENT_BYTES / 8, false, true)] + #[case::over_limit(MAX_ELEMENT_BYTES / 8 + 2, false, false)] + #[case::validity_counts(MAX_ELEMENT_BYTES / 8, true, false)] + fn retained_span_includes_validity_in_byte_estimate( + #[case] rows: u64, + #[case] nullable: bool, + #[case] retained: bool, + ) { + let original = vec![0..rows / 2, rows / 2 + 2..rows]; + let mut ranges = original.clone(); + let actual = preserve_null_span( + &DType::Primitive(PType::I64, nullable.into()), + &[0, rows / 2, rows / 2 + 2, rows], + rows - 2, + &mut ranges, + ); + assert_eq!(actual.is_some(), retained); + if !retained { + assert_eq!(ranges, original); + } + } + + #[rstest] + #[case::sparse(DType::Primitive(PType::I64, false.into()), vec![0..4, 100..104])] + #[case::variable_width(DType::Utf8(true.into()), vec![0..4, 6..10])] + #[case::one_range(DType::Bool(false.into()), vec![0..4])] + #[case::empty(DType::Bool(false.into()), vec![])] + fn unsuitable_null_spans_stay_compact(#[case] dtype: DType, #[case] ranges: Vec>) { + let mut actual = ranges.clone(); + let selected = ranges.iter().map(|range| range.end - range.start).sum(); + assert!(preserve_null_span(&dtype, &[], selected, &mut actual).is_none()); + assert_eq!(actual, ranges); + } + + #[rstest] + fn dense_ranges_use_relative_masks(#[values(0, u64::MAX - 20)] base: u64) -> VortexResult<()> { + let splits = element_splits( + &DType::Primitive(PType::I64, true.into()), + vec![base + 1..base + 5, base + 7..base + 9, base + 12..base + 13], + )?; + assert_eq!(splits.len(), 1); + assert_eq!(splits[0].rows, base + 1..base + 13); + assert_eq!( + splits[0].mask.slices(), + AllOr::Some(&[(0, 4), (6, 8), (11, 12)][..]) + ); + Ok(()) + } + + #[test] + fn sparse_gaps_and_unknown_widths_stay_separate() -> VortexResult<()> { + let ranges = vec![0..2, 100..104, 105..109]; + let splits = element_splits(&DType::Bool(false.into()), ranges.clone())?; + assert_eq!(splits.len(), 2); + assert_eq!(splits[0].rows, 0..2); + assert_eq!(splits[1].rows, 100..109); + assert_eq!(splits[1].mask.true_count(), 8); + let unknown = element_splits(&DType::Utf8(false.into()), ranges.clone())?; + assert_eq!( + unknown + .iter() + .map(|split| split.rows.clone()) + .collect::>(), + ranges + ); + Ok(()) + } + + #[rstest] + #[case::at_limit(MAX_ELEMENT_BYTES / 8, false, 1)] + #[case::over_limit(MAX_ELEMENT_BYTES / 8 + 2, false, 2)] + #[case::validity_counts(MAX_ELEMENT_BYTES / 8, true, 2)] + fn grouping_bounds_the_mask_span( + #[case] rows: u64, + #[case] nullable: bool, + #[case] count: usize, + ) -> VortexResult<()> { + let ranges = (0..rows) + .step_by(1024) + .map(|start| start + u64::from(start != 0)..(start + 1024).min(rows)) + .collect::>(); + let selected = rows - (ranges.len() as u64 - 1); + let splits = element_splits(&DType::Primitive(PType::I64, nullable.into()), ranges)?; + assert_eq!(splits.len(), count); + assert_eq!( + splits + .iter() + .map(|split| split.mask.true_count()) + .sum::(), + usize::try_from(selected)? + ); + Ok(()) + } + + #[test] + fn large_fragments_keep_contiguous_slices() -> VortexResult<()> { + let ranges = vec![0..8192, 8193..16384]; + let splits = element_splits(&DType::Primitive(PType::I64, false.into()), ranges.clone())?; + assert_eq!( + splits + .iter() + .map(|split| split.rows.clone()) + .collect::>(), + ranges + ); + assert!( + splits + .iter() + .all(|split| matches!(split.mask, Mask::AllTrue(_))) + ); + Ok(()) + } + + #[test] + #[allow(clippy::single_range_in_vec_init)] + fn contiguous_ranges_need_no_mask_allocation() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I64, false.into()); + assert!(element_splits(&dtype, vec![])?.is_empty()); + let splits = element_splits(&dtype, vec![0..MAX_ELEMENT_BYTES])?; + assert!(matches!(splits[0].mask, Mask::AllTrue(_))); + Ok(()) + } + + #[rstest] + #[case::grouped(4, 30, true)] + #[case::isolated(2, 30, false)] + #[case::single_worker(4, 1, false)] + fn overlap_requires_grouped_parallel_work( + #[case] selected: usize, + #[case] workers: usize, + #[case] expected: bool, + ) -> VortexResult<()> { + let elements = segment(1024, DType::Primitive(PType::I64, false.into())); + assert_eq!( + can_overlap(&elements, &[0..2, 10..12], selected, workers)?, + expected + ); + Ok(()) + } + + #[rstest] + #[case::within_limit(MAX_ELEMENT_BYTES / 16, false, true)] + #[case::large_chunk(MAX_ELEMENT_BYTES / 8, false, false)] + #[case::validity_counts(MAX_ELEMENT_BYTES / 16, true, false)] + fn charges_whole_chunks_and_validity( + #[case] rows: u64, + #[case] nullable: bool, + #[case] expected: bool, + ) -> VortexResult<()> { + let elements = segment(rows, DType::Primitive(PType::I64, nullable.into())); + assert_eq!(can_overlap(&elements, &[0..2, 10..12], 4, 30)?, expected); + Ok(()) + } + + #[test] + fn skips_untouched_chunks_and_zone_data() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I64, false.into()); + let data = ConcatPlan::try_new( + dtype.clone(), + vec![ + segment(1024, dtype.clone()), + segment(10_000_000, dtype.clone()), + ], + )? + .into_plan(); + let elements = + LegacyStatsPlan::new(data, segment(10_000_000, dtype), 1024, [].into()).into_plan(); + assert!(can_overlap(&elements, &[0..2, 10..12], 4, 30)?); + assert!(!can_overlap(&elements, &[0..2, 1024..1026], 4, 30)?); + Ok(()) + } + + #[test] + fn unknown_width_and_overflow_stay_sequential() -> VortexResult<()> { + let unknown = segment(10, DType::Utf8(false.into())); + let overflow = segment(u64::MAX, DType::Primitive(PType::I64, true.into())); + for plan in [unknown, overflow] { + assert!(!can_overlap(&plan, &[0..2, 5..7], 4, 30)?); + } + Ok(()) + } +} diff --git a/vortex-layout/src/plan/physical/operators/mod.rs b/vortex-layout/src/plan/physical/operators/mod.rs index b0f5937d71d..c4545ec5e19 100644 --- a/vortex-layout/src/plan/physical/operators/mod.rs +++ b/vortex-layout/src/plan/physical/operators/mod.rs @@ -9,6 +9,7 @@ mod concat; mod eval; mod filter; +mod list; mod once; mod pack; mod scan; @@ -18,6 +19,7 @@ mod whole; pub(crate) use concat::*; pub(crate) use eval::*; pub(crate) use filter::*; +pub(crate) use list::*; pub(crate) use once::*; pub(crate) use pack::*; pub(crate) use scan::*; diff --git a/vortex-layout/src/plan/physical/operators/pack.rs b/vortex-layout/src/plan/physical/operators/pack.rs index 4fefa368f05..7a4e83efc40 100644 --- a/vortex-layout/src/plan/physical/operators/pack.rs +++ b/vortex-layout/src/plan/physical/operators/pack.rs @@ -117,7 +117,7 @@ impl Operator for PackSource { } else { Validity::NonNullable }; - let array = assemble(&self.fields, arrays, validity).into_array(); + let array = assemble(&self.fields, arrays, len, validity).into_array(); Ok(Step::Final(array)) } } @@ -138,8 +138,9 @@ impl Operator for WrapStage { let Input::Batch(batch) = input else { return Ok(Step::End); }; + let len = batch.len(); Ok(Step::Final( - assemble(&self.fields, vec![batch], Validity::NonNullable).into_array(), + assemble(&self.fields, vec![batch], len, Validity::NonNullable).into_array(), )) } } @@ -149,8 +150,12 @@ impl Operator for WrapStage { /// The type is not derived or validated per batch: a `Pack` plan checks when it is built that /// each field's plan produces the field's dtype over the struct's rows, and a pack source takes /// the same rows from every field. -fn assemble(fields: &StructFields, arrays: Vec, validity: Validity) -> StructArray { - let len = arrays.first().map_or(0, |array| array.len()); +fn assemble( + fields: &StructFields, + arrays: Vec, + len: usize, + validity: Validity, +) -> StructArray { debug_assert!( arrays.iter().all(|array| array.len() == len) && fields diff --git a/vortex-layout/src/plan/physical/pipeline/children.rs b/vortex-layout/src/plan/physical/pipeline/children.rs new file mode 100644 index 00000000000..974e0afcd74 --- /dev/null +++ b/vortex-layout/src/plan/physical/pipeline/children.rs @@ -0,0 +1,29 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Descriptions of child projections for the scan planner to schedule. + +use std::fmt; + +use super::Split; +use super::ValueCache; +use crate::plan::physical::PlanRef; + +/// A source's child projection, handed to its enclosing planner for execution. +/// Only scans configured by that planner produce these requests. +pub struct ChildProjection { + pub(crate) id: u64, + pub(crate) plan: PlanRef, + pub(crate) splits: Vec, + pub(crate) values: ValueCache, +} + +impl fmt::Debug for ChildProjection { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ChildProjection") + .field("id", &self.id) + .field("plan", &self.plan) + .field("splits", &self.splits) + .finish_non_exhaustive() + } +} diff --git a/vortex-layout/src/plan/physical/pipeline/compile.rs b/vortex-layout/src/plan/physical/pipeline/compile.rs index 74f5296cd19..cf2dcb10184 100644 --- a/vortex-layout/src/plan/physical/pipeline/compile.rs +++ b/vortex-layout/src/plan/physical/pipeline/compile.rs @@ -8,6 +8,7 @@ use std::ops::Range; use rustc_hash::FxHashMap; use smallvec::SmallVec; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -143,6 +144,10 @@ impl Compiler<'_> { &self.core.values } + pub(crate) fn plan_children(&self) -> bool { + self.core.plan_children + } + /// The rows `rows` of `scan`'s segment, filtered by `filter`, read by the split compiling. /// A segment several readers read is decoded once and shared. pub(crate) fn scan( @@ -296,6 +301,30 @@ impl Shares { }) } + /// Counts a single split's readers without lowering chunks its mask excludes. + pub(crate) fn add_selected( + &mut self, + plan: &PlanRef, + rows: Range, + mask: &Mask, + ) -> VortexResult<()> { + if mask.all_false() { + return Ok(()); + } + if mask.all_true() { + return self.add(plan, rows); + } + let at = Reach::Selected { + offset: 0, + rows: rows.clone(), + mask: mask.clone(), + }; + let reaches = &mut self.reaches; + plan.reach(rows, &at, &mut |segment, rows| { + reaches.push((segment, rows)) + }) + } + /// Keeps only the segments more than one reader reads. The rest are read by the one /// pipeline that needs them. pub(crate) fn retain_shared(&mut self) { @@ -350,13 +379,24 @@ pub enum Reach { Offset(u64), /// Every row is read for these scanned rows, as a dictionary's values are for its codes. Fixed(Range), + /// A linear mapping restricted to selected rows in the scan's root domain. + Selected { + /// Root row index of this plan's row zero. + offset: u64, + /// The root row range the mask covers. + rows: Range, + /// Selected positions relative to `rows.start`. + mask: Mask, + }, } impl Reach { /// The scanned rows `rows` are read for. pub fn root(&self, rows: &Range) -> Range { match self { - Reach::Offset(offset) => rows.start + offset..rows.end + offset, + Reach::Offset(offset) | Reach::Selected { offset, .. } => { + rows.start + offset..rows.end + offset + } Reach::Fixed(root) => root.clone(), } } @@ -366,6 +406,11 @@ impl Reach { match self { Reach::Offset(offset) => Reach::Offset(offset + by), Reach::Fixed(root) => Reach::Fixed(root.clone()), + Reach::Selected { offset, rows, mask } => Reach::Selected { + offset: offset + by, + rows: rows.clone(), + mask: mask.clone(), + }, } } @@ -373,6 +418,30 @@ impl Reach { pub fn fixed(&self, rows: &Range) -> Self { Reach::Fixed(self.root(rows)) } + + /// Whether the selection needs any row of this child, before lowering it. + pub(crate) fn selects(&self, local: &Range) -> bool { + let Reach::Selected { offset, rows, mask } = self else { + return true; + }; + let start = (local.start + offset).max(rows.start); + let end = (local.end + offset).min(rows.end); + if start >= end { + return false; + } + let start = usize::try_from(start - rows.start).vortex_expect("offset fits selection mask"); + let end = usize::try_from(end - rows.start).vortex_expect("offset fits selection mask"); + match mask { + Mask::AllTrue(_) => true, + Mask::AllFalse(_) => false, + Mask::Values(values) => match values.cached_indices() { + Some(indices) => indices + .get(indices.partition_point(|&index| index < start)) + .is_some_and(|&index| index < end), + None => values.bit_buffer().count_range(start, end) > 0, + }, + } + } } /// The chunks of `concat` overlapping `rows`: each chunk's index, where it starts, and the rows diff --git a/vortex-layout/src/plan/physical/pipeline/mod.rs b/vortex-layout/src/plan/physical/pipeline/mod.rs index b6037c9e6ea..5dd18e783c1 100644 --- a/vortex-layout/src/plan/physical/pipeline/mod.rs +++ b/vortex-layout/src/plan/physical/pipeline/mod.rs @@ -38,12 +38,14 @@ //! read, and a claimed port is dropped once its source has drained it, so a decoded segment //! lives only as long as a split that may still read it. +mod children; mod compile; mod port; mod scan; pub mod synthetic; mod values; +pub use self::children::ChildProjection; pub use self::compile::Chain; pub use self::compile::Compiler; pub use self::compile::Reach; diff --git a/vortex-layout/src/plan/physical/pipeline/scan.rs b/vortex-layout/src/plan/physical/pipeline/scan.rs index 9d650ebbd26..e5a8ee922e9 100644 --- a/vortex-layout/src/plan/physical/pipeline/scan.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan.rs @@ -31,6 +31,7 @@ use super::port::Slab; use super::port::SourceEnd; use super::values::ValueCache; use crate::plan::physical::PlanRef; +use crate::plan::physical::operator::ChildInbox; use crate::plan::physical::operator::Cx; use crate::plan::physical::operator::InletSet; use crate::plan::physical::operator::Input; @@ -39,6 +40,7 @@ use crate::plan::physical::operator::Source; use crate::plan::physical::operator::SourceBlocked; use crate::plan::physical::operator::SourceReady; use crate::plan::physical::operator::Step; +use crate::plan::physical::pipeline::ChildProjection; use crate::segments::SegmentId; #[cfg(debug_assertions)] @@ -82,9 +84,11 @@ impl Split { pub enum Turn { /// Read this segment and [`deliver`](Scan::deliver) its bytes. Read(ReadRequest), + /// Plan and execute this child projection, then return its arrays to the scan. + Child(ChildProjection), /// The next rows of split `split`, in row order within the split. Output(usize, ArrayRef), - /// Nothing can run until a read is delivered. + /// Nothing can run until a read or child projection completes. Waiting, /// Every split has produced all its rows. Done, @@ -111,6 +115,7 @@ pub(crate) struct Pipeline { /// Reads it asked for that have not arrived, so retiring it cancels them only when any /// are left. pending: usize, + child_inbox: Option>, /// Once its source has finished, the next stage to finish: stages are finished one at a /// time, each when the outlets have room for what it flushes. finishing: Option, @@ -141,6 +146,7 @@ fn drive( session: &VortexSession, exec: &mut ExecutionCtx, reads: &mut SmallVec<[SegmentId; 1]>, + child_requests: &mut Vec, ) -> VortexResult { let Pipeline { source, @@ -153,6 +159,7 @@ fn drive( source_more, touched, changed, + child_inbox, .. } = pipe; let mut cx = Cx { @@ -164,6 +171,8 @@ fn drive( exec, touched, changed, + child_requests, + child_inbox, }; loop { // Closing needs no room, so a pipeline with nothing left to flush closes at once. @@ -191,9 +200,17 @@ fn drive( let ready = source.ready(&mut cx)?; cx.changed.clear(); vortex_ensure!( - cx.reads.is_empty() || ready == SourceReady::Blocked(SourceBlocked::Io), + cx.reads.is_empty() + || matches!( + ready, + SourceReady::Blocked(SourceBlocked::Io | SourceBlocked::Child) + ), "A source that queues reads must block on IO" ); + vortex_ensure!( + cx.child_requests.is_empty() || ready == SourceReady::Blocked(SourceBlocked::Child), + "A source that queues asynchronous work must block on its completion" + ); match ready { SourceReady::Ready => {} SourceReady::Blocked(blocked) => return Ok(Progress::Blocked(blocked)), @@ -281,6 +298,9 @@ pub(crate) struct Core { reads: FxHashMap, SegmentId)>, next_read: u64, new_reads: VecDeque, + child_requests: FxHashMap, + next_child: u64, + new_children: VecDeque, /// Split slots whose stage output has something new, with a flag per slot so a slot is /// queued once. dirty: Vec, @@ -291,6 +311,12 @@ pub(crate) struct Core { pub(crate) split: usize, /// Whole values kept for later scans of the same split. pub(crate) values: ValueCache, + pub(crate) plan_children: bool, +} + +struct PendingChild { + owner: Option, + local_id: u64, } impl Core { @@ -306,11 +332,15 @@ impl Core { reads: FxHashMap::default(), next_read: 0, new_reads: VecDeque::new(), + child_requests: FxHashMap::default(), + next_child: 0, + new_children: VecDeque::new(), dirty: Vec::new(), dirty_flags: Vec::new(), shares: Shares::new(splits), split: usize::MAX, values: ValueCache::default(), + plan_children: false, } } @@ -331,6 +361,7 @@ impl Core { outlets, bytes: VecDeque::new(), pending: 0, + child_inbox: None, finishing: None, touched, changed, @@ -405,16 +436,21 @@ impl Core { .ok_or_else(|| vortex_err!("Pipeline {id} is not live"))?; pipe.queued = false; let mut reads = SmallVec::new(); + let mut child_requests = Vec::new(); let progress = drive( pipe, &mut self.arena, &self.session, &mut self.exec, &mut reads, + &mut child_requests, )?; for segment_id in reads { self.ask(id, segment_id); } + for child in child_requests { + self.ask_child(id, child)?; + } #[cfg(debug_assertions)] self.validate_progress(id, &progress); match &progress { @@ -422,6 +458,16 @@ impl Core { Progress::Blocked(SourceBlocked::Io) if !self.pipeline(id).bytes.is_empty() => { self.make_runnable(id); } + Progress::Blocked(SourceBlocked::Child) + if !self.pipeline(id).bytes.is_empty() + || self + .pipeline(id) + .child_inbox + .as_ref() + .is_some_and(|inbox| !inbox.completed.is_empty()) => + { + self.make_runnable(id); + } // Nothing changes a port while its source's pipeline runs, so the condition the run reported // holds until a neighbour wakes it. Progress::Blocked(blocked) => self.pipeline_mut(id).state = State::Blocked(*blocked), @@ -441,6 +487,52 @@ impl Core { self.pipeline_mut(id).pending += 1; } + fn ask_child(&mut self, owner: PipelineId, child: ChildProjection) -> VortexResult<()> { + vortex_ensure!( + !self + .child_requests + .values() + .any(|pending| pending.owner == Some(owner) && pending.local_id == child.id), + "A source published child {} twice", + child.id + ); + let id = self.next_child; + self.next_child = id + .checked_add(1) + .ok_or_else(|| vortex_err!("Scan child identifiers exhausted"))?; + self.child_requests.insert( + id, + PendingChild { + owner: Some(owner), + local_id: child.id, + }, + ); + self.pipeline_mut(owner) + .child_inbox + .get_or_insert_with(Default::default) + .pending += 1; + self.new_children.push_back(ChildProjection { id, ..child }); + Ok(()) + } + + fn complete_child(&mut self, id: u64, arrays: Option>) -> VortexResult<()> { + let child = self + .child_requests + .remove(&id) + .ok_or_else(|| vortex_err!("Unknown scan child {id}"))?; + if let Some(owner) = child.owner { + let inbox = self + .pipeline_mut(owner) + .child_inbox + .as_mut() + .ok_or_else(|| vortex_err!("Child completed for a source with no child inbox"))?; + inbox.pending -= 1; + inbox.completed.push_back((child.local_id, arrays)); + self.make_runnable(owner); + } + Ok(()) + } + /// Wakes the sources of what pipeline `id` pushed and the sinks of what it took. fn wake_neighbours(&mut self, id: PipelineId) { for index in 0..self.pipeline(id).outlets.len() { @@ -529,6 +621,13 @@ impl Core { if pipe.pending > 0 { self.cancel_reads(id); } + if pipe.child_inbox.is_some_and(|inbox| inbox.pending > 0) { + for child in self.child_requests.values_mut() { + if child.owner == Some(id) { + child.owner = None; + } + } + } self.free_pipelines.push(id); } } @@ -671,15 +770,40 @@ impl Scan { 0, splits.iter().map(|split| split.rows.clone()).collect(), ); - // Only the rows some split reads can share a segment between readers. - let start = splits - .iter() - .map(|split| split.rows.start) - .min() - .unwrap_or(0); - let end = splits.iter().map(|split| split.rows.end).max().unwrap_or(0); match &root { - Root::Plan(plan) => core.shares.add(plan, start..end)?, + Root::Plan(plan) => { + if let [split] = splits.as_slice() { + core.shares + .add_selected(plan, split.rows.clone(), &split.mask)?; + } else if splits + .windows(2) + .all(|pair| pair[0].rows.end <= pair[1].rows.start) + { + // List selections can leave large gaps between element ranges. Keep those + // chunks lazy, merging adjacent ranges so dense scans still visit once. + let mut rows = 0..0; + for split in &splits { + if split.rows.start > rows.end { + if !rows.is_empty() { + core.shares.add(plan, rows.clone())?; + } + rows.start = split.rows.start; + } + rows.end = split.rows.end; + } + if !rows.is_empty() { + core.shares.add(plan, rows)?; + } + } else { + let start = splits + .iter() + .map(|split| split.rows.start) + .min() + .unwrap_or(0); + let end = splits.iter().map(|split| split.rows.end).max().unwrap_or(0); + core.shares.add(plan, start..end)?; + } + } } core.shares.retain_shared(); core.dirty_flags = vec![false; 1]; @@ -714,6 +838,11 @@ impl Scan { self } + pub(crate) fn with_child_planning(mut self, enabled: bool) -> Self { + self.core.plan_children = enabled; + self + } + /// Runs pipelines until the owner has something to do. pub fn step(&mut self) -> VortexResult { loop { @@ -723,6 +852,18 @@ impl Scan { if let Some(read) = self.core.new_reads.pop_front() { return Ok(Turn::Read(read)); } + while let Some(child) = self.core.new_children.pop_front() { + if self + .core + .child_requests + .get(&child.id) + .is_some_and(|pending| pending.owner.is_none()) + { + self.core.child_requests.remove(&child.id); + continue; + } + return Ok(Turn::Child(child)); + } if let Some(slot) = self.core.dirty.pop() { self.core.dirty_flags[slot] = false; self.advance(slot)?; @@ -741,7 +882,7 @@ impl Scan { if self.live == 0 { return Ok(Turn::Done); } - if !self.core.reads.is_empty() { + if !self.core.reads.is_empty() || !self.core.child_requests.is_empty() { return Ok(Turn::Waiting); } vortex_bail!( @@ -756,6 +897,15 @@ impl Scan { self.core.deliver(read, bytes) } + /// Delivers a child projection to its source. + pub(crate) fn complete_child( + &mut self, + id: u64, + arrays: Option>, + ) -> VortexResult<()> { + self.core.complete_child(id, arrays) + } + /// Starts the next split in a free slot. fn activate(&mut self) -> VortexResult<()> { let slot = self diff --git a/vortex-layout/src/plan/physical/pipeline/scan/validate.rs b/vortex-layout/src/plan/physical/pipeline/scan/validate.rs index df404bd34ea..d1749972d65 100644 --- a/vortex-layout/src/plan/physical/pipeline/scan/validate.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan/validate.rs @@ -88,6 +88,12 @@ impl Core { pipe.pending > 0 || !pipe.bytes.is_empty(), "pipeline {id} waits for a read with none in flight or delivered" ), + Progress::Blocked(SourceBlocked::Child) => assert!( + pipe.child_inbox + .as_ref() + .is_some_and(|inbox| inbox.pending > 0 || !inbox.completed.is_empty()), + "pipeline {id} waits for a child with none in flight or completed" + ), Progress::Done => { for (index, &port) in pipe.inlets.iter().enumerate() { assert!( diff --git a/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs b/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs index 19bcb5504db..e8bd51e1ac9 100644 --- a/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs @@ -64,6 +64,7 @@ fn run( events.push(None); arrays.push(array); } + Turn::Child(_) => unreachable!("child planning is disabled"), Turn::Waiting => { if inflight.is_empty() { return Err(vortex_err!("scan waits with no reads in flight")); @@ -193,6 +194,7 @@ fn splits_produce_their_own_rows() -> VortexResult<()> { match scan.step()? { Turn::Read(read) => inflight.push(read), Turn::Output(split, array) => by_split[split].push(array), + Turn::Child(_) => unreachable!("child planning is disabled"), Turn::Waiting => { let read: ReadRequest = inflight.remove(0); scan.deliver(read.id, empty_segment())?; @@ -258,6 +260,7 @@ fn a_source_has_several_reads_in_flight(#[values(1, 3)] deliveries: usize) -> Vo match scan.step()? { Turn::Read(read) => inflight.push(read), Turn::Output(_, array) => arrays.push(array), + Turn::Child(_) => unreachable!("child planning is disabled"), Turn::Waiting => { if first_wait { assert_eq!( @@ -417,6 +420,7 @@ fn finishing_cancels_reads_in_flight() -> VortexResult<()> { match scan.step()? { Turn::Read(read) => inflight.push(read), Turn::Output(split, array) => by_split[split].push(array), + Turn::Child(_) => unreachable!("child planning is disabled"), Turn::Waiting => { let read = inflight.remove(0); scan.deliver(read.id, empty_segment())?; diff --git a/vortex-layout/src/plan/physical/pipeline/tests.rs b/vortex-layout/src/plan/physical/pipeline/tests.rs index 8fc9c67b900..0327d652b3d 100644 --- a/vortex-layout/src/plan/physical/pipeline/tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/tests.rs @@ -6,6 +6,7 @@ use std::collections::VecDeque; use std::ops::Range; +use std::sync::Arc; use rstest::rstest; use vortex_array::ArrayContext; @@ -14,9 +15,12 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::List; +use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::list::ListArraySlotsExt; use vortex_array::assert_arrays_eq; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::StructFields; @@ -25,6 +29,7 @@ use vortex_array::expr::gt; use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_array::serde::SerializeOptions; +use vortex_array::validity::Validity; use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_buffer::ByteBufferMut; @@ -32,7 +37,9 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_mask::Mask; +use vortex_session::registry::CachedId; use vortex_session::registry::ReadContext; +use vortex_utils::parallelism::get_available_parallelism; use super::port::SourceEnd; use super::synthetic::row_dtype; @@ -42,9 +49,11 @@ use crate::OwnedLayoutChildren; use crate::layouts::chunked::ChunkedLayout; use crate::layouts::dict::DictLayout; use crate::layouts::flat::FlatLayout; +use crate::layouts::foreign::new_foreign_layout; use crate::layouts::struct_::StructLayout; use crate::plan::physical::EvalPlan; use crate::plan::physical::Filter; +use crate::plan::physical::ListPackPlan; use crate::plan::physical::PackPlan; use crate::plan::physical::PlanRef; use crate::plan::physical::SegmentScan; @@ -224,6 +233,7 @@ fn drive( splits[split].push(array.clone()); arrays.push(array); } + Turn::Child(_) => unreachable!("child planning is disabled"), Turn::Waiting => { if !batch.is_empty() { events.push(Event::Io(std::mem::take(&mut batch))); @@ -288,6 +298,101 @@ fn delivery(order: Delivery) -> impl FnMut(&[ReadRequest]) -> usize { } } +#[rstest] +fn grouped_list_ranges_preserve_order_and_shared_reads( + #[values(Delivery::Fifo, Delivery::Lifo)] order: Delivery, + #[values(false, true)] shared: bool, +) -> VortexResult<()> { + let mut store = Store::default(); + let elements = Buffer::from_iter(0_i64..20).into_array(); + let offsets = Buffer::from_iter((0_u64..=10).map(|row| row * 2)).into_array(); + let input = + ListArray::try_new(elements.clone(), offsets.clone(), Validity::NonNullable)?.into_array(); + let layout = if shared { + store.flat(&elements)? + } else { + store.chunked(&elements, &[5, 5, 5, 5])? + }; + let plan = ListPackPlan::try_new( + false.into(), + 10, + lower(&layout)?, + lower(&store.flat(&offsets)?)?, + None, + )? + .into_plan(); + let rows = [0_u64, 1, 7, 8]; + let mask = Mask::from_iter((0..10).map(|row| rows.contains(&row))); + let result = run(&store, &plan, 0..10, mask, delivery(order))?; + let actual = ChunkedArray::try_new(result.arrays, input.dtype().clone())?.into_array(); + let expected = input.take(Buffer::from_iter(rows).into_array())?; + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); + if get_available_parallelism().unwrap_or(1) > 1 { + if shared { + let reads = result + .events + .iter() + .filter(|event| **event == Event::Delivered(0)); + assert_eq!(reads.count(), 1); + } else { + assert!( + result.events.iter().any(|event| { + matches!(event, Event::Io(ids) if ids.contains(&0) && ids.contains(&2)) + }), + "{:?}", + result.events + ); + } + } + Ok(()) +} + +#[rstest] +fn dense_nullable_list_ranges_share_one_scan( + #[values(Delivery::Fifo, Delivery::Lifo)] order: Delivery, + #[values(false, true)] partial: bool, +) -> VortexResult<()> { + let mut store = Store::default(); + let elements = Buffer::from_iter(0_i64..20).into_array(); + let offsets = Buffer::from_iter((0_u64..=10).map(|row| row * 2)).into_array(); + let valid = BoolArray::from_iter((0..10).map(|row| row != 2 && row != 7)).into_array(); + let input = ListArray::try_new( + elements.clone(), + offsets.clone(), + Validity::Array(valid.clone()), + )? + .into_array(); + let plan = ListPackPlan::try_new( + true.into(), + 10, + lower(&store.chunked(&elements, &[5, 5, 5, 5])?)?, + lower(&store.flat(&offsets)?)?, + Some(lower(&store.flat(&valid)?)?), + )? + .into_plan(); + let mask = Mask::from_iter((0..10).map(|row| !partial || row != 4)); + let result = run(&store, &plan, 0..10, mask.clone(), delivery(order))?; + if !partial { + assert_eq!(result.arrays[0].as_::().elements().len(), 20); + } + let actual = ChunkedArray::try_new(result.arrays, input.dtype().clone())?.into_array(); + assert_arrays_eq!( + actual, + input.filter(mask)?, + &mut SESSION.create_execution_ctx() + ); + assert!( + result + .events + .iter() + .any(|event| { matches!(event, Event::Io(ids) if (0..4).all(|id| ids.contains(&id))) }), + "{:?}", + result.events + ); + assert_eq!(segments_read(&result.events), vec![0, 1, 2, 3, 4, 5]); + Ok(()) +} + /// Delivers segments in the given order. fn scripted(order: &[u32]) -> impl FnMut(&[ReadRequest]) -> usize + '_ { let mut next = order.iter(); @@ -432,6 +537,135 @@ fn splits_of_one_scan_read_each_segment_once( Ok(()) } +#[rstest] +#[case::gaps(vec![0..2, 8..10, 18..20])] +#[case::overlapping(vec![0..8, 4..12, 10..20])] +#[case::unordered(vec![15..20, 0..4, 7..12])] +#[case::empty(vec![0..0, 4..8, 20..20])] +fn split_ranges_preserve_shared_reads( + #[case] ranges: Vec>, + #[values(1, 3)] active: usize, + #[values(Delivery::Fifo, Delivery::Lifo)] order: Delivery, +) -> VortexResult<()> { + let mut store = Store::default(); + let (plan, expected) = fixture(&mut store)?; + let splits: Vec<_> = ranges.into_iter().map(Split::all).collect(); + let scan = Scan::try_new(SESSION.clone(), plan, splits.clone())?.with_max_active(active); + let run = drive(&store, scan, delivery(order))?; + assert_eq!(run.leftover, (0, 0)); + let reads = segments_read(&run.events); + let mut unique = reads.clone(); + unique.dedup(); + assert_eq!(reads, unique, "each shared segment is read once"); + for (index, split) in splits.iter().enumerate() { + assert_view( + &expected, + &split.rows, + &split.mask, + run.splits.get(index).cloned().unwrap_or_default(), + )?; + } + Ok(()) +} + +#[rstest] +fn chunk_in_gap_between_splits_stays_lazy(#[values(1, 2)] active: usize) -> VortexResult<()> { + static UNKNOWN_ID: CachedId = CachedId::new("test.unread-gap"); + let mut store = Store::default(); + let values = PrimitiveArray::from_iter(0..ROWS as i32).into_array(); + let known = store.flat(&values)?; + let unknown = new_foreign_layout( + *UNKNOWN_ID, + values.dtype().clone(), + ROWS, + Vec::new(), + Vec::new(), + Vec::new(), + ); + let layout = ChunkedLayout::new( + 3 * ROWS, + values.dtype().clone(), + OwnedLayoutChildren::layout_children(vec![Arc::clone(&known), unknown, known]), + ) + .into_layout(); + let scan = Scan::try_new( + SESSION.clone(), + lower(&layout)?, + vec![Split::all(0..2), Split::all(2 * ROWS..2 * ROWS + 2)], + )? + .with_max_active(active); + let run = drive(&store, scan, delivery(Delivery::Lifo))?; + assert_eq!(segments_read(&run.events), vec![0]); + assert_eq!(run.leftover, (0, 0)); + assert_eq!(run.splits.len(), 2); + for pieces in run.splits { + assert_view(&values, &(0..2), &Mask::new_true(2), pieces)?; + } + Ok(()) +} + +#[rstest] +fn chunk_excluded_by_mask_stays_lazy( + #[values(false, true)] nested: bool, + #[values(false, true)] empty: bool, + #[values(false, true)] cached_indices: bool, +) -> VortexResult<()> { + static UNKNOWN_ID: CachedId = CachedId::new("test.masked-gap"); + let mut store = Store::default(); + let values = PrimitiveArray::from_iter(0..ROWS as i32).into_array(); + let known = store.flat(&values)?; + let unknown = new_foreign_layout( + *UNKNOWN_ID, + values.dtype().clone(), + ROWS, + Vec::new(), + Vec::new(), + Vec::new(), + ); + let mut layout = ChunkedLayout::new( + 3 * ROWS, + values.dtype().clone(), + OwnedLayoutChildren::layout_children(vec![Arc::clone(&known), unknown, Arc::clone(&known)]), + ) + .into_layout(); + let offset = if nested { ROWS } else { 0 }; + if nested { + layout = ChunkedLayout::new( + 4 * ROWS, + values.dtype().clone(), + OwnedLayoutChildren::layout_children(vec![known, layout]), + ) + .into_layout(); + } + let rows = offset + 5..offset + 3 * ROWS - 5; + let len = usize::try_from(rows.end - rows.start)?; + let mut mask = if empty { + Mask::new_false(len) + } else { + Mask::from_indices(len, [0, len - 1]) + }; + if !cached_indices { + mask = Mask::from_buffer(mask.to_bit_buffer()); + } + let scan = Scan::try_new(SESSION.clone(), lower(&layout)?, vec![Split { rows, mask }])?; + let run = drive(&store, scan, delivery(Delivery::Lifo))?; + assert_eq!(run.leftover, (0, 0)); + if empty { + assert!(segments_read(&run.events).is_empty()); + assert!(run.splits.iter().all(Vec::is_empty)); + } else { + assert_eq!(segments_read(&run.events), vec![0]); + let expected = PrimitiveArray::from_iter([5, ROWS as i32 - 6]).into_array(); + assert_view( + &expected, + &(0..2), + &Mask::new_true(2), + run.splits[0].clone(), + )?; + } + Ok(()) +} + #[test] fn every_read_is_issued_before_any_is_delivered() -> VortexResult<()> { let mut store = Store::default(); @@ -762,7 +996,11 @@ fn pack_readiness_tracks_changes_and_consumption( let mut exec = SESSION.create_execution_ctx(); let mut touched = InletSet::new(2); let mut changed = InletSet::all(2); + let mut child_requests = Vec::new(); + let mut child_inbox = None; let mut cx = Cx { + child_requests: &mut child_requests, + child_inbox: &mut child_inbox, arena: &mut arena, inlets: &inlets, bytes: &mut bytes, @@ -834,3 +1072,29 @@ fn pack_readiness_tracks_changes_and_consumption( } Ok(()) } + +#[rstest] +fn nullable_struct_without_fields_keeps_its_rows( + #[values(Sel::All, Sel::EveryOther)] selection: Sel, +) -> VortexResult<()> { + let mut store = Store::default(); + let validity = BoolArray::from_iter((0..ROWS).map(|row| row % 3 != 0)).into_array(); + let expected = StructArray::try_new_with_dtype( + Vec::new(), + StructFields::empty(), + ROWS as usize, + Validity::Array(validity.clone()), + )? + .into_array(); + let layout = StructLayout::new(ROWS, expected.dtype().clone(), vec![store.flat(&validity)?]) + .into_layout(); + let mask = selection.mask(ROWS as usize); + let run = run( + &store, + &lower(&layout)?, + 0..ROWS, + mask.clone(), + delivery(Delivery::Fifo), + )?; + assert_view(&expected, &(0..ROWS), &mask, run.arrays) +} diff --git a/vortex-layout/src/plan/physical/plans/concat.rs b/vortex-layout/src/plan/physical/plans/concat.rs index b2e1ffea0f7..377e4bf270c 100644 --- a/vortex-layout/src/plan/physical/plans/concat.rs +++ b/vortex-layout/src/plan/physical/plans/concat.rs @@ -189,6 +189,9 @@ impl PlanVTable for Concat { visit: &mut dyn FnMut(SegmentId, Range), ) -> VortexResult<()> { for (index, start, local) in overlapping(plan, &rows) { + if !at.selects(&local) { + continue; + } plan.child_required(index)?.reach( local.start - start..local.end - start, &at.shift(start), diff --git a/vortex-layout/src/plan/physical/plans/list_pack.rs b/vortex-layout/src/plan/physical/plans/list_pack.rs index bb075c7f4e2..0d459d3537b 100644 --- a/vortex-layout/src/plan/physical/plans/list_pack.rs +++ b/vortex-layout/src/plan/physical/plans/list_pack.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::borrow::Cow; +use std::ops::Range; use std::sync::Arc; use vortex_array::EmptyMetadata; @@ -10,6 +11,8 @@ use vortex_array::dtype::Nullability; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_mask::AllOr; +use vortex_mask::Mask; use vortex_session::registry::CachedId; use crate::plan::physical::Plan; @@ -18,6 +21,12 @@ use crate::plan::physical::PlanId; use crate::plan::physical::PlanParts; use crate::plan::physical::PlanRef; use crate::plan::physical::PlanVTable; +use crate::plan::physical::operators::ListSource; +use crate::plan::physical::operators::UNBOUNDED; +use crate::plan::physical::pipeline::Chain; +use crate::plan::physical::pipeline::Compiler; +use crate::plan::physical::pipeline::Reach; +use crate::segments::SegmentId; const ELEMENTS: usize = 0; const OFFSETS: usize = 1; @@ -124,6 +133,73 @@ impl PlanVTable for ListPack { _ => Cow::Owned(format!("child[{index}]")), } } + + fn compile( + plan: &Plan, + rows: Range, + mask: &Mask, + compiler: &mut Compiler<'_>, + ) -> VortexResult> { + if mask.all_false() { + return Ok(None); + } + let boundaries = offset_mask(mask); + let offsets = compiler + .compile(&plan.offsets()?, rows.start..rows.end + 1, &boundaries)? + .ok_or_else(|| vortex_err!("Selected list offsets produced no rows"))?; + let mut chains = vec![offsets]; + if let Some(validity) = plan.validity()? { + chains.push( + compiler + .compile(&validity, rows, mask)? + .ok_or_else(|| vortex_err!("Selected list validity produced no rows"))?, + ); + } + let source = ListSource::new( + plan.clone(), + mask.clone(), + compiler.values().clone(), + compiler.plan_children(), + ); + Ok(Some(compiler.join_with(chains, |_| UNBOUNDED, source))) + } + + fn reach( + plan: &Plan, + rows: Range, + at: &Reach, + visit: &mut dyn FnMut(SegmentId, Range), + ) -> VortexResult<()> { + // Elements have a different row domain, known only after the selected offsets arrive. + // Their scan records its own segment readers when those ranges have been resolved. + let root_rows = at.root(&rows); + plan.offsets()? + .reach(rows.start..rows.end + 1, &at.fixed(&rows), &mut |id, _| { + visit(id, root_rows.clone()) + })?; + if let Some(validity) = plan.validity()? { + validity.reach(rows, at, visit)?; + } + Ok(()) + } +} + +/// A selected list needs both its starting and ending offsets. +fn offset_mask(mask: &Mask) -> Mask { + match mask.indices() { + AllOr::None => Mask::new_false(mask.len() + 1), + AllOr::All => Mask::new_true(mask.len() + 1), + AllOr::Some(indices) => { + let mut boundaries = Vec::with_capacity(indices.len() * 2); + for &index in indices { + if boundaries.last() != Some(&index) { + boundaries.push(index); + } + boundaries.push(index + 1); + } + Mask::from_indices(mask.len() + 1, boundaries) + } + } } fn validate_children(dtype: &DType, row_count: u64, children: &PlanChildren) -> VortexResult<()> { diff --git a/vortex-layout/src/plan/physical/tests.rs b/vortex-layout/src/plan/physical/tests.rs index da6fa56ccb9..2b74e23d17a 100644 --- a/vortex-layout/src/plan/physical/tests.rs +++ b/vortex-layout/src/plan/physical/tests.rs @@ -4,7 +4,10 @@ use std::fmt; use std::num::NonZeroUsize; use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnRef; @@ -43,6 +46,7 @@ use vortex_session::registry::ReadContext; use super::*; use crate::LayoutBuildContext; +use crate::LayoutChildren; use crate::LayoutEncoding; use crate::LayoutRef; use crate::OwnedLayoutChildren; @@ -150,6 +154,60 @@ fn chunked_plan_exposes_chunks() -> VortexResult<()> { Ok(()) } +#[derive(Clone)] +struct CountedChunkRows { + rows: Arc<[u64]>, + calls: Arc, +} + +impl LayoutChildren for CountedChunkRows { + fn to_arc(&self) -> Arc { + Arc::new(self.clone()) + } + + fn child(&self, _idx: usize, _dtype: &DType) -> VortexResult { + Err(vortex_err!("Lowering must leave chunk children lazy")) + } + + fn child_row_count(&self, idx: usize) -> u64 { + self.calls.fetch_add(1, Ordering::Relaxed); + self.rows[idx] + } + + fn nchildren(&self) -> usize { + self.rows.len() + } +} + +#[rstest] +#[case::normal(vec![3, 2], vec![0, 3])] +#[case::empty_chunks(vec![0, 2, 0, 3, 0], vec![0, 0, 2, 2, 5])] +#[case::empty(vec![], vec![])] +#[case::zero_rows(vec![0], vec![0])] +fn lowering_reuses_validated_chunk_boundaries( + #[case] rows: Vec, + #[case] offsets: Vec, +) -> VortexResult<()> { + let calls = Arc::new(AtomicUsize::new(0)); + let row_count = rows.iter().sum(); + let layout = ChunkedLayout::new( + row_count, + DType::Null, + Arc::new(CountedChunkRows { + rows: rows.into(), + calls: Arc::clone(&calls), + }), + ) + .into_layout(); + calls.store(0, Ordering::Relaxed); + let plan = lower(&layout)?; + assert_eq!(plan.as_::().row_offsets(), offsets); + assert_eq!(plan.row_count(), row_count); + assert_eq!(plan.child_count(), offsets.len()); + assert_eq!(calls.load(Ordering::Relaxed), 0); + Ok(()) +} + #[test] fn chunked_plan_lowers_each_chunk_on_access() -> VortexResult<()> { let dtype = primitive(PType::I32, Nullability::NonNullable); diff --git a/vortex-layout/src/scan/execution/adaptor.rs b/vortex-layout/src/scan/execution/adaptor.rs index fc6053d55ca..edf681b294f 100644 --- a/vortex-layout/src/scan/execution/adaptor.rs +++ b/vortex-layout/src/scan/execution/adaptor.rs @@ -19,6 +19,7 @@ use vortex_io::request::IoIntent; use vortex_io::request::IoResult; use vortex_scan::planning::planner::State; +use crate::plan::physical::pipeline::ChildProjection; use crate::plan::physical::pipeline::ReadId; use crate::plan::physical::pipeline::Scan; use crate::plan::physical::pipeline::Turn; @@ -40,6 +41,8 @@ pub(super) enum AdaptorStep { Yield, /// Requests to publish, each returned once. NeedsIO(IoBatch), + /// A child projection for the enclosing planner to schedule. + Child(ChildProjection), /// Rows the plan produced. Piece(Piece), } @@ -97,6 +100,7 @@ impl ScanAdaptor { Ok(match self.step()? { Turn::Waiting | Turn::Done => AdaptorStep::Yield, Turn::Output(_, array) => AdaptorStep::Piece(self.piece(array)), + Turn::Child(child) => AdaptorStep::Child(child), Turn::Read(read) => { // Publish every read the scan has ready together, so the owner can coalesce them. let mut batch = vec![self.request(read.id, read.segment_id)?]; @@ -125,7 +129,7 @@ impl ScanAdaptor { match &turn { Turn::Done => self.done = true, Turn::Waiting => self.pending = Some(Turn::Waiting), - Turn::Read(_) | Turn::Output(..) => {} + Turn::Read(_) | Turn::Child(_) | Turn::Output(..) => {} } Ok(turn) } @@ -139,6 +143,18 @@ impl ScanAdaptor { } } + pub(super) fn complete_child( + &mut self, + id: u64, + arrays: Option>, + ) -> VortexResult<()> { + self.scan.complete_child(id, arrays)?; + if matches!(self.pending, Some(Turn::Waiting)) { + self.pending = None; + } + Ok(()) + } + fn request(&mut self, read: ReadId, segment_id: SegmentId) -> VortexResult { let location = self.location(segment_id)?; let id = self.next_id; diff --git a/vortex-layout/src/scan/execution/children.rs b/vortex-layout/src/scan/execution/children.rs new file mode 100644 index 00000000000..e57ba8a4816 --- /dev/null +++ b/vortex-layout/src/scan/execution/children.rs @@ -0,0 +1,230 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Query-wide admission and execution of large child projections. +//! +//! Permits cover the estimated decoded bytes of whole touched chunks until their task completes. +//! Retained output, encoded buffers and allocator overhead are outside this estimate; this is not +//! a limit on the query's total heap. Cancellation drops child handles and their IO sessions. + +use std::future::poll_fn; +use std::sync::Arc; +use std::sync::OnceLock; +#[cfg(test)] +use std::sync::atomic::AtomicUsize; +#[cfg(test)] +use std::sync::atomic::Ordering; + +use async_lock::Mutex; +use async_lock::Semaphore; +use async_lock::SemaphoreGuardArc; +use futures::FutureExt; +use futures::StreamExt; +use futures::TryStreamExt; +use futures::future::BoxFuture; +use futures::stream; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::Chunked; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::chunked::ChunkedArrayExt; +use vortex_array::dtype::DType; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_io::request::IoService; +use vortex_io::session::RuntimeSessionExt; +use vortex_scan::planning::driver::Progress; +use vortex_scan::planning::driver::Run; +use vortex_scan::planning::planner::WorkScope; +use vortex_scan::planning::selection::SelectedRows; +use vortex_session::VortexSession; +use vortex_utils::parallelism::get_available_parallelism; + +use super::ScanPlans; +use super::SegmentLocation; +use super::projection::ProjectionPlanner; +use crate::plan::physical::PlanRef; +use crate::plan::physical::pipeline::Split; +use crate::plan::physical::pipeline::ValueCache; + +const BUDGET_BYTES: u64 = 32 * 1024 * 1024; +const PERMIT_BYTES: u64 = 256 * 1024; + +pub(super) mod splits; + +#[cfg(test)] +mod tests; + +#[derive(Debug)] +pub(super) struct ChildWork { + pub(super) split: Split, + /// Estimated bytes of whole touched chunks, including validity. + pub(super) bytes: u64, +} + +struct ByteBudget { + permits: Arc, + acquisition: Mutex<()>, +} + +impl Default for ByteBudget { + fn default() -> Self { + Self { + permits: Arc::new(Semaphore::new((BUDGET_BYTES / PERMIT_BYTES) as usize)), + acquisition: Mutex::new(()), + } + } +} + +impl ByteBudget { + async fn acquire(&self, bytes: u64) -> VortexResult> { + vortex_ensure!( + bytes > 0 && bytes <= BUDGET_BYTES, + "Child work exceeds its admission budget" + ); + // Acquire a whole weight before admitting another claimant, so partial acquisitions + // cannot fill the budget and deadlock one another. Completed tasks release independently. + let _acquisition = self.acquisition.lock().await; + let count = usize::try_from(bytes.div_ceil(PERMIT_BYTES))?; + let mut permits = Vec::with_capacity(count); + for _ in 0..count { + permits.push(self.permits.acquire_arc().await); + } + Ok(permits) + } +} + +pub(in crate::scan) struct Children { + session: VortexSession, + locations: Arc<[SegmentLocation]>, + io: Arc, + workers: usize, + budget: OnceLock, + #[cfg(test)] + admitted: AtomicUsize, +} + +pub(in crate::scan) fn child_executor( + session: VortexSession, + locations: Arc<[SegmentLocation]>, + io: Arc, +) -> Option> { + let workers = get_available_parallelism().unwrap_or(1); + (workers > 1).then(|| { + Arc::new(Children { + session, + locations, + io, + workers, + budget: OnceLock::new(), + #[cfg(test)] + admitted: Default::default(), + }) + }) +} + +impl Children { + #[cfg(test)] + pub(in crate::scan) fn admitted(&self) -> usize { + self.admitted.load(Ordering::Relaxed) + } + + pub(super) fn execute( + self: Arc, + plan: PlanRef, + work: Vec, + values: ValueCache, + ) -> BoxFuture<'static, VortexResult>> { + async move { + let workers = self.workers; + stream::iter(work) + .map(|work| { + let children = Arc::clone(&self); + let plan = plan.clone(); + let values = values.clone(); + async move { + let permits = children + .budget + .get_or_init(ByteBudget::default) + .acquire(work.bytes) + .await?; + #[cfg(test)] + children.admitted.fetch_add(1, Ordering::Relaxed); + let handle = children.session.handle(); + handle + .spawn(async move { + let _permits = permits; + let scope = WorkScope { + file_ordinal: 0, + rows: work.split.rows, + }; + let plans = ScanPlans { + session: children.session.clone(), + locations: Arc::clone(&children.locations), + projection: plan, + row_offset: 0, + values, + // Eligible child plans contain only fixed-width segment leaves. + // They cannot recursively wait for this same byte budget. + children: None, + }; + let root = ProjectionPlanner::new( + plans, + SelectedRows { + scope: scope.clone(), + mask: work.split.mask, + }, + ); + let mut run = Run::new(); + run.admit(Box::new(root), scope, children.io.session()); + let mut result = Vec::new(); + let mut exec = children.session.create_execution_ctx(); + loop { + match run.advance()? { + Progress::Batch(batch) => { + materialize(batch.array, &mut exec, &mut result)?; + } + Progress::RootDone(_) | Progress::Idle => break, + Progress::Waiting => { + poll_fn(|cx| run.poll_completion(cx)).await? + } + } + } + Ok::<_, vortex_error::VortexError>(result) + }) + .await + } + }) + .buffered(workers) + .try_collect::>>() + .await + .map(|arrays| arrays.into_iter().flatten().collect()) + } + .boxed() + } +} + +/// Decode on the child task, preserving chunk boundaries so the parent need only concatenate +/// once. Canonicalizing the whole job would copy its chunks into an intermediate buffer first. +fn materialize( + array: ArrayRef, + exec: &mut ExecutionCtx, + output: &mut Vec, +) -> VortexResult<()> { + if let Some(chunks) = array.as_opt::() { + for chunk in chunks.iter_chunks() { + materialize(chunk.clone(), exec, output)?; + } + return Ok(()); + } + output.push(match array.dtype() { + DType::Primitive(..) => array.execute::(exec)?.into_array(), + DType::Bool(..) => array.execute::(exec)?.into_array(), + _ => vortex_bail!("Child projection is not fixed-width"), + }); + Ok(()) +} diff --git a/vortex-layout/src/scan/execution/children/splits.rs b/vortex-layout/src/scan/execution/children/splits.rs new file mode 100644 index 00000000000..89bf01eff97 --- /dev/null +++ b/vortex-layout/src/scan/execution/children/splits.rs @@ -0,0 +1,152 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Split selected child projections at stored segment boundaries before admitting them. + +use std::ops::Range; + +use rustc_hash::FxHashSet; +use vortex_array::dtype::DType; +use vortex_error::VortexResult; + +use super::ChildWork; +use crate::plan::physical::Concat; +use crate::plan::physical::Filter; +use crate::plan::physical::LegacyStats; +use crate::plan::physical::PlanRef; +use crate::plan::physical::SegmentScan; +use crate::plan::physical::Zoned; +use crate::plan::physical::pipeline::Split; +use crate::plan::physical::pipeline::overlapping; +use crate::segments::SegmentId; + +const MAX_ELEMENT_BYTES: u64 = 8 * 1024 * 1024; +const CHILD_TARGET_BYTES: u64 = 256 * 1024; + +#[cfg(test)] +mod tests; + +/// Divide substantial fixed-width selections at stored segment boundaries. Each segment is +/// owned by at most one task, so independent execution cannot duplicate its IO and decoding. +pub(in crate::scan::execution) fn parallel_splits( + plan: &PlanRef, + splits: &[Split], +) -> VortexResult>> { + let width = match plan.dtype() { + DType::Primitive(ptype, _) => ptype.byte_width() as u64, + DType::Bool(_) => 1, + _ => return Ok(None), + }; + let selected = splits.iter().fold(0_u64, |total, split| { + total.saturating_add((split.mask.true_count() as u64).saturating_mul(width)) + }); + if selected < CHILD_TARGET_BYTES * 2 { + return Ok(None); + } + let mut seen = FxHashSet::default(); + let mut work = Vec::new(); + for split in splits { + let mut leaves = Vec::new(); + if !selected_leaves(plan, split, 0, &mut seen, &mut leaves)? { + return Ok(None); + } + let mut first = 0; + while first < leaves.len() { + let start = leaves[first].0.start; + let mut end = leaves[first].0.end; + let mut bytes = leaves[first].1; + first += 1; + while bytes < CHILD_TARGET_BYTES + && first < leaves.len() + && bytes.saturating_add(leaves[first].1) <= MAX_ELEMENT_BYTES + { + end = leaves[first].0.end; + bytes += leaves[first].1; + first += 1; + } + work.push(ChildWork { + split: Split { + rows: start..end, + mask: split.mask.slice( + usize::try_from(start - split.rows.start)? + ..usize::try_from(end - split.rows.start)?, + ), + }, + bytes, + }); + } + } + Ok((work.len() > 1).then_some(work)) +} + +fn selected_leaves( + plan: &PlanRef, + split: &Split, + offset: u64, + seen: &mut FxHashSet, + leaves: &mut Vec<(Range, u64)>, +) -> VortexResult { + if split.mask.all_false() { + return Ok(true); + } + if let Some(filter) = plan.as_opt::() { + return selected_leaves(&filter.child_plan()?, split, offset, seen, leaves); + } + if let Some(zoned) = plan.as_opt::() { + return match zoned.data_plan()? { + Some(data) => selected_leaves(&data, split, offset, seen, leaves), + None => Ok(false), + }; + } + if let Some(stats) = plan.as_opt::() { + return selected_leaves(&stats.data_plan()?, split, offset, seen, leaves); + } + if let Some(concat) = plan.as_opt::() { + for (index, start, rows) in overlapping(concat, &split.rows) { + let mask = split.mask.slice( + usize::try_from(rows.start - split.rows.start)? + ..usize::try_from(rows.end - split.rows.start)?, + ); + // Do not lower an untouched lazy child, including unknown layouts in holes. + if mask.all_false() { + continue; + } + let child = Split { + rows: rows.start - start..rows.end - start, + mask, + }; + if !selected_leaves( + &concat.child_required(index)?, + &child, + offset + start, + seen, + leaves, + )? { + return Ok(false); + } + } + return Ok(true); + } + let Some(segment) = plan.as_opt::() else { + return Ok(false); + }; + let width = match plan.dtype() { + DType::Primitive(ptype, _) => ptype.byte_width() as u64, + DType::Bool(_) => 1, + _ => return Ok(false), + }; + let validity = if plan.dtype().is_nullable() { + plan.row_count().div_ceil(8) + } else { + 0 + }; + let bytes = plan + .row_count() + .saturating_mul(width) + .saturating_add(validity); + if bytes == 0 || bytes > MAX_ELEMENT_BYTES || !seen.insert(segment.segment_id()) { + return Ok(false); + } + leaves.push((offset + split.rows.start..offset + split.rows.end, bytes)); + Ok(true) +} diff --git a/vortex-layout/src/scan/execution/children/splits/tests.rs b/vortex-layout/src/scan/execution/children/splits/tests.rs new file mode 100644 index 00000000000..d7ba1c8f7e8 --- /dev/null +++ b/vortex-layout/src/scan/execution/children/splits/tests.rs @@ -0,0 +1,139 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::dtype::PType; +use vortex_mask::Mask; +use vortex_session::registry::ReadContext; + +use super::*; +use crate::plan::physical::ConcatPlan; +use crate::plan::physical::SegmentScanPlan; + +fn chunks(count: u32, rows: u64, dtype: DType, duplicate: bool) -> VortexResult { + ConcatPlan::try_new( + dtype.clone(), + (0..count) + .map(|id| { + SegmentScanPlan::new( + dtype.clone(), + rows, + SegmentId::from(if duplicate { 0 } else { id }), + ReadContext::new([]), + None, + ) + .into_plan() + }) + .collect(), + ) + .map(|plan| plan.into_plan()) +} + +#[rstest] +fn contiguous_list_row_splits_at_stored_boundaries( + #[values(false, true)] nullable: bool, +) -> VortexResult<()> { + let rows = 8192; + let plan = chunks( + 32, + rows, + DType::Primitive(PType::I64, nullable.into()), + false, + )?; + let jobs = parallel_splits(&plan, &[Split::all(100..rows * 32 - 100)])? + .expect("large row is parallel"); + assert_eq!(jobs.len(), 8); + assert_eq!(jobs[0].split.rows, 100..rows * 4); + assert_eq!(jobs[7].split.rows, rows * 28..rows * 32 - 100); + assert!(jobs.iter().all(|job| job.split.mask.all_true())); + assert_eq!( + jobs.iter().map(|job| job.bytes).sum::(), + rows * 32 * 8 + if nullable { rows * 4 } else { 0 } + ); + Ok(()) +} + +#[test] +fn masks_preserve_holes_and_do_not_charge_unselected_segments() -> VortexResult<()> { + let rows = 32768; + let plan = chunks(5, rows, DType::Primitive(PType::I64, false.into()), false)?; + let mask = Mask::from_slices( + usize::try_from(rows * 5)?, + vec![(0, 32768), (65536, 98304), (131072, 163840)], + ); + let jobs = parallel_splits( + &plan, + &[Split { + rows: 0..rows * 5, + mask, + }], + )? + .expect("three selected chunks"); + assert_eq!( + jobs.iter() + .map(|job| job.split.rows.clone()) + .collect::>(), + vec![0..rows, rows * 2..rows * 3, rows * 4..rows * 5] + ); + assert_eq!(jobs.iter().map(|job| job.bytes).sum::(), rows * 3 * 8); + Ok(()) +} + +#[test] +fn partial_masks_survive_grouping_and_boundary_clipping() -> VortexResult<()> { + let plan = chunks(32, 8192, DType::Primitive(PType::I64, false.into()), false)?; + let mask = Mask::from_iter((0..200_000).map(|i| i % 7 != 0)); + let split = Split { + rows: 999..200_999, + mask: mask.clone(), + }; + let jobs = parallel_splits(&plan, &[split])?.expect("large selected payload"); + let actual: Vec<_> = jobs + .iter() + .flat_map(|job| { + job.split + .mask + .iter() + .enumerate() + .filter_map(|(i, selected)| selected.then_some(job.split.rows.start + i as u64)) + }) + .collect(); + let expected: Vec<_> = mask + .iter() + .enumerate() + .filter_map(|(i, selected)| selected.then_some(999 + i as u64)) + .collect(); + assert_eq!(actual, expected); + Ok(()) +} + +#[test] +fn repeated_segments_and_repeated_ranges_stay_serial() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I64, false.into()); + let repeated = chunks(2, 65536, dtype.clone(), true)?; + assert!(parallel_splits(&repeated, &[Split::all(0..131072)])?.is_none()); + let distinct = chunks(2, 65536, dtype, false)?; + assert!( + parallel_splits( + &distinct, + &[Split::all(0..32000), Split::all(32001..131072)] + )? + .is_none() + ); + Ok(()) +} + +#[rstest] +#[case::single_segment(1, 131072, DType::Primitive(PType::I64, false.into()))] +#[case::oversized_segment(2, MAX_ELEMENT_BYTES / 8 + 1, DType::Primitive(PType::I64, false.into()))] +#[case::tiny_selection(8, 1024, DType::Primitive(PType::I64, false.into()))] +#[case::variable_width(8, 131072, DType::Utf8(false.into()))] +fn unsuitable_child_work_stays_serial( + #[case] count: u32, + #[case] rows: u64, + #[case] dtype: DType, +) -> VortexResult<()> { + let plan = chunks(count, rows, dtype, false)?; + assert!(parallel_splits(&plan, &[Split::all(0..rows * u64::from(count))])?.is_none()); + Ok(()) +} diff --git a/vortex-layout/src/scan/execution/children/tests.rs b/vortex-layout/src/scan/execution/children/tests.rs new file mode 100644 index 00000000000..680dd781c68 --- /dev/null +++ b/vortex-layout/src/scan/execution/children/tests.rs @@ -0,0 +1,357 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::collections::VecDeque; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; +use std::time::Duration; +use std::time::Instant; + +use futures::executor::block_on; +use futures::task::AtomicWaker; +use parking_lot::Mutex as SyncMutex; +use vortex_array::ArrayContext; +use vortex_array::arrays::ChunkedArray; +use vortex_array::assert_arrays_eq; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::PType; +use vortex_array::serde::SerializeOptions; +use vortex_buffer::Alignment; +use vortex_buffer::Buffer; +use vortex_buffer::ByteBufferMut; +use vortex_error::vortex_err; +use vortex_io::request::Completion; +use vortex_io::request::IoBatch; +use vortex_io::request::IoIntent; +use vortex_io::request::IoOwnerId; +use vortex_io::request::IoResult; +use vortex_io::request::IoSource; +use vortex_io::request::IoTarget; +use vortex_session::registry::ReadContext; + +use super::*; +use crate::plan::physical::ConcatPlan; +use crate::plan::physical::SegmentScanPlan; +use crate::plan::physical::pipeline::Split; +use crate::segments::SegmentId; +use crate::test::new_session; + +#[test] +fn byte_budget_prevents_over_admission_and_cancellation_releases_partial_acquisitions() +-> VortexResult<()> { + let budget = ByteBudget::default(); + let held = block_on(budget.acquire(BUDGET_BYTES * 3 / 4))?; + let mut waiting = Box::pin(budget.acquire(BUDGET_BYTES / 2)); + let mut cx = Context::from_waker(Waker::noop()); + assert!(waiting.poll_unpin(&mut cx).is_pending()); + // The waiting acquisition has taken the remaining quarter. Cancelling it must return + // those permits and the acquisition lock, even while the first job is still running. + drop(waiting); + let mut quarter = Box::pin(budget.acquire(BUDGET_BYTES / 4)); + let Poll::Ready(permits) = quarter.poll_unpin(&mut cx) else { + panic!("cancelled admission leaked permits") + }; + let quarter = permits?; + let mut extra = Box::pin(budget.acquire(1)); + assert!(extra.poll_unpin(&mut cx).is_pending()); + drop(quarter); + let Poll::Ready(permits) = extra.poll_unpin(&mut cx) else { + panic!("released permits did not wake admission") + }; + drop(permits?); + drop(held); + let mut whole = Box::pin(budget.acquire(BUDGET_BYTES)); + assert!(matches!(whole.poll_unpin(&mut cx), Poll::Ready(Ok(_)))); + Ok(()) +} + +#[test] +fn admission_rejects_zero_and_oversized_work() { + let budget = ByteBudget::default(); + assert!(block_on(budget.acquire(0)).is_err()); + assert!(block_on(budget.acquire(BUDGET_BYTES + 1)).is_err()); +} + +#[derive(Default)] +struct Gates { + sources: SyncMutex>>, + bytes: Vec, +} + +impl IoService for Gates { + fn session(&self) -> Arc { + let gate = Arc::new(Gate { + bytes: self.bytes.clone(), + ready: AtomicBool::new(false), + fail: AtomicBool::new(false), + requests: SyncMutex::new(VecDeque::new()), + wake: AtomicWaker::new(), + cleared: AtomicUsize::new(0), + }); + self.sources.lock().push(Arc::clone(&gate)); + gate + } +} + +struct Gate { + bytes: Vec, + ready: AtomicBool, + fail: AtomicBool, + requests: SyncMutex>, + wake: AtomicWaker, + cleared: AtomicUsize, +} + +impl Gate { + fn open(&self) { + self.ready.store(true, Ordering::Release); + self.wake.wake(); + } +} + +impl IoSource for Gate { + fn submit(&self, owner: IoOwnerId, batch: IoBatch) -> VortexResult<()> { + for request in batch + .into_iter() + .filter(|request| request.intent == IoIntent::Fetch) + { + let IoTarget::Range { offset, .. } = request.target else { + return Err(vortex_err!("unexpected size request")); + }; + self.requests.lock().push_back(Completion { + owner, + request: request.request, + result: Ok(IoResult::Bytes( + self.bytes[usize::try_from(offset)?].clone(), + )), + }); + } + Ok(()) + } + + fn poll(&self) -> VortexResult> { + if self.fail.load(Ordering::Acquire) { + return Err(vortex_err!("injected child IO failure")); + } + Ok(if self.ready.load(Ordering::Acquire) { + self.requests.lock().pop_front() + } else { + None + }) + } + + fn poll_completion(&self, cx: &mut Context<'_>) -> Poll> { + self.wake.register(cx.waker()); + match self.poll() { + Ok(Some(completion)) => Poll::Ready(Ok(completion)), + Ok(None) => Poll::Pending, + Err(error) => Poll::Ready(Err(error)), + } + } + + fn wait(&self) -> VortexResult { + Err(vortex_err!("child executor must not block on IO")) + } + fn release(&self, owner: IoOwnerId) { + self.requests + .lock() + .retain(|request| request.owner != owner); + } + fn clear(&self) { + self.requests.lock().clear(); + self.cleared.fetch_add(1, Ordering::Relaxed); + } +} + +type Fixture = (Arc, Arc, PlanRef, Vec); + +#[test] +fn materialization_preserves_chunk_buffers_instead_of_concatenating_twice() -> VortexResult<()> { + let session = new_session(); + let chunks = [ + PrimitiveArray::from_option_iter([Some(1_i64), None]), + PrimitiveArray::from_option_iter([None, Some(4_i64)]), + ]; + let array = ChunkedArray::try_new( + chunks.iter().cloned().map(IntoArray::into_array), + chunks[0].dtype().clone(), + )? + .into_array(); + let mut output = Vec::new(); + materialize(array, &mut session.create_execution_ctx(), &mut output)?; + assert_eq!(output.len(), chunks.len()); + for (actual, expected) in output.into_iter().zip(chunks) { + let actual = actual.execute::(&mut session.create_execution_ctx())?; + assert_eq!( + actual.as_slice::().as_ptr(), + expected.as_slice::().as_ptr() + ); + assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); + } + Ok(()) +} + +fn fixture() -> VortexResult { + let session = new_session().with_tokio(); + let mut service = Gates::default(); + let mut plans = Vec::new(); + for id in 0_u32..3 { + let array = Buffer::from_iter(u64::from(id) * 8..u64::from(id + 1) * 8).into_array(); + let ctx = ArrayContext::empty(); + let buffers = array.serialize( + &ctx, + &session, + &SerializeOptions { + offset: 0, + include_padding: true, + }, + )?; + let mut bytes = ByteBufferMut::empty_aligned(Alignment::new(64)); + for buffer in buffers { + bytes.extend_from_slice(buffer.as_ref()); + } + service.bytes.push(BufferHandle::new_host(bytes.freeze())); + plans.push( + SegmentScanPlan::new( + array.dtype().clone(), + 8, + SegmentId::from(id), + ReadContext::new(ctx.to_ids()), + None, + ) + .into_plan(), + ); + } + let service = Arc::new(service); + let children = Arc::new(Children { + session, + io: Arc::::clone(&service), + workers: 4, + budget: OnceLock::new(), + admitted: Default::default(), + locations: (0..3) + .map(|offset| SegmentLocation { + offset, + length: 0, + alignment: Alignment::new(64), + }) + .collect(), + }); + let plan = ConcatPlan::try_new(DType::Primitive(PType::U64, false.into()), plans)?.into_plan(); + // Inflate the estimate to exercise admission without a 48 MiB test fixture. + let work = (0..3) + .map(|id| ChildWork { + split: Split::all(id * 8..(id + 1) * 8), + bytes: BUDGET_BYTES / 2, + }) + .collect(); + Ok((children, service, plan, work)) +} + +async fn until(mut condition: impl FnMut() -> bool) { + let start = Instant::now(); + while !condition() { + assert!( + start.elapsed() < Duration::from_secs(5), + "child execution did not make progress" + ); + tokio::task::yield_now().await; + } +} + +#[tokio::test] +async fn independent_children_respect_the_shared_budget_and_return_in_order() -> VortexResult<()> { + let (children, service, plan, work) = fixture()?; + let task = tokio::spawn(Arc::clone(&children).execute(plan, work, ValueCache::default())); + until(|| service.sources.lock().len() == 2).await; + let first = Arc::clone(&service.sources.lock()[0]); + let second = Arc::clone(&service.sources.lock()[1]); + assert!(!task.is_finished()); + assert_eq!(first.requests.lock().len(), 1); + assert_eq!(second.requests.lock().len(), 1); + // The later job releases capacity while the first job is still held. The third job + // must start without losing output order or exceeding the two-job byte budget. + second.open(); + until(|| service.sources.lock().len() == 3).await; + assert!(!first.ready.load(Ordering::Relaxed)); + service.sources.lock()[2].open(); + first.open(); + let arrays = task.await.expect("child task panicked")?; + assert_eq!(arrays.len(), 3); + for (index, array) in arrays.into_iter().enumerate() { + let start = u64::try_from(index)? * 8; + assert_arrays_eq!( + array, + Buffer::from_iter(start..start + 8).into_array(), + &mut children.session.create_execution_ctx() + ); + } + assert!( + service + .sources + .lock() + .iter() + .all(|source| source.cleared.load(Ordering::Relaxed) == 1) + ); + Ok(()) +} + +#[rstest::rstest] +#[tokio::test] +async fn cancelled_or_failed_parent_releases_child_io_and_budget( + #[values(false, true)] failure: bool, +) -> VortexResult<()> { + let (children, service, plan, work) = fixture()?; + let task = tokio::spawn(Arc::clone(&children).execute(plan, work, ValueCache::default())); + until(|| service.sources.lock().len() == 2).await; + if failure { + let first = Arc::clone(&service.sources.lock()[0]); + first.fail.store(true, Ordering::Release); + first.open(); + assert!( + task.await + .expect("task panicked") + .unwrap_err() + .to_string() + .contains("injected child IO failure") + ); + } else { + task.abort(); + assert!( + task.await + .expect_err("task should be cancelled") + .is_cancelled() + ); + } + until(|| { + service + .sources + .lock() + .iter() + .all(|source| source.cleared.load(Ordering::Relaxed) == 1) + }) + .await; + assert!( + service + .sources + .lock() + .iter() + .all(|source| source.requests.lock().is_empty()) + ); + let mut all = Box::pin( + children + .budget + .get() + .expect("budget initialized") + .acquire(BUDGET_BYTES), + ); + assert!(matches!( + all.poll_unpin(&mut Context::from_waker(Waker::noop())), + Poll::Ready(Ok(_)) + )); + Ok(()) +} diff --git a/vortex-layout/src/scan/execution/mod.rs b/vortex-layout/src/scan/execution/mod.rs index 9d8b8897ef3..3f919d3a5b5 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -3,23 +3,29 @@ //! Scan execution over physical plans through the planner I/O and compute protocol. //! -//! A `ProjectionPlanner` emits one `ProjectionMorsel` for each fixed-size work chunk. The morsel -//! runs the physical projection plan over the chunk's selected rows through -//! the pipeline [`Scan`](crate::plan::physical::pipeline::Scan). Every segment read is published -//! through a `ScanAdaptor`, so the planning driver can advance other work while reads wait. +//! A `ProjectionPlanner` drives the physical projection over selected rows and schedules +//! independently admitted child projections. Every segment read is published through a +//! `ScanAdaptor`, so the driver advances other work while reads wait. Only the assembled result +//! becomes a morsel, which publishes the selected rows in order. mod adaptor; +mod children; mod projection; use std::sync::Arc; +pub(super) use children::child_executor; use vortex_buffer::Alignment; use vortex_io::request::IoTarget; -use vortex_mask::Mask; +use vortex_scan::planning::next::next_fn; use vortex_scan::planning::planner::Planner; use vortex_scan::planning::planner::WorkScope; +use vortex_scan::planning::selection::SelectedRows; +use vortex_scan::planning::selection::SelectionPlanner; +use vortex_scan::selection::Selection; use vortex_session::VortexSession; +use self::children::Children; use self::projection::ProjectionPlanner; use crate::plan::physical::PlanRef; use crate::plan::physical::pipeline::ValueCache; @@ -60,17 +66,18 @@ pub(super) struct ScanPlans { /// Whole values, such as dictionary values, shared by the splits of this scan. Their plans /// always read the same complete row domain, independently of a split's selection. pub(super) values: ValueCache, + pub(super) children: Option>, } -/// The rows of a split selected for projection. -struct SelectedRows { - /// The split the rows belong to. +/// Selects the rows of a fixed work chunk before handing them to its projection planner. +pub(super) fn plan_projection( + plans: ScanPlans, scope: WorkScope, - /// Which rows of `scope.rows` are selected. - mask: Mask, -} - -/// The work for one split: project the rows of `scope` selected by `mask`. -pub(super) fn plan_projection(plans: ScanPlans, scope: WorkScope, mask: Mask) -> Box { - Box::new(ProjectionPlanner::new(plans, SelectedRows { scope, mask })) + selection: Arc, +) -> Box { + Box::new(SelectionPlanner::new( + scope, + selection, + next_fn(move |selected| Ok(ProjectionPlanner::new(plans.clone(), selected))), + )) } diff --git a/vortex-layout/src/scan/execution/projection.rs b/vortex-layout/src/scan/execution/projection.rs index b6180725c27..1732a6b4b22 100644 --- a/vortex-layout/src/scan/execution/projection.rs +++ b/vortex-layout/src/scan/execution/projection.rs @@ -1,12 +1,18 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; + +use futures::FutureExt; +use futures::channel::oneshot; +use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; use vortex_array::dtype::DType; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_err; use vortex_io::request::IoConsumer; use vortex_io::request::IoRequestId; use vortex_io::request::IoResult; @@ -23,64 +29,29 @@ use crate::scan::execution::SelectedRows; use crate::scan::execution::adaptor::AdaptorStep; use crate::scan::execution::adaptor::Piece; use crate::scan::execution::adaptor::ScanAdaptor; - -/// Hands one fixed-size work chunk to a projection morsel. -pub(super) struct ProjectionPlanner { - work: Option<(ScanPlans, SelectedRows)>, -} - -impl ProjectionPlanner { - /// Creates a planner for `selected`. - pub(super) fn new(plans: ScanPlans, selected: SelectedRows) -> Self { - Self { - work: Some((plans, selected)), - } - } -} - -impl IoConsumer for ProjectionPlanner { - fn set_io_result(&mut self, _request: IoRequestId, _result: IoResult) {} -} - -impl Planner for ProjectionPlanner { - fn state(&self) -> State { - if self.work.is_some() { - State::NeedsCompute - } else { - State::Done - } - } - - fn compute(&mut self) -> VortexResult { - let Some((plans, selected)) = self.work.take() else { - return Ok(PlannerOutput::Done); - }; - if selected.mask.all_false() { - return Ok(PlannerOutput::Done); - } - Ok(PlannerOutput::Morsel( - selected.scope.clone(), - Box::new(ProjectionMorsel::new(plans, selected)), - )) - } -} +use crate::scan::execution::children::Children; +use crate::scan::execution::children::splits::parallel_splits; /// Runs the projection plan over a split's selected rows and returns them as one batch, in row /// order. A selection that projects to no rows finishes without a batch. -struct ProjectionMorsel { +pub(super) struct ProjectionPlanner { plans: Option, dtype: DType, selected: SelectedRows, adaptor: Option, pieces: Vec, done: bool, + children: Option>, + pending: Option<(u64, oneshot::Receiver>)>, } -impl ProjectionMorsel { - /// Creates a morsel projecting `selected`. - fn new(plans: ScanPlans, selected: SelectedRows) -> Self { +impl ProjectionPlanner { + /// Creates a planner projecting `selected`. + pub(super) fn new(plans: ScanPlans, selected: SelectedRows) -> Self { Self { dtype: plans.projection.dtype().clone(), + children: plans.children.clone(), + pending: None, plans: Some(plans), selected, adaptor: None, @@ -89,7 +60,7 @@ impl ProjectionMorsel { } } - fn finish(&mut self) -> VortexResult { + fn finish(&mut self) -> VortexResult { self.done = true; self.pieces.sort_by_key(|piece| piece.rows.start); let mut arrays: Vec<_> = std::mem::take(&mut self.pieces) @@ -98,15 +69,18 @@ impl ProjectionMorsel { .filter(|array| !array.is_empty()) .collect(); let array = match arrays.len() { - 0 => return Ok(MorselOutput::Done), + 0 => return Ok(PlannerOutput::Done), 1 => arrays.remove(0), _ => ChunkedArray::try_new(arrays, self.dtype.clone())?.into_array(), }; - Ok(MorselOutput::Batch(array)) + Ok(PlannerOutput::Morsel( + self.selected.scope.clone(), + Box::new(ProjectedBatch(Some(array))), + )) } } -impl IoConsumer for ProjectionMorsel { +impl IoConsumer for ProjectionPlanner { fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { if let Some(adaptor) = self.adaptor.as_mut() { adaptor.set_io_result(request, result); @@ -114,10 +88,11 @@ impl IoConsumer for ProjectionMorsel { } } -impl Morsel for ProjectionMorsel { +impl Planner for ProjectionPlanner { fn state(&self) -> State { match &self.adaptor { _ if self.done => State::Done, + _ if self.pending.is_some() => State::NeedsCompute, None => State::NeedsCompute, Some(adaptor) => match adaptor.state() { State::Done => State::NeedsCompute, @@ -126,9 +101,19 @@ impl Morsel for ProjectionMorsel { } } - fn compute(&mut self) -> VortexResult { + fn compute(&mut self) -> VortexResult { if self.done { - vortex_bail!("ProjectionMorsel: compute called after Done"); + vortex_bail!("ProjectionPlanner: compute called after Done"); + } + if let Some((id, mut result)) = self.pending.take() { + let arrays = result + .try_recv() + .map_err(|_| vortex_err!("Child projection was cancelled"))? + .ok_or_else(|| vortex_err!("Child projection resumed before completion"))?; + self.adaptor + .as_mut() + .ok_or_else(|| vortex_err!("Child completed before projection started"))? + .complete_child(id, Some(arrays))?; } let Some(adaptor) = self.adaptor.as_mut() else { let plans = self @@ -144,20 +129,67 @@ impl Morsel for ProjectionMorsel { }], )? .with_row_offset(plans.row_offset) - .with_value_cache(plans.values); + .with_value_cache(plans.values) + .with_child_planning(self.children.is_some()); self.adaptor = Some(ScanAdaptor::new(scan, plans.locations)); - return Ok(MorselOutput::Continue); + return Ok(PlannerOutput::Continue); }; if adaptor.state() == State::Done { return self.finish(); } Ok(match adaptor.compute()? { - AdaptorStep::Yield => MorselOutput::Continue, - AdaptorStep::NeedsIO(batch) => MorselOutput::NeedsIO(batch), + AdaptorStep::Yield => PlannerOutput::Continue, + AdaptorStep::NeedsIO(batch) => PlannerOutput::NeedsIO(batch), + AdaptorStep::Child(child) => { + let Some(work) = parallel_splits(&child.plan, &child.splits)? else { + adaptor.complete_child(child.id, None)?; + return Ok(PlannerOutput::Continue); + }; + + let children = Arc::clone( + self.children + .as_ref() + .ok_or_else(|| vortex_err!("Child projection has no scheduler"))?, + ); + let (send, recv) = oneshot::channel(); + self.pending = Some((child.id, recv)); + PlannerOutput::Await( + async move { + let arrays = children.execute(child.plan, work, child.values).await?; + send.send(arrays) + .map_err(|_| vortex_err!("Projection planner was cancelled")) + } + .boxed(), + ) + } AdaptorStep::Piece(piece) => { self.pieces.push(piece); - MorselOutput::Continue + PlannerOutput::Continue } }) } } + +/// Planning resolves the projection; this morsel only publishes the assembled batch. +struct ProjectedBatch(Option); + +impl IoConsumer for ProjectedBatch { + fn set_io_result(&mut self, _request: IoRequestId, _result: IoResult) {} +} + +impl Morsel for ProjectedBatch { + fn state(&self) -> State { + if self.0.is_some() { + State::NeedsCompute + } else { + State::Done + } + } + + fn compute(&mut self) -> VortexResult { + Ok(match self.0.take() { + Some(array) => MorselOutput::Batch(array), + None => MorselOutput::Done, + }) + } +} diff --git a/vortex-layout/src/scan/v2/io.rs b/vortex-layout/src/scan/v2/io.rs index ce565846d7e..2c4ee9c121f 100644 --- a/vortex-layout/src/scan/v2/io.rs +++ b/vortex-layout/src/scan/v2/io.rs @@ -34,25 +34,55 @@ use crate::segments::SegmentFuture; use crate::segments::SegmentId; use crate::segments::SegmentSource; -/// Byte ranges of a file's segments, keyed by `(offset, length)`, back to their segment ids. -pub(super) type SegmentRanges = Arc>; - -pub(super) fn segment_ranges(locations: &[SegmentLocation]) -> SegmentRanges { - let ranges = locations.iter().enumerate().map(|(id, location)| { - let id = SegmentId::from(u32::try_from(id).unwrap_or(u32::MAX)); - ((location.offset, location.length as usize), id) - }); - Arc::new(ranges.collect()) +/// File segments are normally ordered, so range lookup can reuse the location table directly. +/// Custom segment sources can supply arbitrary order and retain indexed lookup. +pub(super) enum SegmentRanges { + Sorted(Arc<[SegmentLocation]>), + Indexed(HashMap<(u64, usize), SegmentId>), +} + +pub(super) fn segment_ranges(locations: Arc<[SegmentLocation]>) -> Arc { + if locations.is_sorted_by_key(|location| (location.offset, location.length)) { + return Arc::new(SegmentRanges::Sorted(locations)); + } + Arc::new(SegmentRanges::Indexed( + locations + .iter() + .enumerate() + .map(|(id, location)| { + let id = SegmentId::from(u32::try_from(id).unwrap_or(u32::MAX)); + ((location.offset, location.length as usize), id) + }) + .collect(), + )) +} + +impl SegmentRanges { + fn get(&self, offset: u64, len: usize) -> Option { + match self { + Self::Sorted(locations) => { + let end = locations.partition_point(|location| { + (location.offset, location.length as usize) <= (offset, len) + }); + let index = end.checked_sub(1)?; + let location = locations[index]; + // Match the indexed path's last-entry rule for duplicate ranges. + (location.offset == offset && location.length as usize == len) + .then(|| SegmentId::from(u32::try_from(index).unwrap_or(u32::MAX))) + } + Self::Indexed(ranges) => ranges.get(&(offset, len)).copied(), + } + } } /// Serves splits from any [`SegmentSource`], by segment id. pub(super) struct SegmentScanIo { segments: Arc, - ranges: SegmentRanges, + ranges: Arc, } impl SegmentScanIo { - pub(super) fn new(segments: Arc, ranges: SegmentRanges) -> Self { + pub(super) fn new(segments: Arc, ranges: Arc) -> Self { Self { segments, ranges } } } @@ -77,7 +107,7 @@ impl IoService for SegmentScanIo { /// fetch consumes it, so a source that coalesces requests can fold it into a nearby read. pub(super) struct SegmentIoSource { segments: Arc, - ranges: SegmentRanges, + ranges: Arc, reads: Mutex>>, ready: Mutex>, /// Reads started by prefetches, consumed by the next fetch of their segment. @@ -90,7 +120,7 @@ pub(super) struct SegmentIoSource { type Prefetched = Shared>>>; impl SegmentIoSource { - pub(super) fn new(segments: Arc, ranges: SegmentRanges) -> Self { + pub(super) fn new(segments: Arc, ranges: Arc) -> Self { Self { segments, ranges, @@ -114,9 +144,9 @@ impl IoSource for SegmentIoSource { request.target ); }; - let segment = *self + let segment = self .ranges - .get(&(offset, len)) + .get(offset, len) .ok_or_else(|| vortex_err!("No segment at bytes {offset}+{len}"))?; if request.intent == IoIntent::Announce { self.announced @@ -234,6 +264,7 @@ mod tests { use futures::channel::oneshot; use futures::future; use parking_lot::Mutex; + use rstest::rstest; use vortex_array::buffer::BufferHandle; use vortex_buffer::Alignment; use vortex_buffer::ByteBuffer; @@ -247,12 +278,47 @@ mod tests { use vortex_io::request::IoTarget; use super::SegmentIoSource; + use super::SegmentRanges; use super::segment_ranges; use crate::scan::execution::SegmentLocation; use crate::segments::SegmentFuture; use crate::segments::SegmentId; use crate::segments::SegmentSource; + #[rstest] + #[case::sorted(vec![(0, 0), (0, 4), (8, 4), (u64::MAX, 0)], true, vec![Some(0), Some(1), Some(2), Some(3)])] + #[case::unordered(vec![(8, 4), (0, 4), (0, 0), (u64::MAX, 0)], false, vec![Some(2), Some(1), Some(0), Some(3)])] + #[case::duplicates(vec![(0, 0), (0, 0), (0, 4), (8, 4), (8, 4)], true, vec![Some(1), Some(2), Some(4), None])] + #[case::empty(vec![], true, vec![None, None, None, None])] + fn range_lookup_preserves_segment_ids( + #[case] entries: Vec<(u64, u32)>, + #[case] sorted: bool, + #[case] expected: Vec>, + ) { + let locations: Arc<[_]> = entries + .into_iter() + .map(|(offset, length)| SegmentLocation { + offset, + length, + alignment: Alignment::none(), + }) + .collect(); + let ranges = segment_ranges(Arc::clone(&locations)); + assert_eq!(matches!(ranges.as_ref(), SegmentRanges::Sorted(_)), sorted); + if let SegmentRanges::Sorted(reused) = ranges.as_ref() { + assert!(Arc::ptr_eq(reused, &locations)); + } + for ((offset, len), expected) in [(0, 0), (0, 4), (8, 4), (u64::MAX, 0)] + .into_iter() + .zip(expected) + { + assert_eq!(ranges.get(offset, len), expected.map(SegmentId::from)); + } + for (offset, len) in [(0, 1), (1, 4), (8, 0), (8, 5)] { + assert_eq!(ranges.get(offset, len), None); + } + } + /// Serves four bytes per segment and counts the reads it is asked for. #[derive(Default)] struct CountingSegments(AtomicUsize); @@ -301,7 +367,7 @@ mod tests { length: 4, alignment: Alignment::none(), }); - let io = SegmentIoSource::new(segments, segment_ranges(&locations)); + let io = SegmentIoSource::new(segments, segment_ranges(Arc::from(locations))); let mut pending = request(IoIntent::Fetch, 1); pending.target = IoTarget::range(4, 4); io.submit(IoOwnerId(0), vec![pending, request(IoIntent::Fetch, 0)])?; @@ -334,7 +400,10 @@ mod tests { length: 4, alignment: Alignment::none(), }; - let io = SegmentIoSource::new(Arc::clone(&segments) as _, segment_ranges(&[location])); + let io = SegmentIoSource::new( + Arc::clone(&segments) as _, + segment_ranges(Arc::from([location])), + ); io.submit(IoOwnerId(0), vec![request(IoIntent::Prefetch, 0)])?; assert_eq!(segments.0.load(Ordering::Relaxed), 1); @@ -362,7 +431,10 @@ mod tests { length: 4, alignment: Alignment::none(), }; - let io = SegmentIoSource::new(Arc::clone(&segments) as _, segment_ranges(&[location])); + let io = SegmentIoSource::new( + Arc::clone(&segments) as _, + segment_ranges(Arc::from([location])), + ); io.submit( IoOwnerId(0), vec![ diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 46cbde987fc..b488b3bca74 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -11,8 +11,10 @@ //! Import this module's [`ScanBuilder`] to select V2 explicitly. It accepts a [`ScanFile`] //! directly and supports projection, row ranges, and selections. Unsupported layouts return an //! error; filters and limits are not part of this builder's API. Work is partitioned into fixed -//! chunks of up to 65,536 rows, independently of the stored layout's chunk boundaries. -//! Full-file projections use smaller 8,192-row multiples when serialized bytes per row are high. +//! chunks independently of the stored layout's chunk boundaries. Dense scans use up to 65,536 +//! rows per chunk, in 8,192-row multiples, targeting 8 MiB of serialized data. Sparse identity +//! projections can use larger aligned chunks within that estimate, targeting four tasks per +//! worker to reduce repeated segment reads. mod io; mod split; @@ -35,8 +37,8 @@ use vortex_array::ArrayRef; use vortex_array::VortexSessionExecute; use vortex_array::dtype::DType; use vortex_array::expr::BoundExpression; +use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; use vortex_io::request::IoService; use vortex_io::session::RuntimeSessionExt; use vortex_scan::planning::planner::WorkScope; @@ -46,18 +48,13 @@ use vortex_session::VortexSession; use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutRef; -use crate::layouts::chunked::Chunked; -use crate::layouts::dict::Dict; -use crate::layouts::flat::Flat; -use crate::layouts::struct_::Struct; -use crate::layouts::zoned::LegacyStats; -use crate::layouts::zoned::Zoned; use crate::plan::physical::lower; use crate::plan::physical::optimize; use crate::plan::physical::plan_row_idx_expression; use crate::scan::arrow::to_record_batch; use crate::scan::execution::ScanPlans; pub use crate::scan::execution::SegmentLocation; +use crate::scan::execution::child_executor; use crate::scan::execution::plan_projection; use crate::scan::v2::io::SegmentScanIo; use crate::scan::v2::io::segment_ranges; @@ -69,7 +66,11 @@ const MAX_ROWS_PER_CHUNK: usize = 8 * ROW_BLOCK_SIZE; const TARGET_BYTES_PER_CHUNK: u64 = 8 * 1024 * 1024; /// Limit estimated bytes per task without changing the writer-aligned row granularity. -fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize { +fn rows_per_chunk( + row_count: u64, + locations: &[SegmentLocation], + sparse_target: Option, +) -> usize { if row_count <= MAX_ROWS_PER_CHUNK as u64 { return MAX_ROWS_PER_CHUNK; } @@ -79,7 +80,15 @@ fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize { .sum(); let bytes_per_row = bytes.div_ceil(row_count).max(1); let blocks = TARGET_BYTES_PER_CHUNK / bytes_per_row / ROW_BLOCK_SIZE as u64; - blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE + let max_blocks = sparse_target.map_or(8, |rows| { + // Grow sparse tasks only within the byte target, keeping power-of-two boundaries + // aligned with the original work chunks and enough tasks to occupy the workers. + let blocks = (rows / ROW_BLOCK_SIZE as u64).min(blocks).max(8); + 1_u64 << blocks.ilog2() + }); + usize::try_from(blocks.clamp(1, max_blocks)) + .vortex_expect("the byte target bounds the block count to at most 1024") + * ROW_BLOCK_SIZE } /// Fixed boundaries in the file's row domain, clipped to the requested range. @@ -109,8 +118,10 @@ pub struct ScanFile { /// Builds a projection scan on the planner-driven V2 executor. /// -/// This builder reads the file's layout and segments directly. Unsupported layouts return an -/// error when the stream is polled. Filters and row limits are not supported. +/// This builder reads the file's layout and segments directly. Layouts are checked as their +/// plans are lowered, so unsupported layouts outside the requested range can remain unread. +/// Unsupported plans return an error when the stream is polled. Filters and row limits are +/// not supported. pub struct ScanBuilder { session: VortexSession, file: ScanFile, @@ -237,34 +248,6 @@ impl ScanBuilder { } } -/// Whether every required layout lowers to plans with a pipeline implementation. -/// -/// List layouts lower to a list assembly, which cannot compile before its offsets are read. -fn supports_layout(layout: &LayoutRef, identity: bool) -> VortexResult { - if !(layout.is::() - || layout.is::() - || layout.is::() - || layout.is::() - || layout.is::() - || layout.is::()) - { - return Ok(false); - } - // An identity projection cannot use zone statistics. Leave their layouts uninitialized. - if identity && (layout.is::() || layout.is::()) { - return match layout.slot(0)? { - Some(data) => supports_layout(&data, identity), - None => Ok(false), - }; - } - for child in layout.children()? { - if !supports_layout(&child, identity)? { - return Ok(false); - } - } - Ok(true) -} - /// A scan planned over its file, ready to run its splits. struct PreparedScan { session: VortexSession, @@ -284,10 +267,6 @@ impl PreparedScan { /// fixed-size row chunks. fn try_new(builder: ScanBuilder) -> VortexResult { let file = builder.file; - vortex_ensure!( - supports_layout(&file.layout, builder.projection.is_root())?, - "V2 scans do not support this file's layout" - ); let root = lower(&file.layout)?; let row_count = root.row_count(); // An identity projection needs no rewrites. Keep the lowered children lazy, including @@ -299,15 +278,29 @@ impl PreparedScan { optimize(plan_row_idx_expression(builder.projection, root)?)? }; let chunk_rows = if identity { - rows_per_chunk(row_count, &file.locations) + let rows = builder.row_range.clone().unwrap_or(0..row_count); + let selected = builder.selection.row_count_in_range(&rows); + let span = rows.end.saturating_sub(rows.start); + let sparse_target = (selected <= span / 16).then(|| { + // Tiny selections cannot occupy the full worker budget. Larger aligned tasks + // let nearby selected rows share their stored segments. + let tasks = (4 * get_available_parallelism().unwrap_or(1)) as u64; + span / tasks.min(selected).max(1) + }); + rows_per_chunk(row_count, &file.locations, sparse_target) } else { MAX_ROWS_PER_CHUNK }; let splits = row_chunks(builder.row_range.unwrap_or(0..row_count), chunk_rows); - let io = Arc::new(SegmentScanIo::new( + let io: Arc = Arc::new(SegmentScanIo::new( file.segments, - segment_ranges(&file.locations), + segment_ranges(Arc::clone(&file.locations)), )); + let children = child_executor( + builder.session.clone(), + Arc::clone(&file.locations), + Arc::clone(&io), + ); Ok(Self { session: builder.session.clone(), plans: ScanPlans { @@ -316,6 +309,7 @@ impl PreparedScan { projection, row_offset: builder.row_offset, values: Default::default(), + children, }, io, splits, @@ -328,22 +322,22 @@ impl PreparedScan { /// The splits with selected rows, each ready to be admitted to a planning run. fn split_plans(&self) -> Vec { + let selection = Arc::new(self.selection.clone()); self.splits .iter() .filter_map(|rows| { - let row_mask = self.selection.row_mask(rows); - if row_mask.mask().all_false() { + if self.selection.row_count_in_range(rows) == 0 { return None; } let scope = WorkScope { file_ordinal: 0, - rows: row_mask.row_range(), + rows: rows.clone(), }; Some(SplitPlan { root: plan_projection( self.plans.clone(), scope.clone(), - row_mask.mask().clone(), + Arc::clone(&selection), ), scope, }) diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index a910172f87b..43bb6392757 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -16,6 +16,7 @@ use vortex_array::ArrayContext; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; @@ -36,6 +37,7 @@ use vortex_error::VortexResult; use vortex_io::session::RuntimeSessionExt; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; +use vortex_session::registry::CachedId; use super::MAX_ROWS_PER_CHUNK; use super::PreparedScan; @@ -45,9 +47,13 @@ use super::row_chunks; use super::rows_per_chunk; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::children::OwnedLayoutChildren; +use crate::layouts::chunked::ChunkedLayout; use crate::layouts::chunked::writer::ChunkedLayoutStrategy; use crate::layouts::dict::DictLayout; use crate::layouts::flat::writer::FlatLayoutStrategy; +use crate::layouts::foreign::new_foreign_layout; +use crate::layouts::list::writer::ListLayoutStrategy; use crate::layouts::row_idx::row_idx; use crate::layouts::table::TableStrategy; use crate::layouts::zoned::writer::ZonedLayoutOptions; @@ -66,6 +72,9 @@ use crate::test::new_session; const CHUNK_ROWS: usize = 1000; const CHUNKS: usize = 4; +mod child_work; +static UNSUPPORTED_ID: CachedId = CachedId::new("test.unsupported"); + /// A table chunk of `CHUNK_ROWS` rows starting at row `start`: an integer, a nullable string and a /// list column. fn table_chunk(start: usize) -> VortexResult { @@ -314,19 +323,29 @@ fn fixed_row_chunks(#[case] rows: Range, #[case] expected: Vec>) } #[rstest] -#[case::empty(0, 0, 65_536)] -#[case::single_chunk(1000, 100_000_000, 65_536)] -#[case::no_bytes(100_000, 0, 65_536)] -#[case::small_rows(1_000_000, 64_000_000, 65_536)] -#[case::large_rows(1_000_000, 384_000_000, 16_384)] -#[case::minimum(1_000_000, 2_000_000_000, 8192)] -fn chunk_byte_budget(#[case] rows: u64, #[case] bytes: u32, #[case] expected: usize) { +#[case::empty(0, 0, None, 65_536)] +#[case::single_chunk(1000, 100_000_000, None, 65_536)] +#[case::no_bytes(100_000, 0, None, 65_536)] +#[case::small_rows(1_000_000, 64_000_000, None, 65_536)] +#[case::large_rows(1_000_000, 384_000_000, None, 16_384)] +#[case::minimum(1_000_000, 2_000_000_000, None, 8192)] +#[case::sparse_small_rows(3_339_715, 49_808_116, Some(834_928), 524_288)] +#[case::sparse_medium_rows(1_000_000, 35_640_884, Some(250_000), 131_072)] +#[case::sparse_more_workers(3_339_715, 49_808_116, Some(27_830), 65_536)] +#[case::sparse_lists(1_000_000, 86_983_648, Some(250_000), 65_536)] +#[case::sparse_wide_rows(1_000_000, 3_354_532_688, Some(250_000), 8192)] +fn chunk_byte_budget( + #[case] rows: u64, + #[case] bytes: u32, + #[case] sparse_target: Option, + #[case] expected: usize, +) { let locations = [SegmentLocation { offset: 0, length: bytes, alignment: Alignment::none(), }]; - assert_eq!(rows_per_chunk(rows, &locations), expected); + assert_eq!(rows_per_chunk(rows, &locations, sparse_target), expected); } #[test] @@ -349,6 +368,46 @@ impl SegmentSource for CountingSegments { } } +#[tokio::test] +async fn tiny_selection_shares_a_stored_segment() -> VortexResult<()> { + let session = new_session().with_tokio(); + let input = PrimitiveArray::from_iter(0_i32..1_048_576).into_array(); + let dtype = input.dtype().clone(); + let segments = Arc::new(TestSegments::default()); + let (mut sequence, eof) = SequenceId::root().split(); + let layout = FlatLayoutStrategy::default() + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + SequentialStreamAdapter::new( + dtype.clone(), + stream::iter([Ok((sequence.advance(), input))]), + ) + .sendable(), + eof, + &session, + ) + .await?; + let counted = Arc::new(CountingSegments { + inner: segments, + reads: AtomicUsize::new(0), + }); + let segments: Arc = Arc::::clone(&counted); + // These rows share one stored segment but fall across the old worker-sized task boundary, + // even on one core. Selection-aware sizing lets one projection share its read. + let scan = builder(&session, &segments, &layout)?.with_row_indices( + StrictSortedBuffer::try_new(Buffer::from(vec![10, 350_000]))?, + ); + let actual = collect(dtype, scan.into_stream()?).await?; + assert_arrays_eq!( + actual, + Buffer::from(vec![10_i32, 350_000]).into_array(), + &mut session.create_execution_ctx() + ); + assert_eq!(counted.reads.load(Ordering::Relaxed), 1); + Ok(()) +} + #[tokio::test] async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { let session = new_session().with_tokio(); @@ -416,19 +475,94 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { Ok(()) } -/// V2 rejects list layouts: their element ranges are only known once their offsets are read. #[tokio::test] -async fn list_layouts_return_an_error() -> VortexResult<()> { +async fn list_layouts_match_v1() -> VortexResult<()> { let session = new_session().with_tokio(); let strategy = TableStrategy::new(chunked_flat(), chunked_flat()).with_list_layout(); let (segments, layout) = write(&session, &strategy).await?; - let builder = builder(&session, &segments, &layout)?; - let error = builder - .into_stream()? - .try_next() - .await - .expect_err("V2 must reject an unsupported layout"); - assert!(error.to_string().contains("V2 scans do not support")); + assert_same_as_v1(&session, || builder(&session, &segments, &layout)).await +} + +#[rstest] +#[case::empty(vec![])] +#[case::null_only(vec![3, 6, 9])] +#[case::contiguous(vec![1, 2])] +#[case::contiguous_with_empty(vec![7, 8])] +#[case::mixed(vec![0, 1, 2, 3, 8, 31, 63])] +#[case::all((0..64).collect())] +#[tokio::test] +async fn selected_list_values_match_input( + #[case] rows: Vec, + #[values(false, true)] nested: bool, + #[values(false, true)] nullable: bool, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let mut offsets = vec![0_u32]; + let mut total = 0; + for row in 0..64 { + total += row % 4; + offsets.push(total); + } + let inner_strategy = ListLayoutStrategy::default() + .with_elements(chunked_flat()) + .with_offsets(chunked_flat()) + .with_validity(chunked_flat()); + let (elements, element_strategy) = if nested { + let elements = ListArray::try_new( + Buffer::from_iter(0..2 * total).into_array(), + Buffer::from_iter((0..=total).map(|row| 2 * row)).into_array(), + Validity::Array(BoolArray::from_iter((0..total).map(|row| row % 3 != 0)).into_array()), + )? + .into_array(); + ( + elements, + Arc::new(inner_strategy) as Arc, + ) + } else { + (Buffer::from_iter(0..total).into_array(), chunked_flat()) + }; + let validity = if nullable { + Validity::Array(BoolArray::from_iter((0..64).map(|row| row % 3 != 0)).into_array()) + } else { + Validity::NonNullable + }; + let input = + ListArray::try_new(elements, Buffer::from(offsets).into_array(), validity)?.into_array(); + let dtype = input.dtype().clone(); + let expected = input.take(Buffer::from(rows.clone()).into_array())?; + let strategy = ListLayoutStrategy::default() + .with_elements(element_strategy) + .with_offsets(chunked_flat()) + .with_validity(chunked_flat()); + let segments = Arc::new(TestSegments::default()); + let (mut id, eof) = SequenceId::root().split(); + let layout = strategy + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + SequentialStreamAdapter::new(dtype.clone(), stream::iter([Ok((id.advance(), input))])) + .sendable(), + eof, + &session, + ) + .await?; + let counted = Arc::new(CountingSegments { + inner: segments, + reads: AtomicUsize::new(0), + }); + let segments: Arc = Arc::::clone(&counted); + let empty = rows.is_empty(); + let null_only = nullable && !empty && rows.iter().all(|row| row % 3 == 0); + let scan = builder(&session, &segments, &layout)? + .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(rows))?); + let actual = collect(dtype, scan.into_stream()?).await?; + assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); + if empty { + assert_eq!(counted.reads.load(Ordering::Relaxed), 0); + } else if null_only { + // Only offsets and validity are needed, even when null lists store nonempty elements. + assert_eq!(counted.reads.load(Ordering::Relaxed), 2); + } Ok(()) } @@ -481,3 +615,72 @@ async fn work_chunks_ignore_layout_boundaries( assert_eq!(arrays[0].len(), CHUNKS * CHUNK_ROWS); Ok(()) } + +#[tokio::test] +async fn unknown_layout_returns_an_error() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, known) = write(&session, chunked_flat().as_ref()).await?; + let layout = new_foreign_layout( + *UNSUPPORTED_ID, + known.dtype().clone(), + known.row_count(), + Vec::new(), + Vec::new(), + Vec::new(), + ); + let error = builder(&session, &segments, &layout)? + .into_stream()? + .try_next() + .await + .expect_err("V2 must reject an unsupported layout"); + assert!( + error + .to_string() + .contains("No physical plan implementation for layout 'test.unsupported'") + ); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn unsupported_chunk_is_checked_when_reached( + #[values(false, true)] selected: bool, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, known) = write(&session, chunked_flat().as_ref()).await?; + let rows = known.row_count(); + let unknown = new_foreign_layout( + *UNSUPPORTED_ID, + known.dtype().clone(), + rows, + Vec::new(), + Vec::new(), + Vec::new(), + ); + let layout = ChunkedLayout::new( + 2 * rows, + known.dtype().clone(), + OwnedLayoutChildren::layout_children(vec![Arc::clone(&known), unknown]), + ) + .into_layout(); + let range = if selected { rows..2 * rows } else { 0..rows }; + let result = collect( + known.dtype().clone(), + builder(&session, &segments, &layout)? + .with_row_range(range) + .into_stream()?, + ) + .await; + if selected { + let error = result.expect_err("a reached unsupported chunk must fail"); + assert!(error.to_string().contains("test.unsupported")); + } else { + let expected = collect( + known.dtype().clone(), + builder(&session, &segments, &known)?.into_stream()?, + ) + .await?; + assert_arrays_eq!(result?, expected, &mut session.create_execution_ctx()); + } + Ok(()) +} diff --git a/vortex-layout/src/scan/v2/tests/child_work.rs b/vortex-layout/src/scan/v2/tests/child_work.rs new file mode 100644 index 00000000000..d26a2d97c8c --- /dev/null +++ b/vortex-layout/src/scan/v2/tests/child_work.rs @@ -0,0 +1,101 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_utils::parallelism::get_available_parallelism; + +use super::*; +use crate::layouts::list::ListLayout; +async fn store_array( + session: &VortexSession, + segments: &Arc, + array: ArrayRef, + chunk_rows: usize, +) -> VortexResult { + let dtype = array.dtype().clone(); + let (mut id, eof) = SequenceId::root().split(); + let chunks = (0..array.len()) + .step_by(chunk_rows) + .map(|start| { + Ok(( + id.advance(), + array.slice(start..(start + chunk_rows).min(array.len()))?, + )) + }) + .collect::>(); + chunked_flat() + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(segments), + SequentialStreamAdapter::new(dtype, stream::iter(chunks)).sendable(), + eof, + session, + ) + .await +} + +#[rstest] +#[case::empty(vec![], false)] +#[case::empty_rows(vec![0, 2], false)] +#[case::null_row(vec![3], false)] +#[case::large_row(vec![1], true)] +#[case::mixed(vec![0, 1, 3], true)] +#[case::all(vec![0, 1, 2, 3], true)] +#[tokio::test] +async fn large_list_children_match_input_and_v1( + #[case] selected: Vec, + #[case] parallel: bool, + #[values(false, true)] nested: bool, + #[values(8192, 262144)] chunk_rows: usize, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let segments = Arc::new(TestSegments::default()); + let primitive = PrimitiveArray::from_option_iter( + (0_i64..262144).map(|value| (value % 7 != 0).then_some(value)), + ) + .into_array(); + let mut elements = primitive.clone(); + let mut layout = store_array(&session, &segments, primitive, chunk_rows).await?; + if nested { + let offsets = Buffer::from_iter((0_u64..=2048).map(|row| row * 128)).into_array(); + let valid = BoolArray::from_iter((0..2048).map(|row| row % 11 != 0)).into_array(); + elements = ListArray::try_new(elements, offsets.clone(), Validity::Array(valid.clone()))? + .into_array(); + layout = ListLayout::new( + elements.dtype().clone(), + layout, + store_array(&session, &segments, offsets, 2049).await?, + Some(store_array(&session, &segments, valid, 2048).await?), + ) + .into_layout(); + } + let half = u64::try_from(elements.len())? / 2; + let offsets = Buffer::from(vec![0_u64, 0, half, half, 2 * half]).into_array(); + let valid = BoolArray::from_iter([true, true, true, false]).into_array(); + let input = + ListArray::try_new(elements, offsets.clone(), Validity::Array(valid.clone()))?.into_array(); + let layout = ListLayout::new( + input.dtype().clone(), + layout, + store_array(&session, &segments, offsets, 5).await?, + Some(store_array(&session, &segments, valid, 4).await?), + ) + .into_layout(); + let segments: Arc = segments; + let make = || { + Ok(builder(&session, &segments, &layout)? + .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(selected.clone()))?)) + }; + let scan = PreparedScan::try_new(make()?)?; + let children = scan.plans.children.clone(); + let actual = collect(input.dtype().clone(), scan.execute_stream()?).await?; + assert_arrays_eq!( + actual, + input.take(Buffer::from(selected.clone()).into_array())?, + &mut session.create_execution_ctx() + ); + assert_eq!( + children.is_some_and(|children| children.admitted() > 1), + parallel && chunk_rows == 8192 && get_available_parallelism().unwrap_or(1) > 1 + ); + assert_same_as_v1(&session, make).await +} diff --git a/vortex-scan/Cargo.toml b/vortex-scan/Cargo.toml index d12c207b049..63352a21566 100644 --- a/vortex-scan/Cargo.toml +++ b/vortex-scan/Cargo.toml @@ -26,6 +26,7 @@ vortex-utils = { workspace = true } async-trait = { workspace = true } futures = { workspace = true } roaring = { workspace = true } +smallvec = { workspace = true } tracing = { workspace = true } [dev-dependencies] diff --git a/vortex-scan/src/planning/driver/mod.rs b/vortex-scan/src/planning/driver/mod.rs index 8cb78099272..024748774fc 100644 --- a/vortex-scan/src/planning/driver/mod.rs +++ b/vortex-scan/src/planning/driver/mod.rs @@ -7,12 +7,12 @@ //! batches to its caller as they are produced. Runnable items are prioritised in DFS preorder, //! with roots ordered by admission and children by emission. Each visit takes at most one CPU //! step. A parent's continuation follows its emitted child's entire subtree. Waiting work is -//! skipped and regains its priority when IO makes it runnable. Published IO is registered +//! skipped and regains its priority when IO or an asynchronous task makes it runnable. Published IO is registered //! immediately and state is inspected again before parking. Output follows emission order, //! not global row order. An earlier runnable branch can starve later branches. //! //! A run never blocks: [`Run::advance`] returns [`Progress::Waiting`] when every live item waits -//! for IO, and the caller decides how to wait, by awaiting [`Run::poll_completion`] or by handing +//! for IO or tasks, and the caller decides how to wait, by awaiting [`Run::poll_completion`] or by handing //! a completion it took from a source to [`Run::complete`]. [`Driver::run`] is the blocking //! convenience for one root. @@ -21,7 +21,12 @@ use std::collections::VecDeque; use std::sync::Arc; use std::task::Context; use std::task::Poll; +use std::task::Waker; +use futures::FutureExt; +use futures::future::BoxFuture; +use smallvec::SmallVec; +use smallvec::smallvec; use vortex_array::ArrayRef; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -83,6 +88,9 @@ enum Output { Morsel(MorselOutput), } +// Root, selection, projection, and continuation fit inline. Deeper compositions can spill. +type Priority = SmallVec<[u64; 4]>; + /// One live planner or morsel. Boxed in the run queue and the parked map, so scheduling moves a /// pointer. struct Work { @@ -91,7 +99,7 @@ struct Work { scope: WorkScope, item: Item, /// Stable child-ordinal path, beginning with the root's admission ordinal. - path: Vec, + path: Priority, /// The parent continuation sorts after all previously emitted child subtrees. next_child: u64, /// Published `Fetch` requests not yet delivered, with the target each must be answered with. @@ -99,7 +107,7 @@ struct Work { } impl Work { - fn priority(&self) -> Vec { + fn priority(&self) -> Priority { let mut priority = self.path.clone(); priority.push(self.next_child); priority @@ -107,7 +115,7 @@ impl Work { // Child k takes the continuation's old key as its path. Its subtree then sorts before // the parent's new continuation at k + 1, even if descendants are discovered after IO. - fn child_path(&mut self) -> Vec { + fn child_path(&mut self) -> Priority { let path = self.priority(); self.next_child += 1; path @@ -146,7 +154,7 @@ impl Work { /// Keeps scheduling order separate from lookup for deliveries to runnable work. #[derive(Default)] struct ReadyQueue { - priorities: BTreeMap, IoOwnerId>, + priorities: BTreeMap, work: HashMap>, } @@ -191,6 +199,8 @@ pub struct Run { next_owner: u64, queue: ReadyQueue, parked: HashMap>, + /// Allocated only for owners that publish asynchronous work. + tasks: HashMap>>, /// Batches and finished roots not yet returned by [`advance`](Self::advance). events: VecDeque, /// Visits left before the sources are polled for completions again. @@ -233,7 +243,7 @@ impl Run { let id = RootId(self.next_root); self.next_root += 1; self.roots.insert(id, Root { io, live: 0 }); - self.spawn(id, scope, Item::Planner(root), vec![id.0]); + self.spawn(id, scope, Item::Planner(root), smallvec![id.0]); id } @@ -245,6 +255,9 @@ impl Run { }; self.queue.cancel(root); self.parked.retain(|_, work| work.root != root); + self.tasks.retain(|owner, _| { + self.queue.work.contains_key(owner) || self.parked.contains_key(owner) + }); cancelled.io.clear(); } @@ -297,10 +310,10 @@ impl Run { let work = parked.get_mut(); work.deliver(request, result)?; // An item that still waits for other requests stays parked. - if work.state() != State::Waiting { + if work.state() != State::Waiting && !self.tasks.contains_key(&owner) { let work = parked.remove(); self.queue.insert(work); - } else if work.outstanding.is_empty() { + } else if work.outstanding.is_empty() && !self.tasks.contains_key(&owner) { vortex_bail!( "work for scope {:?} waits with no outstanding fetch", work.scope @@ -332,6 +345,23 @@ impl Run { } let mut delivered = false; for io in sources { + // A root may now wait solely for child work. IO sources need not accept a poll + // with no fetches outstanding. + if !self.tasks.is_empty() + && !self + .parked + .values() + .chain(self.queue.work.values()) + .any(|work| { + !work.outstanding.is_empty() + && self + .roots + .get(&work.root) + .is_some_and(|root| Arc::ptr_eq(&root.io, &io)) + }) + { + continue; + } match io.poll_completion(cx) { Poll::Ready(Ok(completion)) => { if let Err(error) = self.complete(completion) { @@ -343,6 +373,10 @@ impl Run { Poll::Pending => {} } } + match self.take_tasks(cx) { + Ok(completed) => delivered |= completed, + Err(error) => return Poll::Ready(Err(error)), + } if delivered { Poll::Ready(Ok(())) } else { @@ -352,6 +386,7 @@ impl Run { /// Delivers the completions every root's session has ready, without blocking. fn take_completions(&mut self) -> VortexResult<()> { + self.take_tasks(&mut Context::from_waker(Waker::noop()))?; if self.roots.len() == 1 { // The common case, without collecting the sessions. let Some(io) = self.roots.values().next().map(|root| Arc::clone(&root.io)) else { @@ -375,7 +410,35 @@ impl Run { Ok(()) } - fn spawn(&mut self, root: RootId, scope: WorkScope, item: Item, path: Vec) { + fn take_tasks(&mut self, cx: &mut Context<'_>) -> VortexResult { + if self.tasks.is_empty() { + return Ok(false); + } + let mut completed = Vec::new(); + for (&owner, future) in &mut self.tasks { + if let Poll::Ready(result) = future.poll_unpin(cx) { + result?; + completed.push(owner); + } + } + let delivered = !completed.is_empty(); + for owner in completed { + self.tasks.remove(&owner); + let work = self + .parked + .remove(&owner) + .ok_or_else(|| vortex_err!("future completed for unknown planner {owner:?}"))?; + self.settle(work)?; + } + for work in self.parked.values() { + if work.outstanding.is_empty() && !self.tasks.contains_key(&work.id) { + vortex_bail!("work waits after its last task and fetch completed"); + } + } + Ok(delivered) + } + + fn spawn(&mut self, root: RootId, scope: WorkScope, item: Item, path: Priority) { let id = IoOwnerId(self.next_owner); self.next_owner += 1; if let Some(root) = self.roots.get_mut(&root) { @@ -395,6 +458,7 @@ impl Run { /// Retires a finished item, and its root once nothing descended from it is live. fn retire(&mut self, work: &Work) { + self.tasks.remove(&work.id); let Entry::Occupied(mut root) = self.roots.entry(work.root) else { return; }; @@ -406,9 +470,9 @@ impl Run { } } - /// Parks an item that waits for IO, which is a protocol error with no fetch outstanding. + /// Parks an item that waits for IO or tasks; at least one must be outstanding. fn park(&mut self, work: Box) -> VortexResult<()> { - if work.outstanding.is_empty() { + if work.outstanding.is_empty() && !self.tasks.contains_key(&work.id) { vortex_bail!( "work for scope {:?} waits with no outstanding fetch", work.scope @@ -452,6 +516,12 @@ impl Run { self.submit(&mut work, batch)?; self.settle(work)?; } + Output::Planner(PlannerOutput::Await(future)) => { + self.tasks.insert(work.id, future); + // Await suspends this planner even if its next compute is already prepared. + // Outstanding IO still belongs to it and can be delivered while it is suspended. + self.parked.insert(work.id, work); + } Output::Planner(PlannerOutput::Planner(scope, child)) => { let path = work.child_path(); self.spawn(work.root, scope, Item::Planner(child), path); @@ -541,7 +611,10 @@ impl Driver { match run.advance()? { Progress::Batch(batch) => batches.push(batch), Progress::RootDone(_) | Progress::Idle => return Ok(batches), - Progress::Waiting => run.complete(self.io.wait()?)?, + Progress::Waiting if run.tasks.is_empty() => run.complete(self.io.wait()?)?, + Progress::Waiting => { + futures::executor::block_on(std::future::poll_fn(|cx| run.poll_completion(cx)))? + } } } } diff --git a/vortex-scan/src/planning/driver/tests.rs b/vortex-scan/src/planning/driver/tests.rs index 389ad24372d..d27fa0a5be7 100644 --- a/vortex-scan/src/planning/driver/tests.rs +++ b/vortex-scan/src/planning/driver/tests.rs @@ -22,6 +22,8 @@ use crate::planning::tests::scripted::array_of; use crate::planning::tests::scripted::ctx; use crate::planning::tests::scripted::scope; +mod tasks; + fn request(id: u32, offset: u64) -> IoRequest { IoRequest { intent: IoIntent::Fetch, @@ -141,6 +143,34 @@ fn dfs_finishes_child_subtree_before_parent_continues() -> VortexResult<()> { Ok(()) } +#[rstest] +fn nested_planner_chains_preserve_sibling_order( + #[values(1, 3, 8, 16)] depth: u64, +) -> VortexResult<()> { + let log = Log::default(); + let mut steps = vec![one_batch_morsel(scope(0..1), 1), PlannerStep::Done]; + for row in 1..=depth { + steps = vec![ + PlannerStep::Planner(scope(0..row), steps), + one_batch_morsel(scope(row..row + 1), 1), + PlannerStep::Done, + ]; + } + let root = ScriptedPlanner::boxed("root", steps, &log); + let batches = driver(&source()).run(root)?; + let scopes: Vec<_> = batches.iter().map(|batch| batch.scope.clone()).collect(); + assert_eq!( + scopes, + (0..=depth) + .map(|row| scope(row..row + 1)) + .collect::>() + ); + for batch in batches { + assert_arrays_eq!(batch.array, array_of(1), &mut ctx()); + } + Ok(()) +} + #[test] fn needs_io_delivers_every_request_in_order() -> VortexResult<()> { let log = Log::default(); diff --git a/vortex-scan/src/planning/driver/tests/tasks.rs b/vortex-scan/src/planning/driver/tests/tasks.rs new file mode 100644 index 00000000000..140a06d1a3d --- /dev/null +++ b/vortex-scan/src/planning/driver/tests/tasks.rs @@ -0,0 +1,175 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use futures::channel::oneshot; +use futures::task::ArcWake; +use futures::task::waker; +use vortex_error::vortex_err; + +use super::*; + +struct Awaiting { + tasks: VecDeque>>, + read: Option, + waiting_read: bool, + emitted: bool, +} + +impl IoConsumer for Awaiting { + fn set_io_result(&mut self, _request: IoRequestId, _result: IoResult) { + self.waiting_read = false; + } +} + +impl Planner for Awaiting { + fn state(&self) -> State { + if self.emitted { + State::Done + } else if self.tasks.is_empty() && self.read.is_none() && self.waiting_read { + State::Waiting + } else { + State::NeedsCompute + } + } + + fn compute(&mut self) -> VortexResult { + if let Some(read) = self.read.take() { + return Ok(PlannerOutput::NeedsIO(vec![read])); + } + if let Some(task) = self.tasks.pop_front() { + return Ok(PlannerOutput::Await(task)); + } + self.emitted = true; + Ok(PlannerOutput::Planner( + scope(0..3), + ScriptedPlanner::boxed( + "after task", + vec![one_batch_morsel(scope(0..3), 3), PlannerStep::Done], + &Log::default(), + ), + )) + } +} + +fn awaiting(tasks: Vec>>, read: bool) -> Box { + Box::new(Awaiting { + tasks: tasks.into(), + read: read.then(|| request(0, 0)), + waiting_read: read, + emitted: false, + }) +} + +#[derive(Default)] +struct WakeCount(AtomicUsize); + +impl ArcWake for WakeCount { + fn wake_by_ref(arc_self: &Arc) { + arc_self.0.fetch_add(1, Ordering::Relaxed); + } +} + +#[rstest] +fn completion_wakes_a_planner_and_preserves_its_continuation( + #[values(false, true)] read: bool, +) -> VortexResult<()> { + let (send, recv) = oneshot::channel::<()>(); + let task = async move { recv.await.map_err(|_| vortex_err!("cancelled")) }.boxed(); + let root = awaiting(vec![task], read); + let mut run = Run::new(); + run.admit(root, scope(0..3), source()); + assert!(matches!(run.advance()?, Progress::Waiting)); + let wakes = Arc::new(WakeCount::default()); + let waker = waker(Arc::clone(&wakes)); + if read { + let Poll::Ready(result) = run.poll_completion(&mut Context::from_waker(&waker)) else { + panic!("the canned read should complete") + }; + result?; + assert!(matches!(run.advance()?, Progress::Waiting)); + } + assert!( + run.poll_completion(&mut Context::from_waker(&waker)) + .is_pending() + ); + send.send(()).map_err(|_| vortex_err!("receiver dropped"))?; + assert!(wakes.0.load(Ordering::Relaxed) > 0); + let Progress::Batch(batch) = run.advance()? else { + panic!("missing continuation output") + }; + assert_arrays_eq!(batch.array, array_of(3), &mut ctx()); + assert!(matches!(run.advance()?, Progress::RootDone(_))); + Ok(()) +} + +#[test] +fn completed_future_keeps_waiting_for_outstanding_io() -> VortexResult<()> { + let root = awaiting(vec![async { Ok(()) }.boxed()], true); + let mut run = Run::new(); + run.admit(root, scope(0..3), source()); + assert!(matches!(run.advance()?, Progress::Waiting)); + assert!(run.tasks.is_empty()); + let Poll::Ready(result) = run.poll_completion(&mut Context::from_waker(Waker::noop())) else { + panic!("the outstanding read should complete") + }; + result?; + assert!(matches!(run.advance()?, Progress::Batch(_))); + Ok(()) +} + +#[test] +fn repeated_awaits_preserve_the_planner_continuation() -> VortexResult<()> { + let tasks = (0..3).map(|_| async { Ok(()) }.boxed()).collect(); + let root = awaiting(tasks, false); + assert_eq!(driver(&source()).run(root)?.len(), 1); + Ok(()) +} + +struct DropCount(Arc); + +impl Drop for DropCount { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::Relaxed); + } +} + +#[rstest] +fn cancellation_drops_pending_child_work( + #[values(false, true)] cancel_root: bool, +) -> VortexResult<()> { + let drops = Arc::new(AtomicUsize::new(0)); + let guard = DropCount(Arc::clone(&drops)); + let root = awaiting( + vec![ + async move { + let _guard = guard; + futures::future::pending().await + } + .boxed(), + ], + false, + ); + let mut run = Run::new(); + let root = run.admit(root, scope(0..3), source()); + assert!(matches!(run.advance()?, Progress::Waiting)); + assert_eq!(drops.load(Ordering::Relaxed), 0); + if cancel_root { + run.cancel(root); + assert_eq!(drops.load(Ordering::Relaxed), 1); + } + drop(run); + assert_eq!(drops.load(Ordering::Relaxed), 1); + Ok(()) +} + +#[test] +fn task_failure_fails_the_run() { + let root = awaiting( + vec![async { Err(vortex_err!("child task failed")) }.boxed()], + false, + ); + let mut run = Run::new(); + run.admit(root, scope(0..3), source()); + let error = run.advance().err().expect("child error must fail the run"); + assert!(error.to_string().contains("child task failed")); +} diff --git a/vortex-scan/src/planning/mod.rs b/vortex-scan/src/planning/mod.rs index 3095712f161..2c69e734294 100644 --- a/vortex-scan/src/planning/mod.rs +++ b/vortex-scan/src/planning/mod.rs @@ -10,6 +10,7 @@ pub mod driver; pub mod morsel; pub mod next; pub mod planner; +pub mod selection; #[cfg(test)] mod tests; diff --git a/vortex-scan/src/planning/planner.rs b/vortex-scan/src/planning/planner.rs index 91d875ff317..cf0b71ad569 100644 --- a/vortex-scan/src/planning/planner.rs +++ b/vortex-scan/src/planning/planner.rs @@ -6,6 +6,7 @@ use std::fmt; use std::ops::Range; +use futures::future::BoxFuture; use vortex_error::VortexResult; use vortex_io::request::IoBatch; use vortex_io::request::IoConsumer; @@ -49,6 +50,10 @@ pub enum PlannerOutput { Continue, /// Requests the planner publishes: fetches it needs, and optional hints. NeedsIO(IoBatch), + /// Suspends this planner until the future completes. The future owns its child tasks; + /// dropping the planner or cancelling its root drops the future. The planner retains + /// results through its own channel and resumes without a separate completion callback. + Await(BoxFuture<'static, VortexResult<()>>), /// A child planner, with the scope it is responsible for. Planner(WorkScope, Box), /// A morsel authorised to produce arrays for the scope. @@ -74,6 +79,7 @@ impl fmt::Debug for PlannerOutput { Self::Done => f.write_str("Done"), Self::Continue => f.write_str("Continue"), Self::NeedsIO(batch) => f.debug_tuple("NeedsIO").field(batch).finish(), + Self::Await(_) => f.write_str("Await"), Self::Planner(scope, _) => f.debug_tuple("Planner").field(scope).finish(), Self::Morsel(scope, _) => f.debug_tuple("Morsel").field(scope).finish(), } diff --git a/vortex-scan/src/planning/selection.rs b/vortex-scan/src/planning/selection.rs new file mode 100644 index 00000000000..75b389690bb --- /dev/null +++ b/vortex-scan/src/planning/selection.rs @@ -0,0 +1,210 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A selection stage that hands its surviving rows to another planner. + +use std::sync::Arc; + +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_io::request::IoConsumer; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoResult; +use vortex_mask::Mask; + +use crate::planning::next::Next; +use crate::planning::planner::Planner; +use crate::planning::planner::PlannerOutput; +use crate::planning::planner::State; +use crate::planning::planner::WorkScope; +use crate::selection::Selection; + +/// Selected rows passed between planning stages. +/// +/// The mask is relative to `scope.rows.start`, and its length equals the scope's row count. +/// A stage can narrow the scope while rebasing the mask to preserve absolute row coordinates. +pub struct SelectedRows { + /// The input row domain covering the selected rows, including any unselected gaps. + pub scope: WorkScope, + /// The rows within that domain that the successor should process. + pub mask: Mask, +} + +/// Applies a scan selection to a scope, then constructs the next planner with the surviving rows. +/// +/// The successor may project, filter, or further partition the rows. This stage neither reads +/// data nor decides how those rows are executed. An empty selection creates no successor. +/// Compact selections exclude unselected leading and trailing rows from the successor's scope. +pub struct SelectionPlanner { + scope: Option, + selection: Arc, + next: Option>, +} + +impl SelectionPlanner { + /// Creates an unevaluated selection stage. The mask and successor are built in `compute`. + pub fn new(scope: WorkScope, selection: Arc, next: Next) -> Self { + Self { + scope: Some(scope), + selection, + next: Some(next), + } + } +} + +impl IoConsumer for SelectionPlanner { + fn set_io_result(&mut self, _request: IoRequestId, _result: IoResult) {} +} + +impl Planner for SelectionPlanner { + fn state(&self) -> State { + if self.scope.is_some() { + State::NeedsCompute + } else { + State::Done + } + } + + fn compute(&mut self) -> VortexResult { + let mut scope = self + .scope + .take() + .ok_or_else(|| vortex_err!("SelectionPlanner computed after completion"))?; + let selected = self.selection.row_mask(&scope.rows); + let Some(start) = selected.mask().first() else { + return Ok(PlannerOutput::Done); + }; + let end = selected.mask().last().vortex_expect("nonempty selection") + 1; + let mut mask = selected.mask().clone(); + // Narrow clusters to avoid planning unrelated chunks. Preserve alignment for sparse + // gathers: slicing their encoded inputs at arbitrary boundaries can cost more than it saves. + if end - start <= mask.true_count().saturating_mul(16) { + mask = mask.slice(start..end); + scope.rows = scope.rows.start + start as u64..scope.rows.start + end as u64; + } + let next = self + .next + .take() + .ok_or_else(|| vortex_err!("SelectionPlanner has no successor"))?; + let child = next(SelectedRows { + scope: scope.clone(), + mask, + })?; + Ok(PlannerOutput::Planner(scope, child)) + } +} + +#[cfg(test)] +mod tests { + use parking_lot::Mutex; + use rstest::rstest; + use vortex_buffer::Buffer; + use vortex_error::vortex_bail; + + use super::*; + use crate::planning::next::next_fn; + use crate::strict_sorted_buffer::StrictSortedBuffer; + + struct Finished; + + impl IoConsumer for Finished { + fn set_io_result(&mut self, _request: IoRequestId, _result: IoResult) {} + } + + impl Planner for Finished { + fn state(&self) -> State { + State::Done + } + + fn compute(&mut self) -> VortexResult { + Ok(PlannerOutput::Done) + } + } + + #[rstest] + #[case::all(Selection::All, vec![10, 11, 12, 13, 14])] + #[case::include(Selection::IncludeByIndex(StrictSortedBuffer::try_new(Buffer::from(vec![1, 10, 13, 20]))?), vec![10, 13])] + #[case::exclude(Selection::ExcludeByIndex(StrictSortedBuffer::try_new(Buffer::from(vec![1, 10, 13, 20]))?), vec![11, 12, 14])] + #[case::roaring(Selection::IncludeRoaring([1, 10, 13, 20].into_iter().collect()), vec![10, 13])] + #[case::exclude_roaring(Selection::ExcludeRoaring([1, 10, 13, 20].into_iter().collect()), vec![11, 12, 14])] + #[case::interior(Selection::IncludeByIndex(StrictSortedBuffer::try_new(Buffer::from(vec![11, 12, 13]))?), vec![11, 12, 13])] + #[case::singleton(Selection::IncludeRoaring([12].into_iter().collect()), vec![12])] + #[case::excluded_edges(Selection::ExcludeRoaring([10, 14].into_iter().collect()), vec![11, 12, 13])] + #[case::empty(Selection::IncludeRoaring(Default::default()), vec![])] + fn hands_selected_rows_to_successor( + #[case] selection: Selection, + #[case] expected: Vec, + ) -> VortexResult<()> { + let scope = WorkScope { + file_ordinal: 7, + rows: 10..15, + }; + let received = Arc::new(Mutex::new(None)); + let output = Arc::clone(&received); + let next = next_fn(move |selected: SelectedRows| { + *output.lock() = Some(selected); + Ok(Finished) + }); + let mut planner = SelectionPlanner::new(scope.clone(), Arc::new(selection), next); + assert_eq!(planner.state(), State::NeedsCompute); + assert!(received.lock().is_none()); + let result = planner.compute()?; + assert_eq!(planner.state(), State::Done); + let selected = received.lock().take(); + if expected.is_empty() { + assert!(matches!(result, PlannerOutput::Done)); + assert!(selected.is_none()); + } else { + let PlannerOutput::Planner(child_scope, _) = result else { + vortex_bail!("selection should emit its successor"); + }; + assert_eq!(child_scope.file_ordinal, scope.file_ordinal); + assert_eq!( + child_scope.rows, + expected[0]..expected[expected.len() - 1] + 1 + ); + let selected = selected.ok_or_else(|| vortex_err!("missing successor input"))?; + assert_eq!(selected.scope, child_scope); + assert_eq!( + selected.mask.len() as u64, + child_scope.rows.end - child_scope.rows.start + ); + let rows = selected + .mask + .iter() + .enumerate() + .filter(|(_, selected)| *selected) + .map(|(index, _)| selected.scope.rows.start + index as u64) + .collect::>(); + assert_eq!(rows, expected); + } + assert!(planner.compute().is_err()); + Ok(()) + } + + #[test] + fn scattered_selection_keeps_its_aligned_scope() -> VortexResult<()> { + let scope = WorkScope { + file_ordinal: 0, + rows: 0..65_536, + }; + let expected = scope.clone(); + let next = next_fn(move |selected: SelectedRows| { + assert_eq!(selected.scope, expected); + assert_eq!(selected.mask.len(), 65_536); + assert_eq!( + selected.mask.indices(), + vortex_mask::AllOr::Some(&[100, 50_000][..]) + ); + Ok(Finished) + }); + let selection = + Selection::IncludeByIndex(StrictSortedBuffer::try_new(Buffer::from(vec![ + 100, 50_000, + ]))?); + let mut planner = SelectionPlanner::new(scope, Arc::new(selection), next); + assert!(matches!(planner.compute()?, PlannerOutput::Planner(..))); + Ok(()) + } +} diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index fb961a63dab..9b9138fc46d 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -56,6 +56,7 @@ arrow-array = { workspace = true } arrow-schema = { workspace = true } divan = { workspace = true } fastlanes = { workspace = true } +futures = { workspace = true } mimalloc = { workspace = true } parquet = { workspace = true } rand = { workspace = true } diff --git a/vortex/examples/generate_list_fixture.rs b/vortex/examples/generate_list_fixture.rs new file mode 100644 index 00000000000..83b6fd4173e --- /dev/null +++ b/vortex/examples/generate_list_fixture.rs @@ -0,0 +1,183 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Generate large list rows for selection-scan diagnostics. +//! +//! Usage: `generate_list_fixture OUTPUT balanced|skewed primitive|nested TARGET_BYTES`. +//! Writes a new file and verifies full decoded equality through V1. Element block sizing is +//! independent of outer row count; nullable lists, empty lists, and nullable elements are included. + +use std::collections::BTreeMap; +use std::path::Path; +use std::sync::Arc; + +use anyhow::Result; +use anyhow::ensure; +use futures::stream; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use vortex::VortexSessionDefault; +use vortex::array::IntoArray; +use vortex::array::VortexSessionExecute; +use vortex::array::arrays::BoolArray; +use vortex::array::arrays::ListArray; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::StructArray; +use vortex::array::stream::ArrayStreamAdapter; +use vortex::array::stream::ArrayStreamExt; +use vortex::array::validity::Validity; +use vortex::buffer::Buffer; +use vortex::editions::Edition; +use vortex::editions::EditionDeclaration; +use vortex::editions::EditionId; +use vortex::editions::EditionMember; +use vortex::editions::EditionSessionExt; +use vortex::file::OpenOptionsSessionExt; +use vortex::file::WriteOptionsSessionExt; +use vortex::file::WriteStrategyBuilder; +use vortex::io::session::RuntimeSessionExt; +use vortex::layout::LayoutRef; +use vortex::layout::layouts::table::TableStrategy; +use vortex::session::VortexSession; +use vortex_arrow::ArrowSessionExt; + +const ROWS: u64 = 512; +const INNER_LENGTH: u64 = 128; + +static EDITION: EditionDeclaration = EditionDeclaration { + edition: Edition { + id: EditionId::new("large-list-fixture", 2026, 10, 0), + min_library_version: None, + }, + added: &[EditionMember::layout(&"vortex.list")], +}; + +fn count_layouts(layout: &LayoutRef, counts: &mut BTreeMap) -> Result<()> { + *counts.entry(layout.encoding_id().to_string()).or_default() += 1; + for child in layout.children()? { + count_layouts(&child, counts)?; + } + Ok(()) +} + +#[tokio::main] +async fn main() -> Result<()> { + let args = std::env::args().skip(1).collect::>(); + ensure!( + args.len() == 4, + "Expected OUTPUT balanced|skewed primitive|nested TARGET_BYTES" + ); + ensure!( + matches!(args[1].as_str(), "balanced" | "skewed"), + "Unknown distribution" + ); + ensure!( + matches!(args[2].as_str(), "primitive" | "nested"), + "Unknown element shape" + ); + let target_bytes = args[3].parse::()?; + ensure!(target_bytes > 0, "TARGET_BYTES must be positive"); + let session = VortexSession::default().with_tokio(); + session.register_edition(&EDITION)?; + session.enable_edition(EDITION.edition.id)?; + + let mut offsets = vec![0_u64]; + let mut total = 0_u64; + for row in 0..ROWS { + total += if row % 19 == 0 { + 0 + } else if args[1] == "balanced" { + 8192 + } else if row % 16 == 0 { + 262_144 + } else { + INNER_LENGTH + }; + offsets.push(total); + } + let mut rng = StdRng::seed_from_u64(42); + let mut elements = PrimitiveArray::from_option_iter((0..total).map(|index| { + let value = rng.random::(); + (index % 7 != 0).then_some(value) + })) + .into_array(); + if args[2] == "nested" { + let inner_count = total / INNER_LENGTH; + elements = ListArray::try_new( + elements, + Buffer::from_iter((0..=inner_count).map(|index| index * INNER_LENGTH)).into_array(), + Validity::Array( + BoolArray::from_iter((0..inner_count).map(|index| index % 11 != 0)).into_array(), + ), + )? + .into_array(); + for offset in &mut offsets { + *offset /= INNER_LENGTH; + } + } + let values = ListArray::try_new( + elements, + Buffer::from(offsets).into_array(), + Validity::Array(BoolArray::from_iter((0..ROWS).map(|row| row % 17 != 0)).into_array()), + )? + .into_array(); + let input = StructArray::from_fields(&[ + ("id", Buffer::from_iter(0..ROWS).into_array()), + ("values", values), + ])? + .into_array(); + let elements = WriteStrategyBuilder::from_session(&session) + .with_data_block_target_bytes(Some(target_bytes)) + .build(); + let defaults = WriteStrategyBuilder::from_session(&session).build(); + let strategy = TableStrategy::new(Arc::clone(&defaults), defaults) + .with_list_layout_factory(move |list| Arc::new(list.with_elements(Arc::clone(&elements)))); + let mut output = tokio::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(&args[0]) + .await?; + let summary = session + .write_options() + .with_strategy(Arc::new(strategy)) + .write( + &mut output, + ArrayStreamAdapter::new(input.dtype().clone(), stream::iter([Ok(input.clone())])), + ) + .await?; + drop(output); + let file = session + .open_options() + .open_path(Path::new(&args[0])) + .await?; + let actual = ArrayStreamAdapter::new(file.dtype().clone(), file.scan()?.into_stream()?) + .read_all() + .await?; + let mut ctx = session.create_execution_ctx(); + let arrow = session.arrow(); + let field = arrow.to_arrow_field("", input.dtype())?; + let expected = arrow.execute_arrow(input, Some(&field), &mut ctx)?; + let actual = arrow.execute_arrow(actual, Some(&field), &mut ctx)?; + ensure!( + expected.as_ref() == actual.as_ref(), + "Decoded fixture differs from generated input" + ); + let mut layouts = BTreeMap::new(); + count_layouts(file.footer().layout(), &mut layouts)?; + ensure!( + layouts.get("vortex.list").copied().unwrap_or(0) > 0, + "Expected ListLayout" + ); + println!( + "{}", + serde_json::json!({ + "output": args[0], "distribution": args[1], "shape": args[2], + "target_bytes": target_bytes, "rows": ROWS, "primitive_elements": total, + "max_row_payload_bytes": if args[1] == "balanced" { 8192 * 8 } else { 262_144 * 8 }, + "bytes": summary.size(), "correct": true, "layouts": layouts, + "layout_tree": file.footer().layout().display_tree().to_string(), + }) + ); + Ok(()) +} diff --git a/vortex/examples/rewrite_list_layout.rs b/vortex/examples/rewrite_list_layout.rs new file mode 100644 index 00000000000..60d8f1c3c76 --- /dev/null +++ b/vortex/examples/rewrite_list_layout.rs @@ -0,0 +1,151 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Rewrite identical rows with flat lists or ListLayout for selection-scan experiments. +//! +//! Usage: `rewrite_list_layout INPUT OUTPUT flat|list|list-elements [TARGET_BYTES]`. +//! Creates a new output file, verifies full Arrow equality through V1, and prints layout metadata. +//! `list-elements` applies TARGET_BYTES only to elements; other children keep the default strategy. +//! Compare this mode at different targets to isolate element sizing from outer-column partitioning. + +use std::collections::BTreeMap; +use std::path::Path; +use std::sync::Arc; + +use anyhow::Result; +use anyhow::ensure; +use vortex::VortexSessionDefault; +use vortex::array::ArrayRef; +use vortex::array::VortexSessionExecute; +use vortex::array::stream::ArrayStreamAdapter; +use vortex::array::stream::ArrayStreamExt; +use vortex::editions::Edition; +use vortex::editions::EditionDeclaration; +use vortex::editions::EditionId; +use vortex::editions::EditionMember; +use vortex::editions::EditionSessionExt; +use vortex::file::OpenOptionsSessionExt; +use vortex::file::VortexFile; +use vortex::file::WriteOptionsSessionExt; +use vortex::file::WriteStrategyBuilder; +use vortex::io::session::RuntimeSessionExt; +use vortex::layout::LayoutRef; +use vortex::layout::LayoutStrategy; +use vortex::layout::layouts::table::TableStrategy; +use vortex::session::VortexSession; +use vortex_arrow::ArrowSessionExt; + +// ListLayout is registered for reading but is not in the default write edition. +static EXPERIMENT_EDITION: EditionDeclaration = EditionDeclaration { + edition: Edition { + id: EditionId::new("list-layout-experiment", 2026, 10, 0), + min_library_version: None, + }, + added: &[EditionMember::layout(&"vortex.list")], +}; + +async fn read_v1(file: &VortexFile) -> Result { + Ok(ArrayStreamAdapter::new( + file.dtype().clone(), + file.scan()?.prepare()?.execute_stream(None)?, + ) + .read_all() + .await?) +} + +fn count_layouts(layout: &LayoutRef, counts: &mut BTreeMap) -> Result<()> { + *counts.entry(layout.encoding_id().to_string()).or_default() += 1; + for child in layout.children()? { + count_layouts(&child, counts)?; + } + Ok(()) +} + +#[tokio::main] +async fn main() -> Result<()> { + let args = std::env::args().skip(1).collect::>(); + ensure!( + (3..=4).contains(&args.len()), + "Expected INPUT OUTPUT flat|list|list-elements [TARGET_BYTES]" + ); + ensure!( + matches!(args[2].as_str(), "flat" | "list" | "list-elements"), + "Unknown layout" + ); + let target_bytes = args.get(3).map(|value| value.parse::()).transpose()?; + ensure!(target_bytes != Some(0), "TARGET_BYTES must be positive"); + let session = VortexSession::default().with_tokio(); + session.register_edition(&EXPERIMENT_EDITION)?; + if args[2] != "flat" { + session.enable_edition(EXPERIMENT_EDITION.edition.id)?; + } + let input = session + .open_options() + .open_path(Path::new(&args[0])) + .await?; + let mut strategy = WriteStrategyBuilder::from_session(&session); + if let Some(target_bytes) = target_bytes { + strategy = strategy.with_data_block_target_bytes(Some(target_bytes)); + } + let strategy = if args[2] == "list-elements" { + let elements = strategy.build(); + let defaults = WriteStrategyBuilder::from_session(&session).build(); + Arc::new( + TableStrategy::new(Arc::clone(&defaults), defaults).with_list_layout_factory( + move |list| Arc::new(list.with_elements(Arc::clone(&elements))), + ), + ) as Arc + } else { + strategy.build() + }; + let stream = ArrayStreamAdapter::new( + input.dtype().clone(), + input.scan()?.prepare()?.execute_stream(None)?, + ); + let mut output = tokio::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(&args[1]) + .await?; + let summary = session + .write_options() + .with_strategy(strategy) + .write(&mut output, stream) + .await?; + drop(output); + let rewritten = session + .open_options() + .open_path(Path::new(&args[1])) + .await?; + let mut counts = BTreeMap::new(); + count_layouts(rewritten.footer().layout(), &mut counts)?; + ensure!( + (counts.get("vortex.list").copied().unwrap_or(0) > 0) == (args[2] != "flat"), + "Written layout does not match the requested representation" + ); + ensure!( + input.dtype() == rewritten.dtype(), + "Rewritten schema differs" + ); + let expected = read_v1(&input).await?; + let actual = read_v1(&rewritten).await?; + let mut ctx = session.create_execution_ctx(); + let arrow = session.arrow(); + let target = arrow.to_arrow_field("", input.dtype())?; + let expected = arrow.execute_arrow(expected, Some(&target), &mut ctx)?; + let actual = arrow.execute_arrow(actual, Some(&target), &mut ctx)?; + ensure!( + expected.as_ref() == actual.as_ref(), + "Rewritten values differ" + ); + println!( + "{}", + serde_json::json!({ + "input": args[0], "output": args[1], "representation": args[2], + "target_bytes": target_bytes, "rows": summary.row_count(), + "file_bytes": summary.size(), "correct": true, "layouts": counts, + "layout_tree": rewritten.footer().layout().display_tree().to_string(), + }) + ); + Ok(()) +} diff --git a/vortex/examples/selection_scan.rs b/vortex/examples/selection_scan.rs new file mode 100644 index 00000000000..825122b1e12 --- /dev/null +++ b/vortex/examples/selection_scan.rs @@ -0,0 +1,468 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Diagnose selection scans on an existing Vortex file. +//! +//! Usage: `selection_scan FILE ITERATIONS ROWS uniform|clustered|all cached|reopen memory|disk +//! time|metrics|check|decode|decode-metrics [SEED]`. Set `VORTEX_SCAN_V2=0/1` in separate processes. +//! Timings exclude instrumentation in `time` and `decode` modes. The flag selects the explicit V2 +//! builder here only. `check` compares against a full V1 scan followed by a take. +//! `decode` consumes the selected rows in all columns as Arrow record batches; `decode-metrics` +//! measures allocations while doing the same work. +//! Instrumented disk runs also report physical reads/bytes after coalescing, summed over all +//! iterations and excluding the initial file open (reopen-mode opens are included). +//! Build with `VORTEX_SELECTION_METRICS=1` for the separate allocation-instrumented executable. +//! Use `indices:16,32` in place of the pattern to select exact, strictly increasing row indices; +//! ROWS must match the number supplied. This supports targeted large-row and null-row checks. + +use std::alloc::GlobalAlloc; +use std::alloc::Layout; +use std::hint::black_box; +use std::path::Path; +use std::pin::pin; +use std::sync::Arc; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; +use std::time::Instant; + +use anyhow::Result; +use anyhow::bail; +use anyhow::ensure; +use futures::StreamExt; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use vortex::VortexSessionDefault; +use vortex::array::ArrayRef; +use vortex::array::Canonical; +use vortex::array::IntoArray; +use vortex::array::VortexSessionExecute; +use vortex::array::memory::MemorySessionExt; +use vortex::array::stream::ArrayStreamAdapter; +use vortex::array::stream::ArrayStreamExt; +use vortex::buffer::Buffer; +use vortex::buffer::ByteBuffer; +use vortex::file::OpenOptionsSessionExt; +use vortex::file::VortexFile; +use vortex::io::InstrumentedReadAt; +use vortex::io::session::RuntimeSessionExt; +use vortex::io::std_file::FileReadAt; +use vortex::layout::scan::v2::ScanBuilder as V2ScanBuilder; +use vortex::layout::segments::SegmentFuture; +use vortex::layout::segments::SegmentId; +use vortex::layout::segments::SegmentSource; +use vortex::metrics::DefaultMetricsRegistry; +use vortex::metrics::MetricValue; +use vortex::metrics::MetricsRegistry; +use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; +use vortex::session::VortexSession; +use vortex_arrow::ArrowSessionExt; + +const TRACK_ALLOCATIONS: bool = option_env!("VORTEX_SELECTION_METRICS").is_some(); +static CALLS: AtomicU64 = AtomicU64::new(0); +static BYTES: AtomicU64 = AtomicU64::new(0); +static LIVE: AtomicU64 = AtomicU64::new(0); +static PEAK: AtomicU64 = AtomicU64::new(0); + +struct Allocator; + +#[global_allocator] +static ALLOCATOR: Allocator = Allocator; + +fn allocated(size: usize) { + CALLS.fetch_add(1, Ordering::Relaxed); + BYTES.fetch_add(size as u64, Ordering::Relaxed); + let live = LIVE.fetch_add(size as u64, Ordering::Relaxed) + size as u64; + PEAK.fetch_max(live, Ordering::Relaxed); +} + +// SAFETY: All operations delegate to the same allocator with the original pointer and layout. +unsafe impl GlobalAlloc for Allocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + // SAFETY: The caller supplies a valid allocation layout. + let pointer = unsafe { mimalloc::MiMalloc.alloc(layout) }; + if !pointer.is_null() && TRACK_ALLOCATIONS { + allocated(layout.size()); + } + pointer + } + + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + // SAFETY: The caller supplies a valid allocation layout. + let pointer = unsafe { mimalloc::MiMalloc.alloc_zeroed(layout) }; + if !pointer.is_null() && TRACK_ALLOCATIONS { + allocated(layout.size()); + } + pointer + } + + unsafe fn realloc(&self, pointer: *mut u8, layout: Layout, size: usize) -> *mut u8 { + // SAFETY: The caller guarantees the original allocation and requested size are valid. + let pointer = unsafe { mimalloc::MiMalloc.realloc(pointer, layout, size) }; + if !pointer.is_null() && TRACK_ALLOCATIONS { + LIVE.fetch_sub(layout.size() as u64, Ordering::Relaxed); + allocated(size); + } + pointer + } + + unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) { + if TRACK_ALLOCATIONS { + LIVE.fetch_sub(layout.size() as u64, Ordering::Relaxed); + } + // SAFETY: This allocation came from mimalloc with the same layout. + unsafe { mimalloc::MiMalloc.dealloc(pointer, layout) }; + } +} + +struct CountingSegments { + inner: Arc, + calls: Arc<[AtomicU64]>, +} + +impl SegmentSource for CountingSegments { + fn request(&self, id: SegmentId) -> SegmentFuture { + self.calls[*id as usize].fetch_add(1, Ordering::Relaxed); + self.inner.request(id) + } +} + +async fn open( + session: &VortexSession, + path: &Path, + bytes: Option<&ByteBuffer>, + metrics: Option<&Arc>, +) -> Result { + let options = session.open_options().with_layout_reader_cache(); + Ok(match (bytes, metrics) { + (Some(bytes), _) => options.open_buffer(bytes.clone())?, + (None, Some(metrics)) => { + let read = Arc::new(FileReadAt::open_with_allocator( + path, + session.handle(), + session.allocator(), + )?); + options + .with_metrics_registry(metrics.clone()) + .open_read(InstrumentedReadAt::new(read, metrics.as_ref())) + .await? + } + (None, None) => options.open_path(path).await?, + }) +} + +fn physical_io(metrics: &DefaultMetricsRegistry) -> [u64; 4] { + let mut totals = [0; 4]; + for metric in metrics.snapshot() { + match (metric.name().as_ref(), metric.value()) { + ("vortex.io.read.total_size", MetricValue::Counter(value)) => { + totals[0] += value.value(); + } + ("vortex.io.read.size", MetricValue::Histogram(value)) => { + totals[1] += value.count() as u64; + } + ("io.requests.individual", MetricValue::Counter(value)) => { + totals[2] += value.value(); + } + ("io.requests.coalesced", MetricValue::Counter(value)) => { + totals[3] += value.value(); + } + _ => {} + } + } + totals +} + +fn v2() -> bool { + std::env::var("VORTEX_SCAN_V2").is_ok_and(|value| value == "1") +} + +fn v2_scan( + file: &VortexFile, + indices: &Buffer, + all: bool, + session: &VortexSession, +) -> Result { + let builder = V2ScanBuilder::new(session.clone(), file.scan_file()); + Ok(if all { + builder + } else { + builder.with_row_indices(StrictSortedBuffer::try_new(indices.clone())?) + }) +} + +async fn scan( + file: &VortexFile, + indices: &Buffer, + all: bool, + session: &VortexSession, +) -> Result { + let stream = if v2() { + v2_scan(file, indices, all, session)?.into_stream()? + } else { + let builder = file.scan()?; + let builder = if all { + builder + } else { + builder.with_row_indices(StrictSortedBuffer::try_new(indices.clone())?) + }; + builder.into_stream()?.boxed() + }; + let array = ArrayStreamAdapter::new(file.dtype().clone(), stream) + .read_all() + .await?; + Ok(array + .execute::(&mut session.create_execution_ctx())? + .into_array()) +} + +async fn decode( + file: &VortexFile, + indices: &Buffer, + all: bool, + session: &VortexSession, +) -> Result { + let schema = Arc::new(session.arrow().to_arrow_schema(file.dtype())?); + let stream = if v2() { + v2_scan(file, indices, all, session)? + .into_record_batch_stream(schema)? + .boxed() + } else { + let builder = file.scan()?; + let builder = if all { + builder + } else { + builder.with_row_indices(StrictSortedBuffer::try_new(indices.clone())?) + }; + builder.into_record_batch_stream(schema)?.boxed() + }; + let mut stream = pin!(stream); + let mut rows = 0; + while let Some(batch) = stream.next().await { + rows += batch?.num_rows(); + } + Ok(rows) +} + +fn indices(rows: u64, count: usize, pattern: &str, seed: u64) -> Result> { + if let Some(indices) = pattern.strip_prefix("indices:") { + let selected = if indices.is_empty() { + Vec::new() + } else { + indices + .split(',') + .map(str::parse::) + .collect::, _>>()? + }; + ensure!( + selected.len() == count, + "ROWS must match the explicit selection length" + ); + ensure!( + selected.iter().all(|&index| index < rows), + "Selected row is outside the file" + ); + ensure!( + selected.windows(2).all(|pair| pair[0] < pair[1]), + "Indices must be strictly increasing" + ); + return Ok(Buffer::from(selected)); + } + let mut rng = StdRng::seed_from_u64(seed); + let mut result = Vec::with_capacity(count); + ensure!(count as u64 <= rows, "Selection exceeds the file row count"); + match pattern { + "all" => ensure!(count == 0, "Pass ROWS=0 for an all-row scan"), + "uniform" => { + while result.len() < count { + result.push(rng.random_range(0..rows)); + result.sort_unstable(); + result.dedup(); + } + } + "clustered" => { + while result.len() < count { + let len = (count - result.len()).min(20) as u64; + let start = rng.random_range(0..=rows - len); + result.extend(start..start + len); + result.sort_unstable(); + result.dedup(); + } + } + _ => bail!("Unknown pattern {pattern}"), + } + Ok(Buffer::from(result)) +} + +fn rss() -> Result { + let status = std::fs::read_to_string("/proc/self/status")?; + let get = |name: &str| -> Option { + status + .lines() + .find(|line| line.starts_with(name))? + .split_whitespace() + .nth(1)? + .parse::() + .ok() + .map(|kb| kb * 1024) + }; + Ok(serde_json::json!({"current": get("VmRSS:"), "peak": get("VmHWM:")})) +} + +#[tokio::main] +async fn main() -> Result<()> { + let args = std::env::args().skip(1).collect::>(); + ensure!( + matches!(args.len(), 7 | 8), + "Expected FILE ITERATIONS ROWS uniform|clustered|all cached|reopen memory|disk time|metrics|check|decode|decode-metrics [SEED]" + ); + let iterations = args[1].parse::()?; + ensure!(iterations > 0, "ITERATIONS must be positive"); + let count = args[2].parse::()?; + let cached = match args[4].as_str() { + "cached" => true, + "reopen" => false, + other => bail!("Unknown open mode {other}"), + }; + let metrics = matches!(args[6].as_str(), "metrics" | "decode-metrics"); + ensure!( + matches!( + args[6].as_str(), + "time" | "metrics" | "check" | "decode" | "decode-metrics" + ), + "Unknown measurement mode" + ); + ensure!( + args[6] == "check" || metrics == TRACK_ALLOCATIONS, + "Build with VORTEX_SELECTION_METRICS=1 only for metrics mode" + ); + let session = VortexSession::default().with_tokio(); + let path = Path::new(&args[0]); + let bytes = match args[5].as_str() { + "memory" => Some(ByteBuffer::from(std::fs::read(path)?)), + "disk" => None, + other => bail!("Unknown storage mode {other}"), + }; + let io_metrics = metrics.then(|| Arc::new(DefaultMetricsRegistry::default())); + let file = open(&session, path, bytes.as_ref(), io_metrics.as_ref()).await?; + let seed = args + .get(7) + .map(|seed| seed.parse::()) + .transpose()? + .unwrap_or(42); + let selected = indices(file.row_count(), count, &args[3], seed)?; + let all = args[3] == "all"; + let decode_rows = matches!(args[6].as_str(), "decode" | "decode-metrics"); + let expected_rows = if all { + usize::try_from(file.row_count())? + } else { + selected.len() + }; + if args[6] == "check" { + let expected = ArrayStreamAdapter::new( + file.dtype().clone(), + file.scan()?.prepare()?.execute_stream(None)?, + ) + .read_all() + .await?; + let expected = if all { + expected + } else { + expected.take(selected.clone().into_array())? + }; + let actual = scan(&file, &selected, all, &session).await?; + let mut ctx = session.create_execution_ctx(); + let arrow = session.arrow(); + let target = arrow.to_arrow_field("", file.dtype())?; + let expected = arrow.execute_arrow(expected, Some(&target), &mut ctx)?; + let actual = arrow.execute_arrow(actual, Some(&target), &mut ctx)?; + ensure!( + actual.as_ref() == expected.as_ref(), + "Selected values differ from full-scan take" + ); + println!( + "{}", + serde_json::json!({"correct": true, "rows": actual.len(), "indices": selected.as_slice()}) + ); + return Ok(()); + } + let locations = file.scan_file().locations; + let calls: Arc<[AtomicU64]> = locations.iter().map(|_| AtomicU64::new(0)).collect(); + let instrument = |file: VortexFile| { + if metrics { + let source = Arc::new(CountingSegments { + inner: file.segment_source(), + calls: Arc::clone(&calls), + }); + file.with_segment_source(source) + } else { + file + } + }; + let file = instrument(file); + let initial_io = io_metrics.as_ref().map(|metrics| physical_io(metrics)); + let mut runs = Vec::with_capacity(iterations); + for _ in 0..iterations { + for calls in calls.iter() { + calls.store(0, Ordering::Relaxed); + } + let before_live = LIVE.load(Ordering::Relaxed); + let before_calls = CALLS.load(Ordering::Relaxed); + let before_bytes = BYTES.load(Ordering::Relaxed); + PEAK.store(before_live, Ordering::Relaxed); + let start = Instant::now(); + let reopened = if cached { + None + } else { + Some(instrument( + open(&session, path, bytes.as_ref(), io_metrics.as_ref()).await?, + )) + }; + let input = reopened.as_ref().unwrap_or(&file); + let (array, actual_rows) = if decode_rows { + (None, decode(input, &selected, all, &session).await?) + } else { + let array = scan(input, &selected, all, &session).await?; + let rows = array.len(); + (Some(array), rows) + }; + drop(reopened); + let elapsed = start.elapsed().as_nanos(); + ensure!(actual_rows == expected_rows, "Wrong selected row count"); + let peak_live = PEAK.load(Ordering::Relaxed).saturating_sub(before_live); + let output_live = LIVE.load(Ordering::Relaxed) as i128 - before_live as i128; + drop(black_box(array)); + let retained_live = LIVE.load(Ordering::Relaxed) as i128 - before_live as i128; + runs.push(serde_json::json!({ + "elapsed_ns": elapsed, "instrumented": metrics, + "allocation_calls": CALLS.load(Ordering::Relaxed) - before_calls, + "allocated_bytes": BYTES.load(Ordering::Relaxed) - before_bytes, + "baseline_live_bytes": before_live, + "total_peak_live_bytes": before_live + peak_live, + "post_drop_live_bytes": LIVE.load(Ordering::Relaxed), + "peak_live_bytes": peak_live, "output_live_bytes": output_live, + "retained_live_bytes": retained_live, + "requests": calls.iter().map(|n| n.load(Ordering::Relaxed)).sum::(), + "requested_bytes": calls.iter().zip(locations.iter()).map(|(n, l)| n.load(Ordering::Relaxed) * u64::from(l.length)).sum::(), + "unique_segments": calls.iter().filter(|n| n.load(Ordering::Relaxed) != 0).count(), + "unique_bytes": calls.iter().zip(locations.iter()).filter(|(n, _)| n.load(Ordering::Relaxed) != 0).map(|(_, l)| u64::from(l.length)).sum::(), + })); + } + let physical_io = io_metrics + .as_ref() + .zip(initial_io) + .map(|(metrics, before)| { + let after = physical_io(metrics); + serde_json::json!({ + "bytes": after[0] - before[0], + "reads": after[1] - before[1], + "individual_requests": after[2] - before[2], + "coalesced_requests": after[3] - before[3], + }) + }); + println!( + "{}", + serde_json::json!({"args": args, "indices": selected.as_slice(), "runs": runs, "physical_io": physical_io, "rss": rss()?, "v2": std::env::var("VORTEX_SCAN_V2").ok()}) + ); + Ok(()) +} From a8349ec45b3e26d935ca1c19643da420cac56efc Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 19 Oct 2026 11:01:52 +0100 Subject: [PATCH 2/7] Remove scan experiment scaffolding and always benchmark V2 Use the planner-driven V2 ScanBuilder directly for every Vortex random access measurement, including object-store reads. Remove the reader toggle, its configuration plumbing, V1 documentation, and unused reader-cache setup. Remove the experiment-only regression workflow, driver script, and three fixture/selection examples, along with their futures dev dependency. Keep the planner implementation and regression tests. The previous commit c4fa7818b preserves the diagnostic harnesses and benchmark context. Inspected the resulting diff and searched for stale switch/example references. No tests, builds, or benchmarks were rerun for this cleanup. Signed-off-by: Joe Isaacs --- .github/workflows/scan-v2-regression.yml | 110 ----- Cargo.lock | 1 - benchmarks/random-access-bench/README.md | 11 +- benchmarks/random-access-bench/src/lib.rs | 24 +- benchmarks/random-access-bench/src/main.rs | 4 - scripts/scan-v2-regression.py | 312 -------------- vortex-bench/src/random_access/take.rs | 40 +- vortex/Cargo.toml | 1 - vortex/examples/generate_list_fixture.rs | 183 -------- vortex/examples/rewrite_list_layout.rs | 151 ------- vortex/examples/selection_scan.rs | 468 --------------------- 11 files changed, 17 insertions(+), 1288 deletions(-) delete mode 100644 .github/workflows/scan-v2-regression.yml delete mode 100644 scripts/scan-v2-regression.py delete mode 100644 vortex/examples/generate_list_fixture.rs delete mode 100644 vortex/examples/rewrite_list_layout.rs delete mode 100644 vortex/examples/selection_scan.rs diff --git a/.github/workflows/scan-v2-regression.yml b/.github/workflows/scan-v2-regression.yml deleted file mode 100644 index 5610959b28d..00000000000 --- a/.github/workflows/scan-v2-regression.yml +++ /dev/null @@ -1,110 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright the Vortex contributors - -name: Scan V2 Regression Comparison -run-name: Scan V2 regression comparison on ${{ github.ref_name }} - -on: - push: - branches: ["ji/scan-v2-pr10540"] - workflow_dispatch: { } - -permissions: - contents: read - -concurrency: - group: scan-v2-regression-${{ github.ref }} - cancel-in-progress: true - -jobs: - compare: - if: github.repository == 'vortex-data/vortex' - timeout-minutes: 180 - runs-on: >- - runs-on=${{ github.run_id }}/runner=bench-dedicated/family=c8gd.8xlarge/image=ubuntu24-full-arm64-pre-v2/extras=s3-cache/tag=scan-v2-regression - env: - # Previous random-access branch, before rebasing onto PR #10540. - BASELINE_COMMIT: "ad83fcbfc05d1885c08b7ae5cc57302e412638f3" - CARGO_BUILD_JOBS: "8" - CARGO_TARGET_DIR: ${{ github.workspace }}/target - RUSTFLAGS: "-C target-cpu=native -C force-frame-pointers=yes" - RUST_LOG: warn - steps: - - uses: runs-on/action@v2 - with: - sccache: s3 - - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7 - with: - fetch-depth: 0 - - uses: ./.github/actions/setup-prebuild - with: - repo-token: ${{ secrets.GITHUB_TOKEN }} - enable-sccache: "true" - - uses: ./.github/actions/system-info - - - name: Build current scan and data preparation binaries - shell: bash - run: | - set -euo pipefail - mkdir -p scan-v2-binaries scan-v2-results - env -u VORTEX_SELECTION_METRICS cargo build --locked --profile release_debug \ - -p vortex --example selection_scan --example rewrite_list_layout - cp target/release_debug/examples/selection_scan scan-v2-binaries/current - cp target/release_debug/examples/rewrite_list_layout scan-v2-binaries/rewrite - env -u VORTEX_SELECTION_METRICS cargo build --locked --profile release_debug \ - -p random-access-bench - cp target/release_debug/random-access-bench scan-v2-binaries/standard - - - name: Build previous V2 from the pinned checkpoint - shell: bash - run: | - set -euo pipefail - git worktree add --detach "$RUNNER_TEMP/scan-v2-baseline" "$BASELINE_COMMIT" - ( - cd "$RUNNER_TEMP/scan-v2-baseline" - env -u VORTEX_SELECTION_METRICS cargo build --locked --profile release_debug \ - -p vortex --example selection_scan - ) - cp target/release_debug/examples/selection_scan scan-v2-binaries/baseline - sha256sum scan-v2-binaries/* > scan-v2-results/binaries.sha256 - - - name: Prepare shared datasets and verify ListLayout rewrites - shell: bash - env: - VORTEX_SCAN_V2: "0" - run: | - set -euo pipefail - scan-v2-binaries/standard --prepare-data --formats vortex - mkdir -p scan-v2-data - source_file=vortex-bench/data/random_access/nested_lists/nested_lists.vortex - scan-v2-binaries/rewrite "$source_file" scan-v2-data/flat-control.vortex flat \ - > scan-v2-results/layout-flat.json - scan-v2-binaries/rewrite "$source_file" scan-v2-data/elements-default.vortex list-elements 1048576 \ - > scan-v2-results/layout-1MiB.json - scan-v2-binaries/rewrite "$source_file" scan-v2-data/elements-64k.vortex list-elements 65536 \ - > scan-v2-results/layout-64KiB.json - - - name: Isolate benchmark CPUs - run: sudo bash scripts/setup-benchmark.sh - - - name: Compare Arrow and standard suites and repeat borderline regressions - shell: bash - env: - CANDIDATE_COMMIT: ${{ github.sha }} - run: | - set -euo pipefail - bash scripts/bench-taskset.sh python3 scripts/scan-v2-regression.py \ - --baseline scan-v2-binaries/baseline \ - --candidate scan-v2-binaries/current \ - --standard-benchmark scan-v2-binaries/standard \ - --baseline-commit "$BASELINE_COMMIT" \ - --candidate-commit "$CANDIDATE_COMMIT" \ - | tee scan-v2-results/run.log - - - name: Upload results, raw samples, and layout metadata - if: always() - uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7 - with: - name: scan-v2-regression-${{ github.sha }} - path: scan-v2-results/ - if-no-files-found: warn diff --git a/Cargo.lock b/Cargo.lock index 6f3d8741fbf..736e0146827 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11323,7 +11323,6 @@ dependencies = [ "arrow-schema 59.3.0", "codspeed-divan-compat", "fastlanes 0.7.2", - "futures", "mimalloc", "parquet 59.3.0", "rand 0.10.3", diff --git a/benchmarks/random-access-bench/README.md b/benchmarks/random-access-bench/README.md index cddea03dd97..f845bb8ee11 100644 --- a/benchmarks/random-access-bench/README.md +++ b/benchmarks/random-access-bench/README.md @@ -28,16 +28,7 @@ work in each timed iteration. CI drives the full matrix via cargo run -p random-access-bench --profile release_debug --features lance ``` -Pass `--scan-v2` to use Vortex's planner-driven V2 scan API. The default uses V1. Run each -version separately with the same dataset, pattern, and open-mode options; save their outputs -to separate files because measurement names stay the same: - -```bash -cargo run -p random-access-bench --profile release_debug -- \ - --formats vortex --display-format gh-json -o v1.json -cargo run -p random-access-bench --profile release_debug -- \ - --formats vortex --scan-v2 --display-format gh-json -o v2.json -``` +Vortex measurements always use the planner-driven V2 scan builder. ## Running against S3 diff --git a/benchmarks/random-access-bench/src/lib.rs b/benchmarks/random-access-bench/src/lib.rs index 466e7c980ea..357e688ed53 100644 --- a/benchmarks/random-access-bench/src/lib.rs +++ b/benchmarks/random-access-bench/src/lib.rs @@ -153,11 +153,10 @@ async fn benchmark_random_access( storage: &str, reopen: bool, remote: Option<&RemoteDataDir>, - scan_v2: bool, ) -> Result { let time_limit = Duration::from_secs(time_limit_secs); let mut runs = Vec::new(); - let prepared_accessor = open_accessor(dataset, format, remote, scan_v2).await?; + let prepared_accessor = open_accessor(dataset, format, remote).await?; let cached_accessor = (!reopen).then_some(prepared_accessor); if let Some(accessor) = &cached_accessor { @@ -177,7 +176,7 @@ async fn benchmark_random_access( let arr = if let Some(accessor) = &cached_accessor { accessor.take(indices).await? } else { - open_accessor(dataset, format, remote, scan_v2) + open_accessor(dataset, format, remote) .await? .take(indices) .await? @@ -307,7 +306,6 @@ async fn open_accessor( dataset: &dyn BenchDataset, format: Format, remote: Option<&RemoteDataDir>, - scan_v2: bool, ) -> Result> { let name = format!( "random-access/{}/{}-{}", @@ -325,13 +323,12 @@ async fn open_accessor( } Format::OnDiskVortex | Format::VortexCompact => { let path = dataset_path(dataset, format).await?; - let accessor = match remote { - Some(remote) => { - VortexRandomAccessor::open_object_store(remote, &path, name, format).await? - } - None => VortexRandomAccessor::open(path, name, format).await?, - }; - Ok(Box::new(accessor.with_scan_v2(scan_v2))) + Ok(match remote { + Some(remote) => Box::new( + VortexRandomAccessor::open_object_store(remote, &path, name, format).await?, + ), + None => Box::new(VortexRandomAccessor::open(path, name, format).await?), + }) } Format::Parquet => { let path = dataset_path(dataset, format).await?; @@ -378,8 +375,6 @@ pub struct RunConfig { pub time_limit: u64, /// Whether each benchmark run reuses or reopens its file handle. pub open_mode: OpenMode, - /// Uses the explicit V2 scan API for Vortex formats when enabled. - pub scan_v2: bool, /// Format used to emit measurements. pub display_format: DisplayFormat, /// Optional path for the primary result output. @@ -398,7 +393,6 @@ pub async fn run(config: RunConfig) -> Result<()> { patterns, time_limit, open_mode, - scan_v2, display_format, output_path, ingest_output, @@ -452,7 +446,6 @@ pub async fn run(config: RunConfig) -> Result<()> { storage, reopen, remote, - scan_v2, ) .await?; @@ -482,7 +475,6 @@ pub async fn run(config: RunConfig) -> Result<()> { storage, reopen, remote, - scan_v2, ) .await?; diff --git a/benchmarks/random-access-bench/src/main.rs b/benchmarks/random-access-bench/src/main.rs index da0483a8bc6..3adb7e22beb 100644 --- a/benchmarks/random-access-bench/src/main.rs +++ b/benchmarks/random-access-bench/src/main.rs @@ -88,9 +88,6 @@ struct Args { /// Whether to reopen the file on each iteration, use a cached handle, or run both. #[arg(long, value_enum, default_value_t = OpenMode::Both)] open_mode: OpenMode, - /// Use the explicit V2 scan API for Vortex formats. - #[arg(long)] - scan_v2: bool, /// Read the data files from this remote directory (e.g. `s3://bucket/prefix/`) instead of /// local disk. The directory must mirror the layout of the local benchmark data directory, /// as produced by `--prepare-data`. @@ -127,7 +124,6 @@ async fn main() -> Result<()> { patterns: args.patterns, time_limit: args.time_limit, open_mode: args.open_mode, - scan_v2: args.scan_v2, display_format: args.display_format, output_path: args.output_path, ingest_output: args.ingest_output, diff --git a/scripts/scan-v2-regression.py b/scripts/scan-v2-regression.py deleted file mode 100644 index 17c637c27b0..00000000000 --- a/scripts/scan-v2-regression.py +++ /dev/null @@ -1,312 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright the Vortex contributors - -"""Compare V1, previous V2, and current V2 with full Arrow materialization.""" - -import argparse -import csv -import hashlib -import itertools -import json -import os -import platform -import statistics -import subprocess -from pathlib import Path - - -def digest(path: Path) -> str: - checksum = hashlib.sha256() - with path.open("rb") as source: - for chunk in iter(lambda: source.read(1024 * 1024), b""): - checksum.update(chunk) - return checksum.hexdigest() - - -def change(current: float, baseline: float) -> float: - return 100 * (current / baseline - 1) - - -def measure( - case: dict, - versions: dict[str, tuple[Path, str]], - files: dict[str, Path], - cpus: list[int], - output: Path, - iterations: int, - rounds: int, -) -> dict: - env = {k: v for k, v in os.environ.items() if not k.startswith(("VORTEX_", "TOKIO_"))} - env["RUST_LOG"] = "warn" - seeds = [] - for seed in [17, 42, 73]: - medians = {v: [] for v in versions} - expected_indices = None - for repeat in range(rounds): - order = list(versions)[:: 1 if repeat % 2 == 0 else -1] - for version in order: - binary, flag = versions[version] - stem = "-".join(str(v) for v in (*case.values(), seed, repeat, version)) - command = [ - "taskset", - "-c", - ",".join(map(str, cpus[: case["cores"]])), - str(binary), - str(files[case["dataset"]]), - str(iterations), - "100", - case["pattern"], - case["mode"], - "disk", - "decode", - str(seed), - ] - with (output / f"{stem}.json").open("w") as target, (output / f"{stem}.log").open("w") as log: - subprocess.run( - command, env=dict(env, VORTEX_SCAN_V2=flag), stdout=target, stderr=log, check=True, timeout=180 - ) - result = json.loads((output / f"{stem}.json").read_text()) - if result["v2"] != flag or len(result["runs"]) != iterations: - raise RuntimeError(f"Wrong scan mode or sample count: {stem}") - if any(run["instrumented"] for run in result["runs"]): - raise RuntimeError(f"Instrumented binary used for timing: {stem}") - if expected_indices is None: - expected_indices = result["indices"] - if result["indices"] != expected_indices: - raise RuntimeError(f"Versions selected different rows: {stem}") - samples = result["runs"][iterations // 8 :] - medians[version].append(statistics.median(x["elapsed_ns"] for x in samples) / 1e6) - times = {v: statistics.median(values) for v, values in medians.items()} - seeds.append(dict(seed=seed, **times, rounds=medians)) - row = dict(case, **{v + "_ms": statistics.median(s[v] for s in seeds) for v in versions}, seeds=seeds) - row["vs_old_v2_pct"] = change(row["new_v2_ms"], row["old_v2_ms"]) - if "v1" in versions: - row["vs_v1_pct"] = change(row["new_v2_ms"], row["v1_ms"]) - print(json.dumps({k: v for k, v in row.items() if k != "seeds"}), flush=True) - return row - - -def save(output: Path, metadata: dict, summary: list[dict], confirmations: list[dict]): - (output / "results.json").write_text( - json.dumps(dict(metadata, summary=summary, confirmations=confirmations), indent=2) + "\n" - ) - with (output / "results.csv").open("w") as target: - writer = csv.DictWriter(target, fieldnames=[k for k in summary[0] if k != "seeds"]) - writer.writeheader() - writer.writerows({k: v for k, v in row.items() if k != "seeds"} for row in summary) - - -def report(output: Path, metadata: dict, summary: list[dict], confirmations: list[dict]): - lines = [ - "# Scan V2 regression comparison", - "", - f"Candidate: `{metadata['candidate_commit']}`; previous V2: `{metadata['baseline_commit']}`.", - "", - "100 selected rows; full Arrow decode/drop; warm local disk. " - "Cached reuses file preparation; reopen includes opening/metadata.", - "Three seeds, two alternating rounds of 64 queries, first 8 discarded. Positive percentages mean slower.", - "", - f"Cases above +3% vs previous V2: {sum(r['vs_old_v2_pct'] > 3 for r in summary)}/{len(summary)}; " - f"above +5% vs V1: {sum(r['vs_v1_pct'] > 5 for r in summary)}/{len(summary)}.", - "", - "| Dataset | Cores | Selection | Open | V1 ms | Previous V2 ms | Current V2 ms | vs V1 | vs previous V2 |", - "|---|---:|---|---|---:|---:|---:|---:|---:|", - ] - for r in summary: - lines.append( - f"| {r['dataset']} | {r['cores']} | {r['pattern']} | {r['mode']} | {r['v1_ms']:.4f} | " - f"{r['old_v2_ms']:.4f} | {r['new_v2_ms']:.4f} | {r['vs_v1_pct']:+.1f}% | {r['vs_old_v2_pct']:+.1f}% |" - ) - if confirmations: - lines += [ - "", - "## Longer reruns", - "", - "Rerun when the first-pass aggregate exceeds +2%, or any seed exceeds +3%, versus previous V2.", - "Four alternating rounds of 128 queries, first 16 discarded, same three seeds.", - "", - "| Dataset | Cores | Selection | Open | Previous V2 ms | Current V2 ms | Change |", - "|---|---:|---|---|---:|---:|---:|", - ] - for r in confirmations: - lines.append( - f"| {r['dataset']} | {r['cores']} | {r['pattern']} | {r['mode']} | " - f"{r['old_v2_ms']:.4f} | {r['new_v2_ms']:.4f} | {r['vs_old_v2_pct']:+.1f}% |" - ) - text = "\n".join(lines) + "\n" - (output / "summary.md").write_text(text) - if step_summary := os.environ.get("GITHUB_STEP_SUMMARY"): - with Path(step_summary).open("a") as target: - target.write(text) - - -def compare_standard(binary: Path, cpus: list[int], output: Path): - """Compare the standard suite, including its legacy six-row taxi lookup.""" - output.mkdir(exist_ok=True) - binary = binary.resolve() - env = {k: v for k, v in os.environ.items() if not k.startswith(("VORTEX_", "TOKIO_"))} - env["RUST_LOG"] = "warn" - metadata = dict( - binary=str(binary), - sha256=digest(binary), - workdir=str(Path.cwd()), - cpus=cpus[:30], - rounds=2, - seconds_per_target=5, - seed=42, - ) - metadata["output"] = "standard canonical-array output; one-second cached warmup" - commands, summary = [], [] - for cores in [1, 30]: - samples = {} - for repeat in range(2): - for version in ["v1", "v2"] if repeat == 0 else ["v2", "v1"]: - stem = f"{cores}-{repeat}-{version}" - path = output / f"{stem}.jsonl" - command = [ - "taskset", - "-c", - ",".join(map(str, cpus[:cores])), - str(binary), - "--formats", - "vortex", - "--open-mode", - "both", - "--time-limit", - "5", - "--display-format", - "gh-json", - "-o", - str(path), - ] - if version == "v2": - command.append("--scan-v2") - commands.append(command) - print(f"Standard suite: {stem}", flush=True) - with (output / f"{stem}.log").open("w") as log: - subprocess.run(command, env=env, stdout=log, stderr=log, check=True, timeout=600) - rows = [json.loads(line) for line in path.read_text().splitlines() if line.strip()] - if len(rows) != 18 or len({row["name"] for row in rows}) != 18: - raise RuntimeError(f"Expected all 18 standard targets: {stem}") - for row in rows: - if row["unit"] != "ns": - raise RuntimeError(f"Unexpected standard benchmark unit: {row}") - samples.setdefault(row["name"], {"v1": [], "v2": []})[version].append(row["value"] / 1e6) - for name, rounds in samples.items(): - if any(len(values) != 2 for values in rounds.values()): - raise RuntimeError(f"Unpaired standard benchmark target: {name}") - times = {v: statistics.median(values) for v, values in rounds.items()} - row = dict( - cores=cores, - name=name, - v1_ms=times["v1"], - v2_ms=times["v2"], - vs_v1_pct=change(times["v2"], times["v1"]), - rounds=rounds, - ) - summary.append(row) - print(json.dumps(row), flush=True) - (output / "results.json").write_text( - json.dumps(dict(metadata, commands=commands, summary=summary), indent=2) + "\n" - ) - lines = [ - "## Standard random-access suite", - "", - "V1 versus V2 from the same binary; canonical-array output, seed 42, warm local disk.", - "Two alternating rounds, five seconds per target; positive percentages mean slower.", - "", - f"Cases above +3% vs V1: {sum(r['vs_v1_pct'] > 3 for r in summary)}/{len(summary)}.", - "", - "| Cores | Target | V1 ms | V2 ms | Change |", - "|---:|---|---:|---:|---:|", - ] - lines.extend( - f"| {r['cores']} | {r['name']} | {r['v1_ms']:.4f} | {r['v2_ms']:.4f} | {r['vs_v1_pct']:+.1f}% |" - for r in summary - ) - text = "\n".join(lines) + "\n" - (output / "summary.md").write_text(text) - if step_summary := os.environ.get("GITHUB_STEP_SUMMARY"): - with Path(step_summary).open("a") as target: - target.write(text) - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--baseline", type=Path, required=True) - parser.add_argument("--candidate", type=Path, required=True) - parser.add_argument("--baseline-commit", required=True) - parser.add_argument("--candidate-commit", required=True) - parser.add_argument( - "--standard-benchmark", type=Path, help="Also compare V1/V2 with the standard random-access binary" - ) - parser.add_argument("--data-dir", type=Path, default=Path("vortex-bench/data")) - parser.add_argument("--list-dir", type=Path, default=Path("scan-v2-data")) - parser.add_argument("--output", type=Path, default=Path("scan-v2-results")) - args = parser.parse_args() - args.output.mkdir(parents=True, exist_ok=True) - raw = args.output / "raw" - raw.mkdir(exist_ok=True) - cpus = sorted(os.sched_getaffinity(0)) - if len(cpus) < 30: - raise RuntimeError(f"This comparison requires 30 available CPUs; found {len(cpus)}") - files = { - "taxi": args.data_dir / "taxi/taxi.vortex", - "feature-vectors": args.data_dir / "random_access/feature_vectors/feature_vectors.vortex", - "nested-structs": args.data_dir / "random_access/nested_structs/nested_structs.vortex", - "flat-lists": args.list_dir / "flat-control.vortex", - "list-1MiB": args.list_dir / "elements-default.vortex", - "list-64KiB": args.list_dir / "elements-64k.vortex", - } - versions = { - "v1": (args.candidate.resolve(), "0"), - "old_v2": (args.baseline.resolve(), "1"), - "new_v2": (args.candidate.resolve(), "1"), - } - metadata = dict( - baseline_commit=args.baseline_commit, - candidate_commit=args.candidate_commit, - platform=platform.platform(), - cpus=cpus[:30], - method=dict( - rows=100, - seeds=[17, 42, 73], - iterations=64, - warmups=8, - rounds=2, - output="full Arrow decode and drop", - storage="warm local disk", - aggregation="median of round medians, then median across seeds", - confirmation_iterations=128, - confirmation_warmups=16, - confirmation_rounds=4, - ), - files={k: dict(path=str(p), sha256=digest(p), bytes=p.stat().st_size) for k, p in files.items()}, - versions={k: dict(binary=str(p), sha256=digest(p), scan_v2=f) for k, (p, f) in versions.items()}, - ) - (args.output / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n") - summary = [] - confirmations = [] - for cores, dataset, pattern, mode in itertools.product( - [1, 30], files, ["clustered", "uniform"], ["cached", "reopen"] - ): - case = dict(dataset=dataset, cores=cores, pattern=pattern, mode=mode) - summary.append(measure(case, versions, files, cpus, raw, 64, 2)) - save(args.output, metadata, summary, confirmations) - reruns = args.output / "reruns" - reruns.mkdir(exist_ok=True) - for row in summary: - if row["vs_old_v2_pct"] <= 2 and all(change(s["new_v2"], s["old_v2"]) <= 3 for s in row["seeds"]): - continue - case = {k: row[k] for k in ["dataset", "cores", "pattern", "mode"]} - confirmations.append( - measure(case, {k: v for k, v in versions.items() if k != "v1"}, files, cpus, reruns, 128, 4) - ) - save(args.output, metadata, summary, confirmations) - report(args.output, metadata, summary, confirmations) - if args.standard_benchmark: - compare_standard(args.standard_benchmark, cpus, args.output / "standard") - - -if __name__ == "__main__": - main() diff --git a/vortex-bench/src/random_access/take.rs b/vortex-bench/src/random_access/take.rs index c11285f26f9..2291905825d 100644 --- a/vortex-bench/src/random_access/take.rs +++ b/vortex-bench/src/random_access/take.rs @@ -34,7 +34,7 @@ use vortex::array::stream::ArrayStreamExt; use vortex::buffer::Buffer; use vortex::file::OpenOptionsSessionExt; use vortex::file::VortexFile; -use vortex::layout::scan::v2::ScanBuilder as V2ScanBuilder; +use vortex::layout::scan::v2::ScanBuilder; use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; use vortex::utils::aliases::hash_map::HashMap; @@ -135,7 +135,6 @@ pub struct VortexRandomAccessor { name: String, format: Format, file: VortexFile, - scan_v2: bool, } impl VortexRandomAccessor { @@ -145,16 +144,11 @@ impl VortexRandomAccessor { name: impl Into, format: Format, ) -> anyhow::Result { - let file = SESSION - .open_options() - .with_layout_reader_cache() - .open_path(path.as_ref()) - .await?; + let file = SESSION.open_options().open_path(path.as_ref()).await?; Ok(Self { name: name.into(), format, file, - scan_v2: false, }) } @@ -167,22 +161,14 @@ impl VortexRandomAccessor { ) -> anyhow::Result { let file = SESSION .open_options() - .with_layout_reader_cache() .open_object_store(remote.store(), remote.key(path)?) .await?; Ok(Self { name: name.into(), format, file, - scan_v2: false, }) } - - /// Selects the explicit planner-driven V2 scan API for this benchmark accessor. - pub fn with_scan_v2(mut self, enabled: bool) -> Self { - self.scan_v2 = enabled; - self - } } #[async_trait] @@ -197,22 +183,12 @@ impl RandomAccessor for VortexRandomAccessor { async fn take(&self, indices: &[u64]) -> anyhow::Result { let indices_buf: Buffer = Buffer::from(indices.to_vec()); - let indices = StrictSortedBuffer::try_new(indices_buf)?; - let array = if self.scan_v2 { - let stream = V2ScanBuilder::new(SESSION.clone(), self.file.scan_file()) - .with_row_indices(indices) - .into_stream()?; - ArrayStreamAdapter::new(self.file.dtype().clone(), stream) - .read_all() - .await? - } else { - self.file - .scan()? - .with_row_indices(indices) - .into_array_stream()? - .read_all() - .await? - }; + let stream = ScanBuilder::new(SESSION.clone(), self.file.scan_file()) + .with_row_indices(StrictSortedBuffer::try_new(indices_buf)?) + .into_stream()?; + let array = ArrayStreamAdapter::new(self.file.dtype().clone(), stream) + .read_all() + .await?; // We canonicalize / decompress for equivalence to Arrow's `RecordBatch`es. let mut ctx = SESSION.create_execution_ctx(); diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index 9b9138fc46d..fb961a63dab 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -56,7 +56,6 @@ arrow-array = { workspace = true } arrow-schema = { workspace = true } divan = { workspace = true } fastlanes = { workspace = true } -futures = { workspace = true } mimalloc = { workspace = true } parquet = { workspace = true } rand = { workspace = true } diff --git a/vortex/examples/generate_list_fixture.rs b/vortex/examples/generate_list_fixture.rs deleted file mode 100644 index 83b6fd4173e..00000000000 --- a/vortex/examples/generate_list_fixture.rs +++ /dev/null @@ -1,183 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Generate large list rows for selection-scan diagnostics. -//! -//! Usage: `generate_list_fixture OUTPUT balanced|skewed primitive|nested TARGET_BYTES`. -//! Writes a new file and verifies full decoded equality through V1. Element block sizing is -//! independent of outer row count; nullable lists, empty lists, and nullable elements are included. - -use std::collections::BTreeMap; -use std::path::Path; -use std::sync::Arc; - -use anyhow::Result; -use anyhow::ensure; -use futures::stream; -use rand::RngExt; -use rand::SeedableRng; -use rand::rngs::StdRng; -use vortex::VortexSessionDefault; -use vortex::array::IntoArray; -use vortex::array::VortexSessionExecute; -use vortex::array::arrays::BoolArray; -use vortex::array::arrays::ListArray; -use vortex::array::arrays::PrimitiveArray; -use vortex::array::arrays::StructArray; -use vortex::array::stream::ArrayStreamAdapter; -use vortex::array::stream::ArrayStreamExt; -use vortex::array::validity::Validity; -use vortex::buffer::Buffer; -use vortex::editions::Edition; -use vortex::editions::EditionDeclaration; -use vortex::editions::EditionId; -use vortex::editions::EditionMember; -use vortex::editions::EditionSessionExt; -use vortex::file::OpenOptionsSessionExt; -use vortex::file::WriteOptionsSessionExt; -use vortex::file::WriteStrategyBuilder; -use vortex::io::session::RuntimeSessionExt; -use vortex::layout::LayoutRef; -use vortex::layout::layouts::table::TableStrategy; -use vortex::session::VortexSession; -use vortex_arrow::ArrowSessionExt; - -const ROWS: u64 = 512; -const INNER_LENGTH: u64 = 128; - -static EDITION: EditionDeclaration = EditionDeclaration { - edition: Edition { - id: EditionId::new("large-list-fixture", 2026, 10, 0), - min_library_version: None, - }, - added: &[EditionMember::layout(&"vortex.list")], -}; - -fn count_layouts(layout: &LayoutRef, counts: &mut BTreeMap) -> Result<()> { - *counts.entry(layout.encoding_id().to_string()).or_default() += 1; - for child in layout.children()? { - count_layouts(&child, counts)?; - } - Ok(()) -} - -#[tokio::main] -async fn main() -> Result<()> { - let args = std::env::args().skip(1).collect::>(); - ensure!( - args.len() == 4, - "Expected OUTPUT balanced|skewed primitive|nested TARGET_BYTES" - ); - ensure!( - matches!(args[1].as_str(), "balanced" | "skewed"), - "Unknown distribution" - ); - ensure!( - matches!(args[2].as_str(), "primitive" | "nested"), - "Unknown element shape" - ); - let target_bytes = args[3].parse::()?; - ensure!(target_bytes > 0, "TARGET_BYTES must be positive"); - let session = VortexSession::default().with_tokio(); - session.register_edition(&EDITION)?; - session.enable_edition(EDITION.edition.id)?; - - let mut offsets = vec![0_u64]; - let mut total = 0_u64; - for row in 0..ROWS { - total += if row % 19 == 0 { - 0 - } else if args[1] == "balanced" { - 8192 - } else if row % 16 == 0 { - 262_144 - } else { - INNER_LENGTH - }; - offsets.push(total); - } - let mut rng = StdRng::seed_from_u64(42); - let mut elements = PrimitiveArray::from_option_iter((0..total).map(|index| { - let value = rng.random::(); - (index % 7 != 0).then_some(value) - })) - .into_array(); - if args[2] == "nested" { - let inner_count = total / INNER_LENGTH; - elements = ListArray::try_new( - elements, - Buffer::from_iter((0..=inner_count).map(|index| index * INNER_LENGTH)).into_array(), - Validity::Array( - BoolArray::from_iter((0..inner_count).map(|index| index % 11 != 0)).into_array(), - ), - )? - .into_array(); - for offset in &mut offsets { - *offset /= INNER_LENGTH; - } - } - let values = ListArray::try_new( - elements, - Buffer::from(offsets).into_array(), - Validity::Array(BoolArray::from_iter((0..ROWS).map(|row| row % 17 != 0)).into_array()), - )? - .into_array(); - let input = StructArray::from_fields(&[ - ("id", Buffer::from_iter(0..ROWS).into_array()), - ("values", values), - ])? - .into_array(); - let elements = WriteStrategyBuilder::from_session(&session) - .with_data_block_target_bytes(Some(target_bytes)) - .build(); - let defaults = WriteStrategyBuilder::from_session(&session).build(); - let strategy = TableStrategy::new(Arc::clone(&defaults), defaults) - .with_list_layout_factory(move |list| Arc::new(list.with_elements(Arc::clone(&elements)))); - let mut output = tokio::fs::OpenOptions::new() - .write(true) - .create_new(true) - .open(&args[0]) - .await?; - let summary = session - .write_options() - .with_strategy(Arc::new(strategy)) - .write( - &mut output, - ArrayStreamAdapter::new(input.dtype().clone(), stream::iter([Ok(input.clone())])), - ) - .await?; - drop(output); - let file = session - .open_options() - .open_path(Path::new(&args[0])) - .await?; - let actual = ArrayStreamAdapter::new(file.dtype().clone(), file.scan()?.into_stream()?) - .read_all() - .await?; - let mut ctx = session.create_execution_ctx(); - let arrow = session.arrow(); - let field = arrow.to_arrow_field("", input.dtype())?; - let expected = arrow.execute_arrow(input, Some(&field), &mut ctx)?; - let actual = arrow.execute_arrow(actual, Some(&field), &mut ctx)?; - ensure!( - expected.as_ref() == actual.as_ref(), - "Decoded fixture differs from generated input" - ); - let mut layouts = BTreeMap::new(); - count_layouts(file.footer().layout(), &mut layouts)?; - ensure!( - layouts.get("vortex.list").copied().unwrap_or(0) > 0, - "Expected ListLayout" - ); - println!( - "{}", - serde_json::json!({ - "output": args[0], "distribution": args[1], "shape": args[2], - "target_bytes": target_bytes, "rows": ROWS, "primitive_elements": total, - "max_row_payload_bytes": if args[1] == "balanced" { 8192 * 8 } else { 262_144 * 8 }, - "bytes": summary.size(), "correct": true, "layouts": layouts, - "layout_tree": file.footer().layout().display_tree().to_string(), - }) - ); - Ok(()) -} diff --git a/vortex/examples/rewrite_list_layout.rs b/vortex/examples/rewrite_list_layout.rs deleted file mode 100644 index 60d8f1c3c76..00000000000 --- a/vortex/examples/rewrite_list_layout.rs +++ /dev/null @@ -1,151 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Rewrite identical rows with flat lists or ListLayout for selection-scan experiments. -//! -//! Usage: `rewrite_list_layout INPUT OUTPUT flat|list|list-elements [TARGET_BYTES]`. -//! Creates a new output file, verifies full Arrow equality through V1, and prints layout metadata. -//! `list-elements` applies TARGET_BYTES only to elements; other children keep the default strategy. -//! Compare this mode at different targets to isolate element sizing from outer-column partitioning. - -use std::collections::BTreeMap; -use std::path::Path; -use std::sync::Arc; - -use anyhow::Result; -use anyhow::ensure; -use vortex::VortexSessionDefault; -use vortex::array::ArrayRef; -use vortex::array::VortexSessionExecute; -use vortex::array::stream::ArrayStreamAdapter; -use vortex::array::stream::ArrayStreamExt; -use vortex::editions::Edition; -use vortex::editions::EditionDeclaration; -use vortex::editions::EditionId; -use vortex::editions::EditionMember; -use vortex::editions::EditionSessionExt; -use vortex::file::OpenOptionsSessionExt; -use vortex::file::VortexFile; -use vortex::file::WriteOptionsSessionExt; -use vortex::file::WriteStrategyBuilder; -use vortex::io::session::RuntimeSessionExt; -use vortex::layout::LayoutRef; -use vortex::layout::LayoutStrategy; -use vortex::layout::layouts::table::TableStrategy; -use vortex::session::VortexSession; -use vortex_arrow::ArrowSessionExt; - -// ListLayout is registered for reading but is not in the default write edition. -static EXPERIMENT_EDITION: EditionDeclaration = EditionDeclaration { - edition: Edition { - id: EditionId::new("list-layout-experiment", 2026, 10, 0), - min_library_version: None, - }, - added: &[EditionMember::layout(&"vortex.list")], -}; - -async fn read_v1(file: &VortexFile) -> Result { - Ok(ArrayStreamAdapter::new( - file.dtype().clone(), - file.scan()?.prepare()?.execute_stream(None)?, - ) - .read_all() - .await?) -} - -fn count_layouts(layout: &LayoutRef, counts: &mut BTreeMap) -> Result<()> { - *counts.entry(layout.encoding_id().to_string()).or_default() += 1; - for child in layout.children()? { - count_layouts(&child, counts)?; - } - Ok(()) -} - -#[tokio::main] -async fn main() -> Result<()> { - let args = std::env::args().skip(1).collect::>(); - ensure!( - (3..=4).contains(&args.len()), - "Expected INPUT OUTPUT flat|list|list-elements [TARGET_BYTES]" - ); - ensure!( - matches!(args[2].as_str(), "flat" | "list" | "list-elements"), - "Unknown layout" - ); - let target_bytes = args.get(3).map(|value| value.parse::()).transpose()?; - ensure!(target_bytes != Some(0), "TARGET_BYTES must be positive"); - let session = VortexSession::default().with_tokio(); - session.register_edition(&EXPERIMENT_EDITION)?; - if args[2] != "flat" { - session.enable_edition(EXPERIMENT_EDITION.edition.id)?; - } - let input = session - .open_options() - .open_path(Path::new(&args[0])) - .await?; - let mut strategy = WriteStrategyBuilder::from_session(&session); - if let Some(target_bytes) = target_bytes { - strategy = strategy.with_data_block_target_bytes(Some(target_bytes)); - } - let strategy = if args[2] == "list-elements" { - let elements = strategy.build(); - let defaults = WriteStrategyBuilder::from_session(&session).build(); - Arc::new( - TableStrategy::new(Arc::clone(&defaults), defaults).with_list_layout_factory( - move |list| Arc::new(list.with_elements(Arc::clone(&elements))), - ), - ) as Arc - } else { - strategy.build() - }; - let stream = ArrayStreamAdapter::new( - input.dtype().clone(), - input.scan()?.prepare()?.execute_stream(None)?, - ); - let mut output = tokio::fs::OpenOptions::new() - .write(true) - .create_new(true) - .open(&args[1]) - .await?; - let summary = session - .write_options() - .with_strategy(strategy) - .write(&mut output, stream) - .await?; - drop(output); - let rewritten = session - .open_options() - .open_path(Path::new(&args[1])) - .await?; - let mut counts = BTreeMap::new(); - count_layouts(rewritten.footer().layout(), &mut counts)?; - ensure!( - (counts.get("vortex.list").copied().unwrap_or(0) > 0) == (args[2] != "flat"), - "Written layout does not match the requested representation" - ); - ensure!( - input.dtype() == rewritten.dtype(), - "Rewritten schema differs" - ); - let expected = read_v1(&input).await?; - let actual = read_v1(&rewritten).await?; - let mut ctx = session.create_execution_ctx(); - let arrow = session.arrow(); - let target = arrow.to_arrow_field("", input.dtype())?; - let expected = arrow.execute_arrow(expected, Some(&target), &mut ctx)?; - let actual = arrow.execute_arrow(actual, Some(&target), &mut ctx)?; - ensure!( - expected.as_ref() == actual.as_ref(), - "Rewritten values differ" - ); - println!( - "{}", - serde_json::json!({ - "input": args[0], "output": args[1], "representation": args[2], - "target_bytes": target_bytes, "rows": summary.row_count(), - "file_bytes": summary.size(), "correct": true, "layouts": counts, - "layout_tree": rewritten.footer().layout().display_tree().to_string(), - }) - ); - Ok(()) -} diff --git a/vortex/examples/selection_scan.rs b/vortex/examples/selection_scan.rs deleted file mode 100644 index 825122b1e12..00000000000 --- a/vortex/examples/selection_scan.rs +++ /dev/null @@ -1,468 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Diagnose selection scans on an existing Vortex file. -//! -//! Usage: `selection_scan FILE ITERATIONS ROWS uniform|clustered|all cached|reopen memory|disk -//! time|metrics|check|decode|decode-metrics [SEED]`. Set `VORTEX_SCAN_V2=0/1` in separate processes. -//! Timings exclude instrumentation in `time` and `decode` modes. The flag selects the explicit V2 -//! builder here only. `check` compares against a full V1 scan followed by a take. -//! `decode` consumes the selected rows in all columns as Arrow record batches; `decode-metrics` -//! measures allocations while doing the same work. -//! Instrumented disk runs also report physical reads/bytes after coalescing, summed over all -//! iterations and excluding the initial file open (reopen-mode opens are included). -//! Build with `VORTEX_SELECTION_METRICS=1` for the separate allocation-instrumented executable. -//! Use `indices:16,32` in place of the pattern to select exact, strictly increasing row indices; -//! ROWS must match the number supplied. This supports targeted large-row and null-row checks. - -use std::alloc::GlobalAlloc; -use std::alloc::Layout; -use std::hint::black_box; -use std::path::Path; -use std::pin::pin; -use std::sync::Arc; -use std::sync::atomic::AtomicU64; -use std::sync::atomic::Ordering; -use std::time::Instant; - -use anyhow::Result; -use anyhow::bail; -use anyhow::ensure; -use futures::StreamExt; -use rand::RngExt; -use rand::SeedableRng; -use rand::rngs::StdRng; -use vortex::VortexSessionDefault; -use vortex::array::ArrayRef; -use vortex::array::Canonical; -use vortex::array::IntoArray; -use vortex::array::VortexSessionExecute; -use vortex::array::memory::MemorySessionExt; -use vortex::array::stream::ArrayStreamAdapter; -use vortex::array::stream::ArrayStreamExt; -use vortex::buffer::Buffer; -use vortex::buffer::ByteBuffer; -use vortex::file::OpenOptionsSessionExt; -use vortex::file::VortexFile; -use vortex::io::InstrumentedReadAt; -use vortex::io::session::RuntimeSessionExt; -use vortex::io::std_file::FileReadAt; -use vortex::layout::scan::v2::ScanBuilder as V2ScanBuilder; -use vortex::layout::segments::SegmentFuture; -use vortex::layout::segments::SegmentId; -use vortex::layout::segments::SegmentSource; -use vortex::metrics::DefaultMetricsRegistry; -use vortex::metrics::MetricValue; -use vortex::metrics::MetricsRegistry; -use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; -use vortex::session::VortexSession; -use vortex_arrow::ArrowSessionExt; - -const TRACK_ALLOCATIONS: bool = option_env!("VORTEX_SELECTION_METRICS").is_some(); -static CALLS: AtomicU64 = AtomicU64::new(0); -static BYTES: AtomicU64 = AtomicU64::new(0); -static LIVE: AtomicU64 = AtomicU64::new(0); -static PEAK: AtomicU64 = AtomicU64::new(0); - -struct Allocator; - -#[global_allocator] -static ALLOCATOR: Allocator = Allocator; - -fn allocated(size: usize) { - CALLS.fetch_add(1, Ordering::Relaxed); - BYTES.fetch_add(size as u64, Ordering::Relaxed); - let live = LIVE.fetch_add(size as u64, Ordering::Relaxed) + size as u64; - PEAK.fetch_max(live, Ordering::Relaxed); -} - -// SAFETY: All operations delegate to the same allocator with the original pointer and layout. -unsafe impl GlobalAlloc for Allocator { - unsafe fn alloc(&self, layout: Layout) -> *mut u8 { - // SAFETY: The caller supplies a valid allocation layout. - let pointer = unsafe { mimalloc::MiMalloc.alloc(layout) }; - if !pointer.is_null() && TRACK_ALLOCATIONS { - allocated(layout.size()); - } - pointer - } - - unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { - // SAFETY: The caller supplies a valid allocation layout. - let pointer = unsafe { mimalloc::MiMalloc.alloc_zeroed(layout) }; - if !pointer.is_null() && TRACK_ALLOCATIONS { - allocated(layout.size()); - } - pointer - } - - unsafe fn realloc(&self, pointer: *mut u8, layout: Layout, size: usize) -> *mut u8 { - // SAFETY: The caller guarantees the original allocation and requested size are valid. - let pointer = unsafe { mimalloc::MiMalloc.realloc(pointer, layout, size) }; - if !pointer.is_null() && TRACK_ALLOCATIONS { - LIVE.fetch_sub(layout.size() as u64, Ordering::Relaxed); - allocated(size); - } - pointer - } - - unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) { - if TRACK_ALLOCATIONS { - LIVE.fetch_sub(layout.size() as u64, Ordering::Relaxed); - } - // SAFETY: This allocation came from mimalloc with the same layout. - unsafe { mimalloc::MiMalloc.dealloc(pointer, layout) }; - } -} - -struct CountingSegments { - inner: Arc, - calls: Arc<[AtomicU64]>, -} - -impl SegmentSource for CountingSegments { - fn request(&self, id: SegmentId) -> SegmentFuture { - self.calls[*id as usize].fetch_add(1, Ordering::Relaxed); - self.inner.request(id) - } -} - -async fn open( - session: &VortexSession, - path: &Path, - bytes: Option<&ByteBuffer>, - metrics: Option<&Arc>, -) -> Result { - let options = session.open_options().with_layout_reader_cache(); - Ok(match (bytes, metrics) { - (Some(bytes), _) => options.open_buffer(bytes.clone())?, - (None, Some(metrics)) => { - let read = Arc::new(FileReadAt::open_with_allocator( - path, - session.handle(), - session.allocator(), - )?); - options - .with_metrics_registry(metrics.clone()) - .open_read(InstrumentedReadAt::new(read, metrics.as_ref())) - .await? - } - (None, None) => options.open_path(path).await?, - }) -} - -fn physical_io(metrics: &DefaultMetricsRegistry) -> [u64; 4] { - let mut totals = [0; 4]; - for metric in metrics.snapshot() { - match (metric.name().as_ref(), metric.value()) { - ("vortex.io.read.total_size", MetricValue::Counter(value)) => { - totals[0] += value.value(); - } - ("vortex.io.read.size", MetricValue::Histogram(value)) => { - totals[1] += value.count() as u64; - } - ("io.requests.individual", MetricValue::Counter(value)) => { - totals[2] += value.value(); - } - ("io.requests.coalesced", MetricValue::Counter(value)) => { - totals[3] += value.value(); - } - _ => {} - } - } - totals -} - -fn v2() -> bool { - std::env::var("VORTEX_SCAN_V2").is_ok_and(|value| value == "1") -} - -fn v2_scan( - file: &VortexFile, - indices: &Buffer, - all: bool, - session: &VortexSession, -) -> Result { - let builder = V2ScanBuilder::new(session.clone(), file.scan_file()); - Ok(if all { - builder - } else { - builder.with_row_indices(StrictSortedBuffer::try_new(indices.clone())?) - }) -} - -async fn scan( - file: &VortexFile, - indices: &Buffer, - all: bool, - session: &VortexSession, -) -> Result { - let stream = if v2() { - v2_scan(file, indices, all, session)?.into_stream()? - } else { - let builder = file.scan()?; - let builder = if all { - builder - } else { - builder.with_row_indices(StrictSortedBuffer::try_new(indices.clone())?) - }; - builder.into_stream()?.boxed() - }; - let array = ArrayStreamAdapter::new(file.dtype().clone(), stream) - .read_all() - .await?; - Ok(array - .execute::(&mut session.create_execution_ctx())? - .into_array()) -} - -async fn decode( - file: &VortexFile, - indices: &Buffer, - all: bool, - session: &VortexSession, -) -> Result { - let schema = Arc::new(session.arrow().to_arrow_schema(file.dtype())?); - let stream = if v2() { - v2_scan(file, indices, all, session)? - .into_record_batch_stream(schema)? - .boxed() - } else { - let builder = file.scan()?; - let builder = if all { - builder - } else { - builder.with_row_indices(StrictSortedBuffer::try_new(indices.clone())?) - }; - builder.into_record_batch_stream(schema)?.boxed() - }; - let mut stream = pin!(stream); - let mut rows = 0; - while let Some(batch) = stream.next().await { - rows += batch?.num_rows(); - } - Ok(rows) -} - -fn indices(rows: u64, count: usize, pattern: &str, seed: u64) -> Result> { - if let Some(indices) = pattern.strip_prefix("indices:") { - let selected = if indices.is_empty() { - Vec::new() - } else { - indices - .split(',') - .map(str::parse::) - .collect::, _>>()? - }; - ensure!( - selected.len() == count, - "ROWS must match the explicit selection length" - ); - ensure!( - selected.iter().all(|&index| index < rows), - "Selected row is outside the file" - ); - ensure!( - selected.windows(2).all(|pair| pair[0] < pair[1]), - "Indices must be strictly increasing" - ); - return Ok(Buffer::from(selected)); - } - let mut rng = StdRng::seed_from_u64(seed); - let mut result = Vec::with_capacity(count); - ensure!(count as u64 <= rows, "Selection exceeds the file row count"); - match pattern { - "all" => ensure!(count == 0, "Pass ROWS=0 for an all-row scan"), - "uniform" => { - while result.len() < count { - result.push(rng.random_range(0..rows)); - result.sort_unstable(); - result.dedup(); - } - } - "clustered" => { - while result.len() < count { - let len = (count - result.len()).min(20) as u64; - let start = rng.random_range(0..=rows - len); - result.extend(start..start + len); - result.sort_unstable(); - result.dedup(); - } - } - _ => bail!("Unknown pattern {pattern}"), - } - Ok(Buffer::from(result)) -} - -fn rss() -> Result { - let status = std::fs::read_to_string("/proc/self/status")?; - let get = |name: &str| -> Option { - status - .lines() - .find(|line| line.starts_with(name))? - .split_whitespace() - .nth(1)? - .parse::() - .ok() - .map(|kb| kb * 1024) - }; - Ok(serde_json::json!({"current": get("VmRSS:"), "peak": get("VmHWM:")})) -} - -#[tokio::main] -async fn main() -> Result<()> { - let args = std::env::args().skip(1).collect::>(); - ensure!( - matches!(args.len(), 7 | 8), - "Expected FILE ITERATIONS ROWS uniform|clustered|all cached|reopen memory|disk time|metrics|check|decode|decode-metrics [SEED]" - ); - let iterations = args[1].parse::()?; - ensure!(iterations > 0, "ITERATIONS must be positive"); - let count = args[2].parse::()?; - let cached = match args[4].as_str() { - "cached" => true, - "reopen" => false, - other => bail!("Unknown open mode {other}"), - }; - let metrics = matches!(args[6].as_str(), "metrics" | "decode-metrics"); - ensure!( - matches!( - args[6].as_str(), - "time" | "metrics" | "check" | "decode" | "decode-metrics" - ), - "Unknown measurement mode" - ); - ensure!( - args[6] == "check" || metrics == TRACK_ALLOCATIONS, - "Build with VORTEX_SELECTION_METRICS=1 only for metrics mode" - ); - let session = VortexSession::default().with_tokio(); - let path = Path::new(&args[0]); - let bytes = match args[5].as_str() { - "memory" => Some(ByteBuffer::from(std::fs::read(path)?)), - "disk" => None, - other => bail!("Unknown storage mode {other}"), - }; - let io_metrics = metrics.then(|| Arc::new(DefaultMetricsRegistry::default())); - let file = open(&session, path, bytes.as_ref(), io_metrics.as_ref()).await?; - let seed = args - .get(7) - .map(|seed| seed.parse::()) - .transpose()? - .unwrap_or(42); - let selected = indices(file.row_count(), count, &args[3], seed)?; - let all = args[3] == "all"; - let decode_rows = matches!(args[6].as_str(), "decode" | "decode-metrics"); - let expected_rows = if all { - usize::try_from(file.row_count())? - } else { - selected.len() - }; - if args[6] == "check" { - let expected = ArrayStreamAdapter::new( - file.dtype().clone(), - file.scan()?.prepare()?.execute_stream(None)?, - ) - .read_all() - .await?; - let expected = if all { - expected - } else { - expected.take(selected.clone().into_array())? - }; - let actual = scan(&file, &selected, all, &session).await?; - let mut ctx = session.create_execution_ctx(); - let arrow = session.arrow(); - let target = arrow.to_arrow_field("", file.dtype())?; - let expected = arrow.execute_arrow(expected, Some(&target), &mut ctx)?; - let actual = arrow.execute_arrow(actual, Some(&target), &mut ctx)?; - ensure!( - actual.as_ref() == expected.as_ref(), - "Selected values differ from full-scan take" - ); - println!( - "{}", - serde_json::json!({"correct": true, "rows": actual.len(), "indices": selected.as_slice()}) - ); - return Ok(()); - } - let locations = file.scan_file().locations; - let calls: Arc<[AtomicU64]> = locations.iter().map(|_| AtomicU64::new(0)).collect(); - let instrument = |file: VortexFile| { - if metrics { - let source = Arc::new(CountingSegments { - inner: file.segment_source(), - calls: Arc::clone(&calls), - }); - file.with_segment_source(source) - } else { - file - } - }; - let file = instrument(file); - let initial_io = io_metrics.as_ref().map(|metrics| physical_io(metrics)); - let mut runs = Vec::with_capacity(iterations); - for _ in 0..iterations { - for calls in calls.iter() { - calls.store(0, Ordering::Relaxed); - } - let before_live = LIVE.load(Ordering::Relaxed); - let before_calls = CALLS.load(Ordering::Relaxed); - let before_bytes = BYTES.load(Ordering::Relaxed); - PEAK.store(before_live, Ordering::Relaxed); - let start = Instant::now(); - let reopened = if cached { - None - } else { - Some(instrument( - open(&session, path, bytes.as_ref(), io_metrics.as_ref()).await?, - )) - }; - let input = reopened.as_ref().unwrap_or(&file); - let (array, actual_rows) = if decode_rows { - (None, decode(input, &selected, all, &session).await?) - } else { - let array = scan(input, &selected, all, &session).await?; - let rows = array.len(); - (Some(array), rows) - }; - drop(reopened); - let elapsed = start.elapsed().as_nanos(); - ensure!(actual_rows == expected_rows, "Wrong selected row count"); - let peak_live = PEAK.load(Ordering::Relaxed).saturating_sub(before_live); - let output_live = LIVE.load(Ordering::Relaxed) as i128 - before_live as i128; - drop(black_box(array)); - let retained_live = LIVE.load(Ordering::Relaxed) as i128 - before_live as i128; - runs.push(serde_json::json!({ - "elapsed_ns": elapsed, "instrumented": metrics, - "allocation_calls": CALLS.load(Ordering::Relaxed) - before_calls, - "allocated_bytes": BYTES.load(Ordering::Relaxed) - before_bytes, - "baseline_live_bytes": before_live, - "total_peak_live_bytes": before_live + peak_live, - "post_drop_live_bytes": LIVE.load(Ordering::Relaxed), - "peak_live_bytes": peak_live, "output_live_bytes": output_live, - "retained_live_bytes": retained_live, - "requests": calls.iter().map(|n| n.load(Ordering::Relaxed)).sum::(), - "requested_bytes": calls.iter().zip(locations.iter()).map(|(n, l)| n.load(Ordering::Relaxed) * u64::from(l.length)).sum::(), - "unique_segments": calls.iter().filter(|n| n.load(Ordering::Relaxed) != 0).count(), - "unique_bytes": calls.iter().zip(locations.iter()).filter(|(n, _)| n.load(Ordering::Relaxed) != 0).map(|(_, l)| u64::from(l.length)).sum::(), - })); - } - let physical_io = io_metrics - .as_ref() - .zip(initial_io) - .map(|(metrics, before)| { - let after = physical_io(metrics); - serde_json::json!({ - "bytes": after[0] - before[0], - "reads": after[1] - before[1], - "individual_requests": after[2] - before[2], - "coalesced_requests": after[3] - before[3], - }) - }); - println!( - "{}", - serde_json::json!({"args": args, "indices": selected.as_slice(), "runs": runs, "physical_io": physical_io, "rss": rss()?, "v2": std::env::var("VORTEX_SCAN_V2").ok()}) - ); - Ok(()) -} From d72860ea565f73e97cb246d4639ef5cdd3923f35 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 19 Oct 2026 14:30:50 +0100 Subject: [PATCH 3/7] Benchmark deferred list pipelines locally and in random-access CI Add an opt-in prototype that compiles list element pipelines before offsets and validity resolve, then wakes their sources with the selected ranges. Keep the existing list implementation as the default and comparison control. Restore selection diagnostics and generate primitive, nested and short ListLayout fixtures. The random-access benchmark label now runs reference checks and alternating comparisons, with raw results uploaded as an artifact. Local validation: 72 reference checks and 32 timing cases, plus a reduced-iteration run of the CI comparison harness. This prototype does not yet implement the decoded-work admission budget or shared segment optimization. Signed-off-by: Joe Isaacs --- .github/workflows/pr-bench-runner.yml | 42 ++ Cargo.lock | 1 + scripts/bench-list-pipeline.py | 213 ++++++++ vortex-layout/src/plan/physical/operator.rs | 7 + .../src/plan/physical/operators/list.rs | 3 + .../plan/physical/operators/list/deferred.rs | 366 ++++++++++++++ .../src/plan/physical/pipeline/compile.rs | 6 + .../src/plan/physical/pipeline/scan.rs | 18 + .../plan/physical/pipeline/scan/validate.rs | 3 + .../src/plan/physical/pipeline/tests.rs | 2 + .../src/plan/physical/plans/list_pack.rs | 4 + vortex/Cargo.toml | 1 + vortex/examples/generate_list_fixture.rs | 187 +++++++ vortex/examples/selection_scan.rs | 470 ++++++++++++++++++ 14 files changed, 1323 insertions(+) create mode 100644 scripts/bench-list-pipeline.py create mode 100644 vortex-layout/src/plan/physical/operators/list/deferred.rs create mode 100644 vortex/examples/generate_list_fixture.rs create mode 100644 vortex/examples/selection_scan.rs diff --git a/.github/workflows/pr-bench-runner.yml b/.github/workflows/pr-bench-runner.yml index c7f5846c807..8b9fabb9643 100644 --- a/.github/workflows/pr-bench-runner.yml +++ b/.github/workflows/pr-bench-runner.yml @@ -77,6 +77,21 @@ jobs: target/release_debug/${{ inputs.benchmark_id }} polarsignals-cloud-token: ${{ secrets.POLAR_SIGNALS_API_KEY }} project-id: "e5d846e1-b54c-46e7-9174-8bf055a3af56" + - name: Build list pipeline comparison + if: inputs.benchmark_id == 'random-access-bench' && inputs.remote_data_dir == '' + env: + RUSTFLAGS: "-C target-cpu=native -C force-frame-pointers=yes" + run: | + cargo build --locked --package vortex --profile release_debug \ + --example selection_scan --example generate_list_fixture + - name: Upload list pipeline comparison binaries + if: inputs.benchmark_id == 'random-access-bench' && inputs.remote_data_dir == '' + uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7 + with: + name: list-pipeline-binaries + path: | + target/release_debug/examples/selection_scan + target/release_debug/examples/generate_list_fixture - name: Upload benchmark binary uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7 with: @@ -122,6 +137,13 @@ jobs: shell: bash run: chmod +x target/release_debug/${{ inputs.benchmark_id }} + - name: Download list pipeline comparison binaries + if: inputs.benchmark_id == 'random-access-bench' && inputs.remote_data_dir == '' + uses: actions/download-artifact@9000827ccba6bdab643e8b6fd33ac0654aef8333 # v8 + with: + name: list-pipeline-binaries + path: target/release_debug/examples/ + - name: Setup Polar Signals if: github.event.pull_request.head.repo.fork == false uses: polarsignals/gh-actions-ps-profiling@68ae857e375a826606352016e5b90f01a2a7ff7a # v0.8.1 @@ -179,6 +201,26 @@ jobs: run: | python3 scripts/compress-split.py + - name: Compare current and deferred list pipelines + if: inputs.benchmark_id == 'random-access-bench' && inputs.remote_data_dir == '' + env: + RUSTFLAGS: "-C target-cpu=native -C force-frame-pointers=yes" + run: | + chmod +x target/release_debug/examples/selection_scan \ + target/release_debug/examples/generate_list_fixture + python3 scripts/bench-list-pipeline.py + + - name: Upload list pipeline comparison results + if: always() && inputs.benchmark_id == 'random-access-bench' && inputs.remote_data_dir == '' + uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7 + with: + name: list-pipeline-results + path: | + list-pipeline-results/*.json + list-pipeline-results/*.jsonl + list-pipeline-results/*.log + list-pipeline-results/raw/ + - name: Run ${{ inputs.benchmark_name }} benchmark (per-column-encoder) if: inputs.benchmark_id == 'string-bench' shell: bash diff --git a/Cargo.lock b/Cargo.lock index 736e0146827..6f3d8741fbf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11323,6 +11323,7 @@ dependencies = [ "arrow-schema 59.3.0", "codspeed-divan-compat", "fastlanes 0.7.2", + "futures", "mimalloc", "parquet 59.3.0", "rand 0.10.3", diff --git a/scripts/bench-list-pipeline.py b/scripts/bench-list-pipeline.py new file mode 100644 index 00000000000..0e2edd1e248 --- /dev/null +++ b/scripts/bench-list-pipeline.py @@ -0,0 +1,213 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +"""Compare current and prebuilt list pipelines on identical ListLayout files.""" + +import argparse +import hashlib +import itertools +import json +import os +import platform +import statistics +import subprocess +from pathlib import Path + +ROOT = Path(__file__).resolve().parent.parent +SEEDS = [17, 42, 73] +VERSIONS = ["baseline", "candidate"] + + +def digest(path: Path) -> str: + checksum = hashlib.sha256() + with path.open("rb") as source: + for chunk in iter(lambda: source.read(1024 * 1024), b""): + checksum.update(chunk) + return checksum.hexdigest() + + +class Comparison: + def __init__(self, args: argparse.Namespace) -> None: + self.binaries = { + "baseline": (args.baseline_binary or args.binary).resolve(), + "candidate": args.binary.resolve(), + } + self.results = args.output.resolve() + self.data = self.results / "data" + self.generator = args.generator.resolve() + self.iterations = args.iterations + self.rounds = args.rounds + self.discard = 8 + if self.iterations <= self.discard or self.rounds <= 0: + raise ValueError("Need more than 8 iterations and at least one round") + (self.results / "raw").mkdir(parents=True, exist_ok=True) + self.data.mkdir(parents=True, exist_ok=True) + + def prepare(self) -> dict[str, Path]: + for target_bytes in [65536, 1048576]: + for size, shape in [("large", "primitive"), ("large", "nested"), ("small", "primitive")]: + name = f"{size}-{shape}-{target_bytes}" + path = self.data / f"{name}.vortex" + if path.exists(): + continue + command = [ + str(self.generator), + str(path), + "small" if size == "small" else "skewed", + shape, + str(target_bytes), + ] + result = subprocess.run(command, capture_output=True, text=True, timeout=600) + (self.results / f"layout-{name}.log").write_text(result.stderr) + result.check_returncode() + (self.results / f"layout-{name}.json").write_text(result.stdout) + files = {path.stem: path for path in sorted(self.data.glob("*.vortex"))} + assert len(files) == 6, files + metadata = { + "baseline": "current list implementation, prototype disabled", + "commit": subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=ROOT, text=True).strip(), + "candidate": "prebuilt pipeline prototype; VORTEX_LIST_PIPELINE=1", + "platform": platform.platform(), + "logical_cpu_count": os.cpu_count(), + "affinity": "unrestricted", + "seeds": SEEDS, + "rounds": self.rounds, + "iterations": self.iterations, + "discard": self.discard, + "output": "full Arrow decode and drop", + "storage": "warm local disk", + "rustflags": os.environ.get("RUSTFLAGS", "from Cargo configuration"), + "profile": "release_debug", + "files": { + name: {"path": str(path), "sha256": digest(path), "bytes": path.stat().st_size} + for name, path in files.items() + }, + "binaries": {version: digest(path) for version, path in self.binaries.items()}, + } + self.write_json("metadata.json", metadata) + return files + + def write_json(self, name: str, value: object) -> None: + (self.results / name).write_text(json.dumps(value, indent=2) + "\n") + + def run( + self, + version: str, + path: Path, + count: int, + pattern: str, + mode: str, + seed: int, + check: bool, + stem: str, + ) -> dict: + env = {key: value for key, value in os.environ.items() if not key.startswith(("VORTEX_", "TOKIO_"))} + env.update(VORTEX_SCAN_V2="1", RUST_LOG="warn") + if version == "candidate": + env["VORTEX_LIST_PIPELINE"] = "1" + command = [ + str(self.binaries[version]), + str(path), + "1" if check else str(self.iterations), + str(count), + pattern, + mode, + "disk", + "check" if check else "decode", + str(seed), + ] + with (self.results / "commands.jsonl").open("a") as output: + output.write( + json.dumps({"command": command, "selection_gate": version == "candidate", "VORTEX_SCAN_V2": "1"}) + "\n" + ) + result = subprocess.run(command, env=env, capture_output=True, text=True, timeout=180) + (self.results / "raw" / f"{stem}.log").write_text(result.stderr) + (self.results / "raw" / f"{stem}.json").write_text(result.stdout) + result.check_returncode() + return json.loads(result.stdout) + + def check(self, files: dict[str, Path]) -> int: + checks = [] + selections = [ + (0, "indices:"), + (1, "indices:0"), + (1, "indices:17"), + (1, "indices:16"), + (100, "uniform"), + (100, "clustered"), + ] + for name, path in files.items(): + for count, pattern in selections: + for version in VERSIONS: + stem = f"check-{name}-{count}-{pattern.replace(':', '_')}-{version}" + data = self.run(version, path, count, pattern, "cached", 42, True, stem) + assert data["correct"] + checks.append({"dataset": name, "version": version, "pattern": pattern, "rows": count}) + self.write_json("correctness.json", checks) + print(json.dumps({"correctness": name, "passed": True}), flush=True) + return len(checks) + + def measure(self, files: dict[str, Path]) -> list[dict]: + summary = [] + for name, path in files.items(): + patterns = [(100, "uniform"), (100, "clustered")] + if not name.startswith("small"): + patterns.insert(0, (1, "indices:16")) + for (count, pattern), mode in itertools.product(patterns, ["cached", "reopen"]): + medians: dict[str, list[float]] = {version: [] for version in VERSIONS} + for seed in SEEDS: + expected_indices = None + for repeat in range(self.rounds): + order = VERSIONS if repeat % 2 == 0 else list(reversed(VERSIONS)) + for version in order: + stem = f"{name}-{count}-{pattern.replace(':', '_')}-{mode}-{seed}-{repeat}-{version}" + data = self.run(version, path, count, pattern, mode, seed, False, stem) + assert data["v2"] == "1" + assert len(data["runs"]) == self.iterations + assert not any(run["instrumented"] for run in data["runs"]) + if expected_indices is None: + expected_indices = data["indices"] + assert data["indices"] == expected_indices + median = statistics.median(run["elapsed_ns"] for run in data["runs"][self.discard :]) + medians[version].append(median / 1e6) + row = {"dataset": name, "rows": count, "pattern": pattern, "mode": mode} + row.update({version + "_ms": statistics.median(times) for version, times in medians.items()}) + row["change_pct"] = 100 * (row["candidate_ms"] / row["baseline_ms"] - 1) + row["round_medians_ms"] = medians + summary.append(row) + self.write_json("summary.json", summary) + print(json.dumps({key: value for key, value in row.items() if key != "round_medians_ms"}), flush=True) + return summary + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--binary", type=Path, default=ROOT / "target/release_debug/examples/selection_scan") + parser.add_argument("--baseline-binary", type=Path, help="Defaults to --binary with the prototype disabled") + parser.add_argument("--generator", type=Path, default=ROOT / "target/release_debug/examples/generate_list_fixture") + parser.add_argument("--output", type=Path, default=ROOT / "list-pipeline-results") + parser.add_argument("--iterations", type=int, default=64) + parser.add_argument("--rounds", type=int, default=2) + comparison = Comparison(parser.parse_args()) + files = comparison.prepare() + checks = comparison.check(files) + summary = comparison.measure(files) + if destination := os.environ.get("GITHUB_STEP_SUMMARY"): + with open(destination, "a") as output: + output.write("\n### Deferred list pipeline experiment\n\n") + output.write( + f"{checks} reference checks passed; full Arrow decode on warm local disk. " + "The prototype prebuilds pipelines but does not implement the current decoded-work admission budget. " + "Positive changes are slower.\n\n" + ) + output.write("| Dataset | Rows | Selection | Open mode | Current ms | Prototype ms | Change |\n") + output.write("| --- | ---: | --- | --- | ---: | ---: | ---: |\n") + for row in summary: + output.write( + f"| {row['dataset']} | {row['rows']} | {row['pattern']} | {row['mode']} | " + f"{row['baseline_ms']:.3f} | {row['candidate_ms']:.3f} | {row['change_pct']:+.1f}% |\n" + ) + + +if __name__ == "__main__": + main() diff --git a/vortex-layout/src/plan/physical/operator.rs b/vortex-layout/src/plan/physical/operator.rs index de1b1e35560..12d1d56c648 100644 --- a/vortex-layout/src/plan/physical/operator.rs +++ b/vortex-layout/src/plan/physical/operator.rs @@ -46,6 +46,8 @@ pub enum SourceBlocked { Io, /// Wait for a child projection, possibly alongside pending reads. Child, + /// Wait for a list's runtime element selection. + Selection(usize), /// The inlet needs more batches or its sink to close it. It must be open and have room /// for its sink to push; debug builds check that the sink can still wake this source. Inlet(usize), @@ -162,6 +164,7 @@ pub struct Cx<'a> { pub(super) changed: &'a mut InletSet, pub(super) child_requests: &'a mut Vec, pub(super) child_inbox: &'a mut Option>, + pub(super) resolved_selections: &'a mut Vec, } #[derive(Default)] @@ -171,6 +174,10 @@ pub(super) struct ChildInbox { } impl Cx<'_> { + pub(crate) fn resolve_selection(&mut self, selection: usize) { + self.resolved_selections.push(selection); + } + /// Describes child work for the enclosing planner. Return [`SourceBlocked::Child`] while awaiting its completion. /// Identifiers need only be unique among this source's outstanding child projections. pub(crate) fn request_child(&mut self, child: ChildProjection) { diff --git a/vortex-layout/src/plan/physical/operators/list.rs b/vortex-layout/src/plan/physical/operators/list.rs index 7ccc49dabc8..3ec22b8a423 100644 --- a/vortex-layout/src/plan/physical/operators/list.rs +++ b/vortex-layout/src/plan/physical/operators/list.rs @@ -43,8 +43,11 @@ use crate::plan::physical::pipeline::Turn; use crate::plan::physical::pipeline::ValueCache; use crate::segments::SegmentId; +mod deferred; mod work; +pub(crate) use deferred::compile_deferred_list; + /// A list's offsets and validity arrive before its selected element ranges can be planned. pub(crate) struct ListSource { plan: ListPackPlan, diff --git a/vortex-layout/src/plan/physical/operators/list/deferred.rs b/vortex-layout/src/plan/physical/operators/list/deferred.rs new file mode 100644 index 00000000000..18541516441 --- /dev/null +++ b/vortex-layout/src/plan/physical/operators/list/deferred.rs @@ -0,0 +1,366 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmark prototype: prebuild list pipelines and resolve their selection at runtime. +//! +//! Enabled only by `VORTEX_LIST_PIPELINE`. Supports primitive and nested-list fixture plans; +//! unsupported plans fail explicitly. This experiment does not implement the planner's decoded-work +//! admission budget or shared segment optimization. + +use std::ops::Range; +use std::sync::Arc; +use std::sync::OnceLock; + +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::expr::stats::Precision; +use vortex_array::expr::stats::Stat; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use super::selected_ranges; +use crate::plan::physical::Concat; +use crate::plan::physical::Filter; +use crate::plan::physical::LegacyStats; +use crate::plan::physical::ListPack; +use crate::plan::physical::ListPackPlan; +use crate::plan::physical::PlanRef; +use crate::plan::physical::SegmentScan; +use crate::plan::physical::SegmentScanPlan; +use crate::plan::physical::Zoned; +use crate::plan::physical::operator::Cx; +use crate::plan::physical::operator::Input; +use crate::plan::physical::operator::Operator; +use crate::plan::physical::operator::Source; +use crate::plan::physical::operator::SourceBlocked; +use crate::plan::physical::operator::SourceReady; +use crate::plan::physical::operator::Step; +use crate::plan::physical::operators::ConcatSource; +use crate::plan::physical::operators::ScanSource; +use crate::plan::physical::operators::UNBOUNDED; +use crate::plan::physical::operators::join_inlet; +use crate::plan::physical::pipeline::Chain; +use crate::plan::physical::pipeline::Compiler; + +#[derive(Clone)] +struct Demand { + ranges: Arc>>>, + gate: Option, + domain: Range, + offset: u64, + boundaries: bool, +} + +impl Demand { + fn ranges(&self, len: u64) -> Option>> { + let selected = self.ranges.get()?; + let mut result: Vec> = Vec::new(); + for range in selected { + let start = range.start.max(self.domain.start); + let end = range.end.min(self.domain.end); + if start >= end { + continue; + } + let end = end + u64::from(self.boundaries); + let start = start.max(self.offset); + let end = end.min(self.offset + len); + if start >= end { + continue; + } + let range = start - self.offset..end - self.offset; + if let Some(last) = result.last_mut() + && last.end >= range.start + { + last.end = last.end.max(range.end); + } else { + result.push(range); + } + } + Some(result) + } + + fn blocked(&self) -> VortexResult { + Ok(SourceReady::Blocked(SourceBlocked::Selection( + self.gate + .ok_or_else(|| vortex_err!("Unresolved demand has no producer"))?, + ))) + } + + fn shifted(&self, offset: u64) -> Self { + Self { + offset: self.offset + offset, + ..self.clone() + } + } + + fn boundaries(&self, rows: u64) -> Self { + Self { + domain: self.offset..self.offset + rows, + boundaries: true, + ..self.clone() + } + } +} + +pub(crate) fn compile_deferred_list( + plan: &ListPackPlan, + rows: Range, + mask: &Mask, + compiler: &mut Compiler<'_>, +) -> VortexResult { + let ranges = match mask.slices() { + AllOr::All => vec![rows.clone()], + AllOr::None => Vec::new(), + AllOr::Some(slices) => slices + .iter() + .map(|&(start, end)| rows.start + start as u64..rows.start + end as u64) + .collect(), + }; + let demand = Demand { + ranges: Arc::new(OnceLock::from(ranges)), + gate: None, + domain: 0..plan.row_count(), + offset: 0, + boundaries: false, + }; + compile_list(plan, demand, compiler) +} + +fn compile_list( + plan: &ListPackPlan, + demand: Demand, + compiler: &mut Compiler<'_>, +) -> VortexResult { + let offsets = compile( + &plan.offsets()?, + demand.boundaries(plan.row_count()), + compiler, + )?; + let mut chains = vec![offsets]; + if let Some(validity) = plan.validity()? { + chains.push(compile(&validity, demand.clone(), compiler)?); + } + let elements = plan.elements()?; + let element_demand = Demand { + ranges: Arc::new(OnceLock::new()), + gate: Some(compiler.selection_gate()), + domain: 0..elements.row_count(), + offset: 0, + boundaries: false, + }; + chains.push(compile(&elements, element_demand.clone(), compiler)?); + let source = DeferredList { + plan: plan.clone(), + demand, + elements: element_demand, + assembled: None, + done: false, + }; + Ok(compiler.join_with(chains, |_| UNBOUNDED, source)) +} + +fn compile(plan: &PlanRef, demand: Demand, compiler: &mut Compiler<'_>) -> VortexResult { + if let Some(filter) = plan.as_opt::() { + return compile(&filter.child_plan()?, demand, compiler); + } + if let Some(zoned) = plan.as_opt::() { + let data = zoned + .data_plan()? + .ok_or_else(|| vortex_err!("Zoned data is absent"))?; + return compile(&data, demand, compiler); + } + if let Some(stats) = plan.as_opt::() { + return compile(&stats.data_plan()?, demand, compiler); + } + if let Some(concat) = plan.as_opt::() { + let mut chains = Vec::new(); + for (index, &start) in concat.row_offsets().iter().enumerate() { + chains.push(compile( + &concat.child_required(index)?, + demand.shifted(start), + compiler, + )?); + } + let count = chains.len(); + return Ok(compiler.join(chains, ConcatSource::new(count))); + } + if let Some(list) = plan.as_opt::() { + return compile_list(list, demand, compiler); + } + if let Some(scan) = plan.as_opt::() { + return Ok(Chain::new(DeferredScan { + plan: scan.clone(), + demand, + source: None, + empty: false, + })); + } + vortex_bail!( + "Deferred-list benchmark prototype cannot compile {}", + plan.id() + ) +} + +struct DeferredScan { + plan: SegmentScanPlan, + demand: Demand, + source: Option, + empty: bool, +} + +impl Source for DeferredScan { + fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult { + if self.empty { + return Ok(SourceReady::Ready); + } + if self.source.is_none() { + let Some(ranges) = self.demand.ranges(self.plan.row_count()) else { + return self.demand.blocked(); + }; + let Some(first) = ranges.first() else { + self.empty = true; + return Ok(SourceReady::Ready); + }; + let start = first.start; + let end = ranges + .last() + .ok_or_else(|| vortex_err!("Missing selected range"))? + .end; + let len = usize::try_from(end - start)?; + let mask = if ranges.len() == 1 { + None + } else { + Some(Mask::from_slices( + len, + ranges + .iter() + .map(|range| { + Ok(( + usize::try_from(range.start - start)?, + usize::try_from(range.end - start)?, + )) + }) + .collect::>>()?, + )) + }; + let slice = if start == 0 && end == self.plan.row_count() { + None + } else { + Some(usize::try_from(start)?..usize::try_from(end)?) + }; + self.source = Some(ScanSource::new(self.plan.clone(), slice, mask)); + } + self.source + .as_mut() + .ok_or_else(|| vortex_err!("Missing segment source"))? + .ready(cx) + } +} + +impl Operator for DeferredScan { + fn compute(&mut self, input: Input, cx: &mut Cx<'_>) -> VortexResult { + match self.source.as_mut() { + Some(source) => source.compute(input, cx), + None if self.empty => Ok(Step::End), + None => vortex_bail!("Segment computed before its selection"), + } + } +} + +struct DeferredList { + plan: ListPackPlan, + demand: Demand, + elements: Demand, + assembled: Option<(ArrayRef, Validity)>, + done: bool, +} + +impl Source for DeferredList { + fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult { + if self.done { + return Ok(SourceReady::Ready); + } + let element_inlet = 1 + usize::from(self.plan.dtype().is_nullable()); + if self.assembled.is_none() { + let Some(selected) = self.demand.ranges(self.plan.row_count()) else { + return self.demand.blocked(); + }; + for index in 0..element_inlet { + if !cx.inlet(index).closed() { + return Ok(SourceReady::Blocked(SourceBlocked::Inlet(index))); + } + } + let offsets = join_inlet(cx, 0, self.plan.offsets()?.dtype())? + .cast(DType::Primitive(PType::U64, false.into()))? + .execute::(cx.exec())?; + let validity = if self.plan.dtype().is_nullable() { + Validity::Array(join_inlet(cx, 1, &DType::Bool(false.into()))?) + } else { + Validity::NonNullable + }; + let count = selected + .iter() + .map(|range| range.end - range.start) + .sum::(); + let valid = validity.execute_mask(usize::try_from(count)?, cx.exec())?; + let selected_indices = selected + .iter() + .flat_map(|range| range.clone()) + .map(|index| index as usize); + let (offsets, ranges) = selected_ranges( + selected_indices, + offsets.as_slice::(), + &valid, + self.plan.elements()?.row_count(), + )?; + let last = offsets.last().copied().unwrap_or(0); + let offsets = offsets.into_array(); + let stats = offsets.statistics(); + stats.set(Stat::IsSorted, Precision::Exact(true.into())); + stats.set(Stat::Min, Precision::Exact(0_u64.into())); + stats.set(Stat::Max, Precision::Exact(last.into())); + self.elements + .ranges + .set(ranges) + .map_err(|_| vortex_err!("Selection resolved twice"))?; + cx.resolve_selection( + self.elements + .gate + .ok_or_else(|| vortex_err!("Missing element gate"))?, + ); + self.assembled = Some((offsets, validity)); + } + Ok(if cx.inlet(element_inlet).closed() { + SourceReady::Ready + } else { + SourceReady::Blocked(SourceBlocked::Inlet(element_inlet)) + }) + } +} + +impl Operator for DeferredList { + fn compute(&mut self, _input: Input, cx: &mut Cx<'_>) -> VortexResult { + if self.done { + return Ok(Step::End); + } + self.done = true; + let inlet = 1 + usize::from(self.plan.dtype().is_nullable()); + let elements = join_inlet(cx, inlet, self.plan.elements()?.dtype())?; + let (offsets, validity) = self + .assembled + .take() + .ok_or_else(|| vortex_err!("List selection unresolved"))?; + Ok(Step::Final( + ListArray::try_new(elements, offsets, validity)?.into_array(), + )) + } +} diff --git a/vortex-layout/src/plan/physical/pipeline/compile.rs b/vortex-layout/src/plan/physical/pipeline/compile.rs index cf2dcb10184..7baaa9905a3 100644 --- a/vortex-layout/src/plan/physical/pipeline/compile.rs +++ b/vortex-layout/src/plan/physical/pipeline/compile.rs @@ -81,6 +81,12 @@ pub struct Compiler<'a> { } impl Compiler<'_> { + pub(crate) fn selection_gate(&mut self) -> usize { + let id = self.core.selection_waiters.len(); + self.core.selection_waiters.push(Vec::new()); + id + } + /// Compiles `plan` over `rows` of its domain, restricted to `mask`, into a chain producing /// the rows `mask` selects, in order, or `None` when that is no row. pub fn compile( diff --git a/vortex-layout/src/plan/physical/pipeline/scan.rs b/vortex-layout/src/plan/physical/pipeline/scan.rs index e5a8ee922e9..246680c030f 100644 --- a/vortex-layout/src/plan/physical/pipeline/scan.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan.rs @@ -147,6 +147,7 @@ fn drive( exec: &mut ExecutionCtx, reads: &mut SmallVec<[SegmentId; 1]>, child_requests: &mut Vec, + resolved_selections: &mut Vec, ) -> VortexResult { let Pipeline { source, @@ -173,6 +174,7 @@ fn drive( changed, child_requests, child_inbox, + resolved_selections, }; loop { // Closing needs no room, so a pipeline with nothing left to flush closes at once. @@ -312,6 +314,7 @@ pub(crate) struct Core { /// Whole values kept for later scans of the same split. pub(crate) values: ValueCache, pub(crate) plan_children: bool, + pub(crate) selection_waiters: Vec>, } struct PendingChild { @@ -341,6 +344,7 @@ impl Core { split: usize::MAX, values: ValueCache::default(), plan_children: false, + selection_waiters: Vec::new(), } } @@ -437,6 +441,7 @@ impl Core { pipe.queued = false; let mut reads = SmallVec::new(); let mut child_requests = Vec::new(); + let mut resolved_selections = Vec::new(); let progress = drive( pipe, &mut self.arena, @@ -444,7 +449,20 @@ impl Core { &mut self.exec, &mut reads, &mut child_requests, + &mut resolved_selections, )?; + for selection in resolved_selections { + for waiter in std::mem::take(&mut self.selection_waiters[selection]) { + if self.pipelines[waiter].as_ref().is_some_and(|pipe| { + pipe.state == State::Blocked(SourceBlocked::Selection(selection)) + }) { + self.make_runnable(waiter); + } + } + } + if let Progress::Blocked(SourceBlocked::Selection(selection)) = &progress { + self.selection_waiters[*selection].push(id); + } for segment_id in reads { self.ask(id, segment_id); } diff --git a/vortex-layout/src/plan/physical/pipeline/scan/validate.rs b/vortex-layout/src/plan/physical/pipeline/scan/validate.rs index d1749972d65..e5d2638de56 100644 --- a/vortex-layout/src/plan/physical/pipeline/scan/validate.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan/validate.rs @@ -63,6 +63,9 @@ impl Core { pub(super) fn validate_progress(&self, id: PipelineId, progress: &Progress) { let pipe = self.pipeline(id); match progress { + Progress::Blocked(SourceBlocked::Selection(selection)) => { + assert!(self.selection_waiters[*selection].contains(&id)); + } Progress::Blocked(SourceBlocked::Inlet(index)) => { assert!( *index < pipe.inlets.len(), diff --git a/vortex-layout/src/plan/physical/pipeline/tests.rs b/vortex-layout/src/plan/physical/pipeline/tests.rs index 0327d652b3d..b4a50eb769c 100644 --- a/vortex-layout/src/plan/physical/pipeline/tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/tests.rs @@ -998,7 +998,9 @@ fn pack_readiness_tracks_changes_and_consumption( let mut changed = InletSet::all(2); let mut child_requests = Vec::new(); let mut child_inbox = None; + let mut resolved_selections = Vec::new(); let mut cx = Cx { + resolved_selections: &mut resolved_selections, child_requests: &mut child_requests, child_inbox: &mut child_inbox, arena: &mut arena, diff --git a/vortex-layout/src/plan/physical/plans/list_pack.rs b/vortex-layout/src/plan/physical/plans/list_pack.rs index 0d459d3537b..923b15c7ea4 100644 --- a/vortex-layout/src/plan/physical/plans/list_pack.rs +++ b/vortex-layout/src/plan/physical/plans/list_pack.rs @@ -143,6 +143,10 @@ impl PlanVTable for ListPack { if mask.all_false() { return Ok(None); } + if std::env::var_os("VORTEX_LIST_PIPELINE").is_some() { + return crate::plan::physical::operators::compile_deferred_list(plan, rows, mask, compiler) + .map(Some); + } let boundaries = offset_mask(mask); let offsets = compiler .compile(&plan.offsets()?, rows.start..rows.end + 1, &boundaries)? diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index fb961a63dab..9b9138fc46d 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -56,6 +56,7 @@ arrow-array = { workspace = true } arrow-schema = { workspace = true } divan = { workspace = true } fastlanes = { workspace = true } +futures = { workspace = true } mimalloc = { workspace = true } parquet = { workspace = true } rand = { workspace = true } diff --git a/vortex/examples/generate_list_fixture.rs b/vortex/examples/generate_list_fixture.rs new file mode 100644 index 00000000000..6968df77f7a --- /dev/null +++ b/vortex/examples/generate_list_fixture.rs @@ -0,0 +1,187 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Generate large list rows for selection-scan diagnostics. +//! +//! Usage: `generate_list_fixture OUTPUT balanced|skewed|small primitive|nested TARGET_BYTES`. +//! Writes a new file and verifies full decoded equality through V1. Element block sizing is +//! independent of outer row count; nullable lists, empty lists, and nullable elements are included. + +use std::collections::BTreeMap; +use std::path::Path; +use std::sync::Arc; + +use anyhow::Result; +use anyhow::ensure; +use futures::stream; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use vortex::VortexSessionDefault; +use vortex::array::IntoArray; +use vortex::array::VortexSessionExecute; +use vortex::array::arrays::BoolArray; +use vortex::array::arrays::ListArray; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::StructArray; +use vortex::array::stream::ArrayStreamAdapter; +use vortex::array::stream::ArrayStreamExt; +use vortex::array::validity::Validity; +use vortex::buffer::Buffer; +use vortex::editions::Edition; +use vortex::editions::EditionDeclaration; +use vortex::editions::EditionId; +use vortex::editions::EditionMember; +use vortex::editions::EditionSessionExt; +use vortex::file::OpenOptionsSessionExt; +use vortex::file::WriteOptionsSessionExt; +use vortex::file::WriteStrategyBuilder; +use vortex::io::session::RuntimeSessionExt; +use vortex::layout::LayoutRef; +use vortex::layout::layouts::table::TableStrategy; +use vortex::session::VortexSession; +use vortex_arrow::ArrowSessionExt; + +const INNER_LENGTH: u64 = 128; + +static EDITION: EditionDeclaration = EditionDeclaration { + edition: Edition { + id: EditionId::new("large-list-fixture", 2026, 10, 0), + min_library_version: None, + }, + added: &[EditionMember::layout(&"vortex.list")], +}; + +fn count_layouts(layout: &LayoutRef, counts: &mut BTreeMap) -> Result<()> { + *counts.entry(layout.encoding_id().to_string()).or_default() += 1; + for child in layout.children()? { + count_layouts(&child, counts)?; + } + Ok(()) +} + +#[tokio::main] +async fn main() -> Result<()> { + let args = std::env::args().skip(1).collect::>(); + ensure!( + args.len() == 4, + "Expected OUTPUT balanced|skewed|small primitive|nested TARGET_BYTES" + ); + ensure!( + matches!(args[1].as_str(), "balanced" | "skewed" | "small"), + "Unknown distribution" + ); + ensure!( + matches!(args[2].as_str(), "primitive" | "nested"), + "Unknown element shape" + ); + let target_bytes = args[3].parse::()?; + ensure!(target_bytes > 0, "TARGET_BYTES must be positive"); + let session = VortexSession::default().with_tokio(); + session.register_edition(&EDITION)?; + session.enable_edition(EDITION.edition.id)?; + + let rows = if args[1] == "small" { 1_000_000 } else { 512 }; + let mut offsets = vec![0_u64]; + let mut total = 0_u64; + for row in 0..rows { + total += if row % 19 == 0 { + 0 + } else if args[1] == "small" { + 1 + row % 20 + } else if args[1] == "balanced" { + 8192 + } else if row % 16 == 0 { + 262_144 + } else { + INNER_LENGTH + }; + offsets.push(total); + } + let mut rng = StdRng::seed_from_u64(42); + let mut elements = PrimitiveArray::from_option_iter((0..total).map(|index| { + let value = rng.random::(); + (index % 7 != 0).then_some(value) + })) + .into_array(); + if args[2] == "nested" { + let inner_count = total / INNER_LENGTH; + elements = ListArray::try_new( + elements, + Buffer::from_iter((0..=inner_count).map(|index| index * INNER_LENGTH)).into_array(), + Validity::Array( + BoolArray::from_iter((0..inner_count).map(|index| index % 11 != 0)).into_array(), + ), + )? + .into_array(); + for offset in &mut offsets { + *offset /= INNER_LENGTH; + } + } + let values = ListArray::try_new( + elements, + Buffer::from(offsets).into_array(), + Validity::Array(BoolArray::from_iter((0..rows).map(|row| row % 17 != 0)).into_array()), + )? + .into_array(); + let input = StructArray::from_fields(&[ + ("id", Buffer::from_iter(0..rows).into_array()), + ("values", values), + ])? + .into_array(); + let elements = WriteStrategyBuilder::from_session(&session) + .with_data_block_target_bytes(Some(target_bytes)) + .build(); + let defaults = WriteStrategyBuilder::from_session(&session).build(); + let strategy = TableStrategy::new(Arc::clone(&defaults), defaults) + .with_list_layout_factory(move |list| Arc::new(list.with_elements(Arc::clone(&elements)))); + let mut output = tokio::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(&args[0]) + .await?; + let summary = session + .write_options() + .with_strategy(Arc::new(strategy)) + .write( + &mut output, + ArrayStreamAdapter::new(input.dtype().clone(), stream::iter([Ok(input.clone())])), + ) + .await?; + drop(output); + let file = session + .open_options() + .open_path(Path::new(&args[0])) + .await?; + let actual = ArrayStreamAdapter::new(file.dtype().clone(), file.scan()?.into_stream()?) + .read_all() + .await?; + let mut ctx = session.create_execution_ctx(); + let arrow = session.arrow(); + let field = arrow.to_arrow_field("", input.dtype())?; + let expected = arrow.execute_arrow(input, Some(&field), &mut ctx)?; + let actual = arrow.execute_arrow(actual, Some(&field), &mut ctx)?; + ensure!( + expected.as_ref() == actual.as_ref(), + "Decoded fixture differs from generated input" + ); + let mut layouts = BTreeMap::new(); + count_layouts(file.footer().layout(), &mut layouts)?; + ensure!( + layouts.get("vortex.list").copied().unwrap_or(0) > 0, + "Expected ListLayout" + ); + println!( + "{}", + serde_json::json!({ + "output": args[0], "distribution": args[1], "shape": args[2], + "target_bytes": target_bytes, "rows": rows, "primitive_elements": total, + "max_row_payload_bytes": match args[1].as_str() { + "small" => 20 * 8, "balanced" => 8192 * 8, _ => 262_144 * 8, + }, + "bytes": summary.size(), "correct": true, "layouts": layouts, + "layout_tree": file.footer().layout().display_tree().to_string(), + }) + ); + Ok(()) +} diff --git a/vortex/examples/selection_scan.rs b/vortex/examples/selection_scan.rs new file mode 100644 index 00000000000..57def21b9e0 --- /dev/null +++ b/vortex/examples/selection_scan.rs @@ -0,0 +1,470 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Diagnose selection scans on an existing Vortex file. +//! +//! Usage: `selection_scan FILE ITERATIONS ROWS uniform|clustered|all cached|reopen memory|disk +//! time|metrics|check|decode|decode-metrics [SEED]`. Set `VORTEX_SCAN_V2=0/1` in separate processes. +//! Timings exclude instrumentation in `time` and `decode` modes. The flag selects the explicit V2 +//! builder here only. `check` compares against a full V1 scan followed by a take. +//! `decode` consumes the selected rows in all columns as Arrow record batches; `decode-metrics` +//! measures allocations while doing the same work. +//! Instrumented disk runs also report physical reads/bytes after coalescing, summed over all +//! iterations and excluding the initial file open (reopen-mode opens are included). +//! Build with `VORTEX_SELECTION_METRICS=1` for the separate allocation-instrumented executable. +//! Use `indices:16,32` in place of the pattern to select exact, strictly increasing row indices; +//! ROWS must match the number supplied. This supports targeted large-row and null-row checks. + +use std::alloc::GlobalAlloc; +use std::alloc::Layout; +use std::hint::black_box; +use std::path::Path; +use std::pin::pin; +use std::sync::Arc; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; +use std::time::Instant; + +use anyhow::Result; +use anyhow::bail; +use anyhow::ensure; +use futures::StreamExt; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use vortex::VortexSessionDefault; +use vortex::array::ArrayRef; +use vortex::array::Canonical; +use vortex::array::IntoArray; +use vortex::array::VortexSessionExecute; +use vortex::array::memory::MemorySessionExt; +use vortex::array::stream::ArrayStreamAdapter; +use vortex::array::stream::ArrayStreamExt; +use vortex::buffer::Buffer; +use vortex::buffer::ByteBuffer; +use vortex::file::OpenOptionsSessionExt; +use vortex::file::VortexFile; +use vortex::io::InstrumentedReadAt; +use vortex::io::session::RuntimeSessionExt; +use vortex::io::std_file::FileReadAt; +use vortex::layout::scan::v2::ScanBuilder as V2ScanBuilder; +use vortex::layout::segments::SegmentFuture; +use vortex::layout::segments::SegmentId; +use vortex::layout::segments::SegmentSource; +use vortex::metrics::DefaultMetricsRegistry; +use vortex::metrics::MetricValue; +use vortex::metrics::MetricsRegistry; +use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; +use vortex::session::VortexSession; +use vortex_arrow::ArrowSessionExt; + +const TRACK_ALLOCATIONS: bool = option_env!("VORTEX_SELECTION_METRICS").is_some(); +static CALLS: AtomicU64 = AtomicU64::new(0); +static BYTES: AtomicU64 = AtomicU64::new(0); +static LIVE: AtomicU64 = AtomicU64::new(0); +static PEAK: AtomicU64 = AtomicU64::new(0); + +struct Allocator; + +#[global_allocator] +static ALLOCATOR: Allocator = Allocator; + +fn allocated(size: usize) { + CALLS.fetch_add(1, Ordering::Relaxed); + BYTES.fetch_add(size as u64, Ordering::Relaxed); + let live = LIVE.fetch_add(size as u64, Ordering::Relaxed) + size as u64; + PEAK.fetch_max(live, Ordering::Relaxed); +} + +// SAFETY: All operations delegate to the same allocator with the original pointer and layout. +unsafe impl GlobalAlloc for Allocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + // SAFETY: The caller supplies a valid allocation layout. + let pointer = unsafe { mimalloc::MiMalloc.alloc(layout) }; + if !pointer.is_null() && TRACK_ALLOCATIONS { + allocated(layout.size()); + } + pointer + } + + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + // SAFETY: The caller supplies a valid allocation layout. + let pointer = unsafe { mimalloc::MiMalloc.alloc_zeroed(layout) }; + if !pointer.is_null() && TRACK_ALLOCATIONS { + allocated(layout.size()); + } + pointer + } + + unsafe fn realloc(&self, pointer: *mut u8, layout: Layout, size: usize) -> *mut u8 { + // SAFETY: The caller guarantees the original allocation and requested size are valid. + let pointer = unsafe { mimalloc::MiMalloc.realloc(pointer, layout, size) }; + if !pointer.is_null() && TRACK_ALLOCATIONS { + LIVE.fetch_sub(layout.size() as u64, Ordering::Relaxed); + allocated(size); + } + pointer + } + + unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) { + if TRACK_ALLOCATIONS { + LIVE.fetch_sub(layout.size() as u64, Ordering::Relaxed); + } + // SAFETY: This allocation came from mimalloc with the same layout. + unsafe { mimalloc::MiMalloc.dealloc(pointer, layout) }; + } +} + +struct CountingSegments { + inner: Arc, + calls: Arc<[AtomicU64]>, +} + +impl SegmentSource for CountingSegments { + fn request(&self, id: SegmentId) -> SegmentFuture { + self.calls[*id as usize].fetch_add(1, Ordering::Relaxed); + self.inner.request(id) + } +} + +async fn open( + session: &VortexSession, + path: &Path, + bytes: Option<&ByteBuffer>, + metrics: Option<&Arc>, +) -> Result { + let options = session.open_options().with_layout_reader_cache(); + Ok(match (bytes, metrics) { + (Some(bytes), _) => options.open_buffer(bytes.clone())?, + (None, Some(metrics)) => { + let read = Arc::new(FileReadAt::open_with_allocator( + path, + session.handle(), + session.allocator(), + )?); + options + .with_metrics_registry(metrics.clone()) + .open_read(InstrumentedReadAt::new(read, metrics.as_ref())) + .await? + } + (None, None) => options.open_path(path).await?, + }) +} + +fn physical_io(metrics: &DefaultMetricsRegistry) -> [u64; 4] { + let mut totals = [0; 4]; + for metric in metrics.snapshot() { + match (metric.name().as_ref(), metric.value()) { + ("vortex.io.read.total_size", MetricValue::Counter(value)) => { + totals[0] += value.value(); + } + ("vortex.io.read.size", MetricValue::Histogram(value)) => { + totals[1] += value.count() as u64; + } + ("io.requests.individual", MetricValue::Counter(value)) => { + totals[2] += value.value(); + } + ("io.requests.coalesced", MetricValue::Counter(value)) => { + totals[3] += value.value(); + } + _ => {} + } + } + totals +} + +fn v2() -> bool { + std::env::var("VORTEX_SCAN_V2").is_ok_and(|value| value == "1") +} + +fn v2_scan( + file: &VortexFile, + indices: &Buffer, + all: bool, + session: &VortexSession, +) -> Result { + let builder = V2ScanBuilder::new(session.clone(), file.scan_file()); + Ok(if all { + builder + } else { + builder.with_row_indices(StrictSortedBuffer::try_new(indices.clone())?) + }) +} + +async fn scan( + file: &VortexFile, + indices: &Buffer, + all: bool, + session: &VortexSession, +) -> Result { + let stream = if v2() { + v2_scan(file, indices, all, session)?.into_stream()? + } else { + let builder = file.scan()?; + let builder = if all { + builder + } else { + builder.with_row_indices(StrictSortedBuffer::try_new(indices.clone())?) + }; + builder.into_stream()?.boxed() + }; + let array = ArrayStreamAdapter::new(file.dtype().clone(), stream) + .read_all() + .await?; + Ok(array + .execute::(&mut session.create_execution_ctx())? + .into_array()) +} + +async fn decode( + file: &VortexFile, + indices: &Buffer, + all: bool, + session: &VortexSession, +) -> Result { + let schema = Arc::new(session.arrow().to_arrow_schema(file.dtype())?); + let stream = if v2() { + v2_scan(file, indices, all, session)? + .into_record_batch_stream(schema)? + .boxed() + } else { + let builder = file.scan()?; + let builder = if all { + builder + } else { + builder.with_row_indices(StrictSortedBuffer::try_new(indices.clone())?) + }; + builder.into_record_batch_stream(schema)?.boxed() + }; + let mut stream = pin!(stream); + let mut rows = 0; + while let Some(batch) = stream.next().await { + rows += batch?.num_rows(); + } + Ok(rows) +} + +fn indices(rows: u64, count: usize, pattern: &str, seed: u64) -> Result> { + if let Some(indices) = pattern.strip_prefix("indices:") { + let selected = if indices.is_empty() { + Vec::new() + } else { + indices + .split(',') + .map(str::parse::) + .collect::, _>>()? + }; + ensure!( + selected.len() == count, + "ROWS must match the explicit selection length" + ); + ensure!( + selected.iter().all(|&index| index < rows), + "Selected row is outside the file" + ); + ensure!( + selected.windows(2).all(|pair| pair[0] < pair[1]), + "Indices must be strictly increasing" + ); + return Ok(Buffer::from(selected)); + } + let mut rng = StdRng::seed_from_u64(seed); + let mut result = Vec::with_capacity(count); + ensure!(count as u64 <= rows, "Selection exceeds the file row count"); + match pattern { + "all" => ensure!(count == 0, "Pass ROWS=0 for an all-row scan"), + "uniform" => { + while result.len() < count { + result.push(rng.random_range(0..rows)); + result.sort_unstable(); + result.dedup(); + } + } + "clustered" => { + while result.len() < count { + let len = (count - result.len()).min(20) as u64; + let start = rng.random_range(0..=rows - len); + result.extend(start..start + len); + result.sort_unstable(); + result.dedup(); + } + } + _ => bail!("Unknown pattern {pattern}"), + } + Ok(Buffer::from(result)) +} + +fn rss() -> Result { + let Ok(status) = std::fs::read_to_string("/proc/self/status") else { + return Ok(serde_json::Value::Null); + }; + let get = |name: &str| -> Option { + status + .lines() + .find(|line| line.starts_with(name))? + .split_whitespace() + .nth(1)? + .parse::() + .ok() + .map(|kb| kb * 1024) + }; + Ok(serde_json::json!({"current": get("VmRSS:"), "peak": get("VmHWM:")})) +} + +#[tokio::main] +async fn main() -> Result<()> { + let args = std::env::args().skip(1).collect::>(); + ensure!( + matches!(args.len(), 7 | 8), + "Expected FILE ITERATIONS ROWS uniform|clustered|all cached|reopen memory|disk time|metrics|check|decode|decode-metrics [SEED]" + ); + let iterations = args[1].parse::()?; + ensure!(iterations > 0, "ITERATIONS must be positive"); + let count = args[2].parse::()?; + let cached = match args[4].as_str() { + "cached" => true, + "reopen" => false, + other => bail!("Unknown open mode {other}"), + }; + let metrics = matches!(args[6].as_str(), "metrics" | "decode-metrics"); + ensure!( + matches!( + args[6].as_str(), + "time" | "metrics" | "check" | "decode" | "decode-metrics" + ), + "Unknown measurement mode" + ); + ensure!( + args[6] == "check" || metrics == TRACK_ALLOCATIONS, + "Build with VORTEX_SELECTION_METRICS=1 only for metrics mode" + ); + let session = VortexSession::default().with_tokio(); + let path = Path::new(&args[0]); + let bytes = match args[5].as_str() { + "memory" => Some(ByteBuffer::from(std::fs::read(path)?)), + "disk" => None, + other => bail!("Unknown storage mode {other}"), + }; + let io_metrics = metrics.then(|| Arc::new(DefaultMetricsRegistry::default())); + let file = open(&session, path, bytes.as_ref(), io_metrics.as_ref()).await?; + let seed = args + .get(7) + .map(|seed| seed.parse::()) + .transpose()? + .unwrap_or(42); + let selected = indices(file.row_count(), count, &args[3], seed)?; + let all = args[3] == "all"; + let decode_rows = matches!(args[6].as_str(), "decode" | "decode-metrics"); + let expected_rows = if all { + usize::try_from(file.row_count())? + } else { + selected.len() + }; + if args[6] == "check" { + let expected = ArrayStreamAdapter::new( + file.dtype().clone(), + file.scan()?.prepare()?.execute_stream(None)?, + ) + .read_all() + .await?; + let expected = if all { + expected + } else { + expected.take(selected.clone().into_array())? + }; + let actual = scan(&file, &selected, all, &session).await?; + let mut ctx = session.create_execution_ctx(); + let arrow = session.arrow(); + let target = arrow.to_arrow_field("", file.dtype())?; + let expected = arrow.execute_arrow(expected, Some(&target), &mut ctx)?; + let actual = arrow.execute_arrow(actual, Some(&target), &mut ctx)?; + ensure!( + actual.as_ref() == expected.as_ref(), + "Selected values differ from full-scan take" + ); + println!( + "{}", + serde_json::json!({"correct": true, "rows": actual.len(), "indices": selected.as_slice()}) + ); + return Ok(()); + } + let locations = file.scan_file().locations; + let calls: Arc<[AtomicU64]> = locations.iter().map(|_| AtomicU64::new(0)).collect(); + let instrument = |file: VortexFile| { + if metrics { + let source = Arc::new(CountingSegments { + inner: file.segment_source(), + calls: Arc::clone(&calls), + }); + file.with_segment_source(source) + } else { + file + } + }; + let file = instrument(file); + let initial_io = io_metrics.as_ref().map(|metrics| physical_io(metrics)); + let mut runs = Vec::with_capacity(iterations); + for _ in 0..iterations { + for calls in calls.iter() { + calls.store(0, Ordering::Relaxed); + } + let before_live = LIVE.load(Ordering::Relaxed); + let before_calls = CALLS.load(Ordering::Relaxed); + let before_bytes = BYTES.load(Ordering::Relaxed); + PEAK.store(before_live, Ordering::Relaxed); + let start = Instant::now(); + let reopened = if cached { + None + } else { + Some(instrument( + open(&session, path, bytes.as_ref(), io_metrics.as_ref()).await?, + )) + }; + let input = reopened.as_ref().unwrap_or(&file); + let (array, actual_rows) = if decode_rows { + (None, decode(input, &selected, all, &session).await?) + } else { + let array = scan(input, &selected, all, &session).await?; + let rows = array.len(); + (Some(array), rows) + }; + drop(reopened); + let elapsed = start.elapsed().as_nanos(); + ensure!(actual_rows == expected_rows, "Wrong selected row count"); + let peak_live = PEAK.load(Ordering::Relaxed).saturating_sub(before_live); + let output_live = LIVE.load(Ordering::Relaxed) as i128 - before_live as i128; + drop(black_box(array)); + let retained_live = LIVE.load(Ordering::Relaxed) as i128 - before_live as i128; + runs.push(serde_json::json!({ + "elapsed_ns": elapsed, "instrumented": metrics, + "allocation_calls": CALLS.load(Ordering::Relaxed) - before_calls, + "allocated_bytes": BYTES.load(Ordering::Relaxed) - before_bytes, + "baseline_live_bytes": before_live, + "total_peak_live_bytes": before_live + peak_live, + "post_drop_live_bytes": LIVE.load(Ordering::Relaxed), + "peak_live_bytes": peak_live, "output_live_bytes": output_live, + "retained_live_bytes": retained_live, + "requests": calls.iter().map(|n| n.load(Ordering::Relaxed)).sum::(), + "requested_bytes": calls.iter().zip(locations.iter()).map(|(n, l)| n.load(Ordering::Relaxed) * u64::from(l.length)).sum::(), + "unique_segments": calls.iter().filter(|n| n.load(Ordering::Relaxed) != 0).count(), + "unique_bytes": calls.iter().zip(locations.iter()).filter(|(n, _)| n.load(Ordering::Relaxed) != 0).map(|(_, l)| u64::from(l.length)).sum::(), + })); + } + let physical_io = io_metrics + .as_ref() + .zip(initial_io) + .map(|(metrics, before)| { + let after = physical_io(metrics); + serde_json::json!({ + "bytes": after[0] - before[0], + "reads": after[1] - before[1], + "individual_requests": after[2] - before[2], + "coalesced_requests": after[3] - before[3], + }) + }); + println!( + "{}", + serde_json::json!({"args": args, "indices": selected.as_slice(), "runs": runs, "physical_io": physical_io, "rss": rss()?, "v2": std::env::var("VORTEX_SCAN_V2").ok()}) + ); + Ok(()) +} From e3cbd8b2d10a2c9e69b1323947bb9ea396059bdc Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 19 Oct 2026 15:28:36 +0100 Subject: [PATCH 4/7] Activate only selected list chunks in the pipeline scheduler Signed-off-by: Joe Isaacs --- .github/workflows/pr-bench-runner.yml | 26 +++- vortex-layout/benches/pipeline_scan.rs | 119 ++++++++++++++++++ vortex-layout/src/plan/physical/operator.rs | 14 +++ .../plan/physical/operators/list/deferred.rs | 114 ++++++++++++++--- .../src/plan/physical/pipeline/mod.rs | 8 +- .../src/plan/physical/pipeline/scan.rs | 35 ++++++ .../plan/physical/pipeline/scan/validate.rs | 5 + .../physical/pipeline/scheduling_tests.rs | 44 +++++++ .../src/plan/physical/pipeline/tests.rs | 76 +++++++++++ .../src/plan/physical/plans/list_pack.rs | 4 +- 10 files changed, 417 insertions(+), 28 deletions(-) diff --git a/.github/workflows/pr-bench-runner.yml b/.github/workflows/pr-bench-runner.yml index 8b9fabb9643..006cff056d7 100644 --- a/.github/workflows/pr-bench-runner.yml +++ b/.github/workflows/pr-bench-runner.yml @@ -84,6 +84,18 @@ jobs: run: | cargo build --locked --package vortex --profile release_debug \ --example selection_scan --example generate_list_fixture + cargo bench --locked --package vortex-layout --profile release_debug \ + --bench pipeline_scan --no-run --message-format=json > /tmp/list-micro-build.jsonl + python3 - <<'PY' + import json + import shutil + + with open('/tmp/list-micro-build.jsonl') as source: + for line in source: + artifact = json.loads(line) + if artifact.get('target', {}).get('name') == 'pipeline_scan' and artifact.get('executable'): + shutil.copy2(artifact['executable'], 'target/release_debug/pipeline_scan') + PY - name: Upload list pipeline comparison binaries if: inputs.benchmark_id == 'random-access-bench' && inputs.remote_data_dir == '' uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7 @@ -92,6 +104,7 @@ jobs: path: | target/release_debug/examples/selection_scan target/release_debug/examples/generate_list_fixture + target/release_debug/pipeline_scan - name: Upload benchmark binary uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7 with: @@ -142,7 +155,7 @@ jobs: uses: actions/download-artifact@9000827ccba6bdab643e8b6fd33ac0654aef8333 # v8 with: name: list-pipeline-binaries - path: target/release_debug/examples/ + path: target/release_debug/ - name: Setup Polar Signals if: github.event.pull_request.head.repo.fork == false @@ -207,8 +220,17 @@ jobs: RUSTFLAGS: "-C target-cpu=native -C force-frame-pointers=yes" run: | chmod +x target/release_debug/examples/selection_scan \ - target/release_debug/examples/generate_list_fixture + target/release_debug/examples/generate_list_fixture target/release_debug/pipeline_scan python3 scripts/bench-list-pipeline.py + for round in 0 1; do + versions="current fixed" + if [ "$round" = 1 ]; then versions="fixed current"; fi + for version in $versions; do + if [ "$version" = fixed ]; then export VORTEX_LIST_PIPELINE=1; else unset VORTEX_LIST_PIPELINE; fi + target/release_debug/pipeline_scan --bench list_scan --sample-count 100 --sample-size 100 \ + > "list-pipeline-results/micro-$version-$round.log" + done + done - name: Upload list pipeline comparison results if: always() && inputs.benchmark_id == 'random-access-bench' && inputs.remote_data_dir == '' diff --git a/vortex-layout/benches/pipeline_scan.rs b/vortex-layout/benches/pipeline_scan.rs index e9980aa6f66..d73ba6a0331 100644 --- a/vortex-layout/benches/pipeline_scan.rs +++ b/vortex-layout/benches/pipeline_scan.rs @@ -10,6 +10,7 @@ use std::hint::black_box; use divan::Bencher; use mimalloc::MiMalloc; use vortex_array::ArrayContext; +use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::buffer::BufferHandle; @@ -17,7 +18,11 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::StructFields; use vortex_array::serde::SerializeOptions; use vortex_buffer::Alignment; +use vortex_buffer::Buffer; use vortex_buffer::ByteBufferMut; +use vortex_layout::plan::physical::ConcatPlan; +use vortex_layout::plan::physical::FilterPlan; +use vortex_layout::plan::physical::ListPackPlan; use vortex_layout::plan::physical::PackPlan; use vortex_layout::plan::physical::PlanRef; use vortex_layout::plan::physical::SegmentScanPlan; @@ -26,6 +31,7 @@ use vortex_layout::plan::physical::pipeline::Split; use vortex_layout::plan::physical::pipeline::Turn; use vortex_layout::segments::SegmentId; use vortex_layout::session::LayoutSession; +use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::ReadContext; @@ -131,3 +137,116 @@ fn scan(bencher: Bencher, columns: u32) { } }); } + +struct ListFixture { + session: VortexSession, + plan: PlanRef, + segments: Vec, + mask: Mask, +} + +impl ListFixture { + fn new(chunks: u32, uniform: bool) -> Self { + const LISTS: usize = 65_536; + const ELEMENTS_PER_LIST: usize = 8; + let session = vortex_array::array_session().with::(); + let chunk_len = LISTS * ELEMENTS_PER_LIST / chunks as usize; + let elements = Buffer::from_iter(0..chunk_len as u64).into_array(); + let offsets = Buffer::from_iter((0..=LISTS as u64).map(|row| row * 8)).into_array(); + let serialize = |array: &ArrayRef| { + let context = ArrayContext::empty(); + let buffers = array + .serialize( + &context, + &session, + &SerializeOptions { + offset: 0, + include_padding: true, + }, + ) + .unwrap(); + let mut bytes = ByteBufferMut::empty_aligned(Alignment::new(64)); + for buffer in buffers { + bytes.extend_from_slice(buffer.as_ref()); + } + (BufferHandle::new_host(bytes.freeze()), context) + }; + let (element_bytes, element_context) = serialize(&elements); + let (offset_bytes, offset_context) = serialize(&offsets); + let children = (0..chunks) + .map(|id| { + FilterPlan::new( + SegmentScanPlan::new( + elements.dtype().clone(), + chunk_len as u64, + SegmentId::from(id), + ReadContext::new(element_context.to_ids()), + None, + ) + .into_plan(), + ) + .into_plan() + }) + .collect(); + let elements = ConcatPlan::try_new(elements.dtype().clone(), children).unwrap(); + let offsets = FilterPlan::new( + SegmentScanPlan::new( + offsets.dtype().clone(), + LISTS as u64 + 1, + SegmentId::from(chunks), + ReadContext::new(offset_context.to_ids()), + None, + ) + .into_plan(), + ); + let plan = ListPackPlan::try_new( + Nullability::NonNullable, + LISTS as u64, + elements.into_plan(), + offsets.into_plan(), + None, + ) + .unwrap() + .into_plan(); + let mut segments = vec![element_bytes; chunks as usize]; + segments.push(offset_bytes); + let indices = (0..100).map(|index| if uniform { index * (LISTS / 100) } else { index }); + Self { + session, + plan, + segments, + mask: Mask::from_indices(LISTS, indices), + } + } + + fn scan(&self) { + let split = Split { + rows: 0..self.mask.len() as u64, + mask: self.mask.clone(), + }; + let mut scan = Scan::try_new(self.session.clone(), self.plan.clone(), vec![split]).unwrap(); + loop { + match scan.step().unwrap() { + Turn::Read(read) => { + scan.deliver(read.id, self.segments[*read.segment_id as usize].clone()) + .unwrap(); + } + Turn::Output(_, array) => { + assert_eq!(array.len(), 100); + black_box(array); + } + Turn::Child(_) => unreachable!("microbenchmark has no child planner"), + Turn::Waiting => unreachable!("every read is delivered immediately"), + Turn::Done => break, + } + } + } +} + +/// Fixed selected row count, varying physical chunk count. Compare in separate processes with +/// `VORTEX_LIST_PIPELINE` unset and set to `1`; no disk IO or Arrow conversion is timed. +#[divan::bench(args = [8, 128, 1024], consts = [false, true])] +fn list_scan(bencher: Bencher, chunks: u32) { + let fixture = ListFixture::new(chunks, UNIFORM); + bencher.bench_local(|| fixture.scan()); +} diff --git a/vortex-layout/src/plan/physical/operator.rs b/vortex-layout/src/plan/physical/operator.rs index 12d1d56c648..d74b2f714fb 100644 --- a/vortex-layout/src/plan/physical/operator.rs +++ b/vortex-layout/src/plan/physical/operator.rs @@ -10,9 +10,12 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::buffer::BufferHandle; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_session::VortexSession; use super::pipeline::Arena; +use super::pipeline::Chain; +use super::pipeline::Compiler; use super::pipeline::Inlet; use super::pipeline::PortId; use crate::plan::physical::pipeline::ChildProjection; @@ -35,6 +38,8 @@ pub enum Step { pub enum SourceReady { /// The next compute call can produce output or finish without blocking. Ready, + /// Replace this source with selected work in the same scheduler, through [`Source::activate`]. + Activate, /// Nothing can happen until the condition clears. Blocked(SourceBlocked), } @@ -56,6 +61,15 @@ pub enum SourceBlocked { /// The start of a pipeline: produces batches from its inlets and the segments it asks to have /// read. pub trait Source: Operator { + /// Activates work whose selection was unavailable at initial compilation. + /// + /// Called only after returning [`SourceReady::Activate`]. The source must have no inlets, + /// pending reads, or emitted batches. The resulting chain inherits its downstream stages + /// and outlets, without starting another scan or handing work to the outer planner. + fn activate(&mut self, _compiler: &mut Compiler<'_>) -> VortexResult { + vortex_bail!("Source does not support deferred activation") + } + /// Prepares the next [`compute`](Operator::compute) call, after the driver checks outlets. /// Returning [`SourceReady::Ready`] guarantees that call can proceed without blocking. /// Any input or choice made here can be retained by the source for that call. diff --git a/vortex-layout/src/plan/physical/operators/list/deferred.rs b/vortex-layout/src/plan/physical/operators/list/deferred.rs index 18541516441..03cec976d05 100644 --- a/vortex-layout/src/plan/physical/operators/list/deferred.rs +++ b/vortex-layout/src/plan/physical/operators/list/deferred.rs @@ -1,7 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Benchmark prototype: prebuild list pipelines and resolve their selection at runtime. +//! Benchmark prototype: build list dependencies upfront and resolve selection at runtime. +//! +//! Chunk dispatchers wait for offsets and validity, then activate only selected children in the +//! same scheduler. Unselected chunks need no runtime sources or ports. //! //! Enabled only by `VORTEX_LIST_PIPELINE`. Supports primitive and nested-list fixture plans; //! unsupported plans fail explicitly. This experiment does not implement the planner's decoded-work @@ -29,6 +32,7 @@ use vortex_mask::Mask; use super::selected_ranges; use crate::plan::physical::Concat; +use crate::plan::physical::ConcatPlan; use crate::plan::physical::Filter; use crate::plan::physical::LegacyStats; use crate::plan::physical::ListPack; @@ -50,6 +54,7 @@ use crate::plan::physical::operators::UNBOUNDED; use crate::plan::physical::operators::join_inlet; use crate::plan::physical::pipeline::Chain; use crate::plan::physical::pipeline::Compiler; +use crate::plan::physical::pipeline::overlapping; #[derive(Clone)] struct Demand { @@ -64,7 +69,14 @@ impl Demand { fn ranges(&self, len: u64) -> Option>> { let selected = self.ranges.get()?; let mut result: Vec> = Vec::new(); - for range in selected { + let first = selected.partition_point(|range| { + range.end <= self.domain.start + || range.end.saturating_add(u64::from(self.boundaries)) <= self.offset + }); + for range in &selected[first..] { + if range.start >= self.domain.end || range.start >= self.offset + len { + break; + } let start = range.start.max(self.domain.start); let end = range.end.min(self.domain.end); if start >= end { @@ -118,7 +130,7 @@ pub(crate) fn compile_deferred_list( compiler: &mut Compiler<'_>, ) -> VortexResult { let ranges = match mask.slices() { - AllOr::All => vec![rows.clone()], + AllOr::All => vec![rows], AllOr::None => Vec::new(), AllOr::Some(slices) => slices .iter() @@ -182,16 +194,15 @@ fn compile(plan: &PlanRef, demand: Demand, compiler: &mut Compiler<'_>) -> Vorte return compile(&stats.data_plan()?, demand, compiler); } if let Some(concat) = plan.as_opt::() { - let mut chains = Vec::new(); - for (index, &start) in concat.row_offsets().iter().enumerate() { - chains.push(compile( - &concat.child_required(index)?, - demand.shifted(start), - compiler, - )?); - } - let count = chains.len(); - return Ok(compiler.join(chains, ConcatSource::new(count))); + let mut source = DeferredConcat { + plan: concat.clone(), + demand, + }; + return if source.demand.ranges.get().is_some() { + source.activate(compiler) + } else { + Ok(Chain::new(source)) + }; } if let Some(list) = plan.as_opt::() { return compile_list(list, demand, compiler); @@ -210,6 +221,57 @@ fn compile(plan: &PlanRef, demand: Demand, compiler: &mut Compiler<'_>) -> Vorte ) } +/// Keep one dispatcher while selection is unknown; lower only selected children afterwards. +struct DeferredConcat { + plan: ConcatPlan, + demand: Demand, +} + +impl Source for DeferredConcat { + fn ready(&mut self, _cx: &mut Cx<'_>) -> VortexResult { + if self.demand.ranges.get().is_some() { + Ok(SourceReady::Activate) + } else { + self.demand.blocked() + } + } + + fn activate(&mut self, compiler: &mut Compiler<'_>) -> VortexResult { + let ranges = self + .demand + .ranges(self.plan.row_count()) + .ok_or_else(|| vortex_err!("Chunks activated before their selection"))?; + let mut chains = Vec::new(); + let mut previous = None; + for range in ranges { + for (index, start, _) in overlapping(&self.plan, &range) { + // Several selected ranges can lie in one chunk. Give that chunk one source, + // which applies their combined selection and reads its segment only once. + if previous == Some(index) { + continue; + } + previous = Some(index); + chains.push(compile( + &self.plan.child_required(index)?, + self.demand.shifted(start), + compiler, + )?); + } + } + let count = chains.len(); + if count == 1 { + return Ok(chains.remove(0)); + } + Ok(compiler.join(chains, ConcatSource::new(count))) + } +} + +impl Operator for DeferredConcat { + fn compute(&mut self, _input: Input, _cx: &mut Cx<'_>) -> VortexResult { + vortex_bail!("Chunk dispatcher computed before activation") + } +} + struct DeferredScan { plan: SegmentScanPlan, demand: Demand, @@ -312,18 +374,30 @@ impl Source for DeferredList { .map(|range| range.end - range.start) .sum::(); let valid = validity.execute_mask(usize::try_from(count)?, cx.exec())?; - let selected_indices = selected - .iter() - .flat_map(|range| range.clone()) - .map(|index| index as usize); - let (offsets, ranges) = selected_ranges( + let selected = selected + .into_iter() + .map(|range| Ok(usize::try_from(range.start)?..usize::try_from(range.end)?)) + .collect::>>()?; + let selected_indices = selected.iter().flat_map(|range| range.clone()); + let (mut selected_offsets, mut ranges) = selected_ranges( selected_indices, offsets.as_slice::(), &valid, self.plan.elements()?.row_count(), )?; - let last = offsets.last().copied().unwrap_or(0); - let offsets = offsets.into_array(); + if selected.len() == 1 + && !valid.all_true() + && let Some(retained) = super::work::preserve_null_span( + self.plan.elements()?.dtype(), + offsets.as_slice::(), + selected_offsets.last().copied().unwrap_or(0), + &mut ranges, + ) + { + selected_offsets = retained; + } + let last = selected_offsets.last().copied().unwrap_or(0); + let offsets = selected_offsets.into_array(); let stats = offsets.statistics(); stats.set(Stat::IsSorted, Precision::Exact(true.into())); stats.set(Stat::Min, Precision::Exact(0_u64.into())); diff --git a/vortex-layout/src/plan/physical/pipeline/mod.rs b/vortex-layout/src/plan/physical/pipeline/mod.rs index 5dd18e783c1..87ebcebd91e 100644 --- a/vortex-layout/src/plan/physical/pipeline/mod.rs +++ b/vortex-layout/src/plan/physical/pipeline/mod.rs @@ -20,10 +20,10 @@ //! //! # Masks //! -//! Every mask a scan applies is known when the pipelines that apply it are compiled: the split's -//! selection, and the rows the earlier conjuncts kept. The -//! compiler places the mask in the stage that reads the rows, so chunks a mask selects nothing -//! of are never built and never read. +//! The compiler places known masks in the stages that read the rows. Experimental list sources +//! can instead wait for offsets and validity to resolve their element selection. A chunk +//! dispatcher then activates only overlapping children in the same scheduler, preserving its +//! downstream stages and outlets. Chunks outside that selection are never built or read. //! //! # Sharing //! diff --git a/vortex-layout/src/plan/physical/pipeline/scan.rs b/vortex-layout/src/plan/physical/pipeline/scan.rs index 246680c030f..46813d14b00 100644 --- a/vortex-layout/src/plan/physical/pipeline/scan.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan.rs @@ -23,6 +23,7 @@ use vortex_session::VortexSession; use super::DEFAULT_CAPACITY; use super::compile::Chain; +use super::compile::Compiler; use super::compile::Shares; use super::port::Arena; use super::port::PipelineId; @@ -131,6 +132,8 @@ pub(crate) struct Pipeline { /// What one run of a pipeline came to. enum Progress { + /// The source can now activate its selected chunks in this scheduler. + Activate, /// The source waits for this. Blocked(SourceBlocked), /// An outlet is full. @@ -215,6 +218,7 @@ fn drive( ); match ready { SourceReady::Ready => {} + SourceReady::Activate => return Ok(Progress::Activate), SourceReady::Blocked(blocked) => return Ok(Progress::Blocked(blocked)), } } @@ -472,6 +476,7 @@ impl Core { #[cfg(debug_assertions)] self.validate_progress(id, &progress); match &progress { + Progress::Activate => self.activate(id)?, // A read-ahead source can still have delivered bytes to consume. Progress::Blocked(SourceBlocked::Io) if !self.pipeline(id).bytes.is_empty() => { self.make_runnable(id); @@ -499,6 +504,36 @@ impl Core { Ok(()) } + fn activate(&mut self, id: PipelineId) -> VortexResult<()> { + let mut pipe = self.pipelines[id] + .take() + .ok_or_else(|| vortex_err!("Pipeline {id} is not live"))?; + vortex_ensure!( + pipe.inlets.is_empty() + && pipe.pending == 0 + && pipe.bytes.is_empty() + && pipe.child_inbox.is_none() + && pipe.continuing.is_empty() + && pipe.finishing.is_none(), + "Only an unstarted source without inlets or pending work can activate" + ); + let mut chain = pipe.source.activate(&mut Compiler { core: self })?; + chain.stages.append(&mut pipe.stages); + pipe.source = chain.source; + pipe.stages = chain.stages; + pipe.inlets = chain.inlets; + pipe.touched = InletSet::new(pipe.inlets.len()); + pipe.changed = InletSet::all(pipe.inlets.len()); + for (index, &port) in pipe.inlets.iter().enumerate() { + self.arena.get_mut(port).source = SourceEnd::Pipeline { id, inlet: index }; + } + self.pipelines[id] = Some(pipe); + self.make_runnable(id); + #[cfg(debug_assertions)] + self.validate(); + Ok(()) + } + /// Issues a read of `segment_id`, which pipeline `id`'s source asked for. fn ask(&mut self, id: PipelineId, segment_id: SegmentId) { self.read(id, segment_id); diff --git a/vortex-layout/src/plan/physical/pipeline/scan/validate.rs b/vortex-layout/src/plan/physical/pipeline/scan/validate.rs index e5d2638de56..f63bbca1bb6 100644 --- a/vortex-layout/src/plan/physical/pipeline/scan/validate.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan/validate.rs @@ -63,6 +63,11 @@ impl Core { pub(super) fn validate_progress(&self, id: PipelineId, progress: &Progress) { let pipe = self.pipeline(id); match progress { + Progress::Activate => { + assert!(pipe.inlets.is_empty()); + assert_eq!(pipe.pending, 0); + assert!(pipe.bytes.is_empty()); + } Progress::Blocked(SourceBlocked::Selection(selection)) => { assert!(self.selection_waiters[*selection].contains(&id)); } diff --git a/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs b/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs index e8bd51e1ac9..f756abe6540 100644 --- a/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs @@ -34,6 +34,8 @@ use crate::plan::physical::operator::Source; use crate::plan::physical::operator::SourceBlocked; use crate::plan::physical::operator::SourceReady; use crate::plan::physical::operator::Step; +use crate::plan::physical::operators::ConcatSource; +use crate::plan::physical::operators::SelectStage; use crate::segments::SegmentId; use crate::test::SESSION; @@ -111,6 +113,48 @@ fn segment(id: u32) -> Option { Some(SegmentId::from(id)) } +struct ActivateChunks; + +impl Source for ActivateChunks { + fn ready(&mut self, _cx: &mut Cx<'_>) -> VortexResult { + Ok(SourceReady::Activate) + } + + fn activate(&mut self, compiler: &mut Compiler<'_>) -> VortexResult { + let mut chains = Vec::new(); + for id in 0..2 { + let plan = RowSource::plan(8, vec![4], segment(id), false); + let start = u64::from(id) * 4; + chains.push( + compiler + .compile(&plan, start..start + 4, &Mask::new_true(4))? + .ok_or_else(|| vortex_err!("Selected chunk has no chain"))?, + ); + } + Ok(compiler.join(chains, ConcatSource::new(2))) + } +} + +impl Operator for ActivateChunks { + fn compute(&mut self, _input: Input, _cx: &mut Cx<'_>) -> VortexResult { + Err(vortex_err!("Dispatcher must activate before computing")) + } +} + +#[rstest] +fn activation_preserves_stages_and_orders_new_inlets( + #[values(&[0, 1], &[1, 0])] order: &[u32], +) -> VortexResult<()> { + let plan = SyntheticPlan::new(row_dtype(), 8, Vec::new(), |_, _, _| { + Ok(Some( + Chain::new(ActivateChunks).with(SelectStage::new(Some(1..3), None)), + )) + }) + .into_plan(); + let result = run(&plan, 0..8, Mask::new_true(8), scripted(order))?; + assert_rows(indices([1, 2, 5, 6]), result.arrays) +} + /// A source's pieces come out in order, one array each, after its read if it has one. #[rstest] fn source_pieces_come_out_in_order(#[values(None, Some(7))] read: Option) -> VortexResult<()> { diff --git a/vortex-layout/src/plan/physical/pipeline/tests.rs b/vortex-layout/src/plan/physical/pipeline/tests.rs index b4a50eb769c..40a296c03c1 100644 --- a/vortex-layout/src/plan/physical/pipeline/tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/tests.rs @@ -42,6 +42,7 @@ use vortex_session::registry::ReadContext; use vortex_utils::parallelism::get_available_parallelism; use super::port::SourceEnd; +use super::synthetic::SyntheticPlan; use super::synthetic::row_dtype; use super::*; use crate::LayoutRef; @@ -51,8 +52,10 @@ use crate::layouts::dict::DictLayout; use crate::layouts::flat::FlatLayout; use crate::layouts::foreign::new_foreign_layout; use crate::layouts::struct_::StructLayout; +use crate::plan::physical::ConcatPlan; use crate::plan::physical::EvalPlan; use crate::plan::physical::Filter; +use crate::plan::physical::ListPack; use crate::plan::physical::ListPackPlan; use crate::plan::physical::PackPlan; use crate::plan::physical::PlanRef; @@ -68,6 +71,7 @@ use crate::plan::physical::operator::SourceBlocked; use crate::plan::physical::operator::SourceReady; use crate::plan::physical::operator::Step; use crate::plan::physical::operators::PackSource; +use crate::plan::physical::operators::compile_deferred_list; use crate::plan::physical::optimize; use crate::segments::SegmentId; use crate::test::SESSION; @@ -302,6 +306,7 @@ fn delivery(order: Delivery) -> impl FnMut(&[ReadRequest]) -> usize { fn grouped_list_ranges_preserve_order_and_shared_reads( #[values(Delivery::Fifo, Delivery::Lifo)] order: Delivery, #[values(false, true)] shared: bool, + #[values(false, true)] deferred: bool, ) -> VortexResult<()> { let mut store = Store::default(); let elements = Buffer::from_iter(0_i64..20).into_array(); @@ -321,6 +326,11 @@ fn grouped_list_ranges_preserve_order_and_shared_reads( None, )? .into_plan(); + let plan = if deferred { + deferred_lists(plan) + } else { + plan + }; let rows = [0_u64, 1, 7, 8]; let mask = Mask::from_iter((0..10).map(|row| rows.contains(&row))); let result = run(&store, &plan, 0..10, mask, delivery(order))?; @@ -347,10 +357,71 @@ fn grouped_list_ranges_preserve_order_and_shared_reads( Ok(()) } +fn deferred_lists(plan: PlanRef) -> PlanRef { + SyntheticPlan::new( + plan.dtype().clone(), + plan.row_count(), + Vec::new(), + move |rows, mask, compiler| { + compile_deferred_list(plan.as_::(), rows, mask, compiler).map(Some) + }, + ) + .into_plan() +} + +#[rstest] +fn deferred_list_does_not_lower_unselected_chunks( + #[values(Delivery::Fifo, Delivery::Lifo)] order: Delivery, + #[values(false, true)] empty: bool, +) -> VortexResult<()> { + let mut store = Store::default(); + let elements = Buffer::from_iter(10_i64..12).into_array(); + let selected = lower(&store.flat(&elements)?)?; + let unused = SyntheticPlan::new(elements.dtype().clone(), 2, Vec::new(), |_, _, _| { + Err(vortex_err!("Unselected chunk must not be compiled")) + }) + .into_plan(); + let chunks = ConcatPlan::try_new( + elements.dtype().clone(), + vec![unused.clone(), selected, unused], + )?; + let offsets = Buffer::from(vec![0_u64, 2, 4, 6]).into_array(); + let plan = deferred_lists( + ListPackPlan::try_new( + false.into(), + 3, + chunks.into_plan(), + lower(&store.flat(&offsets)?)?, + None, + )? + .into_plan(), + ); + let mask = Mask::from_iter([false, !empty, false]); + let result = run(&store, &plan, 0..3, mask, delivery(order))?; + let expected = ListArray::try_new( + elements, + Buffer::from(vec![0_u64, 2]).into_array(), + Validity::NonNullable, + )? + .into_array(); + let expected = if empty { + expected.slice(0..0)? + } else { + expected + }; + let actual = ChunkedArray::try_new(result.arrays, expected.dtype().clone())?.into_array(); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); + if !empty { + assert_eq!(segments_read(&result.events), vec![0, 1]); + } + Ok(()) +} + #[rstest] fn dense_nullable_list_ranges_share_one_scan( #[values(Delivery::Fifo, Delivery::Lifo)] order: Delivery, #[values(false, true)] partial: bool, + #[values(false, true)] deferred: bool, ) -> VortexResult<()> { let mut store = Store::default(); let elements = Buffer::from_iter(0_i64..20).into_array(); @@ -370,6 +441,11 @@ fn dense_nullable_list_ranges_share_one_scan( Some(lower(&store.flat(&valid)?)?), )? .into_plan(); + let plan = if deferred { + deferred_lists(plan) + } else { + plan + }; let mask = Mask::from_iter((0..10).map(|row| !partial || row != 4)); let result = run(&store, &plan, 0..10, mask.clone(), delivery(order))?; if !partial { diff --git a/vortex-layout/src/plan/physical/plans/list_pack.rs b/vortex-layout/src/plan/physical/plans/list_pack.rs index 923b15c7ea4..94b0b9ac2d9 100644 --- a/vortex-layout/src/plan/physical/plans/list_pack.rs +++ b/vortex-layout/src/plan/physical/plans/list_pack.rs @@ -22,6 +22,7 @@ use crate::plan::physical::PlanParts; use crate::plan::physical::PlanRef; use crate::plan::physical::PlanVTable; use crate::plan::physical::operators::ListSource; +use crate::plan::physical::operators::compile_deferred_list; use crate::plan::physical::operators::UNBOUNDED; use crate::plan::physical::pipeline::Chain; use crate::plan::physical::pipeline::Compiler; @@ -144,8 +145,7 @@ impl PlanVTable for ListPack { return Ok(None); } if std::env::var_os("VORTEX_LIST_PIPELINE").is_some() { - return crate::plan::physical::operators::compile_deferred_list(plan, rows, mask, compiler) - .map(Some); + return compile_deferred_list(plan, rows, mask, compiler).map(Some); } let boundaries = offset_mask(mask); let offsets = compiler From 23c93e26523d64b3b7de7a3e0cd4a4e9b378fa5a Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 19 Oct 2026 16:27:48 +0100 Subject: [PATCH 5/7] Plan selected list element views outside running pipelines Signed-off-by: Joe Isaacs --- scripts/bench-list-pipeline.py | 10 +- vortex-layout/benches/pipeline_scan.rs | 29 +- vortex-layout/src/plan/physical/operator.rs | 21 - .../src/plan/physical/operators/list.rs | 30 +- .../plan/physical/operators/list/deferred.rs | 440 ------------------ .../src/plan/physical/pipeline/children.rs | 56 ++- .../src/plan/physical/pipeline/compile.rs | 6 - .../src/plan/physical/pipeline/mod.rs | 9 +- .../src/plan/physical/pipeline/scan.rs | 62 +-- .../plan/physical/pipeline/scan/validate.rs | 8 - .../physical/pipeline/scheduling_tests.rs | 44 -- .../src/plan/physical/pipeline/tests.rs | 63 ++- .../src/plan/physical/plans/list_pack.rs | 65 +-- vortex-layout/src/scan/execution/adaptor.rs | 9 +- vortex-layout/src/scan/execution/children.rs | 45 +- .../src/scan/execution/children/tests.rs | 56 +++ .../src/scan/execution/projection.rs | 66 ++- vortex/examples/selection_scan.rs | 2 +- 18 files changed, 331 insertions(+), 690 deletions(-) delete mode 100644 vortex-layout/src/plan/physical/operators/list/deferred.rs diff --git a/scripts/bench-list-pipeline.py b/scripts/bench-list-pipeline.py index 0e2edd1e248..afe36851876 100644 --- a/scripts/bench-list-pipeline.py +++ b/scripts/bench-list-pipeline.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright the Vortex contributors -"""Compare current and prebuilt list pipelines on identical ListLayout files.""" +"""Compare current list scans and planner-owned selected views on identical ListLayout files.""" import argparse import hashlib @@ -66,7 +66,7 @@ def prepare(self) -> dict[str, Path]: metadata = { "baseline": "current list implementation, prototype disabled", "commit": subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=ROOT, text=True).strip(), - "candidate": "prebuilt pipeline prototype; VORTEX_LIST_PIPELINE=1", + "candidate": "planner-owned selected views; VORTEX_LIST_PIPELINE=1", "platform": platform.platform(), "logical_cpu_count": os.cpu_count(), "affinity": "unrestricted", @@ -118,7 +118,7 @@ def run( ] with (self.results / "commands.jsonl").open("a") as output: output.write( - json.dumps({"command": command, "selection_gate": version == "candidate", "VORTEX_SCAN_V2": "1"}) + "\n" + json.dumps({"command": command, "selected_view": version == "candidate", "VORTEX_SCAN_V2": "1"}) + "\n" ) result = subprocess.run(command, env=env, capture_output=True, text=True, timeout=180) (self.results / "raw" / f"{stem}.log").write_text(result.stderr) @@ -197,10 +197,10 @@ def main() -> None: output.write("\n### Deferred list pipeline experiment\n\n") output.write( f"{checks} reference checks passed; full Arrow decode on warm local disk. " - "The prototype prebuilds pipelines but does not implement the current decoded-work admission budget. " + "The candidate builds selected child graphs in the planner and retains admission limits for parallel fixed-width work. " "Positive changes are slower.\n\n" ) - output.write("| Dataset | Rows | Selection | Open mode | Current ms | Prototype ms | Change |\n") + output.write("| Dataset | Rows | Selection | Open mode | Current ms | Selected view ms | Change |\n") output.write("| --- | ---: | --- | --- | ---: | ---: | ---: |\n") for row in summary: output.write( diff --git a/vortex-layout/benches/pipeline_scan.rs b/vortex-layout/benches/pipeline_scan.rs index d73ba6a0331..de952cf6b91 100644 --- a/vortex-layout/benches/pipeline_scan.rs +++ b/vortex-layout/benches/pipeline_scan.rs @@ -210,7 +210,13 @@ impl ListFixture { .into_plan(); let mut segments = vec![element_bytes; chunks as usize]; segments.push(offset_bytes); - let indices = (0..100).map(|index| if uniform { index * (LISTS / 100) } else { index }); + let indices = (0..100).map(|index| { + if uniform { + index * (LISTS / 100) + } else { + index + } + }); Self { session, plan, @@ -224,22 +230,33 @@ impl ListFixture { rows: 0..self.mask.len() as u64, mask: self.mask.clone(), }; - let mut scan = Scan::try_new(self.session.clone(), self.plan.clone(), vec![split]).unwrap(); + let scan = Scan::try_new(self.session.clone(), self.plan.clone(), vec![split]).unwrap(); + let arrays = self.execute(scan); + assert_eq!(arrays.iter().map(|array| array.len()).sum::(), 100); + black_box(arrays); + } + + /// A minimal in-memory planner driving each selected child graph through the same protocol. + fn execute(&self, mut scan: Scan) -> Vec { + let mut output = Vec::new(); loop { match scan.step().unwrap() { Turn::Read(read) => { scan.deliver(read.id, self.segments[*read.segment_id as usize].clone()) .unwrap(); } - Turn::Output(_, array) => { - assert_eq!(array.len(), 100); - black_box(array); + Turn::Output(split, array) => output.push((split, array)), + Turn::Child(child) => { + let (id, child_scan) = child.into_scan(self.session.clone()).unwrap(); + let arrays = self.execute(child_scan); + scan.complete_child(id, Some(arrays)).unwrap(); } - Turn::Child(_) => unreachable!("microbenchmark has no child planner"), Turn::Waiting => unreachable!("every read is delivered immediately"), Turn::Done => break, } } + output.sort_by_key(|(split, _)| *split); + output.into_iter().map(|(_, array)| array).collect() } } diff --git a/vortex-layout/src/plan/physical/operator.rs b/vortex-layout/src/plan/physical/operator.rs index d74b2f714fb..de1b1e35560 100644 --- a/vortex-layout/src/plan/physical/operator.rs +++ b/vortex-layout/src/plan/physical/operator.rs @@ -10,12 +10,9 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::buffer::BufferHandle; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_session::VortexSession; use super::pipeline::Arena; -use super::pipeline::Chain; -use super::pipeline::Compiler; use super::pipeline::Inlet; use super::pipeline::PortId; use crate::plan::physical::pipeline::ChildProjection; @@ -38,8 +35,6 @@ pub enum Step { pub enum SourceReady { /// The next compute call can produce output or finish without blocking. Ready, - /// Replace this source with selected work in the same scheduler, through [`Source::activate`]. - Activate, /// Nothing can happen until the condition clears. Blocked(SourceBlocked), } @@ -51,8 +46,6 @@ pub enum SourceBlocked { Io, /// Wait for a child projection, possibly alongside pending reads. Child, - /// Wait for a list's runtime element selection. - Selection(usize), /// The inlet needs more batches or its sink to close it. It must be open and have room /// for its sink to push; debug builds check that the sink can still wake this source. Inlet(usize), @@ -61,15 +54,6 @@ pub enum SourceBlocked { /// The start of a pipeline: produces batches from its inlets and the segments it asks to have /// read. pub trait Source: Operator { - /// Activates work whose selection was unavailable at initial compilation. - /// - /// Called only after returning [`SourceReady::Activate`]. The source must have no inlets, - /// pending reads, or emitted batches. The resulting chain inherits its downstream stages - /// and outlets, without starting another scan or handing work to the outer planner. - fn activate(&mut self, _compiler: &mut Compiler<'_>) -> VortexResult { - vortex_bail!("Source does not support deferred activation") - } - /// Prepares the next [`compute`](Operator::compute) call, after the driver checks outlets. /// Returning [`SourceReady::Ready`] guarantees that call can proceed without blocking. /// Any input or choice made here can be retained by the source for that call. @@ -178,7 +162,6 @@ pub struct Cx<'a> { pub(super) changed: &'a mut InletSet, pub(super) child_requests: &'a mut Vec, pub(super) child_inbox: &'a mut Option>, - pub(super) resolved_selections: &'a mut Vec, } #[derive(Default)] @@ -188,10 +171,6 @@ pub(super) struct ChildInbox { } impl Cx<'_> { - pub(crate) fn resolve_selection(&mut self, selection: usize) { - self.resolved_selections.push(selection); - } - /// Describes child work for the enclosing planner. Return [`SourceBlocked::Child`] while awaiting its completion. /// Identifiers need only be unique among this source's outstanding child projections. pub(crate) fn request_child(&mut self, child: ChildProjection) { diff --git a/vortex-layout/src/plan/physical/operators/list.rs b/vortex-layout/src/plan/physical/operators/list.rs index 3ec22b8a423..7b94351b2fa 100644 --- a/vortex-layout/src/plan/physical/operators/list.rs +++ b/vortex-layout/src/plan/physical/operators/list.rs @@ -38,22 +38,21 @@ use crate::plan::physical::operator::Step; use crate::plan::physical::pipeline::ChildProjection; use crate::plan::physical::pipeline::ReadId; use crate::plan::physical::pipeline::Scan; +use crate::plan::physical::pipeline::SelectedPlan; use crate::plan::physical::pipeline::Split; use crate::plan::physical::pipeline::Turn; use crate::plan::physical::pipeline::ValueCache; use crate::segments::SegmentId; -mod deferred; mod work; -pub(crate) use deferred::compile_deferred_list; - /// A list's offsets and validity arrive before its selected element ranges can be planned. pub(crate) struct ListSource { plan: ListPackPlan, mask: Mask, values: ValueCache, plan_children: bool, + planner_owned: bool, awaiting_child: bool, serial_splits: Option<(Vec, bool)>, initialized: bool, @@ -73,12 +72,14 @@ impl ListSource { mask: Mask, values: ValueCache, plan_children: bool, + planner_owned: bool, ) -> Self { Self { plan, mask, values, plan_children, + planner_owned, awaiting_child: false, serial_splits: None, initialized: false, @@ -148,8 +149,8 @@ impl ListSource { get_available_parallelism().unwrap_or(1), )?; let splits = work::element_splits(elements.dtype(), ranges)?; - // Avoid handing tiny or variable-width work to the outer planner. The planner - // decides whether the substantial fixed-width ranges can be split safely. + // The baseline only hands off substantial fixed-width work. The selected-view path + // lets the planner own every child graph, including tiny and nested selections. let width = match elements.dtype() { DType::Primitive(ptype, _) => ptype.byte_width() as u64, DType::Bool(_) => 1, @@ -158,13 +159,15 @@ impl ListSource { let bytes = splits.iter().fold(0_u64, |bytes, split| { bytes.saturating_add((split.mask.true_count() as u64).saturating_mul(width)) }); - if self.plan_children && bytes >= 512 * 1024 { - self.serial_splits = Some((splits.clone(), overlap)); + if self.planner_owned || (self.plan_children && bytes >= 512 * 1024) { + if !self.planner_owned { + self.serial_splits = Some((splits.clone(), overlap)); + } self.awaiting_child = true; cx.request_child(ChildProjection { id: 0, - plan: elements, - splits, + view: SelectedPlan::new(elements, splits, if overlap { 2 } else { 1 }), + planner_owned: self.planner_owned, values: self.values.clone(), }); self.initialized = true; @@ -235,15 +238,16 @@ impl Source for ListSource { return Ok(SourceReady::Blocked(SourceBlocked::Child)); }; self.awaiting_child = false; - let (splits, overlap) = self - .serial_splits - .take() - .ok_or_else(|| vortex_err!("Missing deferred list ranges"))?; if let Some(arrays) = arrays { self.chunks.extend(arrays.into_iter().enumerate()); self.finish()?; return Ok(SourceReady::Ready); } + vortex_ensure!(!self.planner_owned, "Planner declined a selected list view"); + let (splits, overlap) = self + .serial_splits + .take() + .ok_or_else(|| vortex_err!("Missing deferred list ranges"))?; self.start_elements(cx, splits, overlap)?; } let scan = self diff --git a/vortex-layout/src/plan/physical/operators/list/deferred.rs b/vortex-layout/src/plan/physical/operators/list/deferred.rs deleted file mode 100644 index 03cec976d05..00000000000 --- a/vortex-layout/src/plan/physical/operators/list/deferred.rs +++ /dev/null @@ -1,440 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Benchmark prototype: build list dependencies upfront and resolve selection at runtime. -//! -//! Chunk dispatchers wait for offsets and validity, then activate only selected children in the -//! same scheduler. Unselected chunks need no runtime sources or ports. -//! -//! Enabled only by `VORTEX_LIST_PIPELINE`. Supports primitive and nested-list fixture plans; -//! unsupported plans fail explicitly. This experiment does not implement the planner's decoded-work -//! admission budget or shared segment optimization. - -use std::ops::Range; -use std::sync::Arc; -use std::sync::OnceLock; - -use vortex_array::ArrayRef; -use vortex_array::IntoArray; -use vortex_array::arrays::ListArray; -use vortex_array::arrays::PrimitiveArray; -use vortex_array::builtins::ArrayBuiltins; -use vortex_array::dtype::DType; -use vortex_array::dtype::PType; -use vortex_array::expr::stats::Precision; -use vortex_array::expr::stats::Stat; -use vortex_array::validity::Validity; -use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_error::vortex_err; -use vortex_mask::AllOr; -use vortex_mask::Mask; - -use super::selected_ranges; -use crate::plan::physical::Concat; -use crate::plan::physical::ConcatPlan; -use crate::plan::physical::Filter; -use crate::plan::physical::LegacyStats; -use crate::plan::physical::ListPack; -use crate::plan::physical::ListPackPlan; -use crate::plan::physical::PlanRef; -use crate::plan::physical::SegmentScan; -use crate::plan::physical::SegmentScanPlan; -use crate::plan::physical::Zoned; -use crate::plan::physical::operator::Cx; -use crate::plan::physical::operator::Input; -use crate::plan::physical::operator::Operator; -use crate::plan::physical::operator::Source; -use crate::plan::physical::operator::SourceBlocked; -use crate::plan::physical::operator::SourceReady; -use crate::plan::physical::operator::Step; -use crate::plan::physical::operators::ConcatSource; -use crate::plan::physical::operators::ScanSource; -use crate::plan::physical::operators::UNBOUNDED; -use crate::plan::physical::operators::join_inlet; -use crate::plan::physical::pipeline::Chain; -use crate::plan::physical::pipeline::Compiler; -use crate::plan::physical::pipeline::overlapping; - -#[derive(Clone)] -struct Demand { - ranges: Arc>>>, - gate: Option, - domain: Range, - offset: u64, - boundaries: bool, -} - -impl Demand { - fn ranges(&self, len: u64) -> Option>> { - let selected = self.ranges.get()?; - let mut result: Vec> = Vec::new(); - let first = selected.partition_point(|range| { - range.end <= self.domain.start - || range.end.saturating_add(u64::from(self.boundaries)) <= self.offset - }); - for range in &selected[first..] { - if range.start >= self.domain.end || range.start >= self.offset + len { - break; - } - let start = range.start.max(self.domain.start); - let end = range.end.min(self.domain.end); - if start >= end { - continue; - } - let end = end + u64::from(self.boundaries); - let start = start.max(self.offset); - let end = end.min(self.offset + len); - if start >= end { - continue; - } - let range = start - self.offset..end - self.offset; - if let Some(last) = result.last_mut() - && last.end >= range.start - { - last.end = last.end.max(range.end); - } else { - result.push(range); - } - } - Some(result) - } - - fn blocked(&self) -> VortexResult { - Ok(SourceReady::Blocked(SourceBlocked::Selection( - self.gate - .ok_or_else(|| vortex_err!("Unresolved demand has no producer"))?, - ))) - } - - fn shifted(&self, offset: u64) -> Self { - Self { - offset: self.offset + offset, - ..self.clone() - } - } - - fn boundaries(&self, rows: u64) -> Self { - Self { - domain: self.offset..self.offset + rows, - boundaries: true, - ..self.clone() - } - } -} - -pub(crate) fn compile_deferred_list( - plan: &ListPackPlan, - rows: Range, - mask: &Mask, - compiler: &mut Compiler<'_>, -) -> VortexResult { - let ranges = match mask.slices() { - AllOr::All => vec![rows], - AllOr::None => Vec::new(), - AllOr::Some(slices) => slices - .iter() - .map(|&(start, end)| rows.start + start as u64..rows.start + end as u64) - .collect(), - }; - let demand = Demand { - ranges: Arc::new(OnceLock::from(ranges)), - gate: None, - domain: 0..plan.row_count(), - offset: 0, - boundaries: false, - }; - compile_list(plan, demand, compiler) -} - -fn compile_list( - plan: &ListPackPlan, - demand: Demand, - compiler: &mut Compiler<'_>, -) -> VortexResult { - let offsets = compile( - &plan.offsets()?, - demand.boundaries(plan.row_count()), - compiler, - )?; - let mut chains = vec![offsets]; - if let Some(validity) = plan.validity()? { - chains.push(compile(&validity, demand.clone(), compiler)?); - } - let elements = plan.elements()?; - let element_demand = Demand { - ranges: Arc::new(OnceLock::new()), - gate: Some(compiler.selection_gate()), - domain: 0..elements.row_count(), - offset: 0, - boundaries: false, - }; - chains.push(compile(&elements, element_demand.clone(), compiler)?); - let source = DeferredList { - plan: plan.clone(), - demand, - elements: element_demand, - assembled: None, - done: false, - }; - Ok(compiler.join_with(chains, |_| UNBOUNDED, source)) -} - -fn compile(plan: &PlanRef, demand: Demand, compiler: &mut Compiler<'_>) -> VortexResult { - if let Some(filter) = plan.as_opt::() { - return compile(&filter.child_plan()?, demand, compiler); - } - if let Some(zoned) = plan.as_opt::() { - let data = zoned - .data_plan()? - .ok_or_else(|| vortex_err!("Zoned data is absent"))?; - return compile(&data, demand, compiler); - } - if let Some(stats) = plan.as_opt::() { - return compile(&stats.data_plan()?, demand, compiler); - } - if let Some(concat) = plan.as_opt::() { - let mut source = DeferredConcat { - plan: concat.clone(), - demand, - }; - return if source.demand.ranges.get().is_some() { - source.activate(compiler) - } else { - Ok(Chain::new(source)) - }; - } - if let Some(list) = plan.as_opt::() { - return compile_list(list, demand, compiler); - } - if let Some(scan) = plan.as_opt::() { - return Ok(Chain::new(DeferredScan { - plan: scan.clone(), - demand, - source: None, - empty: false, - })); - } - vortex_bail!( - "Deferred-list benchmark prototype cannot compile {}", - plan.id() - ) -} - -/// Keep one dispatcher while selection is unknown; lower only selected children afterwards. -struct DeferredConcat { - plan: ConcatPlan, - demand: Demand, -} - -impl Source for DeferredConcat { - fn ready(&mut self, _cx: &mut Cx<'_>) -> VortexResult { - if self.demand.ranges.get().is_some() { - Ok(SourceReady::Activate) - } else { - self.demand.blocked() - } - } - - fn activate(&mut self, compiler: &mut Compiler<'_>) -> VortexResult { - let ranges = self - .demand - .ranges(self.plan.row_count()) - .ok_or_else(|| vortex_err!("Chunks activated before their selection"))?; - let mut chains = Vec::new(); - let mut previous = None; - for range in ranges { - for (index, start, _) in overlapping(&self.plan, &range) { - // Several selected ranges can lie in one chunk. Give that chunk one source, - // which applies their combined selection and reads its segment only once. - if previous == Some(index) { - continue; - } - previous = Some(index); - chains.push(compile( - &self.plan.child_required(index)?, - self.demand.shifted(start), - compiler, - )?); - } - } - let count = chains.len(); - if count == 1 { - return Ok(chains.remove(0)); - } - Ok(compiler.join(chains, ConcatSource::new(count))) - } -} - -impl Operator for DeferredConcat { - fn compute(&mut self, _input: Input, _cx: &mut Cx<'_>) -> VortexResult { - vortex_bail!("Chunk dispatcher computed before activation") - } -} - -struct DeferredScan { - plan: SegmentScanPlan, - demand: Demand, - source: Option, - empty: bool, -} - -impl Source for DeferredScan { - fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult { - if self.empty { - return Ok(SourceReady::Ready); - } - if self.source.is_none() { - let Some(ranges) = self.demand.ranges(self.plan.row_count()) else { - return self.demand.blocked(); - }; - let Some(first) = ranges.first() else { - self.empty = true; - return Ok(SourceReady::Ready); - }; - let start = first.start; - let end = ranges - .last() - .ok_or_else(|| vortex_err!("Missing selected range"))? - .end; - let len = usize::try_from(end - start)?; - let mask = if ranges.len() == 1 { - None - } else { - Some(Mask::from_slices( - len, - ranges - .iter() - .map(|range| { - Ok(( - usize::try_from(range.start - start)?, - usize::try_from(range.end - start)?, - )) - }) - .collect::>>()?, - )) - }; - let slice = if start == 0 && end == self.plan.row_count() { - None - } else { - Some(usize::try_from(start)?..usize::try_from(end)?) - }; - self.source = Some(ScanSource::new(self.plan.clone(), slice, mask)); - } - self.source - .as_mut() - .ok_or_else(|| vortex_err!("Missing segment source"))? - .ready(cx) - } -} - -impl Operator for DeferredScan { - fn compute(&mut self, input: Input, cx: &mut Cx<'_>) -> VortexResult { - match self.source.as_mut() { - Some(source) => source.compute(input, cx), - None if self.empty => Ok(Step::End), - None => vortex_bail!("Segment computed before its selection"), - } - } -} - -struct DeferredList { - plan: ListPackPlan, - demand: Demand, - elements: Demand, - assembled: Option<(ArrayRef, Validity)>, - done: bool, -} - -impl Source for DeferredList { - fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult { - if self.done { - return Ok(SourceReady::Ready); - } - let element_inlet = 1 + usize::from(self.plan.dtype().is_nullable()); - if self.assembled.is_none() { - let Some(selected) = self.demand.ranges(self.plan.row_count()) else { - return self.demand.blocked(); - }; - for index in 0..element_inlet { - if !cx.inlet(index).closed() { - return Ok(SourceReady::Blocked(SourceBlocked::Inlet(index))); - } - } - let offsets = join_inlet(cx, 0, self.plan.offsets()?.dtype())? - .cast(DType::Primitive(PType::U64, false.into()))? - .execute::(cx.exec())?; - let validity = if self.plan.dtype().is_nullable() { - Validity::Array(join_inlet(cx, 1, &DType::Bool(false.into()))?) - } else { - Validity::NonNullable - }; - let count = selected - .iter() - .map(|range| range.end - range.start) - .sum::(); - let valid = validity.execute_mask(usize::try_from(count)?, cx.exec())?; - let selected = selected - .into_iter() - .map(|range| Ok(usize::try_from(range.start)?..usize::try_from(range.end)?)) - .collect::>>()?; - let selected_indices = selected.iter().flat_map(|range| range.clone()); - let (mut selected_offsets, mut ranges) = selected_ranges( - selected_indices, - offsets.as_slice::(), - &valid, - self.plan.elements()?.row_count(), - )?; - if selected.len() == 1 - && !valid.all_true() - && let Some(retained) = super::work::preserve_null_span( - self.plan.elements()?.dtype(), - offsets.as_slice::(), - selected_offsets.last().copied().unwrap_or(0), - &mut ranges, - ) - { - selected_offsets = retained; - } - let last = selected_offsets.last().copied().unwrap_or(0); - let offsets = selected_offsets.into_array(); - let stats = offsets.statistics(); - stats.set(Stat::IsSorted, Precision::Exact(true.into())); - stats.set(Stat::Min, Precision::Exact(0_u64.into())); - stats.set(Stat::Max, Precision::Exact(last.into())); - self.elements - .ranges - .set(ranges) - .map_err(|_| vortex_err!("Selection resolved twice"))?; - cx.resolve_selection( - self.elements - .gate - .ok_or_else(|| vortex_err!("Missing element gate"))?, - ); - self.assembled = Some((offsets, validity)); - } - Ok(if cx.inlet(element_inlet).closed() { - SourceReady::Ready - } else { - SourceReady::Blocked(SourceBlocked::Inlet(element_inlet)) - }) - } -} - -impl Operator for DeferredList { - fn compute(&mut self, _input: Input, cx: &mut Cx<'_>) -> VortexResult { - if self.done { - return Ok(Step::End); - } - self.done = true; - let inlet = 1 + usize::from(self.plan.dtype().is_nullable()); - let elements = join_inlet(cx, inlet, self.plan.elements()?.dtype())?; - let (offsets, validity) = self - .assembled - .take() - .ok_or_else(|| vortex_err!("List selection unresolved"))?; - Ok(Step::Final( - ListArray::try_new(elements, offsets, validity)?.into_array(), - )) - } -} diff --git a/vortex-layout/src/plan/physical/pipeline/children.rs b/vortex-layout/src/plan/physical/pipeline/children.rs index 974e0afcd74..a2685079ed8 100644 --- a/vortex-layout/src/plan/physical/pipeline/children.rs +++ b/vortex-layout/src/plan/physical/pipeline/children.rs @@ -1,29 +1,71 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Descriptions of child projections for the scan planner to schedule. +//! Selected plan views handed to the enclosing planner for execution. use std::fmt; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use super::Scan; use super::Split; use super::ValueCache; use crate::plan::physical::PlanRef; -/// A source's child projection, handed to its enclosing planner for execution. -/// Only scans configured by that planner produce these requests. -pub struct ChildProjection { - pub(crate) id: u64, +/// A shared physical plan and its selected row domains. Constructing the view does not lower +/// children or compile a graph; the planner compiles only children touched by these splits. +#[derive(Debug)] +pub(crate) struct SelectedPlan { pub(crate) plan: PlanRef, pub(crate) splits: Vec, + max_active: usize, +} + +impl SelectedPlan { + pub(crate) fn new(plan: PlanRef, splits: Vec, max_active: usize) -> Self { + Self { + plan, + splits, + max_active, + } + } + + pub(crate) fn into_scan(self, session: VortexSession, values: ValueCache) -> VortexResult { + Ok(Scan::try_new(session, self.plan, self.splits)? + .with_value_cache(values) + .with_max_active(self.max_active)) + } +} + +/// A source's selected child projection, handed to its enclosing planner for execution. +pub struct ChildProjection { + pub(crate) id: u64, + pub(crate) view: SelectedPlan, + /// The planner must execute this view even when parallel subdivision is unsuitable. + pub(crate) planner_owned: bool, pub(crate) values: ValueCache, } +impl ChildProjection { + /// Creates the selected child scan at the planner boundary, returning its completion id. + /// Drive this scan and pass its output to [`Scan::complete_child`] on the requesting scan. + pub fn into_scan(self, session: VortexSession) -> VortexResult<(u64, Scan)> { + Ok(( + self.id, + self.view + .into_scan(session, self.values)? + .with_child_planning(true), + )) + } +} + impl fmt::Debug for ChildProjection { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("ChildProjection") .field("id", &self.id) - .field("plan", &self.plan) - .field("splits", &self.splits) + .field("view", &self.view) + .field("planner_owned", &self.planner_owned) .finish_non_exhaustive() } } diff --git a/vortex-layout/src/plan/physical/pipeline/compile.rs b/vortex-layout/src/plan/physical/pipeline/compile.rs index 7baaa9905a3..cf2dcb10184 100644 --- a/vortex-layout/src/plan/physical/pipeline/compile.rs +++ b/vortex-layout/src/plan/physical/pipeline/compile.rs @@ -81,12 +81,6 @@ pub struct Compiler<'a> { } impl Compiler<'_> { - pub(crate) fn selection_gate(&mut self) -> usize { - let id = self.core.selection_waiters.len(); - self.core.selection_waiters.push(Vec::new()); - id - } - /// Compiles `plan` over `rows` of its domain, restricted to `mask`, into a chain producing /// the rows `mask` selects, in order, or `None` when that is no row. pub fn compile( diff --git a/vortex-layout/src/plan/physical/pipeline/mod.rs b/vortex-layout/src/plan/physical/pipeline/mod.rs index 87ebcebd91e..c6deed171e8 100644 --- a/vortex-layout/src/plan/physical/pipeline/mod.rs +++ b/vortex-layout/src/plan/physical/pipeline/mod.rs @@ -20,10 +20,10 @@ //! //! # Masks //! -//! The compiler places known masks in the stages that read the rows. Experimental list sources -//! can instead wait for offsets and validity to resolve their element selection. A chunk -//! dispatcher then activates only overlapping children in the same scheduler, preserving its -//! downstream stages and outlets. Chunks outside that selection are never built or read. +//! Every mask a scan applies is known when the pipelines that apply it are compiled: the split's +//! selection, and the rows the earlier conjuncts kept. The +//! compiler places the mask in the stage that reads the rows, so chunks a mask selects nothing +//! of are never built and never read. //! //! # Sharing //! @@ -46,6 +46,7 @@ pub mod synthetic; mod values; pub use self::children::ChildProjection; +pub(crate) use self::children::SelectedPlan; pub use self::compile::Chain; pub use self::compile::Compiler; pub use self::compile::Reach; diff --git a/vortex-layout/src/plan/physical/pipeline/scan.rs b/vortex-layout/src/plan/physical/pipeline/scan.rs index 46813d14b00..9ed12d08a5e 100644 --- a/vortex-layout/src/plan/physical/pipeline/scan.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan.rs @@ -23,7 +23,6 @@ use vortex_session::VortexSession; use super::DEFAULT_CAPACITY; use super::compile::Chain; -use super::compile::Compiler; use super::compile::Shares; use super::port::Arena; use super::port::PipelineId; @@ -132,8 +131,6 @@ pub(crate) struct Pipeline { /// What one run of a pipeline came to. enum Progress { - /// The source can now activate its selected chunks in this scheduler. - Activate, /// The source waits for this. Blocked(SourceBlocked), /// An outlet is full. @@ -150,7 +147,6 @@ fn drive( exec: &mut ExecutionCtx, reads: &mut SmallVec<[SegmentId; 1]>, child_requests: &mut Vec, - resolved_selections: &mut Vec, ) -> VortexResult { let Pipeline { source, @@ -177,7 +173,6 @@ fn drive( changed, child_requests, child_inbox, - resolved_selections, }; loop { // Closing needs no room, so a pipeline with nothing left to flush closes at once. @@ -218,7 +213,6 @@ fn drive( ); match ready { SourceReady::Ready => {} - SourceReady::Activate => return Ok(Progress::Activate), SourceReady::Blocked(blocked) => return Ok(Progress::Blocked(blocked)), } } @@ -318,7 +312,6 @@ pub(crate) struct Core { /// Whole values kept for later scans of the same split. pub(crate) values: ValueCache, pub(crate) plan_children: bool, - pub(crate) selection_waiters: Vec>, } struct PendingChild { @@ -348,7 +341,6 @@ impl Core { split: usize::MAX, values: ValueCache::default(), plan_children: false, - selection_waiters: Vec::new(), } } @@ -445,7 +437,6 @@ impl Core { pipe.queued = false; let mut reads = SmallVec::new(); let mut child_requests = Vec::new(); - let mut resolved_selections = Vec::new(); let progress = drive( pipe, &mut self.arena, @@ -453,20 +444,7 @@ impl Core { &mut self.exec, &mut reads, &mut child_requests, - &mut resolved_selections, )?; - for selection in resolved_selections { - for waiter in std::mem::take(&mut self.selection_waiters[selection]) { - if self.pipelines[waiter].as_ref().is_some_and(|pipe| { - pipe.state == State::Blocked(SourceBlocked::Selection(selection)) - }) { - self.make_runnable(waiter); - } - } - } - if let Progress::Blocked(SourceBlocked::Selection(selection)) = &progress { - self.selection_waiters[*selection].push(id); - } for segment_id in reads { self.ask(id, segment_id); } @@ -476,7 +454,6 @@ impl Core { #[cfg(debug_assertions)] self.validate_progress(id, &progress); match &progress { - Progress::Activate => self.activate(id)?, // A read-ahead source can still have delivered bytes to consume. Progress::Blocked(SourceBlocked::Io) if !self.pipeline(id).bytes.is_empty() => { self.make_runnable(id); @@ -504,36 +481,6 @@ impl Core { Ok(()) } - fn activate(&mut self, id: PipelineId) -> VortexResult<()> { - let mut pipe = self.pipelines[id] - .take() - .ok_or_else(|| vortex_err!("Pipeline {id} is not live"))?; - vortex_ensure!( - pipe.inlets.is_empty() - && pipe.pending == 0 - && pipe.bytes.is_empty() - && pipe.child_inbox.is_none() - && pipe.continuing.is_empty() - && pipe.finishing.is_none(), - "Only an unstarted source without inlets or pending work can activate" - ); - let mut chain = pipe.source.activate(&mut Compiler { core: self })?; - chain.stages.append(&mut pipe.stages); - pipe.source = chain.source; - pipe.stages = chain.stages; - pipe.inlets = chain.inlets; - pipe.touched = InletSet::new(pipe.inlets.len()); - pipe.changed = InletSet::all(pipe.inlets.len()); - for (index, &port) in pipe.inlets.iter().enumerate() { - self.arena.get_mut(port).source = SourceEnd::Pipeline { id, inlet: index }; - } - self.pipelines[id] = Some(pipe); - self.make_runnable(id); - #[cfg(debug_assertions)] - self.validate(); - Ok(()) - } - /// Issues a read of `segment_id`, which pipeline `id`'s source asked for. fn ask(&mut self, id: PipelineId, segment_id: SegmentId) { self.read(id, segment_id); @@ -950,12 +897,9 @@ impl Scan { self.core.deliver(read, bytes) } - /// Delivers a child projection to its source. - pub(crate) fn complete_child( - &mut self, - id: u64, - arrays: Option>, - ) -> VortexResult<()> { + /// Delivers a child projection to its source, with arrays in selected-split order. + /// `None` declines optional parallel work; a planner-owned selected view requires `Some`. + pub fn complete_child(&mut self, id: u64, arrays: Option>) -> VortexResult<()> { self.core.complete_child(id, arrays) } diff --git a/vortex-layout/src/plan/physical/pipeline/scan/validate.rs b/vortex-layout/src/plan/physical/pipeline/scan/validate.rs index f63bbca1bb6..d1749972d65 100644 --- a/vortex-layout/src/plan/physical/pipeline/scan/validate.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan/validate.rs @@ -63,14 +63,6 @@ impl Core { pub(super) fn validate_progress(&self, id: PipelineId, progress: &Progress) { let pipe = self.pipeline(id); match progress { - Progress::Activate => { - assert!(pipe.inlets.is_empty()); - assert_eq!(pipe.pending, 0); - assert!(pipe.bytes.is_empty()); - } - Progress::Blocked(SourceBlocked::Selection(selection)) => { - assert!(self.selection_waiters[*selection].contains(&id)); - } Progress::Blocked(SourceBlocked::Inlet(index)) => { assert!( *index < pipe.inlets.len(), diff --git a/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs b/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs index f756abe6540..e8bd51e1ac9 100644 --- a/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs @@ -34,8 +34,6 @@ use crate::plan::physical::operator::Source; use crate::plan::physical::operator::SourceBlocked; use crate::plan::physical::operator::SourceReady; use crate::plan::physical::operator::Step; -use crate::plan::physical::operators::ConcatSource; -use crate::plan::physical::operators::SelectStage; use crate::segments::SegmentId; use crate::test::SESSION; @@ -113,48 +111,6 @@ fn segment(id: u32) -> Option { Some(SegmentId::from(id)) } -struct ActivateChunks; - -impl Source for ActivateChunks { - fn ready(&mut self, _cx: &mut Cx<'_>) -> VortexResult { - Ok(SourceReady::Activate) - } - - fn activate(&mut self, compiler: &mut Compiler<'_>) -> VortexResult { - let mut chains = Vec::new(); - for id in 0..2 { - let plan = RowSource::plan(8, vec![4], segment(id), false); - let start = u64::from(id) * 4; - chains.push( - compiler - .compile(&plan, start..start + 4, &Mask::new_true(4))? - .ok_or_else(|| vortex_err!("Selected chunk has no chain"))?, - ); - } - Ok(compiler.join(chains, ConcatSource::new(2))) - } -} - -impl Operator for ActivateChunks { - fn compute(&mut self, _input: Input, _cx: &mut Cx<'_>) -> VortexResult { - Err(vortex_err!("Dispatcher must activate before computing")) - } -} - -#[rstest] -fn activation_preserves_stages_and_orders_new_inlets( - #[values(&[0, 1], &[1, 0])] order: &[u32], -) -> VortexResult<()> { - let plan = SyntheticPlan::new(row_dtype(), 8, Vec::new(), |_, _, _| { - Ok(Some( - Chain::new(ActivateChunks).with(SelectStage::new(Some(1..3), None)), - )) - }) - .into_plan(); - let result = run(&plan, 0..8, Mask::new_true(8), scripted(order))?; - assert_rows(indices([1, 2, 5, 6]), result.arrays) -} - /// A source's pieces come out in order, one array each, after its read if it has one. #[rstest] fn source_pieces_come_out_in_order(#[values(None, Some(7))] read: Option) -> VortexResult<()> { diff --git a/vortex-layout/src/plan/physical/pipeline/tests.rs b/vortex-layout/src/plan/physical/pipeline/tests.rs index 40a296c03c1..c5460cef9c5 100644 --- a/vortex-layout/src/plan/physical/pipeline/tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/tests.rs @@ -71,7 +71,6 @@ use crate::plan::physical::operator::SourceBlocked; use crate::plan::physical::operator::SourceReady; use crate::plan::physical::operator::Step; use crate::plan::physical::operators::PackSource; -use crate::plan::physical::operators::compile_deferred_list; use crate::plan::physical::optimize; use crate::segments::SegmentId; use crate::test::SESSION; @@ -211,8 +210,16 @@ struct Run { /// Drives a scan the way an owner does, completing reads only while the scan waits. fn drive( store: &Store, - mut scan: Scan, + scan: Scan, mut pick: impl FnMut(&[ReadRequest]) -> usize, +) -> VortexResult { + drive_inner(store, scan, &mut pick) +} + +fn drive_inner( + store: &Store, + mut scan: Scan, + pick: &mut dyn FnMut(&[ReadRequest]) -> usize, ) -> VortexResult { let mut inflight: Vec = Vec::new(); let mut batch: Vec = Vec::new(); @@ -237,7 +244,19 @@ fn drive( splits[split].push(array.clone()); arrays.push(array); } - Turn::Child(_) => unreachable!("child planning is disabled"), + Turn::Child(child) => { + let (id, child_scan) = child.into_scan(SESSION.clone())?; + let child = drive_inner(store, child_scan, pick)?; + assert_eq!(child.leftover, (0, 0)); + events.extend( + child + .events + .into_iter() + .filter(|event| !matches!(event, Event::Piece(_))), + ); + let arrays = child.splits.into_iter().flatten().collect(); + scan.complete_child(id, Some(arrays))?; + } Turn::Waiting => { if !batch.is_empty() { events.push(Event::Io(std::mem::take(&mut batch))); @@ -306,7 +325,7 @@ fn delivery(order: Delivery) -> impl FnMut(&[ReadRequest]) -> usize { fn grouped_list_ranges_preserve_order_and_shared_reads( #[values(Delivery::Fifo, Delivery::Lifo)] order: Delivery, #[values(false, true)] shared: bool, - #[values(false, true)] deferred: bool, + #[values(false, true)] planned: bool, ) -> VortexResult<()> { let mut store = Store::default(); let elements = Buffer::from_iter(0_i64..20).into_array(); @@ -326,11 +345,7 @@ fn grouped_list_ranges_preserve_order_and_shared_reads( None, )? .into_plan(); - let plan = if deferred { - deferred_lists(plan) - } else { - plan - }; + let plan = if planned { selected_list(plan) } else { plan }; let rows = [0_u64, 1, 7, 8]; let mask = Mask::from_iter((0..10).map(|row| rows.contains(&row))); let result = run(&store, &plan, 0..10, mask, delivery(order))?; @@ -357,20 +372,21 @@ fn grouped_list_ranges_preserve_order_and_shared_reads( Ok(()) } -fn deferred_lists(plan: PlanRef) -> PlanRef { +fn selected_list(plan: PlanRef) -> PlanRef { SyntheticPlan::new( plan.dtype().clone(), plan.row_count(), Vec::new(), move |rows, mask, compiler| { - compile_deferred_list(plan.as_::(), rows, mask, compiler).map(Some) + plan.as_::() + .compile_selected(rows, mask, compiler, true) }, ) .into_plan() } #[rstest] -fn deferred_list_does_not_lower_unselected_chunks( +fn selected_list_view_skips_unselected_children( #[values(Delivery::Fifo, Delivery::Lifo)] order: Delivery, #[values(false, true)] empty: bool, ) -> VortexResult<()> { @@ -378,7 +394,7 @@ fn deferred_list_does_not_lower_unselected_chunks( let elements = Buffer::from_iter(10_i64..12).into_array(); let selected = lower(&store.flat(&elements)?)?; let unused = SyntheticPlan::new(elements.dtype().clone(), 2, Vec::new(), |_, _, _| { - Err(vortex_err!("Unselected chunk must not be compiled")) + Err(vortex_err!("Unselected child must not be compiled")) }) .into_plan(); let chunks = ConcatPlan::try_new( @@ -386,7 +402,7 @@ fn deferred_list_does_not_lower_unselected_chunks( vec![unused.clone(), selected, unused], )?; let offsets = Buffer::from(vec![0_u64, 2, 4, 6]).into_array(); - let plan = deferred_lists( + let plan = selected_list( ListPackPlan::try_new( false.into(), 3, @@ -396,8 +412,13 @@ fn deferred_list_does_not_lower_unselected_chunks( )? .into_plan(), ); - let mask = Mask::from_iter([false, !empty, false]); - let result = run(&store, &plan, 0..3, mask, delivery(order))?; + let result = run( + &store, + &plan, + 0..3, + Mask::from_iter([false, !empty, false]), + delivery(order), + )?; let expected = ListArray::try_new( elements, Buffer::from(vec![0_u64, 2]).into_array(), @@ -421,7 +442,7 @@ fn deferred_list_does_not_lower_unselected_chunks( fn dense_nullable_list_ranges_share_one_scan( #[values(Delivery::Fifo, Delivery::Lifo)] order: Delivery, #[values(false, true)] partial: bool, - #[values(false, true)] deferred: bool, + #[values(false, true)] planned: bool, ) -> VortexResult<()> { let mut store = Store::default(); let elements = Buffer::from_iter(0_i64..20).into_array(); @@ -441,11 +462,7 @@ fn dense_nullable_list_ranges_share_one_scan( Some(lower(&store.flat(&valid)?)?), )? .into_plan(); - let plan = if deferred { - deferred_lists(plan) - } else { - plan - }; + let plan = if planned { selected_list(plan) } else { plan }; let mask = Mask::from_iter((0..10).map(|row| !partial || row != 4)); let result = run(&store, &plan, 0..10, mask.clone(), delivery(order))?; if !partial { @@ -1074,9 +1091,7 @@ fn pack_readiness_tracks_changes_and_consumption( let mut changed = InletSet::all(2); let mut child_requests = Vec::new(); let mut child_inbox = None; - let mut resolved_selections = Vec::new(); let mut cx = Cx { - resolved_selections: &mut resolved_selections, child_requests: &mut child_requests, child_inbox: &mut child_inbox, arena: &mut arena, diff --git a/vortex-layout/src/plan/physical/plans/list_pack.rs b/vortex-layout/src/plan/physical/plans/list_pack.rs index 94b0b9ac2d9..4be082730ae 100644 --- a/vortex-layout/src/plan/physical/plans/list_pack.rs +++ b/vortex-layout/src/plan/physical/plans/list_pack.rs @@ -22,7 +22,6 @@ use crate::plan::physical::PlanParts; use crate::plan::physical::PlanRef; use crate::plan::physical::PlanVTable; use crate::plan::physical::operators::ListSource; -use crate::plan::physical::operators::compile_deferred_list; use crate::plan::physical::operators::UNBOUNDED; use crate::plan::physical::pipeline::Chain; use crate::plan::physical::pipeline::Compiler; @@ -88,6 +87,39 @@ impl ListPackPlan { Ok(unsafe { Self::from_children_unchecked(dtype, row_count, children) }) } + pub(crate) fn compile_selected( + &self, + rows: Range, + mask: &Mask, + compiler: &mut Compiler<'_>, + planner_owned: bool, + ) -> VortexResult> { + let plan = self; + if mask.all_false() { + return Ok(None); + } + let boundaries = offset_mask(mask); + let offsets = compiler + .compile(&plan.offsets()?, rows.start..rows.end + 1, &boundaries)? + .ok_or_else(|| vortex_err!("Selected list offsets produced no rows"))?; + let mut chains = vec![offsets]; + if let Some(validity) = plan.validity()? { + chains.push( + compiler + .compile(&validity, rows, mask)? + .ok_or_else(|| vortex_err!("Selected list validity produced no rows"))?, + ); + } + let source = ListSource::new( + plan.clone(), + mask.clone(), + compiler.values().clone(), + compiler.plan_children(), + planner_owned, + ); + Ok(Some(compiler.join_with(chains, |_| UNBOUNDED, source))) + } + /// Returns the plan producing list elements. pub fn elements(&self) -> VortexResult { self.child_required(ELEMENTS) @@ -141,31 +173,12 @@ impl PlanVTable for ListPack { mask: &Mask, compiler: &mut Compiler<'_>, ) -> VortexResult> { - if mask.all_false() { - return Ok(None); - } - if std::env::var_os("VORTEX_LIST_PIPELINE").is_some() { - return compile_deferred_list(plan, rows, mask, compiler).map(Some); - } - let boundaries = offset_mask(mask); - let offsets = compiler - .compile(&plan.offsets()?, rows.start..rows.end + 1, &boundaries)? - .ok_or_else(|| vortex_err!("Selected list offsets produced no rows"))?; - let mut chains = vec![offsets]; - if let Some(validity) = plan.validity()? { - chains.push( - compiler - .compile(&validity, rows, mask)? - .ok_or_else(|| vortex_err!("Selected list validity produced no rows"))?, - ); - } - let source = ListSource::new( - plan.clone(), - mask.clone(), - compiler.values().clone(), - compiler.plan_children(), - ); - Ok(Some(compiler.join_with(chains, |_| UNBOUNDED, source))) + plan.compile_selected( + rows, + mask, + compiler, + std::env::var_os("VORTEX_LIST_PIPELINE").is_some(), + ) } fn reach( diff --git a/vortex-layout/src/scan/execution/adaptor.rs b/vortex-layout/src/scan/execution/adaptor.rs index edf681b294f..dce9beb9743 100644 --- a/vortex-layout/src/scan/execution/adaptor.rs +++ b/vortex-layout/src/scan/execution/adaptor.rs @@ -29,7 +29,9 @@ use crate::segments::SegmentId; /// Rows a running plan produced, in row order. #[derive(Clone)] pub(super) struct Piece { - /// Orders the piece among the outputs of one scan. + /// The selected split that produced this piece. + pub(super) split: usize, + /// Orders the piece within its split. pub(super) rows: Range, /// The selected rows the piece holds. pub(super) array: ArrayRef, @@ -99,7 +101,7 @@ impl ScanAdaptor { } Ok(match self.step()? { Turn::Waiting | Turn::Done => AdaptorStep::Yield, - Turn::Output(_, array) => AdaptorStep::Piece(self.piece(array)), + Turn::Output(split, array) => AdaptorStep::Piece(self.piece(split, array)), Turn::Child(child) => AdaptorStep::Child(child), Turn::Read(read) => { // Publish every read the scan has ready together, so the owner can coalesce them. @@ -134,10 +136,11 @@ impl ScanAdaptor { Ok(turn) } - fn piece(&mut self, array: ArrayRef) -> Piece { + fn piece(&mut self, split: usize, array: ArrayRef) -> Piece { let start = self.next_piece; self.next_piece += 1; Piece { + split, rows: start..start + 1, array, } diff --git a/vortex-layout/src/scan/execution/children.rs b/vortex-layout/src/scan/execution/children.rs index e57ba8a4816..47255bc85a0 100644 --- a/vortex-layout/src/scan/execution/children.rs +++ b/vortex-layout/src/scan/execution/children.rs @@ -48,6 +48,7 @@ use super::ScanPlans; use super::SegmentLocation; use super::projection::ProjectionPlanner; use crate::plan::physical::PlanRef; +use crate::plan::physical::pipeline::SelectedPlan; use crate::plan::physical::pipeline::Split; use crate::plan::physical::pipeline::ValueCache; @@ -114,7 +115,7 @@ pub(in crate::scan) fn child_executor( io: Arc, ) -> Option> { let workers = get_available_parallelism().unwrap_or(1); - (workers > 1).then(|| { + (workers > 1 || std::env::var_os("VORTEX_LIST_PIPELINE").is_some()).then(|| { Arc::new(Children { session, locations, @@ -133,6 +134,48 @@ impl Children { self.admitted.load(Ordering::Relaxed) } + pub(super) fn can_parallelize(&self) -> bool { + self.workers > 1 + } + + /// Execute a selected graph without holding the fixed-width parallel admission budget. + /// Nested lists may ask their enclosing planner for further selected views. + pub(super) fn execute_view( + self: Arc, + view: SelectedPlan, + values: ValueCache, + ) -> BoxFuture<'static, VortexResult>> { + async move { + let rows = view.splits.first().map_or(0, |split| split.rows.start) + ..view.splits.last().map_or(0, |split| split.rows.end); + let scope = WorkScope { + file_ordinal: 0, + rows, + }; + let plans = ScanPlans { + session: self.session.clone(), + locations: Arc::clone(&self.locations), + projection: view.plan.clone(), + row_offset: 0, + values, + children: Some(Arc::clone(&self)), + }; + let root = ProjectionPlanner::from_view(plans, view, scope.clone()); + let mut run = Run::new(); + run.admit(Box::new(root), scope, self.io.session()); + let mut arrays = Vec::new(); + loop { + match run.advance()? { + Progress::Batch(batch) => arrays.push(batch.array), + Progress::RootDone(_) | Progress::Idle => break, + Progress::Waiting => poll_fn(|cx| run.poll_completion(cx)).await?, + } + } + Ok(arrays) + } + .boxed() + } + pub(super) fn execute( self: Arc, plan: PlanRef, diff --git a/vortex-layout/src/scan/execution/children/tests.rs b/vortex-layout/src/scan/execution/children/tests.rs index 680dd781c68..c6c166730a6 100644 --- a/vortex-layout/src/scan/execution/children/tests.rs +++ b/vortex-layout/src/scan/execution/children/tests.rs @@ -300,6 +300,62 @@ async fn independent_children_respect_the_shared_budget_and_return_in_order() -> Ok(()) } +#[tokio::test] +async fn selected_view_preserves_split_order_and_skips_unselected_segments() -> VortexResult<()> { + let (children, service, plan, _) = fixture()?; + let view = SelectedPlan::new(plan, vec![Split::all(1..4), Split::all(17..20)], 2); + let task = tokio::spawn(Arc::clone(&children).execute_view(view, ValueCache::default())); + until(|| { + service + .sources + .lock() + .first() + .is_some_and(|source| source.requests.lock().len() == 2) + }) + .await; + let source = Arc::clone(&service.sources.lock()[0]); + source.requests.lock().make_contiguous().reverse(); + source.open(); + let arrays = task.await.expect("selected view task panicked")?; + let actual = + ChunkedArray::try_new(arrays, DType::Primitive(PType::U64, false.into()))?.into_array(); + let expected = Buffer::from_iter((1_u64..4).chain(17..20)).into_array(); + assert_arrays_eq!(actual, expected, &mut children.session.create_execution_ctx()); + assert_eq!(children.admitted(), 0); + assert!(children.budget.get().is_none()); + assert_eq!(service.sources.lock().len(), 1); + assert_eq!(source.cleared.load(Ordering::Relaxed), 1); + Ok(()) +} + +#[rstest::rstest] +#[tokio::test] +async fn cancelled_or_failed_selected_view_releases_its_io( + #[values(false, true)] failure: bool, +) -> VortexResult<()> { + let (children, service, plan, _) = fixture()?; + let view = SelectedPlan::new(plan, vec![Split::all(1..4)], 1); + let task = tokio::spawn(Arc::clone(&children).execute_view(view, ValueCache::default())); + until(|| !service.sources.lock().is_empty()).await; + let source = Arc::clone(&service.sources.lock()[0]); + if failure { + source.fail.store(true, Ordering::Release); + source.open(); + let error = task + .await + .expect("selected view task panicked") + .expect_err("injected IO failure must propagate"); + assert!(error.to_string().contains("injected child IO failure")); + } else { + task.abort(); + assert!(task.await.expect_err("task must be cancelled").is_cancelled()); + } + until(|| source.cleared.load(Ordering::Relaxed) == 1).await; + assert!(source.requests.lock().is_empty()); + assert!(children.budget.get().is_none()); + Ok(()) +} + #[rstest::rstest] #[tokio::test] async fn cancelled_or_failed_parent_releases_child_io_and_budget( diff --git a/vortex-layout/src/scan/execution/projection.rs b/vortex-layout/src/scan/execution/projection.rs index 1732a6b4b22..f07054487b5 100644 --- a/vortex-layout/src/scan/execution/projection.rs +++ b/vortex-layout/src/scan/execution/projection.rs @@ -21,8 +21,9 @@ use vortex_scan::planning::morsel::MorselOutput; use vortex_scan::planning::planner::Planner; use vortex_scan::planning::planner::PlannerOutput; use vortex_scan::planning::planner::State; +use vortex_scan::planning::planner::WorkScope; -use crate::plan::physical::pipeline::Scan; +use crate::plan::physical::pipeline::SelectedPlan; use crate::plan::physical::pipeline::Split; use crate::scan::execution::ScanPlans; use crate::scan::execution::SelectedRows; @@ -37,7 +38,8 @@ use crate::scan::execution::children::splits::parallel_splits; pub(super) struct ProjectionPlanner { plans: Option, dtype: DType, - selected: SelectedRows, + scope: WorkScope, + view: Option, adaptor: Option, pieces: Vec, done: bool, @@ -48,12 +50,25 @@ pub(super) struct ProjectionPlanner { impl ProjectionPlanner { /// Creates a planner projecting `selected`. pub(super) fn new(plans: ScanPlans, selected: SelectedRows) -> Self { + let view = SelectedPlan::new( + plans.projection.clone(), + vec![Split { + rows: selected.scope.rows.clone(), + mask: selected.mask, + }], + 1, + ); + Self::from_view(plans, view, selected.scope) + } + + pub(super) fn from_view(plans: ScanPlans, view: SelectedPlan, scope: WorkScope) -> Self { Self { dtype: plans.projection.dtype().clone(), children: plans.children.clone(), pending: None, plans: Some(plans), - selected, + scope, + view: Some(view), adaptor: None, pieces: Vec::new(), done: false, @@ -62,7 +77,8 @@ impl ProjectionPlanner { fn finish(&mut self) -> VortexResult { self.done = true; - self.pieces.sort_by_key(|piece| piece.rows.start); + self.pieces + .sort_by_key(|piece| (piece.split, piece.rows.start)); let mut arrays: Vec<_> = std::mem::take(&mut self.pieces) .into_iter() .map(|piece| piece.array) @@ -74,7 +90,7 @@ impl ProjectionPlanner { _ => ChunkedArray::try_new(arrays, self.dtype.clone())?.into_array(), }; Ok(PlannerOutput::Morsel( - self.selected.scope.clone(), + self.scope.clone(), Box::new(ProjectedBatch(Some(array))), )) } @@ -120,17 +136,14 @@ impl Planner for ProjectionPlanner { .plans .take() .vortex_expect("an uncompiled projection still has its plans"); - let scan = Scan::try_new( - plans.session, - plans.projection, - vec![Split { - rows: self.selected.scope.rows.clone(), - mask: self.selected.mask.clone(), - }], - )? - .with_row_offset(plans.row_offset) - .with_value_cache(plans.values) - .with_child_planning(self.children.is_some()); + let view = self + .view + .take() + .vortex_expect("an uncompiled projection has a view"); + let scan = view + .into_scan(plans.session, plans.values)? + .with_row_offset(plans.row_offset) + .with_child_planning(self.children.is_some()); self.adaptor = Some(ScanAdaptor::new(scan, plans.locations)); return Ok(PlannerOutput::Continue); }; @@ -141,21 +154,30 @@ impl Planner for ProjectionPlanner { AdaptorStep::Yield => PlannerOutput::Continue, AdaptorStep::NeedsIO(batch) => PlannerOutput::NeedsIO(batch), AdaptorStep::Child(child) => { - let Some(work) = parallel_splits(&child.plan, &child.splits)? else { - adaptor.complete_child(child.id, None)?; - return Ok(PlannerOutput::Continue); - }; - let children = Arc::clone( self.children .as_ref() .ok_or_else(|| vortex_err!("Child projection has no scheduler"))?, ); + let work = if children.can_parallelize() { + parallel_splits(&child.view.plan, &child.view.splits)? + } else { + None + }; + if work.is_none() && !child.planner_owned { + adaptor.complete_child(child.id, None)?; + return Ok(PlannerOutput::Continue); + } let (send, recv) = oneshot::channel(); self.pending = Some((child.id, recv)); PlannerOutput::Await( async move { - let arrays = children.execute(child.plan, work, child.values).await?; + let arrays = match work { + Some(work) => { + children.execute(child.view.plan, work, child.values).await? + } + None => children.execute_view(child.view, child.values).await?, + }; send.send(arrays) .map_err(|_| vortex_err!("Projection planner was cancelled")) } diff --git a/vortex/examples/selection_scan.rs b/vortex/examples/selection_scan.rs index 57def21b9e0..2942a7b7242 100644 --- a/vortex/examples/selection_scan.rs +++ b/vortex/examples/selection_scan.rs @@ -143,7 +143,7 @@ async fn open( session.allocator(), )?); options - .with_metrics_registry(metrics.clone()) + .with_metrics_registry(Arc::clone(metrics)) .open_read(InstrumentedReadAt::new(read, metrics.as_ref())) .await? } From 1107f53c34e7f9edd26f3ea372c973bf376ef063 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 19 Oct 2026 17:04:55 +0100 Subject: [PATCH 6/7] Fix metrics registry Arc coercion in selection benchmark Signed-off-by: Joe Isaacs --- vortex/examples/selection_scan.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex/examples/selection_scan.rs b/vortex/examples/selection_scan.rs index 2942a7b7242..77c9c14519f 100644 --- a/vortex/examples/selection_scan.rs +++ b/vortex/examples/selection_scan.rs @@ -143,7 +143,7 @@ async fn open( session.allocator(), )?); options - .with_metrics_registry(Arc::clone(metrics)) + .with_metrics_registry(Arc::::clone(metrics)) .open_read(InstrumentedReadAt::new(read, metrics.as_ref())) .await? } From fe25a609313b42a4690445510aeefbd9574f4368 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 19 Oct 2026 17:34:41 +0100 Subject: [PATCH 7/7] Remove custom ListLayout random-access benchmark machinery Signed-off-by: Joe Isaacs --- .github/workflows/pr-bench-runner.yml | 64 --- scripts/bench-list-pipeline.py | 213 ---------- vortex/examples/generate_list_fixture.rs | 187 --------- vortex/examples/selection_scan.rs | 470 ----------------------- 4 files changed, 934 deletions(-) delete mode 100644 scripts/bench-list-pipeline.py delete mode 100644 vortex/examples/generate_list_fixture.rs delete mode 100644 vortex/examples/selection_scan.rs diff --git a/.github/workflows/pr-bench-runner.yml b/.github/workflows/pr-bench-runner.yml index 006cff056d7..c7f5846c807 100644 --- a/.github/workflows/pr-bench-runner.yml +++ b/.github/workflows/pr-bench-runner.yml @@ -77,34 +77,6 @@ jobs: target/release_debug/${{ inputs.benchmark_id }} polarsignals-cloud-token: ${{ secrets.POLAR_SIGNALS_API_KEY }} project-id: "e5d846e1-b54c-46e7-9174-8bf055a3af56" - - name: Build list pipeline comparison - if: inputs.benchmark_id == 'random-access-bench' && inputs.remote_data_dir == '' - env: - RUSTFLAGS: "-C target-cpu=native -C force-frame-pointers=yes" - run: | - cargo build --locked --package vortex --profile release_debug \ - --example selection_scan --example generate_list_fixture - cargo bench --locked --package vortex-layout --profile release_debug \ - --bench pipeline_scan --no-run --message-format=json > /tmp/list-micro-build.jsonl - python3 - <<'PY' - import json - import shutil - - with open('/tmp/list-micro-build.jsonl') as source: - for line in source: - artifact = json.loads(line) - if artifact.get('target', {}).get('name') == 'pipeline_scan' and artifact.get('executable'): - shutil.copy2(artifact['executable'], 'target/release_debug/pipeline_scan') - PY - - name: Upload list pipeline comparison binaries - if: inputs.benchmark_id == 'random-access-bench' && inputs.remote_data_dir == '' - uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7 - with: - name: list-pipeline-binaries - path: | - target/release_debug/examples/selection_scan - target/release_debug/examples/generate_list_fixture - target/release_debug/pipeline_scan - name: Upload benchmark binary uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7 with: @@ -150,13 +122,6 @@ jobs: shell: bash run: chmod +x target/release_debug/${{ inputs.benchmark_id }} - - name: Download list pipeline comparison binaries - if: inputs.benchmark_id == 'random-access-bench' && inputs.remote_data_dir == '' - uses: actions/download-artifact@9000827ccba6bdab643e8b6fd33ac0654aef8333 # v8 - with: - name: list-pipeline-binaries - path: target/release_debug/ - - name: Setup Polar Signals if: github.event.pull_request.head.repo.fork == false uses: polarsignals/gh-actions-ps-profiling@68ae857e375a826606352016e5b90f01a2a7ff7a # v0.8.1 @@ -214,35 +179,6 @@ jobs: run: | python3 scripts/compress-split.py - - name: Compare current and deferred list pipelines - if: inputs.benchmark_id == 'random-access-bench' && inputs.remote_data_dir == '' - env: - RUSTFLAGS: "-C target-cpu=native -C force-frame-pointers=yes" - run: | - chmod +x target/release_debug/examples/selection_scan \ - target/release_debug/examples/generate_list_fixture target/release_debug/pipeline_scan - python3 scripts/bench-list-pipeline.py - for round in 0 1; do - versions="current fixed" - if [ "$round" = 1 ]; then versions="fixed current"; fi - for version in $versions; do - if [ "$version" = fixed ]; then export VORTEX_LIST_PIPELINE=1; else unset VORTEX_LIST_PIPELINE; fi - target/release_debug/pipeline_scan --bench list_scan --sample-count 100 --sample-size 100 \ - > "list-pipeline-results/micro-$version-$round.log" - done - done - - - name: Upload list pipeline comparison results - if: always() && inputs.benchmark_id == 'random-access-bench' && inputs.remote_data_dir == '' - uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7 - with: - name: list-pipeline-results - path: | - list-pipeline-results/*.json - list-pipeline-results/*.jsonl - list-pipeline-results/*.log - list-pipeline-results/raw/ - - name: Run ${{ inputs.benchmark_name }} benchmark (per-column-encoder) if: inputs.benchmark_id == 'string-bench' shell: bash diff --git a/scripts/bench-list-pipeline.py b/scripts/bench-list-pipeline.py deleted file mode 100644 index afe36851876..00000000000 --- a/scripts/bench-list-pipeline.py +++ /dev/null @@ -1,213 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright the Vortex contributors - -"""Compare current list scans and planner-owned selected views on identical ListLayout files.""" - -import argparse -import hashlib -import itertools -import json -import os -import platform -import statistics -import subprocess -from pathlib import Path - -ROOT = Path(__file__).resolve().parent.parent -SEEDS = [17, 42, 73] -VERSIONS = ["baseline", "candidate"] - - -def digest(path: Path) -> str: - checksum = hashlib.sha256() - with path.open("rb") as source: - for chunk in iter(lambda: source.read(1024 * 1024), b""): - checksum.update(chunk) - return checksum.hexdigest() - - -class Comparison: - def __init__(self, args: argparse.Namespace) -> None: - self.binaries = { - "baseline": (args.baseline_binary or args.binary).resolve(), - "candidate": args.binary.resolve(), - } - self.results = args.output.resolve() - self.data = self.results / "data" - self.generator = args.generator.resolve() - self.iterations = args.iterations - self.rounds = args.rounds - self.discard = 8 - if self.iterations <= self.discard or self.rounds <= 0: - raise ValueError("Need more than 8 iterations and at least one round") - (self.results / "raw").mkdir(parents=True, exist_ok=True) - self.data.mkdir(parents=True, exist_ok=True) - - def prepare(self) -> dict[str, Path]: - for target_bytes in [65536, 1048576]: - for size, shape in [("large", "primitive"), ("large", "nested"), ("small", "primitive")]: - name = f"{size}-{shape}-{target_bytes}" - path = self.data / f"{name}.vortex" - if path.exists(): - continue - command = [ - str(self.generator), - str(path), - "small" if size == "small" else "skewed", - shape, - str(target_bytes), - ] - result = subprocess.run(command, capture_output=True, text=True, timeout=600) - (self.results / f"layout-{name}.log").write_text(result.stderr) - result.check_returncode() - (self.results / f"layout-{name}.json").write_text(result.stdout) - files = {path.stem: path for path in sorted(self.data.glob("*.vortex"))} - assert len(files) == 6, files - metadata = { - "baseline": "current list implementation, prototype disabled", - "commit": subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=ROOT, text=True).strip(), - "candidate": "planner-owned selected views; VORTEX_LIST_PIPELINE=1", - "platform": platform.platform(), - "logical_cpu_count": os.cpu_count(), - "affinity": "unrestricted", - "seeds": SEEDS, - "rounds": self.rounds, - "iterations": self.iterations, - "discard": self.discard, - "output": "full Arrow decode and drop", - "storage": "warm local disk", - "rustflags": os.environ.get("RUSTFLAGS", "from Cargo configuration"), - "profile": "release_debug", - "files": { - name: {"path": str(path), "sha256": digest(path), "bytes": path.stat().st_size} - for name, path in files.items() - }, - "binaries": {version: digest(path) for version, path in self.binaries.items()}, - } - self.write_json("metadata.json", metadata) - return files - - def write_json(self, name: str, value: object) -> None: - (self.results / name).write_text(json.dumps(value, indent=2) + "\n") - - def run( - self, - version: str, - path: Path, - count: int, - pattern: str, - mode: str, - seed: int, - check: bool, - stem: str, - ) -> dict: - env = {key: value for key, value in os.environ.items() if not key.startswith(("VORTEX_", "TOKIO_"))} - env.update(VORTEX_SCAN_V2="1", RUST_LOG="warn") - if version == "candidate": - env["VORTEX_LIST_PIPELINE"] = "1" - command = [ - str(self.binaries[version]), - str(path), - "1" if check else str(self.iterations), - str(count), - pattern, - mode, - "disk", - "check" if check else "decode", - str(seed), - ] - with (self.results / "commands.jsonl").open("a") as output: - output.write( - json.dumps({"command": command, "selected_view": version == "candidate", "VORTEX_SCAN_V2": "1"}) + "\n" - ) - result = subprocess.run(command, env=env, capture_output=True, text=True, timeout=180) - (self.results / "raw" / f"{stem}.log").write_text(result.stderr) - (self.results / "raw" / f"{stem}.json").write_text(result.stdout) - result.check_returncode() - return json.loads(result.stdout) - - def check(self, files: dict[str, Path]) -> int: - checks = [] - selections = [ - (0, "indices:"), - (1, "indices:0"), - (1, "indices:17"), - (1, "indices:16"), - (100, "uniform"), - (100, "clustered"), - ] - for name, path in files.items(): - for count, pattern in selections: - for version in VERSIONS: - stem = f"check-{name}-{count}-{pattern.replace(':', '_')}-{version}" - data = self.run(version, path, count, pattern, "cached", 42, True, stem) - assert data["correct"] - checks.append({"dataset": name, "version": version, "pattern": pattern, "rows": count}) - self.write_json("correctness.json", checks) - print(json.dumps({"correctness": name, "passed": True}), flush=True) - return len(checks) - - def measure(self, files: dict[str, Path]) -> list[dict]: - summary = [] - for name, path in files.items(): - patterns = [(100, "uniform"), (100, "clustered")] - if not name.startswith("small"): - patterns.insert(0, (1, "indices:16")) - for (count, pattern), mode in itertools.product(patterns, ["cached", "reopen"]): - medians: dict[str, list[float]] = {version: [] for version in VERSIONS} - for seed in SEEDS: - expected_indices = None - for repeat in range(self.rounds): - order = VERSIONS if repeat % 2 == 0 else list(reversed(VERSIONS)) - for version in order: - stem = f"{name}-{count}-{pattern.replace(':', '_')}-{mode}-{seed}-{repeat}-{version}" - data = self.run(version, path, count, pattern, mode, seed, False, stem) - assert data["v2"] == "1" - assert len(data["runs"]) == self.iterations - assert not any(run["instrumented"] for run in data["runs"]) - if expected_indices is None: - expected_indices = data["indices"] - assert data["indices"] == expected_indices - median = statistics.median(run["elapsed_ns"] for run in data["runs"][self.discard :]) - medians[version].append(median / 1e6) - row = {"dataset": name, "rows": count, "pattern": pattern, "mode": mode} - row.update({version + "_ms": statistics.median(times) for version, times in medians.items()}) - row["change_pct"] = 100 * (row["candidate_ms"] / row["baseline_ms"] - 1) - row["round_medians_ms"] = medians - summary.append(row) - self.write_json("summary.json", summary) - print(json.dumps({key: value for key, value in row.items() if key != "round_medians_ms"}), flush=True) - return summary - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--binary", type=Path, default=ROOT / "target/release_debug/examples/selection_scan") - parser.add_argument("--baseline-binary", type=Path, help="Defaults to --binary with the prototype disabled") - parser.add_argument("--generator", type=Path, default=ROOT / "target/release_debug/examples/generate_list_fixture") - parser.add_argument("--output", type=Path, default=ROOT / "list-pipeline-results") - parser.add_argument("--iterations", type=int, default=64) - parser.add_argument("--rounds", type=int, default=2) - comparison = Comparison(parser.parse_args()) - files = comparison.prepare() - checks = comparison.check(files) - summary = comparison.measure(files) - if destination := os.environ.get("GITHUB_STEP_SUMMARY"): - with open(destination, "a") as output: - output.write("\n### Deferred list pipeline experiment\n\n") - output.write( - f"{checks} reference checks passed; full Arrow decode on warm local disk. " - "The candidate builds selected child graphs in the planner and retains admission limits for parallel fixed-width work. " - "Positive changes are slower.\n\n" - ) - output.write("| Dataset | Rows | Selection | Open mode | Current ms | Selected view ms | Change |\n") - output.write("| --- | ---: | --- | --- | ---: | ---: | ---: |\n") - for row in summary: - output.write( - f"| {row['dataset']} | {row['rows']} | {row['pattern']} | {row['mode']} | " - f"{row['baseline_ms']:.3f} | {row['candidate_ms']:.3f} | {row['change_pct']:+.1f}% |\n" - ) - - -if __name__ == "__main__": - main() diff --git a/vortex/examples/generate_list_fixture.rs b/vortex/examples/generate_list_fixture.rs deleted file mode 100644 index 6968df77f7a..00000000000 --- a/vortex/examples/generate_list_fixture.rs +++ /dev/null @@ -1,187 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Generate large list rows for selection-scan diagnostics. -//! -//! Usage: `generate_list_fixture OUTPUT balanced|skewed|small primitive|nested TARGET_BYTES`. -//! Writes a new file and verifies full decoded equality through V1. Element block sizing is -//! independent of outer row count; nullable lists, empty lists, and nullable elements are included. - -use std::collections::BTreeMap; -use std::path::Path; -use std::sync::Arc; - -use anyhow::Result; -use anyhow::ensure; -use futures::stream; -use rand::RngExt; -use rand::SeedableRng; -use rand::rngs::StdRng; -use vortex::VortexSessionDefault; -use vortex::array::IntoArray; -use vortex::array::VortexSessionExecute; -use vortex::array::arrays::BoolArray; -use vortex::array::arrays::ListArray; -use vortex::array::arrays::PrimitiveArray; -use vortex::array::arrays::StructArray; -use vortex::array::stream::ArrayStreamAdapter; -use vortex::array::stream::ArrayStreamExt; -use vortex::array::validity::Validity; -use vortex::buffer::Buffer; -use vortex::editions::Edition; -use vortex::editions::EditionDeclaration; -use vortex::editions::EditionId; -use vortex::editions::EditionMember; -use vortex::editions::EditionSessionExt; -use vortex::file::OpenOptionsSessionExt; -use vortex::file::WriteOptionsSessionExt; -use vortex::file::WriteStrategyBuilder; -use vortex::io::session::RuntimeSessionExt; -use vortex::layout::LayoutRef; -use vortex::layout::layouts::table::TableStrategy; -use vortex::session::VortexSession; -use vortex_arrow::ArrowSessionExt; - -const INNER_LENGTH: u64 = 128; - -static EDITION: EditionDeclaration = EditionDeclaration { - edition: Edition { - id: EditionId::new("large-list-fixture", 2026, 10, 0), - min_library_version: None, - }, - added: &[EditionMember::layout(&"vortex.list")], -}; - -fn count_layouts(layout: &LayoutRef, counts: &mut BTreeMap) -> Result<()> { - *counts.entry(layout.encoding_id().to_string()).or_default() += 1; - for child in layout.children()? { - count_layouts(&child, counts)?; - } - Ok(()) -} - -#[tokio::main] -async fn main() -> Result<()> { - let args = std::env::args().skip(1).collect::>(); - ensure!( - args.len() == 4, - "Expected OUTPUT balanced|skewed|small primitive|nested TARGET_BYTES" - ); - ensure!( - matches!(args[1].as_str(), "balanced" | "skewed" | "small"), - "Unknown distribution" - ); - ensure!( - matches!(args[2].as_str(), "primitive" | "nested"), - "Unknown element shape" - ); - let target_bytes = args[3].parse::()?; - ensure!(target_bytes > 0, "TARGET_BYTES must be positive"); - let session = VortexSession::default().with_tokio(); - session.register_edition(&EDITION)?; - session.enable_edition(EDITION.edition.id)?; - - let rows = if args[1] == "small" { 1_000_000 } else { 512 }; - let mut offsets = vec![0_u64]; - let mut total = 0_u64; - for row in 0..rows { - total += if row % 19 == 0 { - 0 - } else if args[1] == "small" { - 1 + row % 20 - } else if args[1] == "balanced" { - 8192 - } else if row % 16 == 0 { - 262_144 - } else { - INNER_LENGTH - }; - offsets.push(total); - } - let mut rng = StdRng::seed_from_u64(42); - let mut elements = PrimitiveArray::from_option_iter((0..total).map(|index| { - let value = rng.random::(); - (index % 7 != 0).then_some(value) - })) - .into_array(); - if args[2] == "nested" { - let inner_count = total / INNER_LENGTH; - elements = ListArray::try_new( - elements, - Buffer::from_iter((0..=inner_count).map(|index| index * INNER_LENGTH)).into_array(), - Validity::Array( - BoolArray::from_iter((0..inner_count).map(|index| index % 11 != 0)).into_array(), - ), - )? - .into_array(); - for offset in &mut offsets { - *offset /= INNER_LENGTH; - } - } - let values = ListArray::try_new( - elements, - Buffer::from(offsets).into_array(), - Validity::Array(BoolArray::from_iter((0..rows).map(|row| row % 17 != 0)).into_array()), - )? - .into_array(); - let input = StructArray::from_fields(&[ - ("id", Buffer::from_iter(0..rows).into_array()), - ("values", values), - ])? - .into_array(); - let elements = WriteStrategyBuilder::from_session(&session) - .with_data_block_target_bytes(Some(target_bytes)) - .build(); - let defaults = WriteStrategyBuilder::from_session(&session).build(); - let strategy = TableStrategy::new(Arc::clone(&defaults), defaults) - .with_list_layout_factory(move |list| Arc::new(list.with_elements(Arc::clone(&elements)))); - let mut output = tokio::fs::OpenOptions::new() - .write(true) - .create_new(true) - .open(&args[0]) - .await?; - let summary = session - .write_options() - .with_strategy(Arc::new(strategy)) - .write( - &mut output, - ArrayStreamAdapter::new(input.dtype().clone(), stream::iter([Ok(input.clone())])), - ) - .await?; - drop(output); - let file = session - .open_options() - .open_path(Path::new(&args[0])) - .await?; - let actual = ArrayStreamAdapter::new(file.dtype().clone(), file.scan()?.into_stream()?) - .read_all() - .await?; - let mut ctx = session.create_execution_ctx(); - let arrow = session.arrow(); - let field = arrow.to_arrow_field("", input.dtype())?; - let expected = arrow.execute_arrow(input, Some(&field), &mut ctx)?; - let actual = arrow.execute_arrow(actual, Some(&field), &mut ctx)?; - ensure!( - expected.as_ref() == actual.as_ref(), - "Decoded fixture differs from generated input" - ); - let mut layouts = BTreeMap::new(); - count_layouts(file.footer().layout(), &mut layouts)?; - ensure!( - layouts.get("vortex.list").copied().unwrap_or(0) > 0, - "Expected ListLayout" - ); - println!( - "{}", - serde_json::json!({ - "output": args[0], "distribution": args[1], "shape": args[2], - "target_bytes": target_bytes, "rows": rows, "primitive_elements": total, - "max_row_payload_bytes": match args[1].as_str() { - "small" => 20 * 8, "balanced" => 8192 * 8, _ => 262_144 * 8, - }, - "bytes": summary.size(), "correct": true, "layouts": layouts, - "layout_tree": file.footer().layout().display_tree().to_string(), - }) - ); - Ok(()) -} diff --git a/vortex/examples/selection_scan.rs b/vortex/examples/selection_scan.rs deleted file mode 100644 index 77c9c14519f..00000000000 --- a/vortex/examples/selection_scan.rs +++ /dev/null @@ -1,470 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Diagnose selection scans on an existing Vortex file. -//! -//! Usage: `selection_scan FILE ITERATIONS ROWS uniform|clustered|all cached|reopen memory|disk -//! time|metrics|check|decode|decode-metrics [SEED]`. Set `VORTEX_SCAN_V2=0/1` in separate processes. -//! Timings exclude instrumentation in `time` and `decode` modes. The flag selects the explicit V2 -//! builder here only. `check` compares against a full V1 scan followed by a take. -//! `decode` consumes the selected rows in all columns as Arrow record batches; `decode-metrics` -//! measures allocations while doing the same work. -//! Instrumented disk runs also report physical reads/bytes after coalescing, summed over all -//! iterations and excluding the initial file open (reopen-mode opens are included). -//! Build with `VORTEX_SELECTION_METRICS=1` for the separate allocation-instrumented executable. -//! Use `indices:16,32` in place of the pattern to select exact, strictly increasing row indices; -//! ROWS must match the number supplied. This supports targeted large-row and null-row checks. - -use std::alloc::GlobalAlloc; -use std::alloc::Layout; -use std::hint::black_box; -use std::path::Path; -use std::pin::pin; -use std::sync::Arc; -use std::sync::atomic::AtomicU64; -use std::sync::atomic::Ordering; -use std::time::Instant; - -use anyhow::Result; -use anyhow::bail; -use anyhow::ensure; -use futures::StreamExt; -use rand::RngExt; -use rand::SeedableRng; -use rand::rngs::StdRng; -use vortex::VortexSessionDefault; -use vortex::array::ArrayRef; -use vortex::array::Canonical; -use vortex::array::IntoArray; -use vortex::array::VortexSessionExecute; -use vortex::array::memory::MemorySessionExt; -use vortex::array::stream::ArrayStreamAdapter; -use vortex::array::stream::ArrayStreamExt; -use vortex::buffer::Buffer; -use vortex::buffer::ByteBuffer; -use vortex::file::OpenOptionsSessionExt; -use vortex::file::VortexFile; -use vortex::io::InstrumentedReadAt; -use vortex::io::session::RuntimeSessionExt; -use vortex::io::std_file::FileReadAt; -use vortex::layout::scan::v2::ScanBuilder as V2ScanBuilder; -use vortex::layout::segments::SegmentFuture; -use vortex::layout::segments::SegmentId; -use vortex::layout::segments::SegmentSource; -use vortex::metrics::DefaultMetricsRegistry; -use vortex::metrics::MetricValue; -use vortex::metrics::MetricsRegistry; -use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; -use vortex::session::VortexSession; -use vortex_arrow::ArrowSessionExt; - -const TRACK_ALLOCATIONS: bool = option_env!("VORTEX_SELECTION_METRICS").is_some(); -static CALLS: AtomicU64 = AtomicU64::new(0); -static BYTES: AtomicU64 = AtomicU64::new(0); -static LIVE: AtomicU64 = AtomicU64::new(0); -static PEAK: AtomicU64 = AtomicU64::new(0); - -struct Allocator; - -#[global_allocator] -static ALLOCATOR: Allocator = Allocator; - -fn allocated(size: usize) { - CALLS.fetch_add(1, Ordering::Relaxed); - BYTES.fetch_add(size as u64, Ordering::Relaxed); - let live = LIVE.fetch_add(size as u64, Ordering::Relaxed) + size as u64; - PEAK.fetch_max(live, Ordering::Relaxed); -} - -// SAFETY: All operations delegate to the same allocator with the original pointer and layout. -unsafe impl GlobalAlloc for Allocator { - unsafe fn alloc(&self, layout: Layout) -> *mut u8 { - // SAFETY: The caller supplies a valid allocation layout. - let pointer = unsafe { mimalloc::MiMalloc.alloc(layout) }; - if !pointer.is_null() && TRACK_ALLOCATIONS { - allocated(layout.size()); - } - pointer - } - - unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { - // SAFETY: The caller supplies a valid allocation layout. - let pointer = unsafe { mimalloc::MiMalloc.alloc_zeroed(layout) }; - if !pointer.is_null() && TRACK_ALLOCATIONS { - allocated(layout.size()); - } - pointer - } - - unsafe fn realloc(&self, pointer: *mut u8, layout: Layout, size: usize) -> *mut u8 { - // SAFETY: The caller guarantees the original allocation and requested size are valid. - let pointer = unsafe { mimalloc::MiMalloc.realloc(pointer, layout, size) }; - if !pointer.is_null() && TRACK_ALLOCATIONS { - LIVE.fetch_sub(layout.size() as u64, Ordering::Relaxed); - allocated(size); - } - pointer - } - - unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) { - if TRACK_ALLOCATIONS { - LIVE.fetch_sub(layout.size() as u64, Ordering::Relaxed); - } - // SAFETY: This allocation came from mimalloc with the same layout. - unsafe { mimalloc::MiMalloc.dealloc(pointer, layout) }; - } -} - -struct CountingSegments { - inner: Arc, - calls: Arc<[AtomicU64]>, -} - -impl SegmentSource for CountingSegments { - fn request(&self, id: SegmentId) -> SegmentFuture { - self.calls[*id as usize].fetch_add(1, Ordering::Relaxed); - self.inner.request(id) - } -} - -async fn open( - session: &VortexSession, - path: &Path, - bytes: Option<&ByteBuffer>, - metrics: Option<&Arc>, -) -> Result { - let options = session.open_options().with_layout_reader_cache(); - Ok(match (bytes, metrics) { - (Some(bytes), _) => options.open_buffer(bytes.clone())?, - (None, Some(metrics)) => { - let read = Arc::new(FileReadAt::open_with_allocator( - path, - session.handle(), - session.allocator(), - )?); - options - .with_metrics_registry(Arc::::clone(metrics)) - .open_read(InstrumentedReadAt::new(read, metrics.as_ref())) - .await? - } - (None, None) => options.open_path(path).await?, - }) -} - -fn physical_io(metrics: &DefaultMetricsRegistry) -> [u64; 4] { - let mut totals = [0; 4]; - for metric in metrics.snapshot() { - match (metric.name().as_ref(), metric.value()) { - ("vortex.io.read.total_size", MetricValue::Counter(value)) => { - totals[0] += value.value(); - } - ("vortex.io.read.size", MetricValue::Histogram(value)) => { - totals[1] += value.count() as u64; - } - ("io.requests.individual", MetricValue::Counter(value)) => { - totals[2] += value.value(); - } - ("io.requests.coalesced", MetricValue::Counter(value)) => { - totals[3] += value.value(); - } - _ => {} - } - } - totals -} - -fn v2() -> bool { - std::env::var("VORTEX_SCAN_V2").is_ok_and(|value| value == "1") -} - -fn v2_scan( - file: &VortexFile, - indices: &Buffer, - all: bool, - session: &VortexSession, -) -> Result { - let builder = V2ScanBuilder::new(session.clone(), file.scan_file()); - Ok(if all { - builder - } else { - builder.with_row_indices(StrictSortedBuffer::try_new(indices.clone())?) - }) -} - -async fn scan( - file: &VortexFile, - indices: &Buffer, - all: bool, - session: &VortexSession, -) -> Result { - let stream = if v2() { - v2_scan(file, indices, all, session)?.into_stream()? - } else { - let builder = file.scan()?; - let builder = if all { - builder - } else { - builder.with_row_indices(StrictSortedBuffer::try_new(indices.clone())?) - }; - builder.into_stream()?.boxed() - }; - let array = ArrayStreamAdapter::new(file.dtype().clone(), stream) - .read_all() - .await?; - Ok(array - .execute::(&mut session.create_execution_ctx())? - .into_array()) -} - -async fn decode( - file: &VortexFile, - indices: &Buffer, - all: bool, - session: &VortexSession, -) -> Result { - let schema = Arc::new(session.arrow().to_arrow_schema(file.dtype())?); - let stream = if v2() { - v2_scan(file, indices, all, session)? - .into_record_batch_stream(schema)? - .boxed() - } else { - let builder = file.scan()?; - let builder = if all { - builder - } else { - builder.with_row_indices(StrictSortedBuffer::try_new(indices.clone())?) - }; - builder.into_record_batch_stream(schema)?.boxed() - }; - let mut stream = pin!(stream); - let mut rows = 0; - while let Some(batch) = stream.next().await { - rows += batch?.num_rows(); - } - Ok(rows) -} - -fn indices(rows: u64, count: usize, pattern: &str, seed: u64) -> Result> { - if let Some(indices) = pattern.strip_prefix("indices:") { - let selected = if indices.is_empty() { - Vec::new() - } else { - indices - .split(',') - .map(str::parse::) - .collect::, _>>()? - }; - ensure!( - selected.len() == count, - "ROWS must match the explicit selection length" - ); - ensure!( - selected.iter().all(|&index| index < rows), - "Selected row is outside the file" - ); - ensure!( - selected.windows(2).all(|pair| pair[0] < pair[1]), - "Indices must be strictly increasing" - ); - return Ok(Buffer::from(selected)); - } - let mut rng = StdRng::seed_from_u64(seed); - let mut result = Vec::with_capacity(count); - ensure!(count as u64 <= rows, "Selection exceeds the file row count"); - match pattern { - "all" => ensure!(count == 0, "Pass ROWS=0 for an all-row scan"), - "uniform" => { - while result.len() < count { - result.push(rng.random_range(0..rows)); - result.sort_unstable(); - result.dedup(); - } - } - "clustered" => { - while result.len() < count { - let len = (count - result.len()).min(20) as u64; - let start = rng.random_range(0..=rows - len); - result.extend(start..start + len); - result.sort_unstable(); - result.dedup(); - } - } - _ => bail!("Unknown pattern {pattern}"), - } - Ok(Buffer::from(result)) -} - -fn rss() -> Result { - let Ok(status) = std::fs::read_to_string("/proc/self/status") else { - return Ok(serde_json::Value::Null); - }; - let get = |name: &str| -> Option { - status - .lines() - .find(|line| line.starts_with(name))? - .split_whitespace() - .nth(1)? - .parse::() - .ok() - .map(|kb| kb * 1024) - }; - Ok(serde_json::json!({"current": get("VmRSS:"), "peak": get("VmHWM:")})) -} - -#[tokio::main] -async fn main() -> Result<()> { - let args = std::env::args().skip(1).collect::>(); - ensure!( - matches!(args.len(), 7 | 8), - "Expected FILE ITERATIONS ROWS uniform|clustered|all cached|reopen memory|disk time|metrics|check|decode|decode-metrics [SEED]" - ); - let iterations = args[1].parse::()?; - ensure!(iterations > 0, "ITERATIONS must be positive"); - let count = args[2].parse::()?; - let cached = match args[4].as_str() { - "cached" => true, - "reopen" => false, - other => bail!("Unknown open mode {other}"), - }; - let metrics = matches!(args[6].as_str(), "metrics" | "decode-metrics"); - ensure!( - matches!( - args[6].as_str(), - "time" | "metrics" | "check" | "decode" | "decode-metrics" - ), - "Unknown measurement mode" - ); - ensure!( - args[6] == "check" || metrics == TRACK_ALLOCATIONS, - "Build with VORTEX_SELECTION_METRICS=1 only for metrics mode" - ); - let session = VortexSession::default().with_tokio(); - let path = Path::new(&args[0]); - let bytes = match args[5].as_str() { - "memory" => Some(ByteBuffer::from(std::fs::read(path)?)), - "disk" => None, - other => bail!("Unknown storage mode {other}"), - }; - let io_metrics = metrics.then(|| Arc::new(DefaultMetricsRegistry::default())); - let file = open(&session, path, bytes.as_ref(), io_metrics.as_ref()).await?; - let seed = args - .get(7) - .map(|seed| seed.parse::()) - .transpose()? - .unwrap_or(42); - let selected = indices(file.row_count(), count, &args[3], seed)?; - let all = args[3] == "all"; - let decode_rows = matches!(args[6].as_str(), "decode" | "decode-metrics"); - let expected_rows = if all { - usize::try_from(file.row_count())? - } else { - selected.len() - }; - if args[6] == "check" { - let expected = ArrayStreamAdapter::new( - file.dtype().clone(), - file.scan()?.prepare()?.execute_stream(None)?, - ) - .read_all() - .await?; - let expected = if all { - expected - } else { - expected.take(selected.clone().into_array())? - }; - let actual = scan(&file, &selected, all, &session).await?; - let mut ctx = session.create_execution_ctx(); - let arrow = session.arrow(); - let target = arrow.to_arrow_field("", file.dtype())?; - let expected = arrow.execute_arrow(expected, Some(&target), &mut ctx)?; - let actual = arrow.execute_arrow(actual, Some(&target), &mut ctx)?; - ensure!( - actual.as_ref() == expected.as_ref(), - "Selected values differ from full-scan take" - ); - println!( - "{}", - serde_json::json!({"correct": true, "rows": actual.len(), "indices": selected.as_slice()}) - ); - return Ok(()); - } - let locations = file.scan_file().locations; - let calls: Arc<[AtomicU64]> = locations.iter().map(|_| AtomicU64::new(0)).collect(); - let instrument = |file: VortexFile| { - if metrics { - let source = Arc::new(CountingSegments { - inner: file.segment_source(), - calls: Arc::clone(&calls), - }); - file.with_segment_source(source) - } else { - file - } - }; - let file = instrument(file); - let initial_io = io_metrics.as_ref().map(|metrics| physical_io(metrics)); - let mut runs = Vec::with_capacity(iterations); - for _ in 0..iterations { - for calls in calls.iter() { - calls.store(0, Ordering::Relaxed); - } - let before_live = LIVE.load(Ordering::Relaxed); - let before_calls = CALLS.load(Ordering::Relaxed); - let before_bytes = BYTES.load(Ordering::Relaxed); - PEAK.store(before_live, Ordering::Relaxed); - let start = Instant::now(); - let reopened = if cached { - None - } else { - Some(instrument( - open(&session, path, bytes.as_ref(), io_metrics.as_ref()).await?, - )) - }; - let input = reopened.as_ref().unwrap_or(&file); - let (array, actual_rows) = if decode_rows { - (None, decode(input, &selected, all, &session).await?) - } else { - let array = scan(input, &selected, all, &session).await?; - let rows = array.len(); - (Some(array), rows) - }; - drop(reopened); - let elapsed = start.elapsed().as_nanos(); - ensure!(actual_rows == expected_rows, "Wrong selected row count"); - let peak_live = PEAK.load(Ordering::Relaxed).saturating_sub(before_live); - let output_live = LIVE.load(Ordering::Relaxed) as i128 - before_live as i128; - drop(black_box(array)); - let retained_live = LIVE.load(Ordering::Relaxed) as i128 - before_live as i128; - runs.push(serde_json::json!({ - "elapsed_ns": elapsed, "instrumented": metrics, - "allocation_calls": CALLS.load(Ordering::Relaxed) - before_calls, - "allocated_bytes": BYTES.load(Ordering::Relaxed) - before_bytes, - "baseline_live_bytes": before_live, - "total_peak_live_bytes": before_live + peak_live, - "post_drop_live_bytes": LIVE.load(Ordering::Relaxed), - "peak_live_bytes": peak_live, "output_live_bytes": output_live, - "retained_live_bytes": retained_live, - "requests": calls.iter().map(|n| n.load(Ordering::Relaxed)).sum::(), - "requested_bytes": calls.iter().zip(locations.iter()).map(|(n, l)| n.load(Ordering::Relaxed) * u64::from(l.length)).sum::(), - "unique_segments": calls.iter().filter(|n| n.load(Ordering::Relaxed) != 0).count(), - "unique_bytes": calls.iter().zip(locations.iter()).filter(|(n, _)| n.load(Ordering::Relaxed) != 0).map(|(_, l)| u64::from(l.length)).sum::(), - })); - } - let physical_io = io_metrics - .as_ref() - .zip(initial_io) - .map(|(metrics, before)| { - let after = physical_io(metrics); - serde_json::json!({ - "bytes": after[0] - before[0], - "reads": after[1] - before[1], - "individual_requests": after[2] - before[2], - "coalesced_requests": after[3] - before[3], - }) - }); - println!( - "{}", - serde_json::json!({"args": args, "indices": selected.as_slice(), "runs": runs, "physical_io": physical_io, "rss": rss()?, "v2": std::env::var("VORTEX_SCAN_V2").ok()}) - ); - Ok(()) -}