From c4fa7818b7ec0928ebfb96847083a7ae276f7e39 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 19 Oct 2026 09:43:50 +0100 Subject: [PATCH 01/15] 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 02/15] 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 4eaf31bdfe68e1fe4e53dd67e27c7283602e0050 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 19 Oct 2026 12:51:40 +0000 Subject: [PATCH 03/15] Read segment list elements directly from the list source Once a ListSource has its offsets, it previously compiled a nested Scan over the element plan. When the selected element ranges reach only segment chunks (Concat/Filter/Zoned/LegacyStats over SegmentScan), the list source now issues those reads itself: each touched segment is requested once, all together, decoded on arrival and sliced or masked into the ranges it holds, in list order. Other element plans (dict, struct, nested lists) keep the nested Scan. No operator, pipeline or driver API changes. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_014rtiSjzL4EauVCM7NS84hj --- .../src/plan/physical/operators/list.rs | 31 ++- .../src/plan/physical/operators/list/reads.rs | 180 ++++++++++++++++++ .../src/plan/physical/operators/scan.rs | 4 +- .../src/plan/physical/pipeline/tests.rs | 69 +++++++ 4 files changed, 280 insertions(+), 4 deletions(-) create mode 100644 vortex-layout/src/plan/physical/operators/list/reads.rs diff --git a/vortex-layout/src/plan/physical/operators/list.rs b/vortex-layout/src/plan/physical/operators/list.rs index 7ccc49dabc8..c7b5976774d 100644 --- a/vortex-layout/src/plan/physical/operators/list.rs +++ b/vortex-layout/src/plan/physical/operators/list.rs @@ -25,6 +25,7 @@ use vortex_mask::AllOr; use vortex_mask::Mask; use vortex_utils::parallelism::get_available_parallelism; +use self::reads::ElementReads; use super::join; use super::join_inlet; use crate::plan::physical::ListPackPlan; @@ -43,6 +44,7 @@ use crate::plan::physical::pipeline::Turn; use crate::plan::physical::pipeline::ValueCache; use crate::segments::SegmentId; +mod reads; mod work; /// A list's offsets and validity arrive before its selected element ranges can be planned. @@ -57,6 +59,8 @@ pub(crate) struct ListSource { offsets: Option, validity: Option, elements: Option, + /// Element reads issued by this source when the elements are plain segment chunks. + direct: Option, reads: FxHashMap>, pending_children: usize, chunks: Vec<(usize, ArrayRef)>, @@ -82,6 +86,7 @@ impl ListSource { offsets: None, validity: None, elements: None, + direct: None, reads: FxHashMap::default(), pending_children: 0, chunks: Vec::new(), @@ -174,11 +179,17 @@ impl ListSource { fn start_elements( &mut self, - cx: &Cx<'_>, + cx: &mut Cx<'_>, splits: Vec, overlap: bool, ) -> VortexResult<()> { - let scan = Scan::try_new(cx.session().clone(), self.plan.elements()?, splits)? + let elements = self.plan.elements()?; + if let Some(mut reads) = ElementReads::try_new(&elements, &splits)? { + reads.request(cx); + self.direct = Some(reads); + return Ok(()); + } + let scan = Scan::try_new(cx.session().clone(), elements, splits)? .with_value_cache(self.values.clone()) .with_child_planning(self.plan_children); self.elements = Some(if overlap { @@ -243,6 +254,22 @@ impl Source for ListSource { } self.start_elements(cx, splits, overlap)?; } + if let Some(reads) = self.direct.as_mut() { + let session = cx.session().clone(); + while let Some((segment, bytes)) = cx.take_segment() { + reads.deliver(segment, bytes, &session)?; + } + if !reads.is_done() { + return Ok(SourceReady::Blocked(SourceBlocked::Io)); + } + let reads = self + .direct + .take() + .ok_or_else(|| vortex_err!("Missing list element reads"))?; + self.chunks.extend(reads.finish()?.into_iter().enumerate()); + self.finish()?; + return Ok(SourceReady::Ready); + } let scan = self .elements .as_mut() diff --git a/vortex-layout/src/plan/physical/operators/list/reads.rs b/vortex-layout/src/plan/physical/operators/list/reads.rs new file mode 100644 index 00000000000..e8a2e569926 --- /dev/null +++ b/vortex-layout/src/plan/physical/operators/list/reads.rs @@ -0,0 +1,180 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Reads selected list elements straight from their segments, in the list's own pipeline. +//! +//! Element ranges are known only once the offsets arrive, so they cannot be compiled with the +//! rest of the split. When the elements are plain segment chunks, the list source issues their +//! reads itself instead of compiling a nested scan: each touched segment is read once and decoded +//! on arrival, and the ranges it holds are sliced out of it in list order. + +use std::ops::Range; + +use rustc_hash::FxHashMap; +use vortex_array::ArrayRef; +use vortex_array::buffer::BufferHandle; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +use super::super::scan::decode; +use super::super::scan::select; +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::SegmentScanPlan; +use crate::plan::physical::Zoned; +use crate::plan::physical::operator::Cx; +use crate::plan::physical::pipeline::Split; +use crate::plan::physical::pipeline::overlapping; +use crate::segments::SegmentId; + +/// The selected element rows of one segment chunk. +struct Piece { + slice: Option>, + filter: Option, +} + +/// A requested segment and the pieces, by output position, that read it. +struct Pending { + plan: SegmentScanPlan, + pieces: Vec, +} + +/// The element reads of one list source, from request to the ordered element chunks. +pub(super) struct ElementReads { + pieces: Vec, + output: Vec>, + pending: FxHashMap, + /// Segments in first-use order, so reads are published in element order. + order: Vec, +} + +impl ElementReads { + /// Plans the reads for `splits` of `elements`, or `None` when any selected range reaches a + /// plan other than segment chunks, which then needs a nested scan. + pub(super) fn try_new(elements: &PlanRef, splits: &[Split]) -> VortexResult> { + let mut reads = Self { + pieces: Vec::new(), + output: Vec::new(), + pending: FxHashMap::default(), + order: Vec::new(), + }; + for split in splits { + if !reads.collect(elements, split.rows.clone(), &split.mask)? { + return Ok(None); + } + } + reads.output = (0..reads.pieces.len()).map(|_| None).collect(); + Ok(Some(reads)) + } + + fn collect(&mut self, plan: &PlanRef, rows: Range, mask: &Mask) -> VortexResult { + if mask.all_false() { + return Ok(true); + } + // A filter keeps the rows of the mask it runs with, which this piece applies itself. + if let Some(filter) = plan.as_opt::() { + return self.collect(&filter.child_plan()?, rows, mask); + } + if let Some(zoned) = plan.as_opt::() { + if zoned.is_pruning() { + return Ok(false); + } + return match zoned.data_plan()? { + Some(data) => self.collect(&data, rows, mask), + None => Ok(false), + }; + } + if let Some(stats) = plan.as_opt::() { + return self.collect(&stats.data_plan()?, rows, mask); + } + if let Some(concat) = plan.as_opt::() { + for (index, start, local) in overlapping(concat, &rows) { + let mask = mask.slice( + usize::try_from(local.start - rows.start)? + ..usize::try_from(local.end - rows.start)?, + ); + if !self.collect( + &concat.child_required(index)?, + local.start - start..local.end - start, + &mask, + )? { + return Ok(false); + } + } + return Ok(true); + } + let Some(scan) = plan.as_opt::() else { + return Ok(false); + }; + let slice = if rows.start == 0 && rows.end == scan.row_count() { + None + } else { + Some(usize::try_from(rows.start)?..usize::try_from(rows.end)?) + }; + let filter = (!mask.all_true()).then(|| mask.clone()); + let segment = scan.segment_id(); + let piece = self.pieces.len(); + self.pieces.push(Piece { slice, filter }); + self.pending + .entry(segment) + .or_insert_with(|| { + self.order.push(segment); + Pending { + plan: scan.clone(), + pieces: Vec::new(), + } + }) + .pieces + .push(piece); + Ok(true) + } + + /// Publishes one read per touched segment. + pub(super) fn request(&mut self, cx: &mut Cx<'_>) { + for segment in self.order.drain(..) { + cx.request_read(segment); + } + } + + /// Decodes an arrived segment and narrows it to each piece that reads it. + pub(super) fn deliver( + &mut self, + segment: SegmentId, + bytes: BufferHandle, + session: &VortexSession, + ) -> VortexResult<()> { + let pending = self + .pending + .remove(&segment) + .ok_or_else(|| vortex_err!("Unexpected list element segment {segment}"))?; + let array = decode(&pending.plan, bytes, session)?; + for piece in pending.pieces { + let Piece { slice, filter } = &self.pieces[piece]; + self.output[piece] = Some(select(array.clone(), slice.as_ref(), filter.as_ref())?); + } + Ok(()) + } + + /// Whether every requested segment has arrived. + pub(super) fn is_done(&self) -> bool { + self.pending.is_empty() + } + + /// The element chunks in list order. + pub(super) fn finish(self) -> VortexResult> { + vortex_ensure!( + self.is_done(), + "List elements finished before every read arrived" + ); + self.output + .into_iter() + .map(|array| array.ok_or_else(|| vortex_err!("Missing list element chunk"))) + .collect() + } +} diff --git a/vortex-layout/src/plan/physical/operators/scan.rs b/vortex-layout/src/plan/physical/operators/scan.rs index b6d2c463a81..48c33773007 100644 --- a/vortex-layout/src/plan/physical/operators/scan.rs +++ b/vortex-layout/src/plan/physical/operators/scan.rs @@ -24,7 +24,7 @@ use crate::plan::physical::operator::Step; use crate::segments::SegmentId; /// Decodes the whole of a scan's segment. -fn decode( +pub(super) fn decode( plan: &SegmentScanPlan, segment: BufferHandle, session: &VortexSession, @@ -42,7 +42,7 @@ fn decode( } /// Slices a whole decoded segment to `slice`, then keeps the rows `filter` selects. -fn select( +pub(super) fn select( array: ArrayRef, slice: Option<&Range>, filter: Option<&Mask>, diff --git a/vortex-layout/src/plan/physical/pipeline/tests.rs b/vortex-layout/src/plan/physical/pipeline/tests.rs index 0327d652b3d..cfa3cfdeb37 100644 --- a/vortex-layout/src/plan/physical/pipeline/tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/tests.rs @@ -393,6 +393,75 @@ fn dense_nullable_list_ranges_share_one_scan( Ok(()) } +#[rstest] +fn segment_list_elements_are_read_by_the_list_source( + #[values(Delivery::Fifo, Delivery::Lifo)] order: Delivery, + #[values(false, true)] utf8: bool, +) -> VortexResult<()> { + let mut store = Store::default(); + let elements = if utf8 { + VarBinViewArray::from_iter_str((0..100).map(|value| format!("v{value}"))).into_array() + } else { + Buffer::from_iter(0_i64..100).into_array() + }; + let offsets = Buffer::from_iter((0_u64..=20).map(|row| row * 5)).into_array(); + let input = + ListArray::try_new(elements.clone(), offsets.clone(), Validity::NonNullable)?.into_array(); + let plan = ListPackPlan::try_new( + false.into(), + 20, + lower(&store.chunked(&elements, &[10; 10])?)?, + lower(&store.flat(&offsets)?)?, + None, + )? + .into_plan(); + let rows = [1_u64, 13]; + let mask = Mask::from_iter((0..20).map(|row| rows.contains(&row))); + let result = run(&store, &plan, 0..20, 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()); + // Only the offsets and the two touched element chunks are read, and the element reads are + // published together once the offsets arrive. + assert_eq!(segments_read(&result.events), vec![0, 6, 10]); + assert!( + result.events.contains(&Event::Io(vec![0, 6])), + "{:?}", + result.events + ); + Ok(()) +} + +#[rstest] +fn dict_list_elements_fall_back_to_a_nested_scan( + #[values(Delivery::Fifo, Delivery::Lifo)] order: Delivery, +) -> VortexResult<()> { + let mut store = Store::default(); + let values = VarBinViewArray::from_iter_str(["a", "b", "c"]).into_array(); + let codes = PrimitiveArray::from_iter((0..100).map(|value| (value % 3) as u8)).into_array(); + let elements = VarBinViewArray::from_iter_str((0..100).map(|value| ["a", "b", "c"][value % 3])) + .into_array(); + let offsets = Buffer::from_iter((0_u64..=20).map(|row| row * 5)).into_array(); + let input = ListArray::try_new(elements, offsets.clone(), Validity::NonNullable)?.into_array(); + let layout = + DictLayout::new(store.flat(&values)?, store.chunked(&codes, &[10; 10])?).into_layout(); + let plan = ListPackPlan::try_new( + false.into(), + 20, + lower(&layout)?, + lower(&store.flat(&offsets)?)?, + None, + )? + .into_plan(); + let rows = [1_u64, 2, 13]; + let mask = Mask::from_iter((0..20).map(|row| rows.contains(&row))); + let result = run(&store, &plan, 0..20, 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()); + Ok(()) +} + /// Delivers segments in the given order. fn scripted(order: &[u32]) -> impl FnMut(&[ReadRequest]) -> usize + '_ { let mut next = order.iter(); From 7039c9c1775ca4cb2d5163bfab99e460b3960803 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 19 Oct 2026 13:04:44 +0000 Subject: [PATCH 04/15] Remove parallel child projections and selection-aware reach Drop the child projection path and every core change it needed: PlannerOutput::Await and the driver's task tracking, SourceBlocked::Child, the Cx child inbox, Turn::Child and Scan::complete_child, AdaptorStep::Child, Compiler::plan_children, the Children executor and its split planning, and the ProjectionPlanner/ProjectedBatch conversion. List sources now always run serially, reading segment elements directly or through a nested scan. Also restore Reach, Concat::reach and Shares to develop. Scan::try_new now counts shared readers over the row runs the splits select, merged when adjacent, so chunks in unselected gaps stay lazy without a mask-aware Reach variant. Each run counts as a reader, so a segment reused by distinct chunks is still decoded once. Splits with more than 1024 runs count their whole range. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_014rtiSjzL4EauVCM7NS84hj --- Cargo.lock | 1 - vortex-layout/benches/pipeline_scan.rs | 1 - vortex-layout/src/plan/physical/operator.rs | 22 -- .../src/plan/physical/operators/list.rs | 71 +--- .../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 | 201 +++------- .../plan/physical/pipeline/scan/validate.rs | 6 - .../physical/pipeline/scheduling_tests.rs | 4 - .../src/plan/physical/pipeline/tests.rs | 5 - .../src/plan/physical/plans/concat.rs | 3 - .../src/plan/physical/plans/list_pack.rs | 7 +- 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 | 13 +- .../src/scan/execution/projection.rs | 146 +++---- vortex-layout/src/scan/v2/mod.rs | 9 +- vortex-layout/src/scan/v2/tests.rs | 1 - 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/planner.rs | 6 - 28 files changed, 125 insertions(+), 1775 deletions(-) delete mode 100644 vortex-layout/src/plan/physical/pipeline/children.rs delete mode 100644 vortex-layout/src/scan/execution/children.rs delete mode 100644 vortex-layout/src/scan/execution/children/splits.rs delete mode 100644 vortex-layout/src/scan/execution/children/splits/tests.rs delete mode 100644 vortex-layout/src/scan/execution/children/tests.rs delete mode 100644 vortex-layout/src/scan/v2/tests/child_work.rs delete mode 100644 vortex-scan/src/planning/driver/tests/tasks.rs diff --git a/Cargo.lock b/Cargo.lock index 736e0146827..22aace294d8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -12396,7 +12396,6 @@ dependencies = [ "parking_lot", "roaring", "rstest", - "smallvec", "tracing", "vortex-array", "vortex-buffer", diff --git a/vortex-layout/benches/pipeline_scan.rs b/vortex-layout/benches/pipeline_scan.rs index e9980aa6f66..db4beeaa043 100644 --- a/vortex-layout/benches/pipeline_scan.rs +++ b/vortex-layout/benches/pipeline_scan.rs @@ -124,7 +124,6 @@ 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/plan/physical/operator.rs b/vortex-layout/src/plan/physical/operator.rs index de1b1e35560..96a8d0cfbd7 100644 --- a/vortex-layout/src/plan/physical/operator.rs +++ b/vortex-layout/src/plan/physical/operator.rs @@ -15,7 +15,6 @@ 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. @@ -44,8 +43,6 @@ 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), @@ -160,28 +157,9 @@ 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 index c7b5976774d..f7324824a20 100644 --- a/vortex-layout/src/plan/physical/operators/list.rs +++ b/vortex-layout/src/plan/physical/operators/list.rs @@ -36,7 +36,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::pipeline::ChildProjection; use crate::plan::physical::pipeline::ReadId; use crate::plan::physical::pipeline::Scan; use crate::plan::physical::pipeline::Split; @@ -52,9 +51,6 @@ 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, @@ -62,33 +58,23 @@ pub(crate) struct ListSource { /// Element reads issued by this source when the elements are plain segment chunks. direct: 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 { + pub(crate) fn new(plan: ListPackPlan, mask: Mask, values: ValueCache) -> Self { Self { plan, mask, values, - plan_children, - awaiting_child: false, - serial_splits: None, initialized: false, offsets: None, validity: None, elements: None, direct: None, reads: FxHashMap::default(), - pending_children: 0, chunks: Vec::new(), output: None, done: false, @@ -150,28 +136,6 @@ 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. - 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(()) @@ -190,8 +154,7 @@ impl ListSource { return Ok(()); } let scan = Scan::try_new(cx.session().clone(), elements, splits)? - .with_value_cache(self.values.clone()) - .with_child_planning(self.plan_children); + .with_value_cache(self.values.clone()); self.elements = Some(if overlap { scan.with_max_active(2) } else { @@ -238,22 +201,6 @@ impl Source for ListSource { } 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)?; - } if let Some(reads) = self.direct.as_mut() { let session = cx.session().clone(); while let Some((segment, bytes)) = cx.take_segment() { @@ -282,10 +229,6 @@ impl Source for ListSource { .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) => { @@ -296,16 +239,8 @@ impl Source for ListSource { 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 - })); + return Ok(SourceReady::Blocked(SourceBlocked::Io)); } Turn::Done => { self.finish()?; diff --git a/vortex-layout/src/plan/physical/pipeline/children.rs b/vortex-layout/src/plan/physical/pipeline/children.rs deleted file mode 100644 index 974e0afcd74..00000000000 --- a/vortex-layout/src/plan/physical/pipeline/children.rs +++ /dev/null @@ -1,29 +0,0 @@ -// 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 cf2dcb10184..74f5296cd19 100644 --- a/vortex-layout/src/plan/physical/pipeline/compile.rs +++ b/vortex-layout/src/plan/physical/pipeline/compile.rs @@ -8,7 +8,6 @@ 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; @@ -144,10 +143,6 @@ 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( @@ -301,30 +296,6 @@ 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) { @@ -379,24 +350,13 @@ 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) | Reach::Selected { offset, .. } => { - rows.start + offset..rows.end + offset - } + Reach::Offset(offset) => rows.start + offset..rows.end + offset, Reach::Fixed(root) => root.clone(), } } @@ -406,11 +366,6 @@ 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(), - }, } } @@ -418,30 +373,6 @@ 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 5dd18e783c1..b6037c9e6ea 100644 --- a/vortex-layout/src/plan/physical/pipeline/mod.rs +++ b/vortex-layout/src/plan/physical/pipeline/mod.rs @@ -38,14 +38,12 @@ //! 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 e5a8ee922e9..1d4b01c1bfd 100644 --- a/vortex-layout/src/plan/physical/pipeline/scan.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan.rs @@ -18,6 +18,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; +use vortex_mask::AllOr; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -31,7 +32,6 @@ 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; @@ -40,7 +40,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::pipeline::ChildProjection; use crate::segments::SegmentId; #[cfg(debug_assertions)] @@ -84,11 +83,9 @@ 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 or child projection completes. + /// Nothing can run until a read is delivered. Waiting, /// Every split has produced all its rows. Done, @@ -115,7 +112,6 @@ 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, @@ -146,7 +142,6 @@ fn drive( session: &VortexSession, exec: &mut ExecutionCtx, reads: &mut SmallVec<[SegmentId; 1]>, - child_requests: &mut Vec, ) -> VortexResult { let Pipeline { source, @@ -159,7 +154,6 @@ fn drive( source_more, touched, changed, - child_inbox, .. } = pipe; let mut cx = Cx { @@ -171,8 +165,6 @@ 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. @@ -200,17 +192,9 @@ fn drive( let ready = source.ready(&mut cx)?; cx.changed.clear(); vortex_ensure!( - cx.reads.is_empty() - || matches!( - ready, - SourceReady::Blocked(SourceBlocked::Io | SourceBlocked::Child) - ), + cx.reads.is_empty() || ready == SourceReady::Blocked(SourceBlocked::Io), "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)), @@ -298,9 +282,6 @@ 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, @@ -311,12 +292,6 @@ 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 { @@ -332,15 +307,11 @@ 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, } } @@ -361,7 +332,6 @@ impl Core { outlets, bytes: VecDeque::new(), pending: 0, - child_inbox: None, finishing: None, touched, changed, @@ -436,21 +406,16 @@ 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 { @@ -458,16 +423,6 @@ 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), @@ -487,52 +442,6 @@ 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() { @@ -621,13 +530,6 @@ 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); } } @@ -717,6 +619,39 @@ struct Active { output: Option, } +/// A split selecting more separate runs than this counts its whole row range instead, which +/// bounds the walks over the plan before the scan starts. +const MAX_SELECTED_RUNS: usize = 1024; + +/// The rows the splits select, sorted, with overlapping and adjacent runs merged. +fn selected_runs(splits: &[Split]) -> VortexResult>> { + let mut runs = Vec::new(); + for split in splits { + match split.mask.slices() { + AllOr::None => {} + AllOr::Some(slices) if slices.len() <= MAX_SELECTED_RUNS => { + for &(start, end) in slices { + runs.push( + split.rows.start + u64::try_from(start)? + ..split.rows.start + u64::try_from(end)?, + ); + } + } + AllOr::All | AllOr::Some(_) => runs.push(split.rows.clone()), + } + } + runs.retain(|run| !run.is_empty()); + runs.sort_unstable_by_key(|run| run.start); + let mut merged: Vec> = Vec::with_capacity(runs.len()); + for run in runs { + match merged.last_mut() { + Some(last) if run.start <= last.end => last.end = last.end.max(run.end), + _ => merged.push(run), + } + } + Ok(merged) +} + /// Runs a plan over a list of splits. /// /// The plan runs as one stage under each split's selection. @@ -770,38 +705,14 @@ impl Scan { 0, splits.iter().map(|split| split.rows.clone()).collect(), ); + // Only the rows some split selects can share a segment between readers. Chunks in the + // gaps between them stay lazy: they are never lowered or read. A segment several runs + // reach counts a reader per run, so it is decoded once even when distinct chunks + // reuse it; readers that never come leave their ports unclaimed. match &root { 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)?; + for run in selected_runs(&splits)? { + core.shares.add(plan, run)?; } } } @@ -838,11 +749,6 @@ 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 { @@ -852,18 +758,6 @@ 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)?; @@ -882,7 +776,7 @@ impl Scan { if self.live == 0 { return Ok(Turn::Done); } - if !self.core.reads.is_empty() || !self.core.child_requests.is_empty() { + if !self.core.reads.is_empty() { return Ok(Turn::Waiting); } vortex_bail!( @@ -897,15 +791,6 @@ 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 d1749972d65..df404bd34ea 100644 --- a/vortex-layout/src/plan/physical/pipeline/scan/validate.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan/validate.rs @@ -88,12 +88,6 @@ 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 e8bd51e1ac9..19bcb5504db 100644 --- a/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs @@ -64,7 +64,6 @@ 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")); @@ -194,7 +193,6 @@ 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())?; @@ -260,7 +258,6 @@ 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!( @@ -420,7 +417,6 @@ 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 cfa3cfdeb37..8e673004b26 100644 --- a/vortex-layout/src/plan/physical/pipeline/tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/tests.rs @@ -233,7 +233,6 @@ 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))); @@ -1065,11 +1064,7 @@ 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, diff --git a/vortex-layout/src/plan/physical/plans/concat.rs b/vortex-layout/src/plan/physical/plans/concat.rs index 377e4bf270c..b2e1ffea0f7 100644 --- a/vortex-layout/src/plan/physical/plans/concat.rs +++ b/vortex-layout/src/plan/physical/plans/concat.rs @@ -189,9 +189,6 @@ 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 0d459d3537b..4481180f6ef 100644 --- a/vortex-layout/src/plan/physical/plans/list_pack.rs +++ b/vortex-layout/src/plan/physical/plans/list_pack.rs @@ -155,12 +155,7 @@ impl PlanVTable for ListPack { .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(), - ); + let source = ListSource::new(plan.clone(), mask.clone(), compiler.values().clone()); Ok(Some(compiler.join_with(chains, |_| UNBOUNDED, source))) } diff --git a/vortex-layout/src/scan/execution/adaptor.rs b/vortex-layout/src/scan/execution/adaptor.rs index edf681b294f..fc6053d55ca 100644 --- a/vortex-layout/src/scan/execution/adaptor.rs +++ b/vortex-layout/src/scan/execution/adaptor.rs @@ -19,7 +19,6 @@ 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; @@ -41,8 +40,6 @@ 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), } @@ -100,7 +97,6 @@ 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)?]; @@ -129,7 +125,7 @@ impl ScanAdaptor { match &turn { Turn::Done => self.done = true, Turn::Waiting => self.pending = Some(Turn::Waiting), - Turn::Read(_) | Turn::Child(_) | Turn::Output(..) => {} + Turn::Read(_) | Turn::Output(..) => {} } Ok(turn) } @@ -143,18 +139,6 @@ 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 deleted file mode 100644 index e57ba8a4816..00000000000 --- a/vortex-layout/src/scan/execution/children.rs +++ /dev/null @@ -1,230 +0,0 @@ -// 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 deleted file mode 100644 index 89bf01eff97..00000000000 --- a/vortex-layout/src/scan/execution/children/splits.rs +++ /dev/null @@ -1,152 +0,0 @@ -// 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 deleted file mode 100644 index d7ba1c8f7e8..00000000000 --- a/vortex-layout/src/scan/execution/children/splits/tests.rs +++ /dev/null @@ -1,139 +0,0 @@ -// 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 deleted file mode 100644 index 680dd781c68..00000000000 --- a/vortex-layout/src/scan/execution/children/tests.rs +++ /dev/null @@ -1,357 +0,0 @@ -// 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 3f919d3a5b5..93c48ecae12 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -3,18 +3,17 @@ //! Scan execution over physical plans through the planner I/O and compute protocol. //! -//! 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. +//! A `SelectionPlanner` resolves each fixed-size work chunk's selected rows and hands them to a +//! `ProjectionPlanner`, which emits one `ProjectionMorsel` for them. 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. 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_scan::planning::next::next_fn; @@ -25,7 +24,6 @@ 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; @@ -66,7 +64,6 @@ 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>, } /// Selects the rows of a fixed work chunk before handing them to its projection planner. diff --git a/vortex-layout/src/scan/execution/projection.rs b/vortex-layout/src/scan/execution/projection.rs index 1732a6b4b22..b6180725c27 100644 --- a/vortex-layout/src/scan/execution/projection.rs +++ b/vortex-layout/src/scan/execution/projection.rs @@ -1,18 +1,12 @@ // 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; @@ -29,29 +23,64 @@ use crate::scan::execution::SelectedRows; use crate::scan::execution::adaptor::AdaptorStep; use crate::scan::execution::adaptor::Piece; use crate::scan::execution::adaptor::ScanAdaptor; -use crate::scan::execution::children::Children; -use crate::scan::execution::children::splits::parallel_splits; + +/// 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)), + )) + } +} /// 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. -pub(super) struct ProjectionPlanner { +struct ProjectionMorsel { plans: Option, dtype: DType, selected: SelectedRows, adaptor: Option, pieces: Vec, done: bool, - children: Option>, - pending: Option<(u64, oneshot::Receiver>)>, } -impl ProjectionPlanner { - /// Creates a planner projecting `selected`. - pub(super) fn new(plans: ScanPlans, selected: SelectedRows) -> Self { +impl ProjectionMorsel { + /// Creates a morsel projecting `selected`. + 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, @@ -60,7 +89,7 @@ impl ProjectionPlanner { } } - 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) @@ -69,18 +98,15 @@ impl ProjectionPlanner { .filter(|array| !array.is_empty()) .collect(); let array = match arrays.len() { - 0 => return Ok(PlannerOutput::Done), + 0 => return Ok(MorselOutput::Done), 1 => arrays.remove(0), _ => ChunkedArray::try_new(arrays, self.dtype.clone())?.into_array(), }; - Ok(PlannerOutput::Morsel( - self.selected.scope.clone(), - Box::new(ProjectedBatch(Some(array))), - )) + Ok(MorselOutput::Batch(array)) } } -impl IoConsumer for ProjectionPlanner { +impl IoConsumer for ProjectionMorsel { fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { if let Some(adaptor) = self.adaptor.as_mut() { adaptor.set_io_result(request, result); @@ -88,11 +114,10 @@ impl IoConsumer for ProjectionPlanner { } } -impl Planner for ProjectionPlanner { +impl Morsel for ProjectionMorsel { 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, @@ -101,19 +126,9 @@ impl Planner for ProjectionPlanner { } } - fn compute(&mut self) -> VortexResult { + fn compute(&mut self) -> VortexResult { if self.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))?; + vortex_bail!("ProjectionMorsel: compute called after Done"); } let Some(adaptor) = self.adaptor.as_mut() else { let plans = self @@ -129,67 +144,20 @@ impl Planner for ProjectionPlanner { }], )? .with_row_offset(plans.row_offset) - .with_value_cache(plans.values) - .with_child_planning(self.children.is_some()); + .with_value_cache(plans.values); self.adaptor = Some(ScanAdaptor::new(scan, plans.locations)); - return Ok(PlannerOutput::Continue); + return Ok(MorselOutput::Continue); }; if adaptor.state() == State::Done { return self.finish(); } Ok(match adaptor.compute()? { - 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::Yield => MorselOutput::Continue, + AdaptorStep::NeedsIO(batch) => MorselOutput::NeedsIO(batch), AdaptorStep::Piece(piece) => { self.pieces.push(piece); - PlannerOutput::Continue + MorselOutput::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/mod.rs b/vortex-layout/src/scan/v2/mod.rs index b488b3bca74..51b450dd519 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -54,7 +54,6 @@ 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; @@ -292,15 +291,10 @@ impl PreparedScan { MAX_ROWS_PER_CHUNK }; let splits = row_chunks(builder.row_range.unwrap_or(0..row_count), chunk_rows); - let io: Arc = Arc::new(SegmentScanIo::new( + let io = Arc::new(SegmentScanIo::new( file.segments, 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 { @@ -309,7 +303,6 @@ impl PreparedScan { projection, row_offset: builder.row_offset, values: Default::default(), - children, }, io, splits, diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index 43bb6392757..abecf2785a9 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -72,7 +72,6 @@ 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 diff --git a/vortex-layout/src/scan/v2/tests/child_work.rs b/vortex-layout/src/scan/v2/tests/child_work.rs deleted file mode 100644 index d26a2d97c8c..00000000000 --- a/vortex-layout/src/scan/v2/tests/child_work.rs +++ /dev/null @@ -1,101 +0,0 @@ -// 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 63352a21566..d12c207b049 100644 --- a/vortex-scan/Cargo.toml +++ b/vortex-scan/Cargo.toml @@ -26,7 +26,6 @@ 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 024748774fc..8cb78099272 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 or an asynchronous task makes it runnable. Published IO is registered +//! skipped and regains its priority when IO 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 or tasks, and the caller decides how to wait, by awaiting [`Run::poll_completion`] or by handing +//! for IO, 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,12 +21,7 @@ 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; @@ -88,9 +83,6 @@ 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 { @@ -99,7 +91,7 @@ struct Work { scope: WorkScope, item: Item, /// Stable child-ordinal path, beginning with the root's admission ordinal. - path: Priority, + path: Vec, /// 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. @@ -107,7 +99,7 @@ struct Work { } impl Work { - fn priority(&self) -> Priority { + fn priority(&self) -> Vec { let mut priority = self.path.clone(); priority.push(self.next_child); priority @@ -115,7 +107,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) -> Priority { + fn child_path(&mut self) -> Vec { let path = self.priority(); self.next_child += 1; path @@ -154,7 +146,7 @@ impl Work { /// Keeps scheduling order separate from lookup for deliveries to runnable work. #[derive(Default)] struct ReadyQueue { - priorities: BTreeMap, + priorities: BTreeMap, IoOwnerId>, work: HashMap>, } @@ -199,8 +191,6 @@ 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. @@ -243,7 +233,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), smallvec![id.0]); + self.spawn(id, scope, Item::Planner(root), vec![id.0]); id } @@ -255,9 +245,6 @@ 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(); } @@ -310,10 +297,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 && !self.tasks.contains_key(&owner) { + if work.state() != State::Waiting { let work = parked.remove(); self.queue.insert(work); - } else if work.outstanding.is_empty() && !self.tasks.contains_key(&owner) { + } else if work.outstanding.is_empty() { vortex_bail!( "work for scope {:?} waits with no outstanding fetch", work.scope @@ -345,23 +332,6 @@ 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) { @@ -373,10 +343,6 @@ 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 { @@ -386,7 +352,6 @@ 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 { @@ -410,35 +375,7 @@ impl Run { Ok(()) } - 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) { + fn spawn(&mut self, root: RootId, scope: WorkScope, item: Item, path: Vec) { let id = IoOwnerId(self.next_owner); self.next_owner += 1; if let Some(root) = self.roots.get_mut(&root) { @@ -458,7 +395,6 @@ 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; }; @@ -470,9 +406,9 @@ impl Run { } } - /// Parks an item that waits for IO or tasks; at least one must be outstanding. + /// Parks an item that waits for IO, which is a protocol error with no fetch outstanding. fn park(&mut self, work: Box) -> VortexResult<()> { - if work.outstanding.is_empty() && !self.tasks.contains_key(&work.id) { + if work.outstanding.is_empty() { vortex_bail!( "work for scope {:?} waits with no outstanding fetch", work.scope @@ -516,12 +452,6 @@ 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); @@ -611,10 +541,7 @@ impl Driver { match run.advance()? { Progress::Batch(batch) => batches.push(batch), Progress::RootDone(_) | Progress::Idle => return Ok(batches), - 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)))? - } + Progress::Waiting => run.complete(self.io.wait()?)?, } } } diff --git a/vortex-scan/src/planning/driver/tests.rs b/vortex-scan/src/planning/driver/tests.rs index d27fa0a5be7..389ad24372d 100644 --- a/vortex-scan/src/planning/driver/tests.rs +++ b/vortex-scan/src/planning/driver/tests.rs @@ -22,8 +22,6 @@ 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, @@ -143,34 +141,6 @@ 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 deleted file mode 100644 index 140a06d1a3d..00000000000 --- a/vortex-scan/src/planning/driver/tests/tasks.rs +++ /dev/null @@ -1,175 +0,0 @@ -// 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/planner.rs b/vortex-scan/src/planning/planner.rs index cf0b71ad569..91d875ff317 100644 --- a/vortex-scan/src/planning/planner.rs +++ b/vortex-scan/src/planning/planner.rs @@ -6,7 +6,6 @@ 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; @@ -50,10 +49,6 @@ 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. @@ -79,7 +74,6 @@ 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(), } From 99b4a80ae912987717131b44c480e660dc18a8f0 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 19 Oct 2026 17:29:04 +0000 Subject: [PATCH 05/15] Plan V2 random access only over groups of selected rows Reduce this branch to a minimal change against develop: drop list support, the selection planner and every compiler and generic trait change. V2 split planning now groups each fixed chunk's selected rows. A selected run joins the group before it while the unselected gap is at most one writer row block (8,192 rows), or 1 MiB of serialized data when that spans more rows. Each group becomes its own planner root over just its rows, with the selection's keep mask, so rows in wider gaps get no planner, plan or read. The random-access benchmark reads files with the V2 scan when it supports their layouts and falls back to the default scan otherwise. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_014rtiSjzL4EauVCM7NS84hj --- benchmarks/random-access-bench/README.md | 2 - vortex-bench/src/random_access/take.rs | 51 ++- vortex-layout/src/layouts/chunked/mod.rs | 5 - vortex-layout/src/plan/physical/lower.rs | 12 +- .../src/plan/physical/operators/list.rs | 355 --------------- .../src/plan/physical/operators/list/reads.rs | 180 -------- .../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/operators/scan.rs | 4 +- .../src/plan/physical/pipeline/scan.rs | 51 +-- .../src/plan/physical/pipeline/tests.rs | 328 -------------- .../src/plan/physical/plans/list_pack.rs | 71 --- vortex-layout/src/plan/physical/tests.rs | 58 --- vortex-layout/src/scan/execution/mod.rs | 34 +- vortex-layout/src/scan/v2/io.rs | 108 +---- vortex-layout/src/scan/v2/mod.rs | 166 ++++--- vortex-layout/src/scan/v2/tests.rs | 325 +++++--------- vortex-scan/src/planning/mod.rs | 1 - vortex-scan/src/planning/selection.rs | 210 --------- 20 files changed, 324 insertions(+), 2063 deletions(-) delete mode 100644 vortex-layout/src/plan/physical/operators/list.rs delete mode 100644 vortex-layout/src/plan/physical/operators/list/reads.rs delete mode 100644 vortex-layout/src/plan/physical/operators/list/work.rs delete mode 100644 vortex-scan/src/planning/selection.rs diff --git a/benchmarks/random-access-bench/README.md b/benchmarks/random-access-bench/README.md index f845bb8ee11..73f52f5b219 100644 --- a/benchmarks/random-access-bench/README.md +++ b/benchmarks/random-access-bench/README.md @@ -28,8 +28,6 @@ work in each timed iteration. CI drives the full matrix via cargo run -p random-access-bench --profile release_debug --features lance ``` -Vortex measurements always use the planner-driven V2 scan builder. - ## Running against S3 The same benchmark can read its data from an object store instead of local disk. The remote diff --git a/vortex-bench/src/random_access/take.rs b/vortex-bench/src/random_access/take.rs index 2291905825d..97d1a005778 100644 --- a/vortex-bench/src/random_access/take.rs +++ b/vortex-bench/src/random_access/take.rs @@ -130,11 +130,28 @@ impl RandomAccessor for ArrowIpcRandomAccessor { /// Random accessor for Vortex format files. /// -/// The file handle is opened at construction time and reused across `take()` calls. +/// The file handle is opened at construction time and reused across `take()` calls. Files the +/// V2 scan supports are read with it; others, such as files with list columns, use the default +/// scan. pub struct VortexRandomAccessor { name: String, format: Format, file: VortexFile, + v2: bool, +} + +/// Whether the V2 scan can read `file`. It checks the file's layouts before reading any rows. +async fn supports_v2(file: &VortexFile) -> bool { + let Ok(stream) = ScanBuilder::new(SESSION.clone(), file.scan_file()) + .with_row_range(0..0) + .into_stream() + else { + return false; + }; + ArrayStreamAdapter::new(file.dtype().clone(), stream) + .read_all() + .await + .is_ok() } impl VortexRandomAccessor { @@ -144,10 +161,15 @@ impl VortexRandomAccessor { name: impl Into, format: Format, ) -> anyhow::Result { - let file = SESSION.open_options().open_path(path.as_ref()).await?; + let file = SESSION + .open_options() + .with_layout_reader_cache() + .open_path(path.as_ref()) + .await?; Ok(Self { name: name.into(), format, + v2: supports_v2(&file).await, file, }) } @@ -161,11 +183,13 @@ 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, + v2: supports_v2(&file).await, file, }) } @@ -182,13 +206,22 @@ impl RandomAccessor for VortexRandomAccessor { } async fn take(&self, indices: &[u64]) -> anyhow::Result { - let indices_buf: Buffer = Buffer::from(indices.to_vec()); - 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?; + let indices = StrictSortedBuffer::try_new(Buffer::from(indices.to_vec()))?; + let array = if self.v2 { + let stream = ScanBuilder::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/src/layouts/chunked/mod.rs b/vortex-layout/src/layouts/chunked/mod.rs index 55c268b295e..459429046c2 100644 --- a/vortex-layout/src/layouts/chunked/mod.rs +++ b/vortex-layout/src/layouts/chunked/mod.rs @@ -125,11 +125,6 @@ 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 9f965206544..1ea284cc1ca 100644 --- a/vortex-layout/src/plan/physical/lower.rs +++ b/vortex-layout/src/plan/physical/lower.rs @@ -83,13 +83,21 @@ 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 these checked cumulative offsets end at the parent's row count. + // that their row counts sum to the parent; offsets were computed with checked addition above. Ok(unsafe { ConcatPlan::from_children_unchecked( layout.dtype().clone(), layout.row_count(), - layout.chunk_offsets()[..layout.nchildren()].into(), + row_offsets.into(), lazy_children(layout.to_layout(), (0..layout.nchildren()).collect()), ) }) diff --git a/vortex-layout/src/plan/physical/operators/list.rs b/vortex-layout/src/plan/physical/operators/list.rs deleted file mode 100644 index f7324824a20..00000000000 --- a/vortex-layout/src/plan/physical/operators/list.rs +++ /dev/null @@ -1,355 +0,0 @@ -// 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 self::reads::ElementReads; -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::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 reads; -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, - initialized: bool, - offsets: Option, - validity: Option, - elements: Option, - /// Element reads issued by this source when the elements are plain segment chunks. - direct: Option, - reads: FxHashMap>, - chunks: Vec<(usize, ArrayRef)>, - output: Option, - done: bool, -} - -impl ListSource { - pub(crate) fn new(plan: ListPackPlan, mask: Mask, values: ValueCache) -> Self { - Self { - plan, - mask, - values, - initialized: false, - offsets: None, - validity: None, - elements: None, - direct: None, - reads: FxHashMap::default(), - 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)?; - self.start_elements(cx, splits, overlap)?; - self.initialized = true; - Ok(()) - } - - fn start_elements( - &mut self, - cx: &mut Cx<'_>, - splits: Vec, - overlap: bool, - ) -> VortexResult<()> { - let elements = self.plan.elements()?; - if let Some(mut reads) = ElementReads::try_new(&elements, &splits)? { - reads.request(cx); - self.direct = Some(reads); - return Ok(()); - } - let scan = Scan::try_new(cx.session().clone(), elements, splits)? - .with_value_cache(self.values.clone()); - 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 let Some(reads) = self.direct.as_mut() { - let session = cx.session().clone(); - while let Some((segment, bytes)) = cx.take_segment() { - reads.deliver(segment, bytes, &session)?; - } - if !reads.is_done() { - return Ok(SourceReady::Blocked(SourceBlocked::Io)); - } - let reads = self - .direct - .take() - .ok_or_else(|| vortex_err!("Missing list element reads"))?; - self.chunks.extend(reads.finish()?.into_iter().enumerate()); - self.finish()?; - return Ok(SourceReady::Ready); - } - 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)?; - } - 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::Waiting => { - return Ok(SourceReady::Blocked(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/reads.rs b/vortex-layout/src/plan/physical/operators/list/reads.rs deleted file mode 100644 index e8a2e569926..00000000000 --- a/vortex-layout/src/plan/physical/operators/list/reads.rs +++ /dev/null @@ -1,180 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Reads selected list elements straight from their segments, in the list's own pipeline. -//! -//! Element ranges are known only once the offsets arrive, so they cannot be compiled with the -//! rest of the split. When the elements are plain segment chunks, the list source issues their -//! reads itself instead of compiling a nested scan: each touched segment is read once and decoded -//! on arrival, and the ranges it holds are sliced out of it in list order. - -use std::ops::Range; - -use rustc_hash::FxHashMap; -use vortex_array::ArrayRef; -use vortex_array::buffer::BufferHandle; -use vortex_error::VortexResult; -use vortex_error::vortex_ensure; -use vortex_error::vortex_err; -use vortex_mask::Mask; -use vortex_session::VortexSession; - -use super::super::scan::decode; -use super::super::scan::select; -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::SegmentScanPlan; -use crate::plan::physical::Zoned; -use crate::plan::physical::operator::Cx; -use crate::plan::physical::pipeline::Split; -use crate::plan::physical::pipeline::overlapping; -use crate::segments::SegmentId; - -/// The selected element rows of one segment chunk. -struct Piece { - slice: Option>, - filter: Option, -} - -/// A requested segment and the pieces, by output position, that read it. -struct Pending { - plan: SegmentScanPlan, - pieces: Vec, -} - -/// The element reads of one list source, from request to the ordered element chunks. -pub(super) struct ElementReads { - pieces: Vec, - output: Vec>, - pending: FxHashMap, - /// Segments in first-use order, so reads are published in element order. - order: Vec, -} - -impl ElementReads { - /// Plans the reads for `splits` of `elements`, or `None` when any selected range reaches a - /// plan other than segment chunks, which then needs a nested scan. - pub(super) fn try_new(elements: &PlanRef, splits: &[Split]) -> VortexResult> { - let mut reads = Self { - pieces: Vec::new(), - output: Vec::new(), - pending: FxHashMap::default(), - order: Vec::new(), - }; - for split in splits { - if !reads.collect(elements, split.rows.clone(), &split.mask)? { - return Ok(None); - } - } - reads.output = (0..reads.pieces.len()).map(|_| None).collect(); - Ok(Some(reads)) - } - - fn collect(&mut self, plan: &PlanRef, rows: Range, mask: &Mask) -> VortexResult { - if mask.all_false() { - return Ok(true); - } - // A filter keeps the rows of the mask it runs with, which this piece applies itself. - if let Some(filter) = plan.as_opt::() { - return self.collect(&filter.child_plan()?, rows, mask); - } - if let Some(zoned) = plan.as_opt::() { - if zoned.is_pruning() { - return Ok(false); - } - return match zoned.data_plan()? { - Some(data) => self.collect(&data, rows, mask), - None => Ok(false), - }; - } - if let Some(stats) = plan.as_opt::() { - return self.collect(&stats.data_plan()?, rows, mask); - } - if let Some(concat) = plan.as_opt::() { - for (index, start, local) in overlapping(concat, &rows) { - let mask = mask.slice( - usize::try_from(local.start - rows.start)? - ..usize::try_from(local.end - rows.start)?, - ); - if !self.collect( - &concat.child_required(index)?, - local.start - start..local.end - start, - &mask, - )? { - return Ok(false); - } - } - return Ok(true); - } - let Some(scan) = plan.as_opt::() else { - return Ok(false); - }; - let slice = if rows.start == 0 && rows.end == scan.row_count() { - None - } else { - Some(usize::try_from(rows.start)?..usize::try_from(rows.end)?) - }; - let filter = (!mask.all_true()).then(|| mask.clone()); - let segment = scan.segment_id(); - let piece = self.pieces.len(); - self.pieces.push(Piece { slice, filter }); - self.pending - .entry(segment) - .or_insert_with(|| { - self.order.push(segment); - Pending { - plan: scan.clone(), - pieces: Vec::new(), - } - }) - .pieces - .push(piece); - Ok(true) - } - - /// Publishes one read per touched segment. - pub(super) fn request(&mut self, cx: &mut Cx<'_>) { - for segment in self.order.drain(..) { - cx.request_read(segment); - } - } - - /// Decodes an arrived segment and narrows it to each piece that reads it. - pub(super) fn deliver( - &mut self, - segment: SegmentId, - bytes: BufferHandle, - session: &VortexSession, - ) -> VortexResult<()> { - let pending = self - .pending - .remove(&segment) - .ok_or_else(|| vortex_err!("Unexpected list element segment {segment}"))?; - let array = decode(&pending.plan, bytes, session)?; - for piece in pending.pieces { - let Piece { slice, filter } = &self.pieces[piece]; - self.output[piece] = Some(select(array.clone(), slice.as_ref(), filter.as_ref())?); - } - Ok(()) - } - - /// Whether every requested segment has arrived. - pub(super) fn is_done(&self) -> bool { - self.pending.is_empty() - } - - /// The element chunks in list order. - pub(super) fn finish(self) -> VortexResult> { - vortex_ensure!( - self.is_done(), - "List elements finished before every read arrived" - ); - self.output - .into_iter() - .map(|array| array.ok_or_else(|| vortex_err!("Missing list element chunk"))) - .collect() - } -} diff --git a/vortex-layout/src/plan/physical/operators/list/work.rs b/vortex-layout/src/plan/physical/operators/list/work.rs deleted file mode 100644 index e93595683d6..00000000000 --- a/vortex-layout/src/plan/physical/operators/list/work.rs +++ /dev/null @@ -1,411 +0,0 @@ -// 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 c4545ec5e19..b0f5937d71d 100644 --- a/vortex-layout/src/plan/physical/operators/mod.rs +++ b/vortex-layout/src/plan/physical/operators/mod.rs @@ -9,7 +9,6 @@ mod concat; mod eval; mod filter; -mod list; mod once; mod pack; mod scan; @@ -19,7 +18,6 @@ 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 7a4e83efc40..4fefa368f05 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, len, validity).into_array(); + let array = assemble(&self.fields, arrays, validity).into_array(); Ok(Step::Final(array)) } } @@ -138,9 +138,8 @@ 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], len, Validity::NonNullable).into_array(), + assemble(&self.fields, vec![batch], Validity::NonNullable).into_array(), )) } } @@ -150,12 +149,8 @@ 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, - len: usize, - validity: Validity, -) -> StructArray { +fn assemble(fields: &StructFields, arrays: Vec, validity: Validity) -> StructArray { + let len = arrays.first().map_or(0, |array| array.len()); debug_assert!( arrays.iter().all(|array| array.len() == len) && fields diff --git a/vortex-layout/src/plan/physical/operators/scan.rs b/vortex-layout/src/plan/physical/operators/scan.rs index 48c33773007..b6d2c463a81 100644 --- a/vortex-layout/src/plan/physical/operators/scan.rs +++ b/vortex-layout/src/plan/physical/operators/scan.rs @@ -24,7 +24,7 @@ use crate::plan::physical::operator::Step; use crate::segments::SegmentId; /// Decodes the whole of a scan's segment. -pub(super) fn decode( +fn decode( plan: &SegmentScanPlan, segment: BufferHandle, session: &VortexSession, @@ -42,7 +42,7 @@ pub(super) fn decode( } /// Slices a whole decoded segment to `slice`, then keeps the rows `filter` selects. -pub(super) fn select( +fn select( array: ArrayRef, slice: Option<&Range>, filter: Option<&Mask>, diff --git a/vortex-layout/src/plan/physical/pipeline/scan.rs b/vortex-layout/src/plan/physical/pipeline/scan.rs index 1d4b01c1bfd..9d650ebbd26 100644 --- a/vortex-layout/src/plan/physical/pipeline/scan.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan.rs @@ -18,7 +18,6 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; -use vortex_mask::AllOr; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -619,39 +618,6 @@ struct Active { output: Option, } -/// A split selecting more separate runs than this counts its whole row range instead, which -/// bounds the walks over the plan before the scan starts. -const MAX_SELECTED_RUNS: usize = 1024; - -/// The rows the splits select, sorted, with overlapping and adjacent runs merged. -fn selected_runs(splits: &[Split]) -> VortexResult>> { - let mut runs = Vec::new(); - for split in splits { - match split.mask.slices() { - AllOr::None => {} - AllOr::Some(slices) if slices.len() <= MAX_SELECTED_RUNS => { - for &(start, end) in slices { - runs.push( - split.rows.start + u64::try_from(start)? - ..split.rows.start + u64::try_from(end)?, - ); - } - } - AllOr::All | AllOr::Some(_) => runs.push(split.rows.clone()), - } - } - runs.retain(|run| !run.is_empty()); - runs.sort_unstable_by_key(|run| run.start); - let mut merged: Vec> = Vec::with_capacity(runs.len()); - for run in runs { - match merged.last_mut() { - Some(last) if run.start <= last.end => last.end = last.end.max(run.end), - _ => merged.push(run), - } - } - Ok(merged) -} - /// Runs a plan over a list of splits. /// /// The plan runs as one stage under each split's selection. @@ -705,16 +671,15 @@ impl Scan { 0, splits.iter().map(|split| split.rows.clone()).collect(), ); - // Only the rows some split selects can share a segment between readers. Chunks in the - // gaps between them stay lazy: they are never lowered or read. A segment several runs - // reach counts a reader per run, so it is decoded once even when distinct chunks - // reuse it; readers that never come leave their ports unclaimed. + // 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) => { - for run in selected_runs(&splits)? { - core.shares.add(plan, run)?; - } - } + Root::Plan(plan) => core.shares.add(plan, start..end)?, } core.shares.retain_shared(); core.dirty_flags = vec![false; 1]; diff --git a/vortex-layout/src/plan/physical/pipeline/tests.rs b/vortex-layout/src/plan/physical/pipeline/tests.rs index 8e673004b26..8fc9c67b900 100644 --- a/vortex-layout/src/plan/physical/pipeline/tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/tests.rs @@ -6,7 +6,6 @@ use std::collections::VecDeque; use std::ops::Range; -use std::sync::Arc; use rstest::rstest; use vortex_array::ArrayContext; @@ -15,12 +14,9 @@ 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; @@ -29,7 +25,6 @@ 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; @@ -37,9 +32,7 @@ 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; @@ -49,11 +42,9 @@ 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; @@ -297,170 +288,6 @@ 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(()) -} - -#[rstest] -fn segment_list_elements_are_read_by_the_list_source( - #[values(Delivery::Fifo, Delivery::Lifo)] order: Delivery, - #[values(false, true)] utf8: bool, -) -> VortexResult<()> { - let mut store = Store::default(); - let elements = if utf8 { - VarBinViewArray::from_iter_str((0..100).map(|value| format!("v{value}"))).into_array() - } else { - Buffer::from_iter(0_i64..100).into_array() - }; - let offsets = Buffer::from_iter((0_u64..=20).map(|row| row * 5)).into_array(); - let input = - ListArray::try_new(elements.clone(), offsets.clone(), Validity::NonNullable)?.into_array(); - let plan = ListPackPlan::try_new( - false.into(), - 20, - lower(&store.chunked(&elements, &[10; 10])?)?, - lower(&store.flat(&offsets)?)?, - None, - )? - .into_plan(); - let rows = [1_u64, 13]; - let mask = Mask::from_iter((0..20).map(|row| rows.contains(&row))); - let result = run(&store, &plan, 0..20, 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()); - // Only the offsets and the two touched element chunks are read, and the element reads are - // published together once the offsets arrive. - assert_eq!(segments_read(&result.events), vec![0, 6, 10]); - assert!( - result.events.contains(&Event::Io(vec![0, 6])), - "{:?}", - result.events - ); - Ok(()) -} - -#[rstest] -fn dict_list_elements_fall_back_to_a_nested_scan( - #[values(Delivery::Fifo, Delivery::Lifo)] order: Delivery, -) -> VortexResult<()> { - let mut store = Store::default(); - let values = VarBinViewArray::from_iter_str(["a", "b", "c"]).into_array(); - let codes = PrimitiveArray::from_iter((0..100).map(|value| (value % 3) as u8)).into_array(); - let elements = VarBinViewArray::from_iter_str((0..100).map(|value| ["a", "b", "c"][value % 3])) - .into_array(); - let offsets = Buffer::from_iter((0_u64..=20).map(|row| row * 5)).into_array(); - let input = ListArray::try_new(elements, offsets.clone(), Validity::NonNullable)?.into_array(); - let layout = - DictLayout::new(store.flat(&values)?, store.chunked(&codes, &[10; 10])?).into_layout(); - let plan = ListPackPlan::try_new( - false.into(), - 20, - lower(&layout)?, - lower(&store.flat(&offsets)?)?, - None, - )? - .into_plan(); - let rows = [1_u64, 2, 13]; - let mask = Mask::from_iter((0..20).map(|row| rows.contains(&row))); - let result = run(&store, &plan, 0..20, 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()); - Ok(()) -} - /// Delivers segments in the given order. fn scripted(order: &[u32]) -> impl FnMut(&[ReadRequest]) -> usize + '_ { let mut next = order.iter(); @@ -605,135 +432,6 @@ 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(); @@ -1136,29 +834,3 @@ 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/list_pack.rs b/vortex-layout/src/plan/physical/plans/list_pack.rs index 4481180f6ef..bb075c7f4e2 100644 --- a/vortex-layout/src/plan/physical/plans/list_pack.rs +++ b/vortex-layout/src/plan/physical/plans/list_pack.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::borrow::Cow; -use std::ops::Range; use std::sync::Arc; use vortex_array::EmptyMetadata; @@ -11,8 +10,6 @@ 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; @@ -21,12 +18,6 @@ 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; @@ -133,68 +124,6 @@ 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()); - 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 2b74e23d17a..da6fa56ccb9 100644 --- a/vortex-layout/src/plan/physical/tests.rs +++ b/vortex-layout/src/plan/physical/tests.rs @@ -4,10 +4,7 @@ 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; @@ -46,7 +43,6 @@ use vortex_session::registry::ReadContext; use super::*; use crate::LayoutBuildContext; -use crate::LayoutChildren; use crate::LayoutEncoding; use crate::LayoutRef; use crate::OwnedLayoutChildren; @@ -154,60 +150,6 @@ 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/mod.rs b/vortex-layout/src/scan/execution/mod.rs index 93c48ecae12..9d8b8897ef3 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -3,11 +3,10 @@ //! Scan execution over physical plans through the planner I/O and compute protocol. //! -//! A `SelectionPlanner` resolves each fixed-size work chunk's selected rows and hands them to a -//! `ProjectionPlanner`, which emits one `ProjectionMorsel` for them. 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` 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. mod adaptor; mod projection; @@ -16,12 +15,9 @@ use std::sync::Arc; use vortex_buffer::Alignment; use vortex_io::request::IoTarget; -use vortex_scan::planning::next::next_fn; +use vortex_mask::Mask; 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::projection::ProjectionPlanner; @@ -66,15 +62,15 @@ pub(super) struct ScanPlans { pub(super) values: ValueCache, } -/// Selects the rows of a fixed work chunk before handing them to its projection planner. -pub(super) fn plan_projection( - plans: ScanPlans, +/// The rows of a split selected for projection. +struct SelectedRows { + /// The split the rows belong to. scope: WorkScope, - selection: Arc, -) -> Box { - Box::new(SelectionPlanner::new( - scope, - selection, - next_fn(move |selected| Ok(ProjectionPlanner::new(plans.clone(), selected))), - )) + /// 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 })) } diff --git a/vortex-layout/src/scan/v2/io.rs b/vortex-layout/src/scan/v2/io.rs index 2c4ee9c121f..ce565846d7e 100644 --- a/vortex-layout/src/scan/v2/io.rs +++ b/vortex-layout/src/scan/v2/io.rs @@ -34,55 +34,25 @@ use crate::segments::SegmentFuture; use crate::segments::SegmentId; use crate::segments::SegmentSource; -/// 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(), - } - } +/// 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()) } /// Serves splits from any [`SegmentSource`], by segment id. pub(super) struct SegmentScanIo { segments: Arc, - ranges: Arc, + ranges: SegmentRanges, } impl SegmentScanIo { - pub(super) fn new(segments: Arc, ranges: Arc) -> Self { + pub(super) fn new(segments: Arc, ranges: SegmentRanges) -> Self { Self { segments, ranges } } } @@ -107,7 +77,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: Arc, + ranges: SegmentRanges, reads: Mutex>>, ready: Mutex>, /// Reads started by prefetches, consumed by the next fetch of their segment. @@ -120,7 +90,7 @@ pub(super) struct SegmentIoSource { type Prefetched = Shared>>>; impl SegmentIoSource { - pub(super) fn new(segments: Arc, ranges: Arc) -> Self { + pub(super) fn new(segments: Arc, ranges: SegmentRanges) -> Self { Self { segments, ranges, @@ -144,9 +114,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 @@ -264,7 +234,6 @@ 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; @@ -278,47 +247,12 @@ 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); @@ -367,7 +301,7 @@ mod tests { length: 4, alignment: Alignment::none(), }); - let io = SegmentIoSource::new(segments, segment_ranges(Arc::from(locations))); + let io = SegmentIoSource::new(segments, segment_ranges(&locations)); let mut pending = request(IoIntent::Fetch, 1); pending.target = IoTarget::range(4, 4); io.submit(IoOwnerId(0), vec![pending, request(IoIntent::Fetch, 0)])?; @@ -400,10 +334,7 @@ mod tests { length: 4, alignment: Alignment::none(), }; - let io = SegmentIoSource::new( - Arc::clone(&segments) as _, - segment_ranges(Arc::from([location])), - ); + let io = SegmentIoSource::new(Arc::clone(&segments) as _, segment_ranges(&[location])); io.submit(IoOwnerId(0), vec![request(IoIntent::Prefetch, 0)])?; assert_eq!(segments.0.load(Ordering::Relaxed), 1); @@ -431,10 +362,7 @@ mod tests { length: 4, alignment: Alignment::none(), }; - let io = SegmentIoSource::new( - Arc::clone(&segments) as _, - segment_ranges(Arc::from([location])), - ); + let io = SegmentIoSource::new(Arc::clone(&segments) as _, segment_ranges(&[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 51b450dd519..d6c5fff3cac 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -11,10 +11,8 @@ //! 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 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. +//! 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. mod io; mod split; @@ -37,10 +35,11 @@ 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_mask::AllOr; use vortex_scan::planning::planner::WorkScope; use vortex_scan::selection::Selection; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; @@ -48,6 +47,12 @@ 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; @@ -63,31 +68,47 @@ use crate::segments::SegmentSource; const ROW_BLOCK_SIZE: usize = 8192; const MAX_ROWS_PER_CHUNK: usize = 8 * ROW_BLOCK_SIZE; const TARGET_BYTES_PER_CHUNK: u64 = 8 * 1024 * 1024; +/// Selected rows closer than a writer block are likely to share stored segments. +const MAX_GAP_BYTES: u64 = 1024 * 1024; -/// Limit estimated bytes per task without changing the writer-aligned row granularity. -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; - } +fn bytes_per_row(row_count: u64, locations: &[SegmentLocation]) -> u64 { let bytes: u64 = locations .iter() .map(|location| u64::from(location.length)) .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; - 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 + bytes.div_ceil(row_count.max(1)).max(1) +} + +/// Limit estimated bytes per task without changing the writer-aligned row granularity. +fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize { + if row_count <= MAX_ROWS_PER_CHUNK as u64 { + return MAX_ROWS_PER_CHUNK; + } + let blocks = + TARGET_BYTES_PER_CHUNK / bytes_per_row(row_count, locations) / ROW_BLOCK_SIZE as u64; + blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE +} + +/// The widest unselected gap kept inside one group of selected rows: one row block, or more +/// when a writer block of serialized bytes spans more rows. +fn max_gap_rows(row_count: u64, locations: &[SegmentLocation]) -> usize { + let rows = MAX_GAP_BYTES / bytes_per_row(row_count, locations); + usize::try_from(rows) + .unwrap_or(usize::MAX) + .max(ROW_BLOCK_SIZE) +} + +/// Groups sorted, disjoint selected slices, starting a new group at each gap wider than +/// `max_gap` rows. +fn group_slices(slices: &[(usize, usize)], max_gap: usize) -> Vec> { + let mut groups: Vec> = Vec::new(); + for &(start, end) in slices { + match groups.last_mut() { + Some(group) if start - group.end <= max_gap => group.end = end, + _ => groups.push(start..end), + } + } + groups } /// Fixed boundaries in the file's row domain, clipped to the requested range. @@ -117,10 +138,8 @@ pub struct ScanFile { /// Builds a projection scan on the planner-driven V2 executor. /// -/// 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. +/// 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. pub struct ScanBuilder { session: VortexSession, file: ScanFile, @@ -247,6 +266,34 @@ 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, @@ -256,6 +303,8 @@ struct PreparedScan { /// The scan's splits, over the root row domain, in row order. splits: Vec>, selection: Selection, + /// The widest unselected gap kept inside one split's group of selected rows. + max_gap: usize, map_fn: Arc VortexResult + Send + Sync>, ordered: bool, concurrency: usize, @@ -266,6 +315,10 @@ 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 @@ -277,23 +330,15 @@ impl PreparedScan { optimize(plan_row_idx_expression(builder.projection, root)?)? }; let chunk_rows = if identity { - 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) + rows_per_chunk(row_count, &file.locations) } else { MAX_ROWS_PER_CHUNK }; let splits = row_chunks(builder.row_range.unwrap_or(0..row_count), chunk_rows); + let max_gap = max_gap_rows(row_count, &file.locations); let io = Arc::new(SegmentScanIo::new( file.segments, - segment_ranges(Arc::clone(&file.locations)), + segment_ranges(&file.locations), )); Ok(Self { session: builder.session.clone(), @@ -307,35 +352,40 @@ impl PreparedScan { io, splits, selection: builder.selection, + max_gap, map_fn: builder.map_fn, ordered: builder.ordered, concurrency: builder.concurrency, }) } - /// The splits with selected rows, each ready to be admitted to a planning run. + /// The groups of selected rows, each ready to be admitted to a planning run. + /// + /// Within each fixed chunk, a selected run joins the group before it while the gap between + /// them is at most `max_gap` rows. A group spans only its own rows and keeps the selection's + /// mask over them, so rows in wider gaps get no planner, plan or read. fn split_plans(&self) -> Vec { - let selection = Arc::new(self.selection.clone()); - self.splits - .iter() - .filter_map(|rows| { - if self.selection.row_count_in_range(rows) == 0 { - return None; - } + let mut plans = Vec::new(); + for rows in &self.splits { + let row_mask = self.selection.row_mask(rows); + let mask = row_mask.mask(); + let groups = match mask.slices() { + AllOr::None => continue, + AllOr::All => vec![0..mask.len()], + AllOr::Some(slices) => group_slices(slices, self.max_gap), + }; + for group in groups { let scope = WorkScope { file_ordinal: 0, - rows: rows.clone(), + rows: rows.start + group.start as u64..rows.start + group.end as u64, }; - Some(SplitPlan { - root: plan_projection( - self.plans.clone(), - scope.clone(), - Arc::clone(&selection), - ), + plans.push(SplitPlan { + root: plan_projection(self.plans.clone(), scope.clone(), mask.slice(group)), scope, - }) - }) - .collect() + }); + } + } + plans } /// Runs every split as a task on the session's runtime and yields their mapped batches. diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index abecf2785a9..2ffbadba0be 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -16,7 +16,6 @@ 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; @@ -37,23 +36,20 @@ 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; use super::ScanBuilder; use super::ScanFile; +use super::group_slices; +use super::max_gap_rows; 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; @@ -72,8 +68,6 @@ use crate::test::new_session; const CHUNK_ROWS: usize = 1000; const CHUNKS: usize = 4; -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 { @@ -322,29 +316,19 @@ fn fixed_row_chunks(#[case] rows: Range, #[case] expected: Vec>) } #[rstest] -#[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, -) { +#[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) { let locations = [SegmentLocation { offset: 0, length: bytes, alignment: Alignment::none(), }]; - assert_eq!(rows_per_chunk(rows, &locations, sparse_target), expected); + assert_eq!(rows_per_chunk(rows, &locations), expected); } #[test] @@ -367,46 +351,6 @@ 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(); @@ -474,94 +418,128 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { Ok(()) } -#[tokio::test] -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?; - assert_same_as_v1(&session, || builder(&session, &segments, &layout)).await +#[rstest] +#[case::empty(vec![], 10, vec![])] +#[case::joined(vec![(0, 2), (5, 6), (16, 20)], 10, vec![0..20])] +#[case::at_limit(vec![(0, 2), (12, 13)], 10, vec![0..13])] +#[case::split(vec![(0, 2), (13, 14), (30, 31)], 10, vec![0..2, 13..14, 30..31])] +#[case::adjacent_only(vec![(0, 2), (3, 4)], 0, vec![0..2, 3..4])] +fn selected_slices_group_across_small_gaps( + #[case] slices: Vec<(usize, usize)>, + #[case] max_gap: usize, + #[case] expected: Vec>, +) { + assert_eq!(group_slices(&slices, max_gap), expected); } #[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()); +#[case::narrow_rows(1_000_000, 4_000_000, 262_144)] +#[case::wide_rows(1_000_000, 1_000_000_000, 8192)] +#[case::empty(0, 0, 1_048_576)] +fn gap_budget_covers_a_writer_block( + #[case] rows: u64, + #[case] bytes: u32, + #[case] expected: usize, +) { + let locations = [SegmentLocation { + offset: 0, + length: bytes, + alignment: Alignment::none(), + }]; + assert_eq!(max_gap_rows(rows, &locations), expected); +} + +/// Three stored chunks of 70,000 consecutive integers, read through a counting source. +async fn counted_ints(session: &VortexSession) -> VortexResult<(Arc, LayoutRef)> { + const STORED_CHUNK_ROWS: i32 = 70_000; let segments = Arc::new(TestSegments::default()); - let (mut id, eof) = SequenceId::root().split(); - let layout = strategy + let (mut sequence, eof) = SequenceId::root().split(); + let arrays = (0..3) + .map(|index| { + let start = index * STORED_CHUNK_ROWS; + Ok(( + sequence.advance(), + PrimitiveArray::from_iter(start..start + STORED_CHUNK_ROWS).into_array(), + )) + }) + .collect::>>(); + let dtype = DType::Primitive(vortex_array::dtype::PType::I32, false.into()); + let layout = chunked_flat() .write_stream( ArrayContext::empty().into(), Arc::::clone(&segments), - SequentialStreamAdapter::new(dtype.clone(), stream::iter([Ok((id.advance(), input))])) - .sendable(), + SequentialStreamAdapter::new(dtype, stream::iter(arrays)).sendable(), eof, - &session, + 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((counted, layout)) +} + +#[rstest] +#[case::default_gap(None, vec![5..30_001, 140_100..140_101], 2)] +#[case::narrow_gap(Some(100), vec![5..7, 30_000..30_001, 140_100..140_101], 3)] +#[tokio::test] +async fn selected_rows_plan_only_their_groups( + #[case] max_gap: Option, + #[case] scopes: Vec>, + #[case] reads: usize, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (counted, layout) = counted_ints(&session).await?; + let segments = Arc::clone(&counted) as Arc; + let rows = vec![5_u64, 6, 30_000, 140_100]; + let mut scan = PreparedScan::try_new( + builder(&session, &segments, &layout)? + .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(rows.clone()))?), + )?; + if let Some(max_gap) = max_gap { + scan.max_gap = max_gap; } + let splits = scan.split_plans(); + assert_eq!( + splits + .iter() + .map(|split| split.scope.rows.clone()) + .collect::>(), + scopes + ); + let mut arrays = Vec::new(); + for split in splits { + arrays.extend(split.run(Arc::clone(&scan.io)).await?); + } + // Each group reads only the stored chunks holding its selected rows. Groups split by a + // wide gap inside one stored chunk read it once each. + assert_eq!(counted.reads.load(Ordering::Relaxed), reads); + let expected = rows + .iter() + .map(|&row| i32::try_from(row)) + .collect::, _>>()?; + assert_arrays_eq!( + ChunkedArray::try_new(arrays, layout.dtype().clone())?.into_array(), + expected.into_array(), + &mut session.create_execution_ctx() + ); + 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<()> { + 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")); Ok(()) } @@ -614,72 +592,3 @@ 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-scan/src/planning/mod.rs b/vortex-scan/src/planning/mod.rs index 2c69e734294..3095712f161 100644 --- a/vortex-scan/src/planning/mod.rs +++ b/vortex-scan/src/planning/mod.rs @@ -10,7 +10,6 @@ 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/selection.rs b/vortex-scan/src/planning/selection.rs deleted file mode 100644 index 75b389690bb..00000000000 --- a/vortex-scan/src/planning/selection.rs +++ /dev/null @@ -1,210 +0,0 @@ -// 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(()) - } -} From 435023cbadbd481318436f808eb173902aed04f7 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 19 Oct 2026 18:56:48 +0000 Subject: [PATCH 06/15] Group selected runs by the stored segments they touch Replace the 8,192-row / 1 MiB gap estimate: within each fixed chunk, a selected run joins the previous group when both reach a stored segment in common, so no segment is read by two groups. Local random-access benchmarks against develop's V2 (3 interleaved passes) show uniform access regressing by 8-14% with both this grouping and the gap estimate: extra planner roots split each chunk's reads across tasks. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_014rtiSjzL4EauVCM7NS84hj --- vortex-layout/src/scan/v2/mod.rs | 88 +++++++++++++++++------------- vortex-layout/src/scan/v2/tests.rs | 60 +++++++------------- 2 files changed, 70 insertions(+), 78 deletions(-) diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index d6c5fff3cac..c43e93b9508 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -19,6 +19,7 @@ mod split; #[cfg(test)] mod tests; +use std::cmp::Ordering; use std::ops::Range; use std::sync::Arc; @@ -53,8 +54,10 @@ use crate::layouts::flat::Flat; use crate::layouts::struct_::Struct; use crate::layouts::zoned::LegacyStats; use crate::layouts::zoned::Zoned; +use crate::plan::physical::PlanRef; use crate::plan::physical::lower; use crate::plan::physical::optimize; +use crate::plan::physical::pipeline::Reach; use crate::plan::physical::plan_row_idx_expression; use crate::scan::arrow::to_record_batch; use crate::scan::execution::ScanPlans; @@ -63,52 +66,67 @@ use crate::scan::execution::plan_projection; use crate::scan::v2::io::SegmentScanIo; use crate::scan::v2::io::segment_ranges; use crate::scan::v2::split::SplitPlan; +use crate::segments::SegmentId; use crate::segments::SegmentSource; const ROW_BLOCK_SIZE: usize = 8192; const MAX_ROWS_PER_CHUNK: usize = 8 * ROW_BLOCK_SIZE; const TARGET_BYTES_PER_CHUNK: u64 = 8 * 1024 * 1024; -/// Selected rows closer than a writer block are likely to share stored segments. -const MAX_GAP_BYTES: u64 = 1024 * 1024; - -fn bytes_per_row(row_count: u64, locations: &[SegmentLocation]) -> u64 { - let bytes: u64 = locations - .iter() - .map(|location| u64::from(location.length)) - .sum(); - bytes.div_ceil(row_count.max(1)).max(1) -} /// Limit estimated bytes per task without changing the writer-aligned row granularity. fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize { if row_count <= MAX_ROWS_PER_CHUNK as u64 { return MAX_ROWS_PER_CHUNK; } - let blocks = - TARGET_BYTES_PER_CHUNK / bytes_per_row(row_count, locations) / ROW_BLOCK_SIZE as u64; + let bytes: u64 = locations + .iter() + .map(|location| u64::from(location.length)) + .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 } -/// The widest unselected gap kept inside one group of selected rows: one row block, or more -/// when a writer block of serialized bytes spans more rows. -fn max_gap_rows(row_count: u64, locations: &[SegmentLocation]) -> usize { - let rows = MAX_GAP_BYTES / bytes_per_row(row_count, locations); - usize::try_from(rows) - .unwrap_or(usize::MAX) - .max(ROW_BLOCK_SIZE) -} - -/// Groups sorted, disjoint selected slices, starting a new group at each gap wider than -/// `max_gap` rows. -fn group_slices(slices: &[(usize, usize)], max_gap: usize) -> Vec> { +/// Groups the sorted, disjoint selected `slices` of the rows starting at `offset`: a slice joins +/// the group before it when both read a stored segment in common, so no segment is read by two +/// groups. A column's segments hold contiguous rows, so a slice can share a segment with an +/// earlier slice of its group only if it shares one with the slice just before it. +fn group_slices( + plan: &PlanRef, + offset: u64, + slices: &[(usize, usize)], +) -> VortexResult>> { let mut groups: Vec> = Vec::new(); + let mut previous = Vec::new(); for &(start, end) in slices { + let mut segments = Vec::new(); + plan.reach( + offset + start as u64..offset + end as u64, + &Reach::Offset(0), + &mut |segment, _| segments.push(segment), + )?; + segments.sort_unstable(); + segments.dedup(); match groups.last_mut() { - Some(group) if start - group.end <= max_gap => group.end = end, + Some(group) if shares_any(&previous, &segments) => group.end = end, _ => groups.push(start..end), } + previous = segments; + } + Ok(groups) +} + +/// Whether two sorted segment lists have a segment in common. +fn shares_any(left: &[SegmentId], right: &[SegmentId]) -> bool { + let (mut left, mut right) = (left.iter().peekable(), right.iter().peekable()); + while let (Some(l), Some(r)) = (left.peek(), right.peek()) { + match l.cmp(r) { + Ordering::Less => _ = left.next(), + Ordering::Greater => _ = right.next(), + Ordering::Equal => return true, + } } - groups + false } /// Fixed boundaries in the file's row domain, clipped to the requested range. @@ -303,8 +321,6 @@ struct PreparedScan { /// The scan's splits, over the root row domain, in row order. splits: Vec>, selection: Selection, - /// The widest unselected gap kept inside one split's group of selected rows. - max_gap: usize, map_fn: Arc VortexResult + Send + Sync>, ordered: bool, concurrency: usize, @@ -335,7 +351,6 @@ impl PreparedScan { MAX_ROWS_PER_CHUNK }; let splits = row_chunks(builder.row_range.unwrap_or(0..row_count), chunk_rows); - let max_gap = max_gap_rows(row_count, &file.locations); let io = Arc::new(SegmentScanIo::new( file.segments, segment_ranges(&file.locations), @@ -352,7 +367,6 @@ impl PreparedScan { io, splits, selection: builder.selection, - max_gap, map_fn: builder.map_fn, ordered: builder.ordered, concurrency: builder.concurrency, @@ -361,10 +375,10 @@ impl PreparedScan { /// The groups of selected rows, each ready to be admitted to a planning run. /// - /// Within each fixed chunk, a selected run joins the group before it while the gap between - /// them is at most `max_gap` rows. A group spans only its own rows and keeps the selection's - /// mask over them, so rows in wider gaps get no planner, plan or read. - fn split_plans(&self) -> Vec { + /// Within each fixed chunk, selected runs that read a stored segment in common form one + /// group. A group spans only its own rows and keeps the selection's mask over them, so rows + /// between groups get no planner, plan or read. + fn split_plans(&self) -> VortexResult> { let mut plans = Vec::new(); for rows in &self.splits { let row_mask = self.selection.row_mask(rows); @@ -372,7 +386,7 @@ impl PreparedScan { let groups = match mask.slices() { AllOr::None => continue, AllOr::All => vec![0..mask.len()], - AllOr::Some(slices) => group_slices(slices, self.max_gap), + AllOr::Some(slices) => group_slices(&self.plans.projection, rows.start, slices)?, }; for group in groups { let scope = WorkScope { @@ -385,7 +399,7 @@ impl PreparedScan { }); } } - plans + Ok(plans) } /// Runs every split as a task on the session's runtime and yields their mapped batches. @@ -395,7 +409,7 @@ impl PreparedScan { fn execute_stream(self) -> VortexResult>> { let concurrency = self.concurrency * get_available_parallelism().unwrap_or(1); let handle = self.session.handle(); - let split_plans = self.split_plans(); + let split_plans = self.split_plans()?; let io = self.io; let map_fn = self.map_fn; // Spawned as the buffer pulls them, so at most `concurrency` splits read at once. diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index 2ffbadba0be..df01dee282d 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -41,10 +41,9 @@ use super::MAX_ROWS_PER_CHUNK; use super::PreparedScan; use super::ScanBuilder; use super::ScanFile; -use super::group_slices; -use super::max_gap_rows; use super::row_chunks; use super::rows_per_chunk; +use super::shares_any; use crate::LayoutRef; use crate::LayoutStrategy; use crate::layouts::chunked::writer::ChunkedLayoutStrategy; @@ -387,7 +386,7 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { }); let segments = Arc::clone(&counted) as Arc; let scan = PreparedScan::try_new(builder(&session, &segments, &layout)?)?; - let splits = scan.split_plans(); + let splits = scan.split_plans()?; assert_eq!( splits.len(), (STORED_CHUNKS * STORED_CHUNK_ROWS).div_ceil(MAX_ROWS_PER_CHUNK) @@ -419,34 +418,16 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { } #[rstest] -#[case::empty(vec![], 10, vec![])] -#[case::joined(vec![(0, 2), (5, 6), (16, 20)], 10, vec![0..20])] -#[case::at_limit(vec![(0, 2), (12, 13)], 10, vec![0..13])] -#[case::split(vec![(0, 2), (13, 14), (30, 31)], 10, vec![0..2, 13..14, 30..31])] -#[case::adjacent_only(vec![(0, 2), (3, 4)], 0, vec![0..2, 3..4])] -fn selected_slices_group_across_small_gaps( - #[case] slices: Vec<(usize, usize)>, - #[case] max_gap: usize, - #[case] expected: Vec>, +#[case::empty(vec![], vec![2], false)] +#[case::disjoint(vec![1, 3], vec![2, 4], false)] +#[case::shared(vec![1, 3, 9], vec![2, 9], true)] +fn sorted_segment_lists_intersect( + #[case] left: Vec, + #[case] right: Vec, + #[case] expected: bool, ) { - assert_eq!(group_slices(&slices, max_gap), expected); -} - -#[rstest] -#[case::narrow_rows(1_000_000, 4_000_000, 262_144)] -#[case::wide_rows(1_000_000, 1_000_000_000, 8192)] -#[case::empty(0, 0, 1_048_576)] -fn gap_budget_covers_a_writer_block( - #[case] rows: u64, - #[case] bytes: u32, - #[case] expected: usize, -) { - let locations = [SegmentLocation { - offset: 0, - length: bytes, - alignment: Alignment::none(), - }]; - assert_eq!(max_gap_rows(rows, &locations), expected); + let ids = |ids: Vec| ids.into_iter().map(SegmentId::from).collect::>(); + assert_eq!(shares_any(&ids(left), &ids(right)), expected); } /// Three stored chunks of 70,000 consecutive integers, read through a counting source. @@ -481,26 +462,24 @@ async fn counted_ints(session: &VortexSession) -> VortexResult<(Arc, + #[case] rows: Vec, #[case] scopes: Vec>, #[case] reads: usize, ) -> VortexResult<()> { let session = new_session().with_tokio(); let (counted, layout) = counted_ints(&session).await?; let segments = Arc::clone(&counted) as Arc; - let rows = vec![5_u64, 6, 30_000, 140_100]; - let mut scan = PreparedScan::try_new( + let scan = PreparedScan::try_new( builder(&session, &segments, &layout)? .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(rows.clone()))?), )?; - if let Some(max_gap) = max_gap { - scan.max_gap = max_gap; - } - let splits = scan.split_plans(); + let splits = scan.split_plans()?; assert_eq!( splits .iter() @@ -512,8 +491,7 @@ async fn selected_rows_plan_only_their_groups( for split in splits { arrays.extend(split.run(Arc::clone(&scan.io)).await?); } - // Each group reads only the stored chunks holding its selected rows. Groups split by a - // wide gap inside one stored chunk read it once each. + // Each group reads only the stored chunks holding its selected rows. assert_eq!(counted.reads.load(Ordering::Relaxed), reads); let expected = rows .iter() From a15cf39bcb984f3fd40e2b7b2d5e9e770b24ebc4 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 19 Oct 2026 19:00:51 +0000 Subject: [PATCH 07/15] Plan each V2 work chunk only over its selected span Revert the selected-run grouping to develop's one planner root per fixed chunk, and narrow that root to the rows from its first selected row to its last, with the selection mask sliced to match. Develop's compiler already builds nothing for chunks the mask does not touch, so rows outside the span get no plan and no read, without splitting a chunk's reads across tasks. Local random-access benchmarks against develop's V2, 3 interleaved passes each (median, 4-core VM, run-to-run spread up to 20%): - this change: geomean -0.7%, no consistent regression; - grouping by an 8,192-row / 1 MiB gap: geomean +0.0%, uniform +8-11%; - grouping by touched segments: geomean +2.8%, uniform +8-14%. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_014rtiSjzL4EauVCM7NS84hj --- vortex-layout/src/scan/v2/mod.rs | 87 ++++++------------------------ vortex-layout/src/scan/v2/tests.rs | 28 +++------- 2 files changed, 22 insertions(+), 93 deletions(-) diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index c43e93b9508..d7724f41544 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -19,7 +19,6 @@ mod split; #[cfg(test)] mod tests; -use std::cmp::Ordering; use std::ops::Range; use std::sync::Arc; @@ -40,7 +39,6 @@ use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_io::request::IoService; use vortex_io::session::RuntimeSessionExt; -use vortex_mask::AllOr; use vortex_scan::planning::planner::WorkScope; use vortex_scan::selection::Selection; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; @@ -54,10 +52,8 @@ use crate::layouts::flat::Flat; use crate::layouts::struct_::Struct; use crate::layouts::zoned::LegacyStats; use crate::layouts::zoned::Zoned; -use crate::plan::physical::PlanRef; use crate::plan::physical::lower; use crate::plan::physical::optimize; -use crate::plan::physical::pipeline::Reach; use crate::plan::physical::plan_row_idx_expression; use crate::scan::arrow::to_record_batch; use crate::scan::execution::ScanPlans; @@ -66,7 +62,6 @@ use crate::scan::execution::plan_projection; use crate::scan::v2::io::SegmentScanIo; use crate::scan::v2::io::segment_ranges; use crate::scan::v2::split::SplitPlan; -use crate::segments::SegmentId; use crate::segments::SegmentSource; const ROW_BLOCK_SIZE: usize = 8192; @@ -87,48 +82,6 @@ fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize { blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE } -/// Groups the sorted, disjoint selected `slices` of the rows starting at `offset`: a slice joins -/// the group before it when both read a stored segment in common, so no segment is read by two -/// groups. A column's segments hold contiguous rows, so a slice can share a segment with an -/// earlier slice of its group only if it shares one with the slice just before it. -fn group_slices( - plan: &PlanRef, - offset: u64, - slices: &[(usize, usize)], -) -> VortexResult>> { - let mut groups: Vec> = Vec::new(); - let mut previous = Vec::new(); - for &(start, end) in slices { - let mut segments = Vec::new(); - plan.reach( - offset + start as u64..offset + end as u64, - &Reach::Offset(0), - &mut |segment, _| segments.push(segment), - )?; - segments.sort_unstable(); - segments.dedup(); - match groups.last_mut() { - Some(group) if shares_any(&previous, &segments) => group.end = end, - _ => groups.push(start..end), - } - previous = segments; - } - Ok(groups) -} - -/// Whether two sorted segment lists have a segment in common. -fn shares_any(left: &[SegmentId], right: &[SegmentId]) -> bool { - let (mut left, mut right) = (left.iter().peekable(), right.iter().peekable()); - while let (Some(l), Some(r)) = (left.peek(), right.peek()) { - match l.cmp(r) { - Ordering::Less => _ = left.next(), - Ordering::Greater => _ = right.next(), - Ordering::Equal => return true, - } - } - false -} - /// Fixed boundaries in the file's row domain, clipped to the requested range. fn row_chunks(rows: Range, rows_per_chunk: usize) -> Vec> { if rows.is_empty() { @@ -373,33 +326,25 @@ impl PreparedScan { }) } - /// The groups of selected rows, each ready to be admitted to a planning run. - /// - /// Within each fixed chunk, selected runs that read a stored segment in common form one - /// group. A group spans only its own rows and keeps the selection's mask over them, so rows - /// between groups get no planner, plan or read. - fn split_plans(&self) -> VortexResult> { - let mut plans = Vec::new(); - for rows in &self.splits { - let row_mask = self.selection.row_mask(rows); - let mask = row_mask.mask(); - let groups = match mask.slices() { - AllOr::None => continue, - AllOr::All => vec![0..mask.len()], - AllOr::Some(slices) => group_slices(&self.plans.projection, rows.start, slices)?, - }; - for group in groups { + /// The splits with selected rows, each ready to be admitted to a planning run. + fn split_plans(&self) -> Vec { + self.splits + .iter() + .filter_map(|rows| { + let row_mask = self.selection.row_mask(rows); + let mask = row_mask.mask(); + // Only the span from the first selected row to the last is planned. + let span = mask.first()?..mask.last()? + 1; let scope = WorkScope { file_ordinal: 0, - rows: rows.start + group.start as u64..rows.start + group.end as u64, + rows: rows.start + span.start as u64..rows.start + span.end as u64, }; - plans.push(SplitPlan { - root: plan_projection(self.plans.clone(), scope.clone(), mask.slice(group)), + Some(SplitPlan { + root: plan_projection(self.plans.clone(), scope.clone(), mask.slice(span)), scope, - }); - } - } - Ok(plans) + }) + }) + .collect() } /// Runs every split as a task on the session's runtime and yields their mapped batches. @@ -409,7 +354,7 @@ impl PreparedScan { fn execute_stream(self) -> VortexResult>> { let concurrency = self.concurrency * get_available_parallelism().unwrap_or(1); let handle = self.session.handle(); - let split_plans = self.split_plans()?; + let split_plans = self.split_plans(); let io = self.io; let map_fn = self.map_fn; // Spawned as the buffer pulls them, so at most `concurrency` splits read at once. diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index df01dee282d..c7115a213c1 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -43,7 +43,6 @@ use super::ScanBuilder; use super::ScanFile; use super::row_chunks; use super::rows_per_chunk; -use super::shares_any; use crate::LayoutRef; use crate::LayoutStrategy; use crate::layouts::chunked::writer::ChunkedLayoutStrategy; @@ -386,7 +385,7 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { }); let segments = Arc::clone(&counted) as Arc; let scan = PreparedScan::try_new(builder(&session, &segments, &layout)?)?; - let splits = scan.split_plans()?; + let splits = scan.split_plans(); assert_eq!( splits.len(), (STORED_CHUNKS * STORED_CHUNK_ROWS).div_ceil(MAX_ROWS_PER_CHUNK) @@ -417,19 +416,6 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { Ok(()) } -#[rstest] -#[case::empty(vec![], vec![2], false)] -#[case::disjoint(vec![1, 3], vec![2, 4], false)] -#[case::shared(vec![1, 3, 9], vec![2, 9], true)] -fn sorted_segment_lists_intersect( - #[case] left: Vec, - #[case] right: Vec, - #[case] expected: bool, -) { - let ids = |ids: Vec| ids.into_iter().map(SegmentId::from).collect::>(); - assert_eq!(shares_any(&ids(left), &ids(right)), expected); -} - /// Three stored chunks of 70,000 consecutive integers, read through a counting source. async fn counted_ints(session: &VortexSession) -> VortexResult<(Arc, LayoutRef)> { const STORED_CHUNK_ROWS: i32 = 70_000; @@ -462,12 +448,10 @@ async fn counted_ints(session: &VortexSession) -> VortexResult<(Arc, #[case] scopes: Vec>, #[case] reads: usize, @@ -479,7 +463,7 @@ async fn selected_rows_plan_only_their_groups( builder(&session, &segments, &layout)? .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(rows.clone()))?), )?; - let splits = scan.split_plans()?; + let splits = scan.split_plans(); assert_eq!( splits .iter() @@ -491,7 +475,7 @@ async fn selected_rows_plan_only_their_groups( for split in splits { arrays.extend(split.run(Arc::clone(&scan.io)).await?); } - // Each group reads only the stored chunks holding its selected rows. + // Each work chunk reads only the stored chunks holding its selected rows. assert_eq!(counted.reads.load(Ordering::Relaxed), reads); let expected = rows .iter() From 37d689f448b38413ab4bf91bac0300000aa540d1 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 19 Oct 2026 23:38:20 +0000 Subject: [PATCH 08/15] Plan V2 random access by the segments a selection reads V2 takes were slower than the default scan on CI. Micro benchmarks of each random-access query, over files in the page cache, traced it to work that does not depend on the selected rows: - Each scan checked support by walking the file's whole layout tree, deserializing every chunk: 245 us on feature vectors before any row. Lowering already fails a layout without a pipeline implementation when a split plans it, so the walk is removed. No list layout is written by default, and V2 neither writes nor reads one. - Fixed 64K-row work chunks cut taxi's 178K-row stored chunks into three, and each split read and decoded the same segments again: 1,002 segment requests for 337 segments on taxi/uniform. Selected runs that read a stored segment in common now share a planner, across chunk boundaries, so every segment is read once. - A planner's single hull split made the pipeline reach, and so lower, every chunk between selected runs: 441 lowered layouts for 9 segments on feature-vectors/correlated. Runs reading exactly the same segments share a masked split; others get splits of their own, all active at once, and the pipeline's shared-segment count reaches each split's rows instead of their hull. 441 lowered layouts become 14. - Lowering a chunked layout reused none of the chunk offsets the layout already holds, reading every chunk's row count again. The random-access benchmark reads every Vortex file with V2, without the support probe or the default scan's layout reader cache. A v2_take divan benchmark runs each query with both scans. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_014rtiSjzL4EauVCM7NS84hj --- Cargo.lock | 2 + benchmarks/random-access-bench/Cargo.toml | 8 + benchmarks/random-access-bench/README.md | 13 ++ .../random-access-bench/benches/v2_take.rs | 126 ++++++++++++++ benchmarks/random-access-bench/src/lib.rs | 2 +- vortex-bench/src/random_access/take.rs | 51 +----- vortex-layout/src/layouts/chunked/mod.rs | 5 + vortex-layout/src/plan/physical/children.rs | 2 +- vortex-layout/src/plan/physical/lower.rs | 28 ++- .../src/plan/physical/pipeline/scan.rs | 15 +- vortex-layout/src/scan/execution/adaptor.rs | 7 +- vortex-layout/src/scan/execution/mod.rs | 19 ++- .../src/scan/execution/projection.rs | 24 ++- vortex-layout/src/scan/v2/mod.rs | 160 +++++++++++------- vortex-layout/src/scan/v2/tests.rs | 24 ++- 15 files changed, 333 insertions(+), 153 deletions(-) create mode 100644 benchmarks/random-access-bench/benches/v2_take.rs diff --git a/Cargo.lock b/Cargo.lock index 22aace294d8..d64d2915fe5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9055,6 +9055,7 @@ version = "0.1.0" dependencies = [ "anyhow", "clap", + "codspeed-divan-compat", "indicatif", "lance-bench", "mimalloc", @@ -9063,6 +9064,7 @@ dependencies = [ "tabled", "tokio", "url", + "vortex", "vortex-bench", ] diff --git a/benchmarks/random-access-bench/Cargo.toml b/benchmarks/random-access-bench/Cargo.toml index cf359fc3993..5a502600232 100644 --- a/benchmarks/random-access-bench/Cargo.toml +++ b/benchmarks/random-access-bench/Cargo.toml @@ -26,6 +26,10 @@ tokio = { workspace = true, features = ["full"] } url = { workspace = true } vortex-bench = { workspace = true } +[dev-dependencies] +divan = { workspace = true } +vortex = { workspace = true } + [features] lance = ["dep:lance-bench"] @@ -33,5 +37,9 @@ lance = ["dep:lance-bench"] name = "random-access-bench" test = false +[[bench]] +name = "v2_take" +harness = false + [lints] workspace = true diff --git a/benchmarks/random-access-bench/README.md b/benchmarks/random-access-bench/README.md index 73f52f5b219..65d573d6100 100644 --- a/benchmarks/random-access-bench/README.md +++ b/benchmarks/random-access-bench/README.md @@ -22,12 +22,25 @@ untimed warm-up, then reuses the open file handle. Reopen mode includes file ope work in each timed iteration. CI drives the full matrix via [`scripts/random-access-split.py`](../../scripts/random-access-split.py). +Vortex measurements use the planner-driven V2 scan builder. + ## Running locally ```bash cargo run -p random-access-bench --profile release_debug --features lance ``` +## Micro benchmarks + +The `v2_take` divan benchmark runs each Vortex query with the default and V2 scans over the +files in the page cache, plus the per-scan cost of a V2 scan that selects no rows. Write the +files first: + +```bash +cargo run -p random-access-bench --profile release_debug -- --prepare-data --formats vortex +cargo bench -p random-access-bench --bench v2_take +``` + ## 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/benches/v2_take.rs b/benchmarks/random-access-bench/benches/v2_take.rs new file mode 100644 index 00000000000..e1c00b1995b --- /dev/null +++ b/benchmarks/random-access-bench/benches/v2_take.rs @@ -0,0 +1,126 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Each random-access query with the default and V2 scans, over files in the page cache. +//! +//! Run `random-access-bench --prepare-data --formats vortex` first to write the files. + +#![expect(clippy::unwrap_used)] + +use std::hint::black_box; +use std::sync::LazyLock; + +use divan::Bencher; +use mimalloc::MiMalloc; +use random_access_bench::AccessPattern; +use random_access_bench::generate_indices; +use tokio::runtime::Runtime; +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; +use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; +use vortex_bench::Format; +use vortex_bench::SESSION; +use vortex_bench::datasets::feature_vectors::FeatureVectorsData; +use vortex_bench::datasets::nested_structs::NestedStructsData; +use vortex_bench::datasets::taxi_data::TaxiData; +use vortex_bench::random_access::BenchDataset; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +static RUNTIME: LazyLock = LazyLock::new(|| Runtime::new().unwrap()); + +/// The queries V2 can scan: nested lists use list layouts, which V2 does not support. +const QUERIES: &[&str] = &[ + "taxi/correlated", + "taxi/uniform", + "feature-vectors/correlated", + "feature-vectors/uniform", + "nested-structs/correlated", + "nested-structs/uniform", +]; + +fn main() { + divan::main(); +} + +struct Query { + file: VortexFile, + indices: Vec, +} + +impl Query { + fn open(query: &str) -> Self { + let (dataset, pattern) = query.split_once('/').unwrap(); + let dataset: Box = match dataset { + "taxi" => Box::new(TaxiData), + "feature-vectors" => Box::new(FeatureVectorsData), + _ => Box::new(NestedStructsData), + }; + let pattern = match pattern { + "correlated" => AccessPattern::Correlated, + _ => AccessPattern::Uniform, + }; + let indices = generate_indices(dataset.as_ref(), pattern); + let file = RUNTIME.block_on(async { + let path = dataset.path(Format::OnDiskVortex).await.unwrap(); + SESSION + .open_options() + .with_layout_reader_cache() + .open_path(path) + .await + .unwrap() + }); + Self { file, indices } + } + + fn indices(&self) -> StrictSortedBuffer { + StrictSortedBuffer::try_new(Buffer::from(self.indices.clone())).unwrap() + } + + fn take_v1(&self) { + let array = RUNTIME.block_on( + self.file + .scan() + .unwrap() + .with_row_indices(self.indices()) + .into_array_stream() + .unwrap() + .read_all(), + ); + black_box(array.unwrap()); + } + + fn scan_v2(&self, builder: impl FnOnce(ScanBuilder) -> ScanBuilder) { + let stream = builder(ScanBuilder::new(SESSION.clone(), self.file.scan_file())) + .into_stream() + .unwrap(); + let array = + RUNTIME.block_on(ArrayStreamAdapter::new(self.file.dtype().clone(), stream).read_all()); + black_box(array.unwrap()); + } +} + +/// The default scan, the target V2 should beat. +#[divan::bench(args = QUERIES)] +fn take_v1(bencher: Bencher, query: &str) { + let query = Query::open(query); + bencher.bench_local(|| query.take_v1()); +} + +#[divan::bench(args = QUERIES)] +fn take_v2(bencher: Bencher, query: &str) { + let query = Query::open(query); + bencher.bench_local(|| query.scan_v2(|scan| scan.with_row_indices(query.indices()))); +} + +/// A V2 scan that selects no rows: its per-scan cost before any split is planned. +#[divan::bench(args = QUERIES)] +fn setup_v2(bencher: Bencher, query: &str) { + let query = Query::open(query); + bencher.bench_local(|| query.scan_v2(|scan| scan.with_row_range(0..0))); +} diff --git a/benchmarks/random-access-bench/src/lib.rs b/benchmarks/random-access-bench/src/lib.rs index 357e688ed53..8cf6bb590d2 100644 --- a/benchmarks/random-access-bench/src/lib.rs +++ b/benchmarks/random-access-bench/src/lib.rs @@ -72,7 +72,7 @@ const POISSON_EXPECTED_COUNT: usize = 100; const CACHED_WARMUP_DURATION: Duration = Duration::from_secs(1); /// Generate indices for the given dataset and access pattern. -fn generate_indices(dataset: &dyn BenchDataset, pattern: AccessPattern) -> Vec { +pub fn generate_indices(dataset: &dyn BenchDataset, pattern: AccessPattern) -> Vec { let row_count = dataset.row_count(); let mut rng = StdRng::seed_from_u64(42); diff --git a/vortex-bench/src/random_access/take.rs b/vortex-bench/src/random_access/take.rs index 97d1a005778..095c4d4d5c5 100644 --- a/vortex-bench/src/random_access/take.rs +++ b/vortex-bench/src/random_access/take.rs @@ -128,30 +128,13 @@ impl RandomAccessor for ArrowIpcRandomAccessor { } } -/// Random accessor for Vortex format files. +/// Random accessor for Vortex format files, read with the V2 scan. /// -/// The file handle is opened at construction time and reused across `take()` calls. Files the -/// V2 scan supports are read with it; others, such as files with list columns, use the default -/// scan. +/// The file handle is opened at construction time and reused across `take()` calls. pub struct VortexRandomAccessor { name: String, format: Format, file: VortexFile, - v2: bool, -} - -/// Whether the V2 scan can read `file`. It checks the file's layouts before reading any rows. -async fn supports_v2(file: &VortexFile) -> bool { - let Ok(stream) = ScanBuilder::new(SESSION.clone(), file.scan_file()) - .with_row_range(0..0) - .into_stream() - else { - return false; - }; - ArrayStreamAdapter::new(file.dtype().clone(), stream) - .read_all() - .await - .is_ok() } impl VortexRandomAccessor { @@ -161,15 +144,10 @@ 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, - v2: supports_v2(&file).await, file, }) } @@ -183,13 +161,11 @@ 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, - v2: supports_v2(&file).await, file, }) } @@ -207,21 +183,12 @@ impl RandomAccessor for VortexRandomAccessor { async fn take(&self, indices: &[u64]) -> anyhow::Result { let indices = StrictSortedBuffer::try_new(Buffer::from(indices.to_vec()))?; - let array = if self.v2 { - let stream = ScanBuilder::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(indices) + .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-layout/src/layouts/chunked/mod.rs b/vortex-layout/src/layouts/chunked/mod.rs index 459429046c2..ecf614f092d 100644 --- a/vortex-layout/src/layouts/chunked/mod.rs +++ b/vortex-layout/src/layouts/chunked/mod.rs @@ -117,6 +117,11 @@ impl Layout { .into_typed() } + /// The first row of each chunk, followed by the layout's row count. + pub fn chunk_offsets(&self) -> &[u64] { + &self.data().chunk_offsets + } + /// Rebuild this layout with owned children. pub fn with_children(&self, children: Vec) -> Self { Self::new( diff --git a/vortex-layout/src/plan/physical/children.rs b/vortex-layout/src/plan/physical/children.rs index bfa1541b87c..21a1f8e73a1 100644 --- a/vortex-layout/src/plan/physical/children.rs +++ b/vortex-layout/src/plan/physical/children.rs @@ -30,7 +30,7 @@ impl PlanChildren { ) -> Self { Self { initializer: Some(Arc::new(initializer)), - cache: (0..len).map(|_| OnceCell::new()).collect::>().into(), + cache: (0..len).map(|_| OnceCell::new()).collect(), } } diff --git a/vortex-layout/src/plan/physical/lower.rs b/vortex-layout/src/plan/physical/lower.rs index 1ea284cc1ca..bba7266093c 100644 --- a/vortex-layout/src/plan/physical/lower.rs +++ b/vortex-layout/src/plan/physical/lower.rs @@ -83,22 +83,16 @@ 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"))?; - } + let nchunks = layout.nchildren(); + let chunks = layout.to_layout(); // 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. + // computes the chunk offsets with checked addition, ending at the parent's row count. Ok(unsafe { ConcatPlan::from_children_unchecked( layout.dtype().clone(), layout.row_count(), - row_offsets.into(), - lazy_children(layout.to_layout(), (0..layout.nchildren()).collect()), + layout.chunk_offsets()[..nchunks].into(), + PlanChildren::lazy(nchunks, move |index| lower_slot(&chunks, index)), ) }) } @@ -159,13 +153,17 @@ fn lazy_children(layout: LayoutRef, slots: Vec) -> PlanChildren { .get(index) .copied() .ok_or_else(|| vortex_err!("Missing plan child slot {index}"))?; - let child = layout - .slot(slot)? - .ok_or_else(|| vortex_err!("Layout child slot {slot} is absent"))?; - lower(&child) + lower_slot(&layout, slot) }) } +fn lower_slot(layout: &LayoutRef, slot: usize) -> VortexResult { + let child = layout + .slot(slot)? + .ok_or_else(|| vortex_err!("Layout child slot {slot} is absent"))?; + lower(&child) +} + fn lower_zoned(layout: &LayoutRef) -> VortexResult { // Zoned and legacy stats layouts share a child shape: transparent data, auxiliary zones. let metadata = if let Some(layout) = layout.as_opt::() { diff --git a/vortex-layout/src/plan/physical/pipeline/scan.rs b/vortex-layout/src/plan/physical/pipeline/scan.rs index 9d650ebbd26..8e9726fcd3a 100644 --- a/vortex-layout/src/plan/physical/pipeline/scan.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan.rs @@ -671,15 +671,14 @@ 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); + // Only the rows some split reads can share a segment between readers, so the rows + // between splits are never visited. match &root { - Root::Plan(plan) => core.shares.add(plan, start..end)?, + Root::Plan(plan) => { + for split in &splits { + core.shares.add(plan, split.rows.clone())?; + } + } } core.shares.retain_shared(); core.dirty_flags = vec![false; 1]; diff --git a/vortex-layout/src/scan/execution/adaptor.rs b/vortex-layout/src/scan/execution/adaptor.rs index fc6053d55ca..543a06f4ec1 100644 --- a/vortex-layout/src/scan/execution/adaptor.rs +++ b/vortex-layout/src/scan/execution/adaptor.rs @@ -28,6 +28,8 @@ use crate::segments::SegmentId; /// Rows a running plan produced, in row order. #[derive(Clone)] pub(super) struct Piece { + /// The scan split that produced the piece. Splits are in row order. + pub(super) split: usize, /// Orders the piece among the outputs of one scan. pub(super) rows: Range, /// The selected rows the piece holds. @@ -96,7 +98,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::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)?]; @@ -130,10 +132,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/mod.rs b/vortex-layout/src/scan/execution/mod.rs index 9d8b8897ef3..4c054c090aa 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -15,13 +15,13 @@ use std::sync::Arc; use vortex_buffer::Alignment; use vortex_io::request::IoTarget; -use vortex_mask::Mask; use vortex_scan::planning::planner::Planner; use vortex_scan::planning::planner::WorkScope; use vortex_session::VortexSession; use self::projection::ProjectionPlanner; use crate::plan::physical::PlanRef; +use crate::plan::physical::pipeline::Split; use crate::plan::physical::pipeline::ValueCache; /// Where a segment's bytes live in its source. @@ -66,11 +66,18 @@ pub(super) struct ScanPlans { struct SelectedRows { /// The split the rows belong to. scope: WorkScope, - /// Which rows of `scope.rows` are selected. - mask: Mask, + /// The pipeline splits producing the selected rows of `scope.rows`, in row order. + splits: Vec, } -/// 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 })) +/// The work for one split: project the rows of `scope` that `splits` select. +pub(super) fn plan_projection( + plans: ScanPlans, + scope: WorkScope, + splits: Vec, +) -> Box { + Box::new(ProjectionPlanner::new( + plans, + SelectedRows { scope, splits }, + )) } diff --git a/vortex-layout/src/scan/execution/projection.rs b/vortex-layout/src/scan/execution/projection.rs index b6180725c27..3f60aa482a4 100644 --- a/vortex-layout/src/scan/execution/projection.rs +++ b/vortex-layout/src/scan/execution/projection.rs @@ -17,7 +17,6 @@ use vortex_scan::planning::planner::PlannerOutput; use vortex_scan::planning::planner::State; use crate::plan::physical::pipeline::Scan; -use crate::plan::physical::pipeline::Split; use crate::scan::execution::ScanPlans; use crate::scan::execution::SelectedRows; use crate::scan::execution::adaptor::AdaptorStep; @@ -55,7 +54,7 @@ impl Planner for ProjectionPlanner { let Some((plans, selected)) = self.work.take() else { return Ok(PlannerOutput::Done); }; - if selected.mask.all_false() { + if selected.splits.is_empty() { return Ok(PlannerOutput::Done); } Ok(PlannerOutput::Morsel( @@ -91,7 +90,9 @@ impl ProjectionMorsel { fn finish(&mut self) -> VortexResult { self.done = true; - self.pieces.sort_by_key(|piece| piece.rows.start); + // Splits run at once, so their pieces arrive interleaved. + 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) @@ -135,16 +136,13 @@ impl Morsel for ProjectionMorsel { .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); + let splits = std::mem::take(&mut self.selected.splits); + // Every split is active at once, so all of their reads are published together. + let active = splits.len(); + let scan = Scan::try_new(plans.session, plans.projection, splits)? + .with_max_active(active) + .with_row_offset(plans.row_offset) + .with_value_cache(plans.values); self.adaptor = Some(ScanAdaptor::new(scan, plans.locations)); return Ok(MorselOutput::Continue); }; diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index d7724f41544..a1d23936d0c 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -9,16 +9,19 @@ //! file's [`SegmentSource`]. The layout reader is never asked to evaluate anything. //! //! 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. +//! directly and supports projection, row ranges, and selections. Layouts without a pipeline +//! implementation 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. A partial selection is instead planned by the stored segments its rows read, so each +//! segment is read once and rows between selected runs are never visited. mod io; mod split; #[cfg(test)] mod tests; +use std::cmp::Ordering; use std::ops::Range; use std::sync::Arc; @@ -36,9 +39,9 @@ use vortex_array::VortexSessionExecute; use vortex_array::dtype::DType; use vortex_array::expr::BoundExpression; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; use vortex_io::request::IoService; use vortex_io::session::RuntimeSessionExt; +use vortex_mask::AllOr; use vortex_scan::planning::planner::WorkScope; use vortex_scan::selection::Selection; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; @@ -46,14 +49,10 @@ 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::pipeline::Reach; +use crate::plan::physical::pipeline::Split; use crate::plan::physical::plan_row_idx_expression; use crate::scan::arrow::to_record_batch; use crate::scan::execution::ScanPlans; @@ -62,11 +61,15 @@ use crate::scan::execution::plan_projection; use crate::scan::v2::io::SegmentScanIo; use crate::scan::v2::io::segment_ranges; use crate::scan::v2::split::SplitPlan; +use crate::segments::SegmentId; use crate::segments::SegmentSource; const ROW_BLOCK_SIZE: usize = 8192; const MAX_ROWS_PER_CHUNK: usize = 8 * ROW_BLOCK_SIZE; const TARGET_BYTES_PER_CHUNK: u64 = 8 * 1024 * 1024; +/// A chunk selected in more runs than this is planned whole: finding the segments of each run +/// would cost more than the chunks it could skip. +const MAX_RUNS: usize = 64; /// Limit estimated bytes per task without changing the writer-aligned row granularity. fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize { @@ -82,6 +85,19 @@ fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize { blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE } +/// Whether two sorted segment lists have a segment in common. +fn shares_any(left: &[SegmentId], right: &[SegmentId]) -> bool { + let (mut left, mut right) = (left.iter().peekable(), right.iter().peekable()); + while let (Some(l), Some(r)) = (left.peek(), right.peek()) { + match l.cmp(r) { + Ordering::Less => _ = left.next(), + Ordering::Greater => _ = right.next(), + Ordering::Equal => return true, + } + } + false +} + /// Fixed boundaries in the file's row domain, clipped to the requested range. fn row_chunks(rows: Range, rows_per_chunk: usize) -> Vec> { if rows.is_empty() { @@ -109,8 +125,9 @@ 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. A layout with no pipeline +/// implementation returns an error when the first split that reads it is planned. Filters and row +/// limits are not supported. pub struct ScanBuilder { session: VortexSession, file: ScanFile, @@ -237,34 +254,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 +273,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 @@ -326,25 +311,84 @@ impl PreparedScan { }) } - /// The splits with selected rows, each ready to be admitted to a planning run. - fn split_plans(&self) -> Vec { - self.splits - .iter() - .filter_map(|rows| { - let row_mask = self.selection.row_mask(rows); - let mask = row_mask.mask(); - // Only the span from the first selected row to the last is planned. - let span = mask.first()?..mask.last()? + 1; + /// The rows each planner projects, as the pipeline splits that produce them, in row order. + /// + /// A fully selected chunk, or one selected in more than [`MAX_RUNS`] runs, is one split of + /// its own planner. Otherwise each selected run joins the planner before it, even across + /// chunk boundaries, while the two read a stored segment in common, so no stored segment is + /// read or decoded by two planners. Within a planner, a run extends the split before it when + /// both read exactly the same segments: every row between them then lies in those segments' + /// chunks, so the split visits no chunk the selection skips. + /// + /// A column's segments hold contiguous rows, so a run can only share a segment with an + /// earlier run of its planner if it shares one with the run just before it. + fn split_rows(&self) -> VortexResult>>> { + let mut planners: Vec>> = Vec::new(); + // The segments the last selected run reads, or none once its planner is closed. + let mut previous = Vec::new(); + for rows in &self.splits { + let row_mask = self.selection.row_mask(rows); + let mask = row_mask.mask(); + let slices = match mask.slices() { + AllOr::None => continue, + AllOr::Some(slices) if slices.len() <= MAX_RUNS => slices, + AllOr::All | AllOr::Some(_) => { + let span = mask.first().unwrap_or(0)..mask.last().map_or(0, |last| last + 1); + let split = rows.start + span.start as u64..rows.start + span.end as u64; + planners.push(vec![split]); + previous.clear(); + continue; + } + }; + for &(start, end) in slices { + let run = rows.start + start as u64..rows.start + end as u64; + let mut segments = Vec::new(); + self.plans.projection.reach( + run.clone(), + &Reach::Offset(0), + &mut |segment, _| segments.push(segment), + )?; + segments.sort_unstable(); + segments.dedup(); + match planners.last_mut() { + Some(splits) if segments == previous && !segments.is_empty() => { + if let Some(split) = splits.last_mut() { + split.end = run.end; + } + } + Some(splits) if shares_any(&previous, &segments) => splits.push(run), + _ => planners.push(vec![run]), + } + previous = segments; + } + } + Ok(planners) + } + + /// The planners with selected rows, each ready to be admitted to a planning run. + fn split_plans(&self) -> VortexResult> { + Ok(self + .split_rows()? + .into_iter() + .filter_map(|ranges| { + let rows = ranges.first()?.start..ranges.last()?.end; + let splits = ranges + .into_iter() + .map(|rows| Split { + mask: self.selection.row_mask(&rows).mask().clone(), + rows, + }) + .collect(); let scope = WorkScope { file_ordinal: 0, - rows: rows.start + span.start as u64..rows.start + span.end as u64, + rows, }; Some(SplitPlan { - root: plan_projection(self.plans.clone(), scope.clone(), mask.slice(span)), + root: plan_projection(self.plans.clone(), scope.clone(), splits), scope, }) }) - .collect() + .collect()) } /// Runs every split as a task on the session's runtime and yields their mapped batches. @@ -354,7 +398,7 @@ impl PreparedScan { fn execute_stream(self) -> VortexResult>> { let concurrency = self.concurrency * get_available_parallelism().unwrap_or(1); let handle = self.session.handle(); - let split_plans = self.split_plans(); + let split_plans = self.split_plans()?; let io = self.io; let map_fn = self.map_fn; // Spawned as the buffer pulls them, so at most `concurrency` splits read at once. diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index c7115a213c1..a91e7dd94e4 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -385,7 +385,7 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { }); let segments = Arc::clone(&counted) as Arc; let scan = PreparedScan::try_new(builder(&session, &segments, &layout)?)?; - let splits = scan.split_plans(); + let splits = scan.split_plans()?; assert_eq!( splits.len(), (STORED_CHUNKS * STORED_CHUNK_ROWS).div_ceil(MAX_ROWS_PER_CHUNK) @@ -449,9 +449,12 @@ async fn counted_ints(session: &VortexSession) -> VortexResult<(Arc, #[case] scopes: Vec>, #[case] reads: usize, @@ -463,7 +466,13 @@ async fn work_chunks_plan_only_their_selected_span( builder(&session, &segments, &layout)? .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(rows.clone()))?), )?; - let splits = scan.split_plans(); + // Runs reading the same stored chunk are one pipeline split, which visits no other chunk. + assert_eq!( + scan.split_rows()?.concat(), + scopes, + "each planner here reads one stored chunk" + ); + let splits = scan.split_plans()?; assert_eq!( splits .iter() @@ -475,7 +484,7 @@ async fn work_chunks_plan_only_their_selected_span( for split in splits { arrays.extend(split.run(Arc::clone(&scan.io)).await?); } - // Each work chunk reads only the stored chunks holding its selected rows. + // Splits read only the stored chunks holding their selected rows, each chunk once. assert_eq!(counted.reads.load(Ordering::Relaxed), reads); let expected = rows .iter() @@ -489,7 +498,8 @@ async fn work_chunks_plan_only_their_selected_span( Ok(()) } -/// V2 rejects list layouts: their element ranges are only known once their offsets are read. +/// V2 rejects list layouts when it plans them: their element ranges are only known once their +/// offsets are read. #[tokio::test] async fn list_layouts_return_an_error() -> VortexResult<()> { let session = new_session().with_tokio(); @@ -501,7 +511,7 @@ async fn list_layouts_return_an_error() -> VortexResult<()> { .try_next() .await .expect_err("V2 must reject an unsupported layout"); - assert!(error.to_string().contains("V2 scans do not support")); + assert!(error.to_string().contains("has no pipeline implementation")); Ok(()) } From 38cc71cc992f8669de701c8428e6b1a7e4a844d2 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 19 Oct 2026 23:53:54 +0000 Subject: [PATCH 09/15] Find V2 segment reads by search when locations are sorted Each V2 scan built a hash map from every segment's byte range to its id, so the IO source could answer range requests. Writers lay segments out in offset order, so the file's location table is searched in place instead, falling back to the map when it is not sorted. An empty scan of the feature-vectors file, 3,917 segments, drops from 85 to 59 us. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_014rtiSjzL4EauVCM7NS84hj --- vortex-layout/src/scan/v2/io.rs | 87 ++++++++++++++++++++++++++------ vortex-layout/src/scan/v2/mod.rs | 4 +- 2 files changed, 73 insertions(+), 18 deletions(-) diff --git a/vortex-layout/src/scan/v2/io.rs b/vortex-layout/src/scan/v2/io.rs index ce565846d7e..c62cc670dc7 100644 --- a/vortex-layout/src/scan/v2/io.rs +++ b/vortex-layout/src/scan/v2/io.rs @@ -34,15 +34,41 @@ 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()) +/// Finds a file's segment by its byte range, to answer a range request with the segment. +#[derive(Clone)] +pub(super) enum SegmentRanges { + /// Every location, indexed by segment id and in offset order, as writers lay segments out. + /// Searched in place, so a scan builds nothing per segment. + Sorted(Arc<[SegmentLocation]>), + /// Segment ids keyed by `(offset, length)`, for locations in any other order. + Keyed(Arc>), +} + +impl SegmentRanges { + pub(super) fn new(locations: Arc<[SegmentLocation]>) -> Self { + if locations.is_sorted_by_key(|location| location.offset) { + return Self::Sorted(locations); + } + 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) + }); + Self::Keyed(Arc::new(ranges.collect())) + } + + fn find(&self, offset: u64, len: usize) -> Option { + match self { + Self::Sorted(locations) => { + let first = locations.partition_point(|location| location.offset < offset); + let index = locations[first..] + .iter() + .take_while(|location| location.offset == offset) + .position(|location| location.length as usize == len)?; + u32::try_from(first + index).ok().map(SegmentId::from) + } + Self::Keyed(ranges) => ranges.get(&(offset, len)).copied(), + } + } } /// Serves splits from any [`SegmentSource`], by segment id. @@ -61,7 +87,7 @@ impl IoService for SegmentScanIo { fn session(&self) -> Arc { Arc::new(SegmentIoSource::new( Arc::clone(&self.segments), - Arc::clone(&self.ranges), + self.ranges.clone(), )) } } @@ -114,9 +140,9 @@ impl IoSource for SegmentIoSource { request.target ); }; - let segment = *self + let segment = self .ranges - .get(&(offset, len)) + .find(offset, len) .ok_or_else(|| vortex_err!("No segment at bytes {offset}+{len}"))?; if request.intent == IoIntent::Announce { self.announced @@ -247,7 +273,7 @@ mod tests { use vortex_io::request::IoTarget; use super::SegmentIoSource; - use super::segment_ranges; + use super::SegmentRanges; use crate::scan::execution::SegmentLocation; use crate::segments::SegmentFuture; use crate::segments::SegmentId; @@ -265,6 +291,29 @@ mod tests { } } + fn location(offset: u64, length: u32) -> SegmentLocation { + SegmentLocation { + offset, + length, + alignment: Alignment::none(), + } + } + + /// Ranges are found by search when locations are in offset order, where a zero-length + /// segment can share its offset with the next, and by key otherwise. + #[test] + fn segment_ranges_find_each_segment() { + let sorted = SegmentRanges::new([location(0, 4), location(4, 0), location(4, 8)].into()); + assert!(matches!(sorted, SegmentRanges::Sorted(_))); + let unsorted = SegmentRanges::new([location(4, 0), location(4, 8), location(0, 4)].into()); + assert!(matches!(unsorted, SegmentRanges::Keyed(_))); + for (ranges, ids) in [(sorted, [0, 1, 2]), (unsorted, [2, 0, 1])] { + let found = [(0, 4), (4, 0), (4, 8)].map(|(offset, len)| ranges.find(offset, len)); + assert_eq!(found, ids.map(|id| Some(SegmentId::from(id)))); + assert_eq!(ranges.find(4, 4), None); + } + } + fn request(intent: IoIntent, id: u32) -> IoRequest { IoRequest { intent, @@ -301,7 +350,7 @@ mod tests { length: 4, alignment: Alignment::none(), }); - let io = SegmentIoSource::new(segments, segment_ranges(&locations)); + let io = SegmentIoSource::new(segments, SegmentRanges::new(locations.into())); 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 +383,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 _, + SegmentRanges::new([location].into()), + ); io.submit(IoOwnerId(0), vec![request(IoIntent::Prefetch, 0)])?; assert_eq!(segments.0.load(Ordering::Relaxed), 1); @@ -362,7 +414,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 _, + SegmentRanges::new([location].into()), + ); io.submit( IoOwnerId(0), vec![ diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index a1d23936d0c..b2de13f8c43 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -58,8 +58,8 @@ use crate::scan::arrow::to_record_batch; use crate::scan::execution::ScanPlans; pub use crate::scan::execution::SegmentLocation; use crate::scan::execution::plan_projection; +use crate::scan::v2::io::SegmentRanges; use crate::scan::v2::io::SegmentScanIo; -use crate::scan::v2::io::segment_ranges; use crate::scan::v2::split::SplitPlan; use crate::segments::SegmentId; use crate::segments::SegmentSource; @@ -291,7 +291,7 @@ impl PreparedScan { let splits = row_chunks(builder.row_range.unwrap_or(0..row_count), chunk_rows); let io = Arc::new(SegmentScanIo::new( file.segments, - segment_ranges(&file.locations), + SegmentRanges::new(Arc::clone(&file.locations)), )); Ok(Self { session: builder.session.clone(), From 85613cffb333cc163f6b7f093ea2b70582524b13 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 20 Oct 2026 00:07:30 +0000 Subject: [PATCH 10/15] Pin tpchgen-cli and use its parquet subcommand tpchgen-cli 4.0.0, released 2026-10-19, removed the --format and --parquet-compression flags the SQL logic test data script passed to an unpinned `uvx tpchgen-cli`, so the job failed before any test ran. Pin 4.0.0 and use `parquet --compression='ZSTD(3)'`, which writes ZSTD column chunks as before. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_014rtiSjzL4EauVCM7NS84hj --- AGENTS.md | 2 +- vortex-sqllogictest/slt/tpch/generate_data.sh | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index c92cb4d2ff9..f38f7998935 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -223,7 +223,7 @@ These CI checks are the ones most often missed when adding files rather than edi default to Snappy. - parquet-rs: `Compression::ZSTD(ZstdLevel::try_new(3)?)` - DuckDB: `(FORMAT parquet, COMPRESSION zstd, COMPRESSION_LEVEL 3)` - - tpchgen-cli: `--parquet-compression='ZSTD(3)'` + - tpchgen-cli: `tpchgen-cli parquet --compression='ZSTD(3)'` - Benchmarks that deliberately compare Parquet codecs, such as the GPU compress benchmark's `--codec`, are exempt. diff --git a/vortex-sqllogictest/slt/tpch/generate_data.sh b/vortex-sqllogictest/slt/tpch/generate_data.sh index 29a0beda69a..9779d3d083c 100755 --- a/vortex-sqllogictest/slt/tpch/generate_data.sh +++ b/vortex-sqllogictest/slt/tpch/generate_data.sh @@ -20,7 +20,8 @@ DATA_DIR="${SCRIPT_DIR}/data" mkdir -p "${DATA_DIR}" echo "Generating TPC-H data (SF=0.1)..." -uvx tpchgen-cli -s 0.1 --format=parquet --parquet-compression='ZSTD(3)' --output-dir "${DATA_DIR}/" +# Pinned: tpchgen-cli 4.0.0 replaced --format with a subcommand, breaking the unpinned call. +uvx tpchgen-cli@4.0.0 parquet -s 0.1 --compression='ZSTD(3)' --output-dir "${DATA_DIR}/" for f in "${DATA_DIR}"/*.parquet; do echo "Converting $(basename "$f") to Vortex..." From 2dab8888e71c9019eee266ca2583618f1da93749 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 20 Oct 2026 00:32:15 +0000 Subject: [PATCH 11/15] Bound V2 planners and share segment reads between them CI showed taxi random access regressing: a column whose segments span most of the file made every selected run share a segment, so the whole take became one planner on one thread. A planner now grows only while it spans at most 256K rows, so taxi/uniform runs as 11 planners. A segment that several planners still reach is read once: the scan's IO counts each segment's planners and hands every planner the same read, dropping it once the last has taken it. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_014rtiSjzL4EauVCM7NS84hj --- vortex-layout/src/scan/v2/io.rs | 99 ++++++++++++++++++++++++++++-- vortex-layout/src/scan/v2/mod.rs | 58 +++++++++++++---- vortex-layout/src/scan/v2/tests.rs | 6 +- 3 files changed, 142 insertions(+), 21 deletions(-) diff --git a/vortex-layout/src/scan/v2/io.rs b/vortex-layout/src/scan/v2/io.rs index c62cc670dc7..b0f954c2c5b 100644 --- a/vortex-layout/src/scan/v2/io.rs +++ b/vortex-layout/src/scan/v2/io.rs @@ -75,11 +75,31 @@ impl SegmentRanges { pub(super) struct SegmentScanIo { segments: Arc, ranges: SegmentRanges, + shared: SharedReads, } impl SegmentScanIo { pub(super) fn new(segments: Arc, ranges: SegmentRanges) -> Self { - Self { segments, ranges } + Self { + segments, + ranges, + shared: SharedReads::default(), + } + } + + /// Reads each segment in `readers` once for as many sessions as it names, rather than once + /// per session. + pub(super) fn share(&self, readers: HashMap) { + let mut shared = self.shared.lock(); + for (segment, readers) in readers { + shared.insert( + segment, + SharedRead { + readers, + read: None, + }, + ); + } } } @@ -88,10 +108,21 @@ impl IoService for SegmentScanIo { Arc::new(SegmentIoSource::new( Arc::clone(&self.segments), self.ranges.clone(), + Arc::clone(&self.shared), )) } } +/// Reads several sessions of one scan make of the same segment, by segment id. +type SharedReads = Arc>>; + +/// A segment read for several sessions: started by the first and dropped once each has taken it. +struct SharedRead { + /// Sessions still to take the read. + readers: usize, + read: Option, +} + /// Serves a split's range requests from a segment source, without blocking. /// /// Each requested range names one segment. [`submit`](IoSource::submit) starts its read at once, @@ -104,6 +135,7 @@ impl IoService for SegmentScanIo { pub(super) struct SegmentIoSource { segments: Arc, ranges: SegmentRanges, + shared: SharedReads, reads: Mutex>>, ready: Mutex>, /// Reads started by prefetches, consumed by the next fetch of their segment. @@ -116,16 +148,41 @@ pub(super) struct SegmentIoSource { type Prefetched = Shared>>>; impl SegmentIoSource { - pub(super) fn new(segments: Arc, ranges: SegmentRanges) -> Self { + fn new(segments: Arc, ranges: SegmentRanges, shared: SharedReads) -> Self { Self { segments, ranges, + shared, reads: Mutex::new(FuturesUnordered::new()), ready: Mutex::default(), prefetched: Mutex::default(), announced: Mutex::default(), } } + + /// Requests `segment`, from the read other sessions share when the scan shares it. + fn request(&self, segment: SegmentId) -> SegmentFuture { + let mut shared = self.shared.lock(); + let Some(entry) = shared.get_mut(&segment) else { + return self.segments.request(segment); + }; + let read = entry + .read + .get_or_insert_with(|| { + self.segments + .request(segment) + .map_err(Arc::new) + .boxed() + .shared() + }) + .clone(); + entry.readers = entry.readers.saturating_sub(1); + if entry.readers == 0 { + shared.remove(&segment); + } + read.map_err(move |err| vortex_err!("shared read of segment {segment} failed: {err}")) + .boxed() + } } impl IoSource for SegmentIoSource { @@ -149,7 +206,7 @@ impl IoSource for SegmentIoSource { .lock() .entry(segment) .or_default() - .push_back(self.segments.request(segment)); + .push_back(self.request(segment)); continue; } if request.intent == IoIntent::Prefetch { @@ -159,7 +216,7 @@ impl IoSource for SegmentIoSource { .lock() .get_mut(&segment) .and_then(VecDeque::pop_front) - .unwrap_or_else(|| self.segments.request(segment)) + .unwrap_or_else(|| self.request(segment)) .map_err(Arc::new) .boxed() .shared(); @@ -183,7 +240,7 @@ impl IoSource for SegmentIoSource { .lock() .get_mut(&segment) .and_then(VecDeque::pop_front) - .unwrap_or_else(|| self.segments.request(segment)), + .unwrap_or_else(|| self.request(segment)), }; let id = request.request; // In-memory and cached segments need no allocated asynchronous task. A pending @@ -269,11 +326,15 @@ mod tests { use vortex_io::request::IoOwnerId; use vortex_io::request::IoRequest; use vortex_io::request::IoRequestId; + use vortex_io::request::IoService; use vortex_io::request::IoSource; use vortex_io::request::IoTarget; + use vortex_utils::aliases::hash_map::HashMap; use super::SegmentIoSource; use super::SegmentRanges; + use super::SegmentScanIo; + use super::SharedReads; use crate::scan::execution::SegmentLocation; use crate::segments::SegmentFuture; use crate::segments::SegmentId; @@ -350,7 +411,11 @@ mod tests { length: 4, alignment: Alignment::none(), }); - let io = SegmentIoSource::new(segments, SegmentRanges::new(locations.into())); + let io = SegmentIoSource::new( + segments, + SegmentRanges::new(locations.into()), + SharedReads::default(), + ); let mut pending = request(IoIntent::Fetch, 1); pending.target = IoTarget::range(4, 4); io.submit(IoOwnerId(0), vec![pending, request(IoIntent::Fetch, 0)])?; @@ -373,6 +438,26 @@ mod tests { Ok(()) } + /// Sessions of a scan that share a segment read it once, and a later session reads it again + /// once every expected reader has taken it. + #[tokio::test] + async fn sessions_share_a_segment_read() -> VortexResult<()> { + let segments = Arc::new(CountingSegments::default()); + let io = SegmentScanIo::new( + Arc::clone(&segments) as _, + SegmentRanges::new([location(0, 4)].into()), + ); + io.share(HashMap::from_iter([(SegmentId::from(0), 2)])); + for reads in [1, 1, 2] { + let session = io.session(); + session.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, 0)])?; + let completion = std::future::poll_fn(|cx| session.poll_completion(cx)).await?; + assert!(completion.result.is_ok()); + assert_eq!(segments.0.load(Ordering::Relaxed), reads); + } + Ok(()) + } + /// A fetch of a prefetched segment waits on the prefetch's read rather than reading again, /// and only the fetch is delivered. #[tokio::test] @@ -386,6 +471,7 @@ mod tests { let io = SegmentIoSource::new( Arc::clone(&segments) as _, SegmentRanges::new([location].into()), + SharedReads::default(), ); io.submit(IoOwnerId(0), vec![request(IoIntent::Prefetch, 0)])?; @@ -417,6 +503,7 @@ mod tests { let io = SegmentIoSource::new( Arc::clone(&segments) as _, SegmentRanges::new([location].into()), + SharedReads::default(), ); io.submit( IoOwnerId(0), diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index b2de13f8c43..26133f26bab 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -39,13 +39,13 @@ use vortex_array::VortexSessionExecute; use vortex_array::dtype::DType; use vortex_array::expr::BoundExpression; use vortex_error::VortexResult; -use vortex_io::request::IoService; use vortex_io::session::RuntimeSessionExt; use vortex_mask::AllOr; use vortex_scan::planning::planner::WorkScope; use vortex_scan::selection::Selection; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; +use vortex_utils::aliases::hash_map::HashMap; use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutRef; @@ -70,6 +70,9 @@ const TARGET_BYTES_PER_CHUNK: u64 = 8 * 1024 * 1024; /// A chunk selected in more runs than this is planned whole: finding the segments of each run /// would cost more than the chunks it could skip. const MAX_RUNS: usize = 64; +/// A planner grows to share stored segments only while it spans at most this many rows, so a +/// segment spanning much of the file cannot join every selected run into one task. +const MAX_PLANNER_ROWS: u64 = 4 * MAX_ROWS_PER_CHUNK as u64; /// Limit estimated bytes per task without changing the writer-aligned row granularity. fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize { @@ -85,6 +88,9 @@ fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize { blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE } +/// The rows of one planner's pipeline splits, in row order. +type PlannerSplits = Vec>; + /// Whether two sorted segment lists have a segment in common. fn shares_any(left: &[SegmentId], right: &[SegmentId]) -> bool { let (mut left, mut right) = (left.iter().peekable(), right.iter().peekable()); @@ -259,7 +265,7 @@ struct PreparedScan { session: VortexSession, plans: ScanPlans, /// Serves the splits' reads. - io: Arc, + io: Arc, /// The scan's splits, over the root row domain, in row order. splits: Vec>, selection: Selection, @@ -315,17 +321,31 @@ impl PreparedScan { /// /// A fully selected chunk, or one selected in more than [`MAX_RUNS`] runs, is one split of /// its own planner. Otherwise each selected run joins the planner before it, even across - /// chunk boundaries, while the two read a stored segment in common, so no stored segment is - /// read or decoded by two planners. Within a planner, a run extends the split before it when + /// chunk boundaries, while the two read a stored segment in common and the planner spans at + /// most [`MAX_PLANNER_ROWS`], so a stored segment is decoded by one planner unless it spans + /// more rows than that. A segment several planners read is still read once: the scan's IO + /// shares the read between them. Within a planner, a run extends the split before it when /// both read exactly the same segments: every row between them then lies in those segments' /// chunks, so the split visits no chunk the selection skips. /// /// A column's segments hold contiguous rows, so a run can only share a segment with an /// earlier run of its planner if it shares one with the run just before it. - fn split_rows(&self) -> VortexResult>>> { - let mut planners: Vec>> = Vec::new(); + /// + /// Also returns how many planners read each segment that more than one of them reads. + fn split_rows(&self) -> VortexResult<(Vec, HashMap)> { + let mut planners: Vec = Vec::new(); // The segments the last selected run reads, or none once its planner is closed. let mut previous = Vec::new(); + // The segments the last planner reads, and how many planners read each segment. + let mut reached: Vec = Vec::new(); + let mut readers: HashMap = HashMap::default(); + let mut close = |reached: &mut Vec| { + reached.sort_unstable(); + reached.dedup(); + for segment in reached.drain(..) { + *readers.entry(segment).or_default() += 1; + } + }; for rows in &self.splits { let row_mask = self.selection.row_mask(rows); let mask = row_mask.mask(); @@ -337,6 +357,7 @@ impl PreparedScan { let split = rows.start + span.start as u64..rows.start + span.end as u64; planners.push(vec![split]); previous.clear(); + close(&mut reached); continue; } }; @@ -350,25 +371,38 @@ impl PreparedScan { )?; segments.sort_unstable(); segments.dedup(); + let joins = planners.last().is_some_and(|splits| { + splits + .first() + .is_some_and(|first| run.end - first.start <= MAX_PLANNER_ROWS) + && shares_any(&previous, &segments) + }); match planners.last_mut() { - Some(splits) if segments == previous && !segments.is_empty() => { + Some(splits) if joins && segments == previous => { if let Some(split) = splits.last_mut() { split.end = run.end; } } - Some(splits) if shares_any(&previous, &segments) => splits.push(run), - _ => planners.push(vec![run]), + Some(splits) if joins => splits.push(run), + _ => { + close(&mut reached); + planners.push(vec![run]); + } } + reached.extend_from_slice(&segments); previous = segments; } } - Ok(planners) + close(&mut reached); + readers.retain(|_, count| *count > 1); + Ok((planners, readers)) } /// The planners with selected rows, each ready to be admitted to a planning run. fn split_plans(&self) -> VortexResult> { - Ok(self - .split_rows()? + let (planners, shared) = self.split_rows()?; + self.io.share(shared); + Ok(planners .into_iter() .filter_map(|ranges| { let rows = ranges.first()?.start..ranges.last()?.end; diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index a91e7dd94e4..05597ac68a7 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -393,7 +393,7 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { for (index, split) in splits.into_iter().enumerate() { let rows = split.scope.rows.clone(); let before = counted.reads.load(Ordering::Relaxed); - let arrays = split.run(Arc::clone(&scan.io)).await?; + let arrays = split.run(Arc::clone(&scan.io) as _).await?; let reads = counted.reads.load(Ordering::Relaxed) - before; let code_chunks = (rows.end - 1) / STORED_CHUNK_ROWS as u64 - rows.start / STORED_CHUNK_ROWS as u64 + 1; @@ -468,7 +468,7 @@ async fn splits_read_each_stored_chunk_once( )?; // Runs reading the same stored chunk are one pipeline split, which visits no other chunk. assert_eq!( - scan.split_rows()?.concat(), + scan.split_rows()?.0.concat(), scopes, "each planner here reads one stored chunk" ); @@ -482,7 +482,7 @@ async fn splits_read_each_stored_chunk_once( ); let mut arrays = Vec::new(); for split in splits { - arrays.extend(split.run(Arc::clone(&scan.io)).await?); + arrays.extend(split.run(Arc::clone(&scan.io) as _).await?); } // Splits read only the stored chunks holding their selected rows, each chunk once. assert_eq!(counted.reads.load(Ordering::Relaxed), reads); From 1816ba7c6ad3927776aabbf4bf8e4160665ed51d Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 20 Oct 2026 07:52:52 +0000 Subject: [PATCH 12/15] Spawn the first V2 planner tasks when the scan is planned Planner tasks were spawned only when the stream's consumer first polled it, after the scan was handed back from the CPU pool. The first `concurrency` tasks now start on the CPU thread right after planning; the rest still start as earlier ones finish. In memory, taxi takes start their tasks 20 to 34 us sooner. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_014rtiSjzL4EauVCM7NS84hj --- vortex-layout/src/scan/v2/mod.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 26133f26bab..074f10b8d99 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -436,7 +436,7 @@ impl PreparedScan { let io = self.io; let map_fn = self.map_fn; // Spawned as the buffer pulls them, so at most `concurrency` splits read at once. - let tasks = stream::iter(split_plans).map(move |split| { + let mut tasks = split_plans.into_iter().map(move |split| { let io = Arc::clone(&io); let map_fn = Arc::clone(&map_fn); handle.spawn(async move { @@ -448,6 +448,9 @@ impl PreparedScan { .collect::>>() }) }); + // The first are spawned now, so they run while the stream is handed back to its consumer. + let started: Vec<_> = tasks.by_ref().take(concurrency).collect(); + let tasks = stream::iter(started).chain(stream::iter(tasks)); let batches = if self.ordered { tasks.buffered(concurrency).boxed() } else { From 72ba81994f0031521d9890e9e0ca44f1fe20820a Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 20 Oct 2026 08:14:17 +0000 Subject: [PATCH 13/15] Plan V2 planners lazily and spawn each as soon as it is planned The scan planned every planner, reaching every selected run, before it spawned any. Planners are now yielded one at a time: a planner is complete once the run after it is known, so it is spawned while the rest are still planned. Shared reads are registered between neighbouring planners as each closes, since a segment two planners read is read by planners next to each other. In memory, taxi/uniform's first planner task starts at about 370 us instead of 830 us. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_014rtiSjzL4EauVCM7NS84hj --- vortex-layout/src/scan/v2/io.rs | 24 +-- vortex-layout/src/scan/v2/mod.rs | 324 ++++++++++++++++++----------- vortex-layout/src/scan/v2/tests.rs | 12 +- 3 files changed, 226 insertions(+), 134 deletions(-) diff --git a/vortex-layout/src/scan/v2/io.rs b/vortex-layout/src/scan/v2/io.rs index b0f954c2c5b..dfa413080a2 100644 --- a/vortex-layout/src/scan/v2/io.rs +++ b/vortex-layout/src/scan/v2/io.rs @@ -87,18 +87,19 @@ impl SegmentScanIo { } } - /// Reads each segment in `readers` once for as many sessions as it names, rather than once - /// per session. - pub(super) fn share(&self, readers: HashMap) { + /// Expects one more session, the next planner's, to read each of `segments`, which the + /// session of the planner handed out before it reads too. Each segment is read once for all + /// the sessions expected to read it. + pub(super) fn share_with_next(&self, segments: impl IntoIterator) { let mut shared = self.shared.lock(); - for (segment, readers) in readers { - shared.insert( - segment, - SharedRead { - readers, + for segment in segments { + shared + .entry(segment) + .or_insert(SharedRead { + readers: 1, read: None, - }, - ); + }) + .readers += 1; } } } @@ -329,7 +330,6 @@ mod tests { use vortex_io::request::IoService; use vortex_io::request::IoSource; use vortex_io::request::IoTarget; - use vortex_utils::aliases::hash_map::HashMap; use super::SegmentIoSource; use super::SegmentRanges; @@ -447,7 +447,7 @@ mod tests { Arc::clone(&segments) as _, SegmentRanges::new([location(0, 4)].into()), ); - io.share(HashMap::from_iter([(SegmentId::from(0), 2)])); + io.share_with_next([SegmentId::from(0)]); for reads in [1, 1, 2] { let session = io.session(); session.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, 0)])?; diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 074f10b8d99..69c7d29e3f9 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -22,6 +22,7 @@ mod split; mod tests; use std::cmp::Ordering; +use std::collections::VecDeque; use std::ops::Range; use std::sync::Arc; @@ -29,9 +30,12 @@ use arrow_array::RecordBatch; use arrow_schema::ArrowError; use arrow_schema::Field; use arrow_schema::SchemaRef; +use futures::FutureExt; use futures::Stream; use futures::StreamExt; use futures::TryStreamExt; +use futures::future; +use futures::future::BoxFuture; use futures::stream; use futures::stream::BoxStream; use vortex_array::ArrayRef; @@ -45,7 +49,6 @@ use vortex_scan::planning::planner::WorkScope; use vortex_scan::selection::Selection; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; -use vortex_utils::aliases::hash_map::HashMap; use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutRef; @@ -317,112 +320,19 @@ impl PreparedScan { }) } - /// The rows each planner projects, as the pipeline splits that produce them, in row order. - /// - /// A fully selected chunk, or one selected in more than [`MAX_RUNS`] runs, is one split of - /// its own planner. Otherwise each selected run joins the planner before it, even across - /// chunk boundaries, while the two read a stored segment in common and the planner spans at - /// most [`MAX_PLANNER_ROWS`], so a stored segment is decoded by one planner unless it spans - /// more rows than that. A segment several planners read is still read once: the scan's IO - /// shares the read between them. Within a planner, a run extends the split before it when - /// both read exactly the same segments: every row between them then lies in those segments' - /// chunks, so the split visits no chunk the selection skips. - /// - /// A column's segments hold contiguous rows, so a run can only share a segment with an - /// earlier run of its planner if it shares one with the run just before it. - /// - /// Also returns how many planners read each segment that more than one of them reads. - fn split_rows(&self) -> VortexResult<(Vec, HashMap)> { - let mut planners: Vec = Vec::new(); - // The segments the last selected run reads, or none once its planner is closed. - let mut previous = Vec::new(); - // The segments the last planner reads, and how many planners read each segment. - let mut reached: Vec = Vec::new(); - let mut readers: HashMap = HashMap::default(); - let mut close = |reached: &mut Vec| { - reached.sort_unstable(); - reached.dedup(); - for segment in reached.drain(..) { - *readers.entry(segment).or_default() += 1; - } - }; - for rows in &self.splits { - let row_mask = self.selection.row_mask(rows); - let mask = row_mask.mask(); - let slices = match mask.slices() { - AllOr::None => continue, - AllOr::Some(slices) if slices.len() <= MAX_RUNS => slices, - AllOr::All | AllOr::Some(_) => { - let span = mask.first().unwrap_or(0)..mask.last().map_or(0, |last| last + 1); - let split = rows.start + span.start as u64..rows.start + span.end as u64; - planners.push(vec![split]); - previous.clear(); - close(&mut reached); - continue; - } - }; - for &(start, end) in slices { - let run = rows.start + start as u64..rows.start + end as u64; - let mut segments = Vec::new(); - self.plans.projection.reach( - run.clone(), - &Reach::Offset(0), - &mut |segment, _| segments.push(segment), - )?; - segments.sort_unstable(); - segments.dedup(); - let joins = planners.last().is_some_and(|splits| { - splits - .first() - .is_some_and(|first| run.end - first.start <= MAX_PLANNER_ROWS) - && shares_any(&previous, &segments) - }); - match planners.last_mut() { - Some(splits) if joins && segments == previous => { - if let Some(split) = splits.last_mut() { - split.end = run.end; - } - } - Some(splits) if joins => splits.push(run), - _ => { - close(&mut reached); - planners.push(vec![run]); - } - } - reached.extend_from_slice(&segments); - previous = segments; - } + /// Plans the scan's planners one at a time, in row order. + fn planners(&self) -> Planners { + Planners { + plans: self.plans.clone(), + selection: self.selection.clone(), + io: Arc::clone(&self.io), + chunks: self.splits.clone().into_iter(), + runs: VecDeque::new(), + whole: None, + splits: Vec::new(), + last: Vec::new(), + reached: Vec::new(), } - close(&mut reached); - readers.retain(|_, count| *count > 1); - Ok((planners, readers)) - } - - /// The planners with selected rows, each ready to be admitted to a planning run. - fn split_plans(&self) -> VortexResult> { - let (planners, shared) = self.split_rows()?; - self.io.share(shared); - Ok(planners - .into_iter() - .filter_map(|ranges| { - let rows = ranges.first()?.start..ranges.last()?.end; - let splits = ranges - .into_iter() - .map(|rows| Split { - mask: self.selection.row_mask(&rows).mask().clone(), - rows, - }) - .collect(); - let scope = WorkScope { - file_ordinal: 0, - rows, - }; - Some(SplitPlan { - root: plan_projection(self.plans.clone(), scope.clone(), splits), - scope, - }) - }) - .collect()) } /// Runs every split as a task on the session's runtime and yields their mapped batches. @@ -432,23 +342,31 @@ impl PreparedScan { fn execute_stream(self) -> VortexResult>> { let concurrency = self.concurrency * get_available_parallelism().unwrap_or(1); let handle = self.session.handle(); - let split_plans = self.split_plans()?; + let planners = self.planners(); let io = self.io; let map_fn = self.map_fn; - // Spawned as the buffer pulls them, so at most `concurrency` splits read at once. - let mut tasks = split_plans.into_iter().map(move |split| { + // Each planner is spawned once planned, as the buffer pulls it, so at most `concurrency` + // run at once. + let mut tasks = planners.map(move |split| -> BoxFuture<'static, VortexResult>> { + let split = match split { + Ok(split) => split, + Err(error) => return future::ready(Err(error)).boxed(), + }; let io = Arc::clone(&io); let map_fn = Arc::clone(&map_fn); - handle.spawn(async move { - split - .run(io) - .await? - .into_iter() - .map(|array| map_fn(array)) - .collect::>>() - }) + handle + .spawn(async move { + split + .run(io) + .await? + .into_iter() + .map(|array| map_fn(array)) + .collect::>>() + }) + .boxed() }); - // The first are spawned now, so they run while the stream is handed back to its consumer. + // The first are planned and spawned now, each before the next is planned, so they run + // while the stream is handed back to its consumer. let started: Vec<_> = tasks.by_ref().take(concurrency).collect(); let tasks = stream::iter(started).chain(stream::iter(tasks)); let batches = if self.ordered { @@ -462,3 +380,173 @@ impl PreparedScan { .boxed()) } } + +/// What the next fixed chunk holds for planning. +enum NextRun { + /// A selected run to group into a planner. + Run(Range), + /// The selected span of a chunk planned whole. + Whole(Range), + /// Every chunk has been visited. + Done, +} + +/// Plans a scan's planners one at a time, in row order, so each can run on a task of its own as +/// soon as it is planned, while the rest are still being planned. +/// +/// A fully selected chunk, or one selected in more than [`MAX_RUNS`] runs, is one split of its +/// own planner. Otherwise each selected run joins the planner before it, even across chunk +/// boundaries, while the two read a stored segment in common and the planner spans at most +/// [`MAX_PLANNER_ROWS`], so a stored segment is decoded by one planner unless it spans more rows +/// than that. Within a planner, a run extends the split before it when both read exactly the +/// same segments: every row between them then lies in those segments' chunks, so the split visits +/// no chunk the selection skips. +/// +/// A column's segments hold contiguous rows, so a run can only share a segment with an earlier +/// run of its planner if it shares one with the run just before it, and a segment several +/// planners read is read by neighbouring planners. So a planner is complete once the run after it +/// is known, and the segments it shares with that run are all it shares with later planners: the +/// scan's IO reads each of them once for both. +struct Planners { + plans: ScanPlans, + selection: Selection, + io: Arc, + /// The fixed chunks not yet visited. + chunks: std::vec::IntoIter>, + /// The selected runs of the last chunk visited, not yet grouped. + runs: VecDeque>, + /// A chunk planned whole, waiting for the planner before it to be handed out. + whole: Option>, + /// The pipeline splits of the planner being grouped. + splits: PlannerSplits, + /// The segments the last run of the planner reads, and those all its runs read. + last: Vec, + reached: Vec, +} + +impl Planners { + /// The next planner, ready to be admitted to a planning run. + fn next_plan(&mut self) -> VortexResult> { + let Some(ranges) = self.next_splits()? else { + return Ok(None); + }; + let rows = ranges[0].start..ranges[ranges.len() - 1].end; + let splits = ranges + .into_iter() + .map(|rows| Split { + mask: self.selection.row_mask(&rows).mask().clone(), + rows, + }) + .collect(); + let scope = WorkScope { + file_ordinal: 0, + rows, + }; + Ok(Some(SplitPlan { + root: plan_projection(self.plans.clone(), scope.clone(), splits), + scope, + })) + } + + /// The pipeline splits of the next planner, in row order. + fn next_splits(&mut self) -> VortexResult> { + if let Some(whole) = self.whole.take() { + return Ok(Some(vec![whole])); + } + loop { + match self.next_run() { + NextRun::Done => { + return Ok((!self.splits.is_empty()).then(|| self.close(&[]))); + } + NextRun::Whole(span) => { + if self.splits.is_empty() { + return Ok(Some(vec![span])); + } + self.whole = Some(span); + return Ok(Some(self.close(&[]))); + } + NextRun::Run(run) => { + let mut segments = Vec::new(); + self.plans.projection.reach( + run.clone(), + &Reach::Offset(0), + &mut |segment, _| segments.push(segment), + )?; + segments.sort_unstable(); + segments.dedup(); + let joins = self.splits.first().is_some_and(|first| { + run.end - first.start <= MAX_PLANNER_ROWS + && shares_any(&self.last, &segments) + }); + if joins { + match self.splits.last_mut() { + Some(split) if segments == self.last => split.end = run.end, + _ => self.splits.push(run), + } + self.reached.extend_from_slice(&segments); + self.last = segments; + continue; + } + let planner = (!self.splits.is_empty()).then(|| self.close(&segments)); + self.splits.push(run); + self.reached.extend_from_slice(&segments); + self.last = segments; + if planner.is_some() { + return Ok(planner); + } + } + } + } + } + + /// Hands out the planner being grouped. The segments it shares with `next`, those the run + /// after it reads, are read once for it and the planner that run starts. + fn close(&mut self, next: &[SegmentId]) -> PlannerSplits { + self.reached.sort_unstable(); + self.reached.dedup(); + self.io.share_with_next( + self.reached + .iter() + .filter(|segment| next.binary_search(segment).is_ok()) + .copied(), + ); + self.reached.clear(); + self.last.clear(); + std::mem::take(&mut self.splits) + } + + fn next_run(&mut self) -> NextRun { + loop { + if let Some(run) = self.runs.pop_front() { + return NextRun::Run(run); + } + let Some(rows) = self.chunks.next() else { + return NextRun::Done; + }; + let row_mask = self.selection.row_mask(&rows); + let mask = row_mask.mask(); + match mask.slices() { + AllOr::None => {} + AllOr::Some(slices) if slices.len() <= MAX_RUNS => self.runs.extend( + slices + .iter() + .map(|&(start, end)| rows.start + start as u64..rows.start + end as u64), + ), + AllOr::All | AllOr::Some(_) => { + let span = mask.first().unwrap_or(0)..mask.last().map_or(0, |last| last + 1); + return NextRun::Whole( + rows.start + span.start as u64..rows.start + span.end as u64, + ); + } + } + } + } +} + +impl Iterator for Planners { + type Item = VortexResult; + + fn next(&mut self) -> Option { + self.next_plan().transpose() + } +} diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index 05597ac68a7..ef7540487cb 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -385,7 +385,7 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { }); let segments = Arc::clone(&counted) as Arc; let scan = PreparedScan::try_new(builder(&session, &segments, &layout)?)?; - let splits = scan.split_plans()?; + let splits = scan.planners().collect::>>()?; assert_eq!( splits.len(), (STORED_CHUNKS * STORED_CHUNK_ROWS).div_ceil(MAX_ROWS_PER_CHUNK) @@ -467,12 +467,16 @@ async fn splits_read_each_stored_chunk_once( .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(rows.clone()))?), )?; // Runs reading the same stored chunk are one pipeline split, which visits no other chunk. + let mut planners = scan.planners(); + let mut pipeline_splits = Vec::new(); + while let Some(splits) = planners.next_splits()? { + pipeline_splits.extend(splits); + } assert_eq!( - scan.split_rows()?.0.concat(), - scopes, + pipeline_splits, scopes, "each planner here reads one stored chunk" ); - let splits = scan.split_plans()?; + let splits = scan.planners().collect::>>()?; assert_eq!( splits .iter() From aec7bad7112546c2df71752c6d767093ea4a44f6 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 20 Oct 2026 08:58:17 +0000 Subject: [PATCH 14/15] Run V2 units on a pool fed by a single file planner A V2 scan now has one root planner, the file planner, which plans the scan on its first step and hands out one unit per step after that. The driver gains distributing roots, whose children are handed back to the caller instead of run, and can admit a handed-off planner or morsel as a root. A new pool in vortex-scan runs those units. Any spare worker steps the file planner, and each unit it hands out waits in a shared queue until a worker claims it and admits it to its own run. Everything a unit emits stays on that worker's run. A worker steps the file planner before its own work whenever fewer units wait than workers are parked, spawns another worker when a unit has no parked worker to take it, and claims more units while the ones it holds wait for IO. Output is reordered by hand-out order when the scan is ordered, and the file planner stops while the consumer falls behind. This replaces the task per unit and the CPU-pool hop that prepared the scan before any unit was planned. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_014rtiSjzL4EauVCM7NS84hj --- vortex-layout/src/scan/v2/mod.rs | 226 +++++----- vortex-layout/src/scan/v2/split.rs | 47 -- vortex-layout/src/scan/v2/tests.rs | 105 +++-- vortex-scan/Cargo.toml | 2 +- vortex-scan/src/planning/driver/mod.rs | 111 ++++- vortex-scan/src/planning/driver/tests.rs | 56 +++ vortex-scan/src/planning/mod.rs | 1 + vortex-scan/src/planning/pool.rs | 519 +++++++++++++++++++++++ vortex-scan/src/planning/pool/tests.rs | 163 +++++++ 9 files changed, 1026 insertions(+), 204 deletions(-) delete mode 100644 vortex-layout/src/scan/v2/split.rs create mode 100644 vortex-scan/src/planning/pool.rs create mode 100644 vortex-scan/src/planning/pool/tests.rs diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 69c7d29e3f9..e430468ad4b 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -4,9 +4,11 @@ //! The planner-driven scan executor. //! //! A V2 scan lowers the file's layout once into a physical plan and runs every split of it -//! through the planning protocol: the [`Run`](vortex_scan::planning::driver::Run) driver advances -//! the split's planners and morsels, and every segment read is a protocol request served from the -//! file's [`SegmentSource`]. The layout reader is never asked to evaluate anything. +//! through the planning protocol. A single file planner hands out the scan's units one at a time +//! to a [`Pool`] of workers, each advancing the planners and morsels of the units it claimed on a +//! [`Run`](vortex_scan::planning::driver::Run) of its own, and every segment read is a protocol +//! request served from the file's [`SegmentSource`]. The layout reader is never asked to evaluate +//! anything. //! //! Import this module's [`ScanBuilder`] to select V2 explicitly. It accepts a [`ScanFile`] //! directly and supports projection, row ranges, and selections. Layouts without a pipeline @@ -17,7 +19,6 @@ //! segment is read once and rows between selected runs are never visited. mod io; -mod split; #[cfg(test)] mod tests; @@ -30,12 +31,10 @@ use arrow_array::RecordBatch; use arrow_schema::ArrowError; use arrow_schema::Field; use arrow_schema::SchemaRef; -use futures::FutureExt; use futures::Stream; use futures::StreamExt; use futures::TryStreamExt; use futures::future; -use futures::future::BoxFuture; use futures::stream; use futures::stream::BoxStream; use vortex_array::ArrayRef; @@ -43,9 +42,18 @@ use vortex_array::VortexSessionExecute; use vortex_array::dtype::DType; use vortex_array::expr::BoundExpression; use vortex_error::VortexResult; +use vortex_io::request::IoConsumer; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoResult; use vortex_io::session::RuntimeSessionExt; use vortex_mask::AllOr; +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 vortex_scan::planning::pool::Pool; +use vortex_scan::planning::pool::PoolOptions; +use vortex_scan::planning::pool::Spawn; use vortex_scan::selection::Selection; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; @@ -63,7 +71,6 @@ pub use crate::scan::execution::SegmentLocation; use crate::scan::execution::plan_projection; use crate::scan::v2::io::SegmentRanges; use crate::scan::v2::io::SegmentScanIo; -use crate::scan::v2::split::SplitPlan; use crate::segments::SegmentId; use crate::segments::SegmentSource; @@ -249,135 +256,138 @@ impl ScanBuilder { } } - /// Runs the scan on the session's runtime, preparing it on the first poll. - /// Preparation errors are returned as stream items. + /// Runs the scan on a [`Pool`] of workers on the session's runtime, starting it on the first + /// poll. Errors, including those preparing the scan, are returned as stream items. + /// + /// The scan's [`FilePlanner`] hands out its units one at a time, and each is run by the next + /// spare worker while the file planner plans the next. Up to `concurrency` units run at once + /// per worker thread, and their batches are yielded in row order when the scan is ordered. pub fn into_stream(self) -> VortexResult>> { let handle = self.session.handle(); - Ok(stream::once(async move { - handle - .spawn_cpu(move || PreparedScan::try_new(self)?.execute_stream()) - .await - }) - .try_flatten() + let spawn: Spawn = Arc::new(move |worker| handle.spawn(worker).detach()); + let io = Arc::new(SegmentScanIo::new( + Arc::clone(&self.file.segments), + SegmentRanges::new(Arc::clone(&self.file.locations)), + )); + let options = PoolOptions { + max_workers: get_available_parallelism().unwrap_or(1), + units_per_worker: self.concurrency, + ordered: self.ordered, + }; + let map = self.map_fn; + let root = FilePlanner { + scan: Some(ScanRequest { + session: self.session, + file: self.file, + projection: self.projection, + row_range: self.row_range, + selection: self.selection, + row_offset: self.row_offset, + }), + io: Arc::clone(&io), + planners: None, + }; + let scope = WorkScope { + file_ordinal: 0, + rows: 0..0, + }; + Ok(stream::once(future::lazy(move |_| { + Pool::stream(Box::new(root), scope, io, map, spawn, options) + })) + .flatten() .boxed()) } } -/// A scan planned over its file, ready to run its splits. -struct PreparedScan { +/// What a scan reads, before it is planned. +struct ScanRequest { session: VortexSession, - plans: ScanPlans, - /// Serves the splits' reads. - io: Arc, - /// The scan's splits, over the root row domain, in row order. - splits: Vec>, + file: ScanFile, + projection: BoundExpression, + row_range: Option>, selection: Selection, - map_fn: Arc VortexResult + Send + Sync>, - ordered: bool, - concurrency: usize, + row_offset: u64, } -impl PreparedScan { +impl ScanRequest { /// Lowers the file's layout, plans the projection, and partitions the selected range into - /// fixed-size row chunks. - fn try_new(builder: ScanBuilder) -> VortexResult { - let file = builder.file; + /// fixed-size row chunks, to plan units over. + fn plan(self, io: Arc) -> VortexResult { + let file = self.file; let root = lower(&file.layout)?; let row_count = root.row_count(); // An identity projection needs no rewrites. Keep the lowered children lazy, including // statistics that a projection-only scan never reads. - let identity = builder.projection.is_root() && builder.projection.dtype() == root.dtype(); + let identity = self.projection.is_root() && self.projection.dtype() == root.dtype(); let projection = if identity { root } else { - optimize(plan_row_idx_expression(builder.projection, root)?)? + optimize(plan_row_idx_expression(self.projection, root)?)? }; let chunk_rows = if identity { rows_per_chunk(row_count, &file.locations) } 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( - file.segments, - SegmentRanges::new(Arc::clone(&file.locations)), - )); - Ok(Self { - session: builder.session.clone(), + let chunks = row_chunks(self.row_range.unwrap_or(0..row_count), chunk_rows); + Ok(Planners { plans: ScanPlans { - session: builder.session, + session: self.session, locations: file.locations, projection, - row_offset: builder.row_offset, + row_offset: self.row_offset, values: Default::default(), }, + selection: self.selection, io, - splits, - selection: builder.selection, - map_fn: builder.map_fn, - ordered: builder.ordered, - concurrency: builder.concurrency, - }) - } - - /// Plans the scan's planners one at a time, in row order. - fn planners(&self) -> Planners { - Planners { - plans: self.plans.clone(), - selection: self.selection.clone(), - io: Arc::clone(&self.io), - chunks: self.splits.clone().into_iter(), + chunks: chunks.into_iter(), runs: VecDeque::new(), whole: None, splits: Vec::new(), last: Vec::new(), reached: Vec::new(), - } + }) } +} - /// Runs every split as a task on the session's runtime and yields their mapped batches. - /// - /// Like the default executor, up to `concurrency` splits per thread run at once, and their - /// batches are yielded in row order when the scan is ordered. - fn execute_stream(self) -> VortexResult>> { - let concurrency = self.concurrency * get_available_parallelism().unwrap_or(1); - let handle = self.session.handle(); - let planners = self.planners(); - let io = self.io; - let map_fn = self.map_fn; - // Each planner is spawned once planned, as the buffer pulls it, so at most `concurrency` - // run at once. - let mut tasks = planners.map(move |split| -> BoxFuture<'static, VortexResult>> { - let split = match split { - Ok(split) => split, - Err(error) => return future::ready(Err(error)).boxed(), - }; - let io = Arc::clone(&io); - let map_fn = Arc::clone(&map_fn); - handle - .spawn(async move { - split - .run(io) - .await? - .into_iter() - .map(|array| map_fn(array)) - .collect::>>() - }) - .boxed() - }); - // The first are planned and spawned now, each before the next is planned, so they run - // while the stream is handed back to its consumer. - let started: Vec<_> = tasks.by_ref().take(concurrency).collect(); - let tasks = stream::iter(started).chain(stream::iter(tasks)); - let batches = if self.ordered { - tasks.buffered(concurrency).boxed() +/// The scan's root planner: plans the scan on its first step, and hands out one unit per step +/// after that, in row order. Each unit is a planner that projects its splits. +struct FilePlanner { + /// The scan, until it is planned. + scan: Option, + io: Arc, + /// The scan's units, once it is planned, until every one has been handed out. + planners: Option, +} + +impl IoConsumer for FilePlanner { + // The file planner publishes no requests. + fn set_io_result(&mut self, _request: IoRequestId, _result: IoResult) {} +} + +impl Planner for FilePlanner { + fn state(&self) -> State { + if self.scan.is_some() || self.planners.is_some() { + State::NeedsCompute } else { - tasks.buffer_unordered(concurrency).boxed() + State::Done + } + } + + fn compute(&mut self) -> VortexResult { + if let Some(scan) = self.scan.take() { + self.planners = Some(scan.plan(Arc::clone(&self.io))?); + } + let Some(planners) = &mut self.planners else { + return Ok(PlannerOutput::Done); }; - Ok(batches - .map_ok(|batches| stream::iter(batches.into_iter().map(Ok))) - .try_flatten() - .boxed()) + match planners.next_plan()? { + Some((scope, root)) => Ok(PlannerOutput::Planner(scope, root)), + None => { + self.planners = None; + Ok(PlannerOutput::Done) + } + } } } @@ -391,8 +401,8 @@ enum NextRun { Done, } -/// Plans a scan's planners one at a time, in row order, so each can run on a task of its own as -/// soon as it is planned, while the rest are still being planned. +/// Plans a scan's units one at a time, in row order, so each can be handed to a worker as soon +/// as it is planned, while the rest are still being planned. /// /// A fully selected chunk, or one selected in more than [`MAX_RUNS`] runs, is one split of its /// own planner. Otherwise each selected run joins the planner before it, even across chunk @@ -425,8 +435,8 @@ struct Planners { } impl Planners { - /// The next planner, ready to be admitted to a planning run. - fn next_plan(&mut self) -> VortexResult> { + /// The next unit, with the rows it covers. + fn next_plan(&mut self) -> VortexResult)>> { let Some(ranges) = self.next_splits()? else { return Ok(None); }; @@ -442,10 +452,8 @@ impl Planners { file_ordinal: 0, rows, }; - Ok(Some(SplitPlan { - root: plan_projection(self.plans.clone(), scope.clone(), splits), - scope, - })) + let root = plan_projection(self.plans.clone(), scope.clone(), splits); + Ok(Some((scope, root))) } /// The pipeline splits of the next planner, in row order. @@ -542,11 +550,3 @@ impl Planners { } } } - -impl Iterator for Planners { - type Item = VortexResult; - - fn next(&mut self) -> Option { - self.next_plan().transpose() - } -} diff --git a/vortex-layout/src/scan/v2/split.rs b/vortex-layout/src/scan/v2/split.rs deleted file mode 100644 index b9371735d39..00000000000 --- a/vortex-layout/src/scan/v2/split.rs +++ /dev/null @@ -1,47 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::future; -use std::sync::Arc; - -use vortex_array::ArrayRef; -use vortex_error::VortexResult; -use vortex_io::request::IoService; -use vortex_scan::planning::driver::Progress; -use vortex_scan::planning::driver::Run; -use vortex_scan::planning::planner::Planner; -use vortex_scan::planning::planner::WorkScope; - -/// One split of a scan, ready to be admitted to a planning run. -/// -/// The split's planner projects the rows of `scope`, and its morsels emit one batch per -/// projection split with selected rows, in whatever order their reads arrive. -pub(super) struct SplitPlan { - /// The rows of the scan's root row domain the split covers. - pub(super) scope: WorkScope, - /// The split's first planner. - pub(super) root: Box, -} - -impl SplitPlan { - /// Runs the split on a planning run of its own, reading through a session of `io`, and - /// returns its batches in row order. - /// - /// The run advances inside this future, on whichever thread polls it, and the future awaits - /// its reads between steps. Nothing is handed to another thread. - pub(super) async fn run(self, io: Arc) -> VortexResult> { - let mut run = Run::new(); - run.admit(self.root, self.scope, io.session()); - let mut batches = Vec::new(); - loop { - match run.advance()? { - Progress::Batch(batch) => batches.push(batch), - Progress::RootDone(_) | Progress::Idle => break, - Progress::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, - } - } - // Morsels finish in whatever order their reads arrive. - batches.sort_by_key(|batch| batch.scope.rows.start); - Ok(batches.into_iter().map(|batch| batch.array).collect()) - } -} diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index ef7540487cb..b098d6089bb 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -33,14 +33,21 @@ use vortex_arrow::ArrowSessionExt; use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_error::VortexResult; +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::Planner; +use vortex_scan::planning::planner::WorkScope; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; use super::MAX_ROWS_PER_CHUNK; -use super::PreparedScan; +use super::Planners; use super::ScanBuilder; use super::ScanFile; +use super::io::SegmentRanges; +use super::io::SegmentScanIo; use super::row_chunks; use super::rows_per_chunk; use crate::LayoutRef; @@ -58,6 +65,51 @@ use crate::segments::SegmentFuture; use crate::segments::SegmentId; use crate::segments::SegmentSource; use crate::segments::TestSegments; + +/// Plans `builder`'s scan, returning its units' planner and the IO they read through. +fn plan(builder: ScanBuilder) -> VortexResult<(Planners, Arc)> { + let io = Arc::new(SegmentScanIo::new( + Arc::clone(&builder.file.segments), + SegmentRanges::new(Arc::clone(&builder.file.locations)), + )); + let scan = super::ScanRequest { + session: builder.session, + file: builder.file, + projection: builder.projection, + row_range: builder.row_range, + selection: builder.selection, + row_offset: builder.row_offset, + }; + Ok((scan.plan(Arc::clone(&io))?, io)) +} + +/// Every unit of a planned scan, in row order. +fn units(planners: &mut Planners) -> VortexResult)>> { + let mut units = Vec::new(); + while let Some(unit) = planners.next_plan()? { + units.push(unit); + } + Ok(units) +} + +/// Runs one unit on a run of its own and returns its batches in row order. +async fn run_unit( + (scope, root): (WorkScope, Box), + io: &SegmentScanIo, +) -> VortexResult> { + let mut run = Run::new(); + run.admit(root, scope, io.session()); + let mut batches = Vec::new(); + loop { + match run.advance()? { + Progress::Batch(batch) => batches.push(batch), + Progress::Waiting => std::future::poll_fn(|cx| run.poll_completion(cx)).await?, + _ => break, + } + } + batches.sort_by_key(|batch| batch.scope.rows.start); + Ok(batches.into_iter().map(|batch| batch.array).collect()) +} use crate::sequence::SequenceId; use crate::sequence::SequentialStreamAdapter; use crate::sequence::SequentialStreamExt as _; @@ -384,16 +436,16 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { reads: AtomicUsize::new(0), }); let segments = Arc::clone(&counted) as Arc; - let scan = PreparedScan::try_new(builder(&session, &segments, &layout)?)?; - let splits = scan.planners().collect::>>()?; + let (mut planners, io) = plan(builder(&session, &segments, &layout)?)?; + let splits = units(&mut planners)?; assert_eq!( splits.len(), (STORED_CHUNKS * STORED_CHUNK_ROWS).div_ceil(MAX_ROWS_PER_CHUNK) ); for (index, split) in splits.into_iter().enumerate() { - let rows = split.scope.rows.clone(); + let rows = split.0.rows.clone(); let before = counted.reads.load(Ordering::Relaxed); - let arrays = split.run(Arc::clone(&scan.io) as _).await?; + let arrays = run_unit(split, &io).await?; let reads = counted.reads.load(Ordering::Relaxed) - before; let code_chunks = (rows.end - 1) / STORED_CHUNK_ROWS as u64 - rows.start / STORED_CHUNK_ROWS as u64 + 1; @@ -462,12 +514,14 @@ async fn splits_read_each_stored_chunk_once( let session = new_session().with_tokio(); let (counted, layout) = counted_ints(&session).await?; let segments = Arc::clone(&counted) as Arc; - let scan = PreparedScan::try_new( - builder(&session, &segments, &layout)? - .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(rows.clone()))?), - )?; + let scan = || { + plan( + builder(&session, &segments, &layout)? + .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(rows.clone()))?), + ) + }; // Runs reading the same stored chunk are one pipeline split, which visits no other chunk. - let mut planners = scan.planners(); + let (mut planners, _) = scan()?; let mut pipeline_splits = Vec::new(); while let Some(splits) = planners.next_splits()? { pipeline_splits.extend(splits); @@ -476,17 +530,18 @@ async fn splits_read_each_stored_chunk_once( pipeline_splits, scopes, "each planner here reads one stored chunk" ); - let splits = scan.planners().collect::>>()?; + let (mut planners, io) = scan()?; + let splits = units(&mut planners)?; assert_eq!( splits .iter() - .map(|split| split.scope.rows.clone()) + .map(|(scope, _)| scope.rows.clone()) .collect::>(), scopes ); let mut arrays = Vec::new(); for split in splits { - arrays.extend(split.run(Arc::clone(&scan.io) as _).await?); + arrays.extend(run_unit(split, &io).await?); } // Splits read only the stored chunks holding their selected rows, each chunk once. assert_eq!(counted.reads.load(Ordering::Relaxed), reads); @@ -523,8 +578,8 @@ async fn list_layouts_return_an_error() -> VortexResult<()> { async fn unordered_returns_every_row() -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, table().as_ref()).await?; - let scan = PreparedScan::try_new(builder(&session, &segments, &layout)?.with_ordered(false))?; - let chunks: Vec = scan.execute_stream()?.try_collect().await?; + let scan = builder(&session, &segments, &layout)?.with_ordered(false); + let chunks: Vec = scan.into_stream()?.try_collect().await?; let rows: usize = chunks.iter().map(|chunk| chunk.len()).sum(); assert_eq!(rows, CHUNKS * CHUNK_ROWS); Ok(()) @@ -535,14 +590,12 @@ async fn primitive_column() -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, table().as_ref()).await?; let dtype = layout.dtype().clone(); - let scan = PreparedScan::try_new( - builder(&session, &segments, &layout)? - .with_projection(get_item("ints", root()).bind(&dtype)?) - .with_row_range(990..1010), - )?; + let scan = builder(&session, &segments, &layout)? + .with_projection(get_item("ints", root()).bind(&dtype)?) + .with_row_range(990..1010); let actual = collect( DType::Primitive(vortex_array::dtype::PType::I32, false.into()), - scan.execute_stream()?, + scan.into_stream()?, ) .await?; assert_arrays_eq!( @@ -561,9 +614,13 @@ async fn work_chunks_ignore_layout_boundaries( ) -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, strategy.get().as_ref()).await?; - let scan = PreparedScan::try_new(builder(&session, &segments, &layout)?)?; - assert_eq!(scan.splits, vec![0..(CHUNKS * CHUNK_ROWS) as u64]); - let arrays: Vec<_> = scan.execute_stream()?.try_collect().await?; + let (planners, _) = plan(builder(&session, &segments, &layout)?)?; + assert_eq!( + planners.chunks.collect::>(), + vec![0..(CHUNKS * CHUNK_ROWS) as u64] + ); + let scan = builder(&session, &segments, &layout)?; + let arrays: Vec<_> = scan.into_stream()?.try_collect().await?; assert_eq!(arrays.len(), 1); assert_eq!(arrays[0].len(), CHUNKS * CHUNK_ROWS); Ok(()) diff --git a/vortex-scan/Cargo.toml b/vortex-scan/Cargo.toml index d12c207b049..a54e2ad63a9 100644 --- a/vortex-scan/Cargo.toml +++ b/vortex-scan/Cargo.toml @@ -25,11 +25,11 @@ vortex-utils = { workspace = true } async-trait = { workspace = true } futures = { workspace = true } +parking_lot = { workspace = true } roaring = { workspace = true } tracing = { workspace = true } [dev-dependencies] -parking_lot = { workspace = true } rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } diff --git a/vortex-scan/src/planning/driver/mod.rs b/vortex-scan/src/planning/driver/mod.rs index 8cb78099272..329ddf964b6 100644 --- a/vortex-scan/src/planning/driver/mod.rs +++ b/vortex-scan/src/planning/driver/mod.rs @@ -15,6 +15,11 @@ //! for IO, 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. +//! +//! A root admitted with [`Run::admit_distributing`] runs nothing it emits: each child is handed +//! back as a [`Progress::Handoff`], for another run to admit with [`Run::admit_child`]. A +//! [`Pool`](crate::planning::pool::Pool) uses this to hand a file planner's units to whichever +//! worker is spare. use std::collections::BTreeMap; use std::collections::VecDeque; @@ -59,10 +64,30 @@ pub struct Batch { pub array: ArrayRef, } +/// A planner or morsel, run as the root of a run. +pub enum Child { + /// A planner. + Planner(Box), + /// A morsel. + Morsel(Box), +} + +/// A child a distributing root emitted, for another run to admit. +pub struct Handoff { + /// The distributing root that emitted the child. + pub root: RootId, + /// The scope the child covers. + pub scope: WorkScope, + /// The child, which has not been visited. + pub child: Child, +} + /// What a call to [`Run::advance`] left behind. pub enum Progress { /// A morsel's batch. One morsel's batches arrive in the order it emitted them. Batch(Batch), + /// A child of a root admitted with [`Run::admit_distributing`], in emission order. + Handoff(Handoff), /// Every item descended from the root has retired, and its IO session was cleared. RootDone(RootId), /// Every live item waits for IO. Await [`Run::poll_completion`], or pass the next completion @@ -72,11 +97,6 @@ pub enum Progress { Idle, } -enum Item { - Planner(Box), - Morsel(Box), -} - /// What one item's `compute()` produced. enum Output { Planner(PlannerOutput), @@ -89,7 +109,7 @@ struct Work { id: IoOwnerId, root: RootId, scope: WorkScope, - item: Item, + item: Child, /// Stable child-ordinal path, beginning with the root's admission ordinal. path: Vec, /// The parent continuation sorts after all previously emitted child subtrees. @@ -115,15 +135,15 @@ impl Work { fn state(&self) -> State { match &self.item { - Item::Planner(planner) => planner.state(), - Item::Morsel(morsel) => morsel.state(), + Child::Planner(planner) => planner.state(), + Child::Morsel(morsel) => morsel.state(), } } fn consumer(&mut self) -> &mut dyn IoConsumer { match &mut self.item { - Item::Planner(planner) => &mut **planner, - Item::Morsel(morsel) => &mut **morsel, + Child::Planner(planner) => &mut **planner, + Child::Morsel(morsel) => &mut **morsel, } } @@ -179,6 +199,8 @@ const IO_POLL_INTERVAL: usize = 64; struct Root { io: Arc, live: usize, + /// Whether the children the root emits are handed off instead of run. + distributes: bool, } /// A set of root planners advanced by its caller on one thread at a time. @@ -229,11 +251,44 @@ impl Run { root: Box, scope: WorkScope, io: Arc, + ) -> RootId { + self.admit_child(Child::Planner(root), scope, io) + } + + /// Admits a planner or morsel, such as one handed off by another run, as a root. + pub fn admit_child(&mut self, child: Child, scope: WorkScope, io: Arc) -> RootId { + self.admit_root(child, scope, io, false) + } + + /// Admits `root` without running what it emits: each child it emits is handed back by + /// [`advance`](Self::advance) as a [`Progress::Handoff`], in emission order. + pub fn admit_distributing( + &mut self, + root: Box, + scope: WorkScope, + io: Arc, + ) -> RootId { + self.admit_root(Child::Planner(root), scope, io, true) + } + + fn admit_root( + &mut self, + root: Child, + scope: WorkScope, + io: Arc, + distributes: bool, ) -> RootId { 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.roots.insert( + id, + Root { + io, + live: 0, + distributes, + }, + ); + self.spawn(id, scope, root, vec![id.0]); id } @@ -375,7 +430,7 @@ impl Run { Ok(()) } - fn spawn(&mut self, root: RootId, scope: WorkScope, item: Item, path: Vec) { + fn spawn(&mut self, root: RootId, scope: WorkScope, item: Child, path: Vec) { let id = IoOwnerId(self.next_owner); self.next_owner += 1; if let Some(root) = self.roots.get_mut(&root) { @@ -431,11 +486,11 @@ impl Run { /// Visits one item: computes once if it is ready, and settles it otherwise. fn visit(&mut self, mut work: Box) -> VortexResult<()> { let output = match &mut work.item { - Item::Planner(planner) => match planner.state() { + Child::Planner(planner) => match planner.state() { State::NeedsCompute => Output::Planner(planner.compute()?), _ => return self.settle(work), }, - Item::Morsel(morsel) => match morsel.state() { + Child::Morsel(morsel) => match morsel.state() { State::NeedsCompute => Output::Morsel(morsel.compute()?), _ => return self.settle(work), }, @@ -453,13 +508,11 @@ impl Run { self.settle(work)?; } Output::Planner(PlannerOutput::Planner(scope, child)) => { - let path = work.child_path(); - self.spawn(work.root, scope, Item::Planner(child), path); + self.emit(&mut work, scope, Child::Planner(child)); self.queue.insert(work); } Output::Planner(PlannerOutput::Morsel(scope, morsel)) => { - let path = work.child_path(); - self.spawn(work.root, scope, Item::Morsel(morsel), path); + self.emit(&mut work, scope, Child::Morsel(morsel)); self.queue.insert(work); } Output::Morsel(MorselOutput::Batch(array)) => { @@ -477,6 +530,25 @@ impl Run { Ok(()) } + /// Spawns a child `work` emitted, or hands it off if `work` is a distributing root. + fn emit(&mut self, work: &mut Work, scope: WorkScope, child: Child) { + let distributes = work.path.len() == 1 + && self + .roots + .get(&work.root) + .is_some_and(|root| root.distributes); + if distributes { + self.events.push_back(Progress::Handoff(Handoff { + root: work.root, + scope, + child, + })); + } else { + let path = work.child_path(); + self.spawn(work.root, scope, child, path); + } + } + /// Registers the requests an item published with its root's session. fn submit(&mut self, work: &mut Work, batch: IoBatch) -> VortexResult<()> { if batch.is_empty() { @@ -540,6 +612,7 @@ impl Driver { loop { match run.advance()? { Progress::Batch(batch) => batches.push(batch), + Progress::Handoff(_) => vortex_bail!("a driver root cannot hand off children"), Progress::RootDone(_) | Progress::Idle => return Ok(batches), Progress::Waiting => run.complete(self.io.wait()?)?, } diff --git a/vortex-scan/src/planning/driver/tests.rs b/vortex-scan/src/planning/driver/tests.rs index 389ad24372d..7437fe92336 100644 --- a/vortex-scan/src/planning/driver/tests.rs +++ b/vortex-scan/src/planning/driver/tests.rs @@ -503,6 +503,9 @@ fn drain(run: &mut Run) -> VortexResult> { events.push(match run.advance()? { Progress::Batch(batch) => format!("batch {:?} {:?}", batch.root, batch.scope.rows), Progress::RootDone(root) => format!("done {root:?}"), + Progress::Handoff(handoff) => { + format!("handoff {:?} {:?}", handoff.root, handoff.scope.rows) + } Progress::Waiting => { events.push("waiting".to_string()); return Ok(events); @@ -515,6 +518,59 @@ fn drain(run: &mut Run) -> VortexResult> { } } +/// A distributing root hands off the planners and morsels it emits, unvisited and in emission +/// order, and another run admits each as a root of its own. +#[test] +fn a_distributing_root_hands_off_its_children() -> VortexResult<()> { + let log = Log::default(); + let file = ScriptedPlanner::boxed( + "file", + vec![ + PlannerStep::Planner( + scope(0..1), + vec![one_batch_morsel(scope(0..1), 1), PlannerStep::Done], + ), + one_batch_morsel(scope(1..3), 2), + PlannerStep::Done, + ], + &log, + ); + let source = source(); + let mut distributing = Run::new(); + distributing.admit_distributing(file, scope(0..3), Arc::clone(&source) as Arc); + let mut handoffs = Vec::new(); + loop { + match distributing.advance()? { + Progress::Handoff(handoff) => handoffs.push(handoff), + Progress::RootDone(_) => break, + _ => vortex_bail!("expected a handoff or the root to finish"), + } + } + let scopes: Vec<_> = handoffs.iter().map(|h| h.scope.rows.clone()).collect(); + assert_eq!(scopes, vec![0..1, 1..3]); + assert_eq!(log.events(), vec!["compute file"; 3]); + + let mut run = Run::new(); + for handoff in handoffs { + run.admit_child( + handoff.child, + handoff.scope, + Arc::clone(&source) as Arc, + ); + } + assert_eq!( + drain(&mut run)?, + vec![ + "batch RootId(0) 0..1", + "done RootId(0)", + "batch RootId(1) 1..3", + "done RootId(1)", + "idle", + ] + ); + Ok(()) +} + /// Roots run side by side: a root with ready work finishes while another waits for its read, and /// each root reads through, and finally clears, its own session. #[test] diff --git a/vortex-scan/src/planning/mod.rs b/vortex-scan/src/planning/mod.rs index 3095712f161..3ac13f9c334 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 pool; #[cfg(test)] mod tests; diff --git a/vortex-scan/src/planning/pool.rs b/vortex-scan/src/planning/pool.rs new file mode 100644 index 00000000000..5c8577ae124 --- /dev/null +++ b/vortex-scan/src/planning/pool.rs @@ -0,0 +1,519 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A pool of workers that run the units a single root planner hands out. +//! +//! The root, such as a file planner, is admitted with [`Run::admit_distributing`] to a run of its +//! own, and each step of it hands out at most one unit. Any worker may step it, and the units it +//! hands out wait in a shared queue until a spare worker claims one and admits it to its own +//! [`Run`]. Everything a unit emits stays on the run that claimed it, so a unit's planners and +//! morsels, and the arrays they decode, stay with one worker. +//! +//! A worker steps the root before its own work whenever fewer units wait than workers are +//! parked, so a unit is handed out as soon as some worker can take it, while the root goes on +//! to plan the next. Workers are spawned as units are handed out with no worker parked to take +//! them, up to a limit, and a worker claims more units while the ones it holds wait for IO. +//! +//! At most a fixed number of units are handed out and not yet consumed, so the root stops +//! planning while the consumer falls behind. + +use std::collections::BTreeMap; +use std::collections::VecDeque; +use std::future; +use std::pin::Pin; +use std::sync::Arc; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; + +use futures::Stream; +use futures::future::BoxFuture; +use parking_lot::Mutex; +use vortex_array::ArrayRef; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_io::request::IoService; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::planning::driver::Child; +use crate::planning::driver::Progress; +use crate::planning::driver::RootId; +use crate::planning::driver::Run; +use crate::planning::planner::Planner; +use crate::planning::planner::WorkScope; + +/// Runs a worker future, typically by spawning it on a runtime. +pub type Spawn = Arc) + Send + Sync>; + +/// Maps each batch to the pool's output, on the worker that produced it. +pub type MapFn = Arc VortexResult + Send + Sync>; + +/// How a [`Pool`] runs its units. +#[derive(Clone, Copy, Debug)] +pub struct PoolOptions { + /// The most workers the pool spawns. + pub max_workers: usize, + /// The most units one worker runs at once. + pub units_per_worker: usize, + /// Whether output follows the order the root handed out its units in. Otherwise each batch + /// is yielded as soon as it is produced. + pub ordered: bool, +} + +/// Runs the units a root planner hands out on a pool of workers. See the [module](self) docs. +pub struct Pool; + +impl Pool { + /// Admits `root` and returns the stream of its units' mapped batches. + /// + /// The first worker is spawned now. Within a unit, ordered output is sorted by the start of + /// each batch's scope, since a unit's morsels finish in whatever order their reads arrive. + /// Dropping the stream stops the workers. + pub fn stream( + root: Box, + scope: WorkScope, + io: Arc, + map: MapFn, + spawn: Spawn, + options: PoolOptions, + ) -> PoolStream { + let max_workers = options.max_workers.max(1); + let units_per_worker = options.units_per_worker.max(1); + let mut file = Run::new(); + file.admit_distributing(root, scope, io.session()); + let shared = Arc::new(Shared { + io, + map, + spawn, + max_workers, + units_per_worker, + max_in_flight: max_workers * units_per_worker, + state: Mutex::new(State { + file: Some(file), + file_waiting: false, + file_done: false, + pending: VecDeque::new(), + next_seq: 0, + in_flight: 0, + workers: 0, + live_workers: 0, + parked: Vec::new(), + stopped: false, + events: VecDeque::new(), + consumer: None, + }), + }); + shared.spawn_worker(&mut shared.state.lock()); + PoolStream { + shared, + ordered: options.ordered, + units: None, + next_unit: 0, + buffered: BTreeMap::new(), + ready: VecDeque::new(), + } + } +} + +/// A unit handed out by the root and not yet claimed. +struct Unit { + /// The unit's position in the order the root handed out its units. + seq: u64, + scope: WorkScope, + child: Child, +} + +/// What workers report to the stream. +enum Event { + /// A mapped batch of a unit, with the start of the scope of the morsel that produced it. + Batch { + seq: u64, + start: u64, + value: T, + }, + /// Every batch of the unit has been reported. + UnitDone(u64), + /// The root has finished, having handed out this many units. + RootDone(u64), + Error(VortexError), +} + +struct Shared { + io: Arc, + map: MapFn, + spawn: Spawn, + max_workers: usize, + units_per_worker: usize, + max_in_flight: usize, + state: Mutex>, +} + +struct State { + /// The root's run, while no worker steps or polls it. + file: Option, + /// Whether the root waits for IO. + file_waiting: bool, + file_done: bool, + pending: VecDeque, + next_seq: u64, + /// Units handed out and not yet consumed. + in_flight: usize, + /// Workers spawned, and those that have not exited. + workers: usize, + live_workers: usize, + /// Wakers of parked workers. + parked: Vec, + /// Set by an error or by dropping the stream. + stopped: bool, + events: VecDeque>, + consumer: Option, +} + +impl State { + /// Whether the root's run can be stepped now. + fn file_ready(&self, max_in_flight: usize) -> bool { + self.file.is_some() && !self.file_waiting && self.in_flight < max_in_flight + } + + fn wake_workers(&mut self) { + for waker in self.parked.drain(..) { + waker.wake(); + } + } + + fn report(&mut self, event: Event) { + self.events.push_back(event); + if let Some(consumer) = self.consumer.take() { + consumer.wake(); + } + } +} + +impl Shared { + fn spawn_worker(self: &Arc, state: &mut State) { + state.workers += 1; + state.live_workers += 1; + let shared = Arc::clone(self); + (self.spawn)(Box::pin(async move { + let exit = WorkerExit(Arc::clone(&shared)); + if let Err(error) = Worker::new(&shared).run().await { + shared.fail(error); + } + drop(exit); + })); + } + + fn fail(&self, error: VortexError) { + let mut state = self.state.lock(); + if !state.stopped { + state.stopped = true; + state.report(Event::Error(error)); + } + state.wake_workers(); + } + + /// Takes the root's run to step it: whenever it is ready if `always`, and otherwise only + /// while fewer units wait than workers are parked. + fn take_file(&self, always: bool) -> Option { + let mut state = self.state.lock(); + if state.stopped || !state.file_ready(self.max_in_flight) { + return None; + } + if !always && state.pending.len() > state.parked.len() { + return None; + } + state.file.take() + } + + /// Steps the root's run until it hands out a unit, waits, or finishes. + fn step_file(self: &Arc, mut file: Run) -> VortexResult<()> { + let progress = file.advance(); + let mut state = self.state.lock(); + match progress? { + Progress::Handoff(handoff) => { + let seq = state.next_seq; + state.next_seq += 1; + state.in_flight += 1; + state.pending.push_back(Unit { + seq, + scope: handoff.scope, + child: handoff.child, + }); + state.file = Some(file); + if state.pending.len() > state.parked.len() && state.workers < self.max_workers { + self.spawn_worker(&mut state); + } + } + Progress::Waiting => { + state.file_waiting = true; + state.file = Some(file); + } + Progress::RootDone(_) | Progress::Idle => { + state.file_done = true; + let units = state.next_seq; + state.report(Event::RootDone(units)); + } + Progress::Batch(_) => vortex_bail!("the pool's root produced a batch"), + } + state.wake_workers(); + Ok(()) + } + + fn claim(&self) -> Option { + let mut state = self.state.lock(); + if state.stopped { + return None; + } + state.pending.pop_front() + } +} + +impl Shared { + fn release(&self, units: usize) { + let mut state = self.state.lock(); + state.in_flight -= units; + state.wake_workers(); + } +} + +/// Reports a worker's exit, including by panic, so the stream ends once every worker has. +struct WorkerExit(Arc>); + +impl Drop for WorkerExit { + fn drop(&mut self) { + let mut state = self.0.state.lock(); + state.live_workers -= 1; + if let Some(consumer) = state.consumer.take() { + consumer.wake(); + } + } +} + +/// One worker's run, holding the units it claimed. +struct Worker<'a, T> { + shared: &'a Arc>, + run: Run, + seqs: HashMap, +} + +impl<'a, T: Send + 'static> Worker<'a, T> { + fn new(shared: &'a Arc>) -> Self { + Self { + shared, + run: Run::new(), + seqs: HashMap::default(), + } + } + + async fn run(mut self) -> VortexResult<()> { + loop { + if self.shared.state.lock().stopped { + return Ok(()); + } + if let Some(file) = self.shared.take_file(false) { + self.shared.step_file(file)?; + continue; + } + if self.run.live_roots() > 0 { + match self.run.advance()? { + Progress::Batch(batch) => { + let seq = self.seq(batch.root)?; + let value = (self.shared.map)(batch.array)?; + self.shared.state.lock().report(Event::Batch { + seq, + start: batch.scope.rows.start, + value, + }); + continue; + } + Progress::RootDone(root) => { + let seq = self.seq(root)?; + self.seqs.remove(&root); + self.shared.state.lock().report(Event::UnitDone(seq)); + continue; + } + Progress::Handoff(_) => vortex_bail!("a unit cannot hand off children"), + Progress::Waiting | Progress::Idle => {} + } + } + if self.run.live_roots() < self.shared.units_per_worker + && let Some(unit) = self.shared.claim() + { + let root = self + .run + .admit_child(unit.child, unit.scope, self.shared.io.session()); + self.seqs.insert(root, unit.seq); + continue; + } + if let Some(file) = self.shared.take_file(true) { + self.shared.step_file(file)?; + continue; + } + if self.run.live_roots() == 0 { + let state = self.shared.state.lock(); + if state.file_done && state.pending.is_empty() { + return Ok(()); + } + } + future::poll_fn(|cx| self.park(cx)).await?; + } + } + + fn seq(&self, root: RootId) -> VortexResult { + self.seqs + .get(&root) + .copied() + .ok_or_else(|| vortex_err!("no unit for {root:?}")) + } + + /// Waits until there is something to do: a unit to claim, the root to step, a delivery to + /// the worker's own run or to the root's, or the pool to stop or finish. + fn park(&mut self, cx: &mut Context<'_>) -> Poll> { + let poll_file = { + let mut state = self.shared.state.lock(); + let can_claim = + !state.pending.is_empty() && self.run.live_roots() < self.shared.units_per_worker; + let finished = state.file_done && state.pending.is_empty(); + if state.stopped + || can_claim + || state.file_ready(self.shared.max_in_flight) + || (finished && self.run.live_roots() == 0) + { + return Poll::Ready(Ok(())); + } + if !state.parked.iter().any(|waker| waker.will_wake(cx.waker())) { + state.parked.push(cx.waker().clone()); + } + if state.file_waiting { + state.file.take() + } else { + None + } + }; + if let Some(mut file) = poll_file { + let delivered = file.poll_completion(cx); + let mut state = self.shared.state.lock(); + state.file = Some(file); + if delivered.is_ready() { + state.file_waiting = false; + state.wake_workers(); + return delivered; + } + } + if self.run.live_roots() > 0 { + return self.run.poll_completion(cx); + } + Poll::Pending + } +} + +/// The mapped batches of a [`Pool`]'s units. +pub struct PoolStream { + shared: Arc>, + ordered: bool, + /// How many units the root handed out, once it has finished. + units: Option, + /// Units fully consumed, in order, when ordered; otherwise units finished. + next_unit: u64, + /// Batches of units not yet yielded, and whether each unit is done. + buffered: BTreeMap, bool)>, + ready: VecDeque, +} + +impl PoolStream { + fn finished(&self) -> bool { + self.units == Some(self.next_unit) + } + + fn handle(&mut self, event: Event) -> VortexResult<()> { + match event { + Event::Batch { seq, start, value } => { + if self.ordered { + self.buffered.entry(seq).or_default().0.push((start, value)); + } else { + self.ready.push_back(value); + } + } + Event::UnitDone(seq) => { + if self.ordered { + self.buffered.entry(seq).or_default().1 = true; + let mut released = 0; + while let Some(entry) = self.buffered.first_entry() + && *entry.key() == self.next_unit + && entry.get().1 + { + let (mut batches, _) = entry.remove(); + batches.sort_by_key(|(start, _)| *start); + self.ready + .extend(batches.into_iter().map(|(_, value)| value)); + self.next_unit += 1; + released += 1; + } + if released > 0 { + self.shared.release(released); + } + } else { + self.next_unit += 1; + self.shared.release(1); + } + } + Event::RootDone(units) => self.units = Some(units), + Event::Error(error) => return Err(error), + } + Ok(()) + } +} + +// Nothing in the stream is pinned. +impl Unpin for PoolStream {} + +impl Stream for PoolStream { + type Item = VortexResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = &mut *self; + loop { + if let Some(value) = this.ready.pop_front() { + return Poll::Ready(Some(Ok(value))); + } + if this.finished() { + return Poll::Ready(None); + } + let events = { + let mut state = this.shared.state.lock(); + if state.events.is_empty() { + if state.stopped { + return Poll::Ready(None); + } + if state.live_workers == 0 { + state.stopped = true; + return Poll::Ready(Some(Err(vortex_err!( + "every pool worker exited before the scan finished" + )))); + } + state.consumer = Some(cx.waker().clone()); + return Poll::Pending; + } + std::mem::take(&mut state.events) + }; + for event in events { + if let Err(error) = this.handle(event) { + return Poll::Ready(Some(Err(error))); + } + } + } + } +} + +impl Drop for PoolStream { + fn drop(&mut self) { + let mut state = self.shared.state.lock(); + state.stopped = true; + state.wake_workers(); + } +} + +#[cfg(test)] +mod tests; diff --git a/vortex-scan/src/planning/pool/tests.rs b/vortex-scan/src/planning/pool/tests.rs new file mode 100644 index 00000000000..ef1400224f2 --- /dev/null +++ b/vortex-scan/src/planning/pool/tests.rs @@ -0,0 +1,163 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::future::Future; +use std::pin::pin; +use std::sync::Arc; +use std::task::Wake; +use std::thread; +use std::thread::Thread; + +use futures::StreamExt; +use rstest::rstest; +use vortex_buffer::buffer; +use vortex_io::request::IoIntent; +use vortex_io::request::IoRequest; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoSource; +use vortex_io::request::IoTarget; + +use super::*; +use crate::planning::morsel::MorselOutput; +use crate::planning::tests::scripted::Log; +use crate::planning::tests::scripted::MorselStep; +use crate::planning::tests::scripted::PlannerStep; +use crate::planning::tests::scripted::RecordingIoSource; +use crate::planning::tests::scripted::ScriptedPlanner; +use crate::planning::tests::scripted::array_of; +use crate::planning::tests::scripted::scope; + +struct ThreadWaker(Thread); + +impl Wake for ThreadWaker { + fn wake(self: Arc) { + self.0.unpark(); + } +} + +fn block_on(future: F) -> F::Output { + let mut future = pin!(future); + let waker = Waker::from(Arc::new(ThreadWaker(thread::current()))); + let mut cx = Context::from_waker(&waker); + loop { + if let Poll::Ready(output) = future.as_mut().poll(&mut cx) { + return output; + } + thread::park(); + } +} + +/// Runs each worker on a thread of its own. +fn threads() -> Spawn { + Arc::new(|worker| { + thread::spawn(move || block_on(worker)); + }) +} + +/// Gives every session a source of its own, with the bytes at offset 0. +struct Sessions; + +impl IoService for Sessions { + fn session(&self) -> Arc { + let source = RecordingIoSource::default(); + source.canned_bytes(0, buffer![0u8, 1, 2, 3].into_byte_buffer()); + Arc::new(source) + } +} + +fn lengths() -> MapFn { + Arc::new(|array: ArrayRef| Ok(array.len())) +} + +/// A unit that reads a segment, then emits a morsel with `rows` rows over `start..start + 1`. +fn unit(start: u64, rows: u64) -> PlannerStep { + PlannerStep::Planner( + scope(start..start + 1), + vec![ + PlannerStep::Io(vec![IoRequest { + intent: IoIntent::Fetch, + request: IoRequestId(0), + target: IoTarget::range(0, 4), + }]), + PlannerStep::Morsel( + scope(start..start + 1), + vec![ + MorselStep::Compute(MorselOutput::Batch(array_of(rows))), + MorselStep::Compute(MorselOutput::Done), + ], + ), + PlannerStep::Done, + ], + ) +} + +fn run(root: Vec, options: PoolOptions) -> Vec> { + let log = Log::default(); + let stream = Pool::stream( + ScriptedPlanner::boxed("file", root, &log), + scope(0..0), + Arc::new(Sessions), + lengths(), + threads(), + options, + ); + block_on(stream.collect()) +} + +fn options(max_workers: usize, ordered: bool) -> PoolOptions { + PoolOptions { + max_workers, + units_per_worker: 2, + ordered, + } +} + +#[rstest] +#[case::one_worker(1)] +#[case::several_workers(3)] +fn ordered_output_follows_hand_out_order(#[case] workers: usize) -> VortexResult<()> { + let mut root: Vec<_> = (1..=8).map(|rows| unit(rows, rows)).collect(); + root.push(PlannerStep::Done); + let lengths = run(root, options(workers, true)) + .into_iter() + .collect::>>()?; + assert_eq!(lengths, (1..=8).collect::>()); + Ok(()) +} + +#[test] +fn unordered_output_yields_every_batch() -> VortexResult<()> { + let mut root: Vec<_> = (1..=8).map(|rows| unit(rows, rows)).collect(); + root.push(PlannerStep::Done); + let mut lengths = run(root, options(3, false)) + .into_iter() + .collect::>>()?; + lengths.sort_unstable(); + assert_eq!(lengths, (1..=8).collect::>()); + Ok(()) +} + +#[test] +fn a_root_with_no_units_ends_the_stream() { + assert!(run(vec![PlannerStep::Done], options(2, true)).is_empty()); +} + +#[test] +fn a_failing_unit_ends_the_stream_with_its_error() { + let root = vec![ + unit(0, 1), + PlannerStep::Planner(scope(1..2), vec![PlannerStep::Wait]), + PlannerStep::Done, + ]; + let outputs = run(root, options(2, true)); + let error = outputs + .iter() + .find_map(|output| output.as_ref().err()) + .map(|error| error.to_string()); + assert!( + error + .as_deref() + .is_some_and(|message| message.contains("no outstanding fetch")), + "{error:?}" + ); +} From 2d5cd2baa52ba8e3c467a7f2a43b9a17e60de0e4 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 20 Oct 2026 09:29:16 +0000 Subject: [PATCH 15/15] Don't link the private file planner from into_stream's docs Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_014rtiSjzL4EauVCM7NS84hj --- vortex-layout/src/scan/v2/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index e430468ad4b..4f8c6f5dc0e 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -259,7 +259,7 @@ impl ScanBuilder { /// Runs the scan on a [`Pool`] of workers on the session's runtime, starting it on the first /// poll. Errors, including those preparing the scan, are returned as stream items. /// - /// The scan's [`FilePlanner`] hands out its units one at a time, and each is run by the next + /// The scan's file planner hands out its units one at a time, and each is run by the next /// spare worker while the file planner plans the next. Up to `concurrency` units run at once /// per worker thread, and their batches are yielded in row order when the scan is ordered. pub fn into_stream(self) -> VortexResult>> {