diff --git a/.agents/skills/bench-performance/SKILL.md b/.agents/skills/bench-performance/SKILL.md index 821f0e022ee..c42a1271ddb 100644 --- a/.agents/skills/bench-performance/SKILL.md +++ b/.agents/skills/bench-performance/SKILL.md @@ -114,7 +114,7 @@ Supported diagnostics: - `--formats parquet,vortex,vortex-compact,lance,arrow`; - `--queries 6`, `--exclude-queries 1,2`, `--iterations N`, `--display-format gh-json`; -- `--hide-progress-bar`, `-o /private/tmp/out.jsonl`, `--gh-json-v3 /private/tmp/out.v3.jsonl`; +- `--hide-progress-bar`, `-o /private/tmp/out.jsonl`, `--ingest-jsonl /private/tmp/out.ingest.jsonl`; - `--verbose`, `--tracing`, `--track-memory`, `--runner NAME`, `--opt key=value`; - `--explain` prints query plans instead of timing; - `--show-metrics` prints Vortex execution-plan metrics after a timed run. @@ -149,7 +149,7 @@ Supported diagnostics: - `--reuse` keeps one DuckDB connection across iterations, useful with Samply to keep work on the same threads; - common flags: `--queries`, `--exclude-queries`, `--iterations`, `--display-format`, - `--hide-progress-bar`, `-o`, `--gh-json-v3`, `--track-memory`, `--verbose`, `--tracing`, + `--hide-progress-bar`, `-o`, `--ingest-jsonl`, `--track-memory`, `--verbose`, `--tracing`, `--runner`, `--opt`, `--explain`. Example: @@ -171,7 +171,7 @@ Important differences: - no `--formats`: the binary always generates/registers Lance data and reports `datafusion:lance`; - no `--explain` or `--show-metrics` path today; - common flags: `--queries`, `--exclude-queries`, `--iterations`, `--display-format`, - `--hide-progress-bar`, `-o`, `--gh-json-v3`, `--track-memory`, `--verbose`, `--tracing`, + `--hide-progress-bar`, `-o`, `--ingest-jsonl`, `--track-memory`, `--verbose`, `--tracing`, `--runner`, `--opt`; - `--threads` is parsed but currently not wired into the Lance/DataFusion session. diff --git a/.agents/skills/pr-review/SKILL.md b/.agents/skills/pr-review/SKILL.md deleted file mode 100644 index e46d5813031..00000000000 --- a/.agents/skills/pr-review/SKILL.md +++ /dev/null @@ -1,93 +0,0 @@ ---- -name: pr-review -description: Review Vortex pull requests for correctness, Rust soundness, performance, API compatibility, and test coverage. Use when reviewing PRs, reviewing code changes, or when the user mentions "/pr-review". ---- - -# Vortex PR Review Skill - -Review Vortex changes for issues that CI may miss: semantic correctness, Rust soundness, -zero-copy and alignment invariants, performance, API compatibility, and missing regression -coverage. - -## Usage Modes - -### GitHub Actions Mode - -When invoked by a PR automation, the prompt may already include PR metadata, issue body, -comments, and changed files. In this mode, use git commands for the diff and history: - -```bash -git diff origin/...HEAD -git diff --stat origin/...HEAD -git log origin/..HEAD --oneline -``` - -If the base ref is missing, fetch only that ref: - -```bash -git fetch origin --depth=1 -``` - -Do not refetch data already provided in the prompt. Use the supplied comments and metadata as -review context. - -### Local CLI Mode - -When the user provides a PR number or URL, use `gh` to fetch the PR metadata, comments, and diff: - -```bash -gh pr view --json title,body,author,baseRefName,headRefName,files,additions,deletions,commits -gh pr view --json comments,reviews -gh pr diff -``` - -If `gh` cannot connect to `api.github.com` because of sandbox networking, rerun with escalated -network permissions. - -## Review Workflow - -1. Read `AGENTS.md` and any nested `AGENTS.md` for project conventions. -2. Identify the change intent from the PR title, body, commits, and tests. -3. Group changed files by area: arrays, encodings, buffers, file/layout, integrations, bindings, - docs, or CI. -4. Trace changed behavior through callers, trait implementations, dtype/nullability handling, - validity masks, and tests. -5. Focus findings on actionable defects. Avoid commenting on formatting or issues already covered - by clippy, rustfmt, or generated API checks. -6. Scope verification to the change. Rust/API changes need Rust checks; docs-only, agent-only, - symlink-only, and metadata-only changes should use targeted validation such as Markdown review, - `ls`, `find`, `git status`, or relevant config linters. - -## Review Areas - -| Area | Focus | -| ------------------ | ---------------------------------------------------------------------------------------------------------------------------------------------------- | -| Correctness | Length and dtype invariants, nullability, validity masks, offset math, canonicalization, boundary conditions, empty arrays, scalar vs array behavior | -| Rust soundness | `unsafe` blocks, aliasing, lifetimes, alignment, FFI boundaries, panic safety, ownership of buffers and arrays | -| Compression and IO | Encoding metadata, statistics, layout evolution, file compatibility, scan projection/filter behavior, async IO edge cases | -| Performance | Unnecessary copies, lost zero-copy behavior, avoidable allocations, poor cache locality, quadratic loops, excessive dynamic dispatch in hot paths | -| Error handling | Correct `vortex_err!` and `vortex_bail!` usage, useful messages, no accidental panics on user data | -| API compatibility | Public API docs, feature flags, crate boundaries, Python/Java binding impacts | -| Tests | Regression coverage, edge cases, parameterized cases with `rstest`, use of `assert_arrays_eq!`, docs doctests when docs change | -| Verification scope | Avoid requesting or running expensive workspace checks when the PR only changes docs, agent files, symlinks, or metadata | - -## Output - -Lead with findings, ordered by severity. For each finding include: - -- File and line reference. -- Why the issue is a real bug or material risk. -- A concrete fix or verification path when possible. - -Use inline review comments when the environment supports them and a precise changed line is the -best place for the feedback. Keep broad design feedback in the summary. - -If no issues are found, say so explicitly and mention any residual risk or tests not run. - -## Principles - -- Review the code that changed, but inspect enough surrounding code to validate invariants. -- Do not infer causation from commit messages alone. Verify with code, tests, or logs. -- Do not ask for broad rewrites when a narrow fix would address the risk. -- Do not downgrade a correctness or soundness issue to a nit because it is inconvenient. -- Be specific and proportionate. diff --git a/.agents/skills/query/SKILL.md b/.agents/skills/query/SKILL.md deleted file mode 100644 index 2a07f3ef312..00000000000 --- a/.agents/skills/query/SKILL.md +++ /dev/null @@ -1,36 +0,0 @@ ---- -name: query -description: Answer questions about the Vortex codebase or pull requests. Use when asked a question via "/query" or when the user wants to understand code, architecture, behavior, or implementation details. ---- - -# Vortex Query Skill - -Answer questions about the Vortex project, its pull requests, and its implementation. - -## Key Context - -- Vortex is a Rust workspace for columnar arrays, compression encodings, file IO, and scan - integrations. -- `vortex-array` defines the core array traits, dtype system, canonical arrays, and base - encodings. -- `vortex-buffer` owns aligned zero-copy buffers. -- `vortex-file` and `vortex-layout` implement file and layout reading. -- `encodings/*` contains specialized compressed encodings. -- Python, Java, DuckDB, and DataFusion integrations live in their own workspace areas. - -## Workflow - -1. Read `AGENTS.md` and any closer scoped `AGENTS.md` before relying on conventions. -2. Use `rg` and targeted file reads to identify the relevant crate, module, and tests. -3. If the question is about a PR, inspect the diff and comments before answering. -4. If the question is about behavior, trace the implementation through public entry points, - encoding-specific implementations, and tests. -5. Answer with concrete file paths and line numbers when they help. - -## Answering Guidelines - -- Separate confirmed facts from inference. -- Prefer precise code references over broad descriptions. -- Mention important uncertainty and describe what would verify it. -- Do not invent architecture. If the repository does not answer the question, say what you - checked and what is still missing. diff --git a/.agents/skills/samply/scripts/profile_inverted_tree.py b/.agents/skills/samply/scripts/profile_inverted_tree.py index f7ad40cf37e..3f12800eee7 100644 --- a/.agents/skills/samply/scripts/profile_inverted_tree.py +++ b/.agents/skills/samply/scripts/profile_inverted_tree.py @@ -6,6 +6,7 @@ import argparse import collections import re +from collections.abc import Callable from pathlib import Path from typing import Any @@ -36,7 +37,7 @@ def thread_matches(thread: dict[str, Any], thread_re: re.Pattern[str] | None, ti return bool(thread_re.search(str(thread.get("name", ""))) or thread_re.search(str(thread.get("tid", "")))) -def frame_labels(thread: dict[str, Any], symbols: dict[str, str]) -> Any: +def frame_labels(thread: dict[str, Any], symbols: dict[str, str]) -> Callable[[int], str]: cache: dict[int, str] = {} def label(frame_index: int) -> str: diff --git a/.agents/skills/samply/scripts/profile_summary.py b/.agents/skills/samply/scripts/profile_summary.py index 53ec1f0f007..1106f2b5e32 100644 --- a/.agents/skills/samply/scripts/profile_summary.py +++ b/.agents/skills/samply/scripts/profile_summary.py @@ -22,14 +22,14 @@ def load_profile(path: Path) -> dict[str, Any]: return json.load(f) -def value_at(table: dict[str, Any], key: str, index: int, default: Any = None) -> Any: +def value_at(table: dict[str, Any], key: str, index: int, default: object = None) -> object: values = table.get(key) if not isinstance(values, list) or index < 0 or index >= len(values): return default return values[index] -def decode_string(thread: dict[str, Any], value: Any) -> str: +def decode_string(thread: dict[str, Any], value: object) -> str: strings = thread.get("stringArray") or [] if isinstance(value, int) and 0 <= value < len(strings): return str(strings[value]) @@ -38,7 +38,7 @@ def decode_string(thread: dict[str, Any], value: Any) -> str: return str(value) -def resource_name(thread: dict[str, Any], resource_index: Any) -> str: +def resource_name(thread: dict[str, Any], resource_index: object) -> str: if not isinstance(resource_index, int): return "" resource_table = thread.get("resourceTable") or {} @@ -72,7 +72,7 @@ def frame_label(thread: dict[str, Any], frame_index: int, symbols: dict[str, str return label -def expand_stack(thread: dict[str, Any], stack_index: Any) -> list[int]: +def expand_stack(thread: dict[str, Any], stack_index: object) -> list[int]: if not isinstance(stack_index, int): return [] diff --git a/.cargo/config.toml b/.cargo/config.toml index 466422a80b1..b40ec53c4c7 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -18,6 +18,19 @@ timeout = 3600 # Environment variables for PyO3. Ensures reproducible builds and avoids spurious recompilations. [env] PYO3_ENVIRONMENT_SIGNATURE = { value = "cpython-3.11-64bit", force = true } +# Read by benchmarks marked `#[cpu_features]` (see the `vortex-bench-support` crate), which +# expands to code reading both with `env!`. VORTEX_BENCH_VARIANT names the run: the benchmark +# is skipped when it is "simulation", and runs otherwise. VORTEX_BENCH_PREFIX and +# VORTEX_BENCH_SUFFIX wrap the benchmark's name, so a CI leg can label the numbers it produced: +# the prefix is the `::` component the leg's benchmark filter selects on, and the suffix is on +# the leaf name, which is the part CodSpeed reports under. +# +# The values below are what a plain `cargo bench` wants: nothing skipped, names left bare. The +# per-feature-set jobs in `.github/workflows/codspeed.yml` set their own. None of them is `force`d, +# so the value those jobs export wins over the one here. +VORTEX_BENCH_VARIANT = "local" +VORTEX_BENCH_PREFIX = "" +VORTEX_BENCH_SUFFIX = "" [cache.global-clean] # Anything older than this duration will be deleted in the source cache. diff --git a/.github/actions/setup-duckdb/action.yml b/.github/actions/setup-duckdb/action.yml new file mode 100644 index 00000000000..51d920304e9 --- /dev/null +++ b/.github/actions/setup-duckdb/action.yml @@ -0,0 +1,26 @@ +name: "Setup DuckDB" +description: "Download the DuckDB CLI and put it on PATH" +inputs: + duckdb_version: + description: "Version of the DuckDB CLI" + default: "1.5.5" +runs: + using: "composite" + steps: + - name: Download DuckDB + shell: bash + run: | + case "$(uname -m)" in + x86_64) duckdb_arch=amd64 ;; + aarch64|arm64) duckdb_arch=arm64 ;; + *) + echo "Unsupported architecture: $(uname -m)" >&2 + exit 1 + ;; + esac + + wget -qO- \ + "https://github.com/duckdb/duckdb/releases/download/v${{ inputs.duckdb_version }}/duckdb_cli-linux-${duckdb_arch}.zip" \ + | funzip > duckdb + chmod +x duckdb + echo "$PWD" >> "$GITHUB_PATH" diff --git a/.github/actions/setup-prebuild/action.yml b/.github/actions/setup-prebuild/action.yml index 4fec5d85c0e..93a5e4c8d29 100644 --- a/.github/actions/setup-prebuild/action.yml +++ b/.github/actions/setup-prebuild/action.yml @@ -1,5 +1,5 @@ name: "Setup Prebuild" -description: "Prebuild AMIs — uses pre-installed tools, falls back to full setup-rust for forks" +description: "Use pre-installed AMI tools and install any missing requirements" inputs: repo-token: @@ -16,6 +16,10 @@ inputs: targets: description: "optional targets override (e.g. wasm32-unknown-unknown)" required: false + protoc-version: + description: "protoc version expected in the prebuilt image" + default: "29.3" + required: false enable-sccache: description: "Should sccache be enabled." required: false @@ -24,7 +28,122 @@ inputs: runs: using: "composite" steps: - # Prebuild path: tools already installed, just configure sccache + - name: Resolve toolchain config + if: github.repository == 'vortex-data/vortex' + id: toolchain-config + shell: bash + run: | + DEFAULT_VERSION=$(grep channel rust-toolchain.toml | awk -F'"' '{print $2}') + TOOLCHAIN="${TOOLCHAIN_OVERRIDE:-$DEFAULT_VERSION}" + echo "toolchain=$TOOLCHAIN" >> "$GITHUB_OUTPUT" + env: + TOOLCHAIN_OVERRIDE: ${{ inputs.toolchain }} + + - name: Check for rustup + if: github.repository == 'vortex-data/vortex' + id: check-rustup + shell: bash + run: echo "exists=$(command -v rustup &> /dev/null && echo 'true' || echo 'false')" >> "$GITHUB_OUTPUT" + + - name: Install Rust + if: github.repository == 'vortex-data/vortex' && steps.check-rustup.outputs.exists != 'true' + uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable + with: + toolchain: ${{ steps.toolchain-config.outputs.toolchain }} + targets: ${{ inputs.targets }} + components: ${{ inputs.components }} + + - name: Reconcile prebuilt Rust installation + if: github.repository == 'vortex-data/vortex' && steps.check-rustup.outputs.exists == 'true' + shell: bash + run: | + TOOLCHAIN_FOUND=false + for INSTALLED in $(rustup toolchain list | awk '{print $1}'); do + case "$INSTALLED" in + "$TOOLCHAIN"|"$TOOLCHAIN"-*) TOOLCHAIN_FOUND=true ;; + esac + done + if [ "$TOOLCHAIN_FOUND" != true ]; then + rustup toolchain install "$TOOLCHAIN" --profile minimal + fi + + TOOLCHAIN_HOST=$(rustc +"$TOOLCHAIN" --print host-tuple) + INSTALLED_COMPONENTS=$(rustup component list --toolchain "$TOOLCHAIN" --installed) + MISSING_COMPONENTS=() + IFS=',' read -ra REQUESTED_COMPONENTS <<< "$COMPONENTS" + for COMPONENT in "${REQUESTED_COMPONENTS[@]}"; do + COMPONENT=${COMPONENT// /} + [ -z "$COMPONENT" ] && continue + if [ "$COMPONENT" = llvm-tools ]; then + COMPONENT=llvm-tools-preview + fi + if ! grep -Fqx "$COMPONENT" <<< "$INSTALLED_COMPONENTS" && \ + ! grep -Fqx "$COMPONENT-$TOOLCHAIN_HOST" <<< "$INSTALLED_COMPONENTS"; then + MISSING_COMPONENTS+=("$COMPONENT") + fi + done + if [ ${#MISSING_COMPONENTS[@]} -gt 0 ]; then + rustup component add --toolchain "$TOOLCHAIN" "${MISSING_COMPONENTS[@]}" + fi + + INSTALLED_TARGETS=$(rustup target list --toolchain "$TOOLCHAIN" --installed) + MISSING_TARGETS=() + IFS=',' read -ra REQUESTED_TARGETS <<< "$TARGETS" + for TARGET in "${REQUESTED_TARGETS[@]}"; do + TARGET=${TARGET// /} + [ -z "$TARGET" ] && continue + if ! grep -Fqx "$TARGET" <<< "$INSTALLED_TARGETS"; then + MISSING_TARGETS+=("$TARGET") + fi + done + if [ ${#MISSING_TARGETS[@]} -gt 0 ]; then + rustup target add --toolchain "$TOOLCHAIN" "${MISSING_TARGETS[@]}" + fi + env: + TOOLCHAIN: ${{ steps.toolchain-config.outputs.toolchain }} + COMPONENTS: ${{ inputs.components }} + TARGETS: ${{ inputs.targets }} + + - name: Reconcile mold installation + if: github.repository == 'vortex-data/vortex' && runner.os == 'Linux' + shell: bash + run: | + if ! command -v mold >/dev/null; then + sudo apt-get update + sudo apt-get install -y mold + fi + + - name: Check protoc installation + if: github.repository == 'vortex-data/vortex' && runner.os == 'Linux' + id: check-protoc + shell: bash + run: | + EXPECTED="libprotoc $PROTOC_VERSION" + ACTUAL=$(protoc --version 2>/dev/null || true) + echo "install=$([ "$ACTUAL" = "$EXPECTED" ] && echo false || echo true)" >> "$GITHUB_OUTPUT" + env: + PROTOC_VERSION: ${{ inputs.protoc-version }} + + - name: Install protoc + if: steps.check-protoc.outputs.install == 'true' + uses: ./.github/actions/setup-protoc + with: + protoc_version: ${{ inputs.protoc-version }} + + - name: Check sccache installation + if: github.repository == 'vortex-data/vortex' && inputs.enable-sccache == 'true' + id: check-sccache + shell: bash + run: echo "installed=$(command -v sccache >/dev/null && echo true || echo false)" >> "$GITHUB_OUTPUT" + + - name: Install sccache + if: >- + github.repository == 'vortex-data/vortex' && + inputs.enable-sccache == 'true' && + steps.check-sccache.outputs.installed != 'true' + uses: mozilla-actions/sccache-action@9e7fa8a12102821edf02ca5dbea1acd0f89a2696 # v0.0.10 + + # Prebuild path: configure the bundled (or fallback-installed) sccache. - name: Configure sccache timeout if: github.repository == 'vortex-data/vortex' && inputs.enable-sccache == 'true' shell: bash @@ -32,10 +151,14 @@ runs: mkdir -p ~/.config/sccache echo 'server_startup_timeout_ms = 60000' > ~/.config/sccache/config - - name: Pre-start sccache server + # Applies to every CMake invocation in opted-in upstream jobs. + - name: Enable compiler caching if: github.repository == 'vortex-data/vortex' && inputs.enable-sccache == 'true' shell: bash - run: sccache --start-server + run: | + sccache --start-server + echo "CMAKE_C_COMPILER_LAUNCHER=sccache" >> "$GITHUB_ENV" + echo "CMAKE_CXX_COMPILER_LAUNCHER=sccache" >> "$GITHUB_ENV" # Fallback path: full setup for forks - name: Full Rust setup diff --git a/.github/actions/setup-rust/action.yml b/.github/actions/setup-rust/action.yml index 345a2d08f7e..e0e04bb39dc 100644 --- a/.github/actions/setup-rust/action.yml +++ b/.github/actions/setup-rust/action.yml @@ -36,7 +36,7 @@ runs: - name: Install Mold if: runner.os == 'Linux' - uses: rui314/setup-mold@9c9c13bf4c3f1adef0cc596abc155580bcb04444 # v1 + uses: rui314/setup-mold@7e4f20ad28a2e8ca6fd0892ccf72e2abb706b9c3 # v1 - name: Check for rustup id: check-rustup @@ -45,7 +45,7 @@ runs: - name: Rust Toolchain id: rust-toolchain - uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable + uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable if: steps.check-rustup.outputs.exists != 'true' with: toolchain: "${{ steps.toolchain-config.outputs.toolchain }}" diff --git a/.github/actions/system-info/action.yml b/.github/actions/system-info/action.yml index 7ca1d99b7ea..f3641d2f9dc 100644 --- a/.github/actions/system-info/action.yml +++ b/.github/actions/system-info/action.yml @@ -10,10 +10,13 @@ runs: echo "=== CPU Model ===" lscpu echo "=== /proc/cpuinfo summary ===" - grep -m1 "model name" /proc/cpuinfo + # x86 exposes "model name"; aarch64 has no such line (lscpu above covers it), and + # `shell: bash` runs with `set -e`, so a missing match must not fail the step. + grep -m1 "model name" /proc/cpuinfo || true nproc echo "=== CPU Flags ===" - grep -m1 "flags" /proc/cpuinfo + # x86 exposes "flags"; aarch64 exposes "Features". + grep -m1 -E "^(flags|Features)" /proc/cpuinfo || true echo "=== Memory ===" free -h diff --git a/.github/actions/upload-parca-debuginfo/action.yml b/.github/actions/upload-parca-debuginfo/action.yml new file mode 100644 index 00000000000..40369a8227e --- /dev/null +++ b/.github/actions/upload-parca-debuginfo/action.yml @@ -0,0 +1,117 @@ +name: "Pre-upload benchmark debuginfo to Polar Signals" +description: "Install parca-debuginfo and upload built benchmark binaries to Polar Signals Cloud" + +inputs: + paths: + description: "Newline-delimited paths to binaries whose debuginfo should be uploaded" + required: true + polarsignals-cloud-token: + description: "Polar Signals service account token" + required: true + project-id: + description: "Polar Signals Cloud project ID" + required: true + store-address: + description: "Parca-compatible debuginfo gRPC endpoint" + required: false + default: "grpc.polarsignals.com:443" + version: + description: "parca-debuginfo version to install" + required: false + default: "0.13.0" + force: + description: "Replace existing debuginfo for the same build ID" + required: false + default: "false" + +runs: + using: "composite" + steps: + - name: Download and verify parca-debuginfo CLI + shell: bash + env: + VERSION: ${{ inputs.version }} + run: | + set -Eeu -o pipefail + + base_url="https://github.com/parca-dev/parca-debuginfo/releases/download/v${VERSION}" + asset="parca-debuginfo_${VERSION}_$(uname -s)_$(uname -m)" + checksums="$(mktemp)" + curl -fsSL "${base_url}/checksums.txt" -o "$checksums" + + mkdir -p "$RUNNER_TEMP/parca-debuginfo" + cd "$RUNNER_TEMP/parca-debuginfo" + + if curl -fsSLO "${base_url}/${asset}"; then + shasum --ignore-missing -a 256 --check "$checksums" + chmod +x "$asset" + mv "$asset" parca-debuginfo + else + archive="${asset}.tar.gz" + curl -fsSLO "${base_url}/${archive}" + shasum --ignore-missing -a 256 --check "$checksums" + tar -xzf "$archive" parca-debuginfo + chmod +x parca-debuginfo + fi + + echo "$PWD" >> "$GITHUB_PATH" + + - name: Upload benchmark binary debuginfo to Polar Signals + shell: bash + env: + PATHS: ${{ inputs.paths }} + POLARSIGNALS_CLOUD_TOKEN: ${{ inputs.polarsignals-cloud-token }} + PROJECT_ID: ${{ inputs.project-id }} + STORE_ADDRESS: ${{ inputs.store-address }} + FORCE: ${{ inputs.force }} + run: | + set -Eeu -o pipefail + + token_file="$(mktemp)" + upload_log="$(mktemp)" + trap 'rm -f "$token_file" "$upload_log"' EXIT + printf "%s" "$POLARSIGNALS_CLOUD_TOKEN" > "$token_file" + + uploaded=0 + while IFS= read -r binary; do + binary="${binary#"${binary%%[![:space:]]*}"}" + binary="${binary%"${binary##*[![:space:]]}"}" + [[ -z "$binary" ]] && continue + + if [[ ! -f "$binary" ]]; then + echo "::error::Debuginfo upload target does not exist: $binary" + exit 1 + fi + + cmd=( + parca-debuginfo + upload + "--store-address=$STORE_ADDRESS" + "--bearer-token-file=$token_file" + "--grpc-headers=projectID=$PROJECT_ID" + "$binary" + ) + if [[ "$FORCE" == "true" ]]; then + cmd+=(--force) + fi + + set +e + "${cmd[@]}" > "$upload_log" 2>&1 + status=$? + set -e + cat "$upload_log" + + if [[ "$status" -ne 0 ]]; then + if grep -Eiq "upload id mismatch|already exists|AlreadyExists|previous upload is still in-progress|not stale yet|only stale uploads can be retried" "$upload_log"; then + echo "::notice::Debuginfo upload already exists or is in progress for $binary; continuing" + else + exit "$status" + fi + fi + uploaded=$((uploaded + 1)) + done <<< "$PATHS" + + if [[ "$uploaded" -eq 0 ]]; then + echo "::error::No debuginfo upload targets were provided" + exit 1 + fi diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index ddefd50673a..8c6f83c93ed 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -1,23 +1,21 @@ -## Rationale for this change +## Summary + +- Tracking Issue: #000 -- Closes: #000 +## Changes -## What changes are included in this PR? -## What APIs are changed? Are there any user-facing changes? - diff --git a/.github/scripts/close_fixed_fuzzer_issues.py b/.github/scripts/close_fixed_fuzzer_issues.py index 48513a3ceb2..1ac32e448d7 100755 --- a/.github/scripts/close_fixed_fuzzer_issues.py +++ b/.github/scripts/close_fixed_fuzzer_issues.py @@ -30,7 +30,7 @@ class FuzzerIssue: body: str -def run(cmd: list[str], **kwargs) -> subprocess.CompletedProcess: +def run(cmd: list[str], **kwargs: object) -> subprocess.CompletedProcess: """Run a command, printing it for visibility.""" print(f" $ {' '.join(cmd)}", flush=True) return subprocess.run(cmd, **kwargs) diff --git a/.github/scripts/fuzz_report/cli.py b/.github/scripts/fuzz_report/cli.py index 2bab8494872..0e05dc4da1c 100644 --- a/.github/scripts/fuzz_report/cli.py +++ b/.github/scripts/fuzz_report/cli.py @@ -41,7 +41,7 @@ def _write_github_output(key: str, value: str) -> None: f.write(f"{key}={value}\n") -def _run_gh(cmd: list[str], *, retries: int = 1, **kwargs) -> subprocess.CompletedProcess: +def _run_gh(cmd: list[str], *, retries: int = 1, **kwargs: object) -> subprocess.CompletedProcess: """Run a gh CLI command with logging and retry on failure. Prints the command before execution and surfaces stderr on failure. diff --git a/.github/scripts/fuzz_report/template.py b/.github/scripts/fuzz_report/template.py index 273527169f0..65b6ef269ec 100644 --- a/.github/scripts/fuzz_report/template.py +++ b/.github/scripts/fuzz_report/template.py @@ -9,10 +9,10 @@ class _NotSetUndefined(Undefined): """Jinja2 undefined that renders as '(not set)' and is falsy.""" - def __str__(self): + def __str__(self) -> str: return "(not set)" - def __bool__(self): + def __bool__(self) -> bool: return False diff --git a/.github/scripts/fuzz_report/tests/conftest.py b/.github/scripts/fuzz_report/tests/conftest.py index ddc61f12cd7..7cc8744c86c 100644 --- a/.github/scripts/fuzz_report/tests/conftest.py +++ b/.github/scripts/fuzz_report/tests/conftest.py @@ -2,20 +2,21 @@ import json import tempfile +from collections.abc import Iterator from pathlib import Path import pytest @pytest.fixture -def temp_dir(): +def temp_dir() -> Iterator[Path]: """Create a temporary directory.""" with tempfile.TemporaryDirectory() as d: yield Path(d) @pytest.fixture -def sample_log_content(): +def sample_log_content() -> str: """Sample fuzzer log content with index out of bounds crash.""" return """ Running: cargo +nightly fuzz run file_io @@ -37,7 +38,7 @@ def sample_log_content(): @pytest.fixture -def sample_issues(): +def sample_issues() -> list[dict[str, object]]: """Sample existing issues for dedup testing.""" return [ { @@ -56,7 +57,7 @@ def sample_issues(): @pytest.fixture -def issues_file(sample_issues, temp_dir): +def issues_file(sample_issues, temp_dir) -> Path: """Write sample issues to a temporary JSON file.""" path = temp_dir / "issues.json" path.write_text(json.dumps(sample_issues)) diff --git a/.github/scripts/fuzz_report/tests/test_dedup.py b/.github/scripts/fuzz_report/tests/test_dedup.py index e62e593e502..926faa5ad52 100644 --- a/.github/scripts/fuzz_report/tests/test_dedup.py +++ b/.github/scripts/fuzz_report/tests/test_dedup.py @@ -2,6 +2,7 @@ import json import tempfile +from collections.abc import Iterator from pathlib import Path import pytest @@ -43,7 +44,7 @@ @pytest.fixture -def issues_file(): +def issues_file() -> Iterator[str]: """Create a temporary issues JSON file.""" with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: json.dump(EXISTING_ISSUES, f) @@ -52,7 +53,7 @@ def issues_file(): Path(f.name).unlink() -def _make_crash_info(**overrides) -> CrashInfo: +def _make_crash_info(**overrides: object) -> CrashInfo: """Helper to create a CrashInfo with defaults.""" defaults = { "panic_location": "brand/new/file.rs:999", @@ -73,7 +74,7 @@ def _make_crash_info(**overrides) -> CrashInfo: class TestCheckSeedHash: - def test_match_found(self): + def test_match_found(self) -> None: result = check_seed_hash( "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", EXISTING_ISSUES, @@ -83,7 +84,7 @@ def test_match_found(self): assert result.confidence == "exact" assert result.issue_number == 100 - def test_no_match(self): + def test_no_match(self) -> None: result = check_seed_hash( "0000000000000000000000000000000000000000000000000000000000000000", EXISTING_ISSUES, @@ -91,13 +92,13 @@ def test_no_match(self): assert result.duplicate is False assert result.check == "seed_hash" - def test_unknown_hash(self): + def test_unknown_hash(self) -> None: result = check_seed_hash("unknown", EXISTING_ISSUES) assert result.duplicate is False class TestCheckPanicLocation: - def test_match_found(self): + def test_match_found(self) -> None: result = check_panic_location( "vortex-array/src/compute/slice.rs:142", EXISTING_ISSUES, @@ -107,17 +108,17 @@ def test_match_found(self): assert result.confidence == "high" assert result.issue_number == 100 - def test_partial_match(self): + def test_partial_match(self) -> None: result = check_panic_location("slice.rs:142", EXISTING_ISSUES) assert result.duplicate is True - def test_no_match(self): + def test_no_match(self) -> None: result = check_panic_location("other/file.rs:999", EXISTING_ISSUES) assert result.duplicate is False class TestCheckStackTrace: - def test_match_found(self): + def test_match_found(self) -> None: result = check_stack_trace( "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb", EXISTING_ISSUES, @@ -126,7 +127,7 @@ def test_match_found(self): assert result.check == "stack_trace" assert result.confidence == "high" - def test_no_match(self): + def test_no_match(self) -> None: result = check_stack_trace( "0000000000000000000000000000000000000000000000000000000000000000", EXISTING_ISSUES, @@ -135,7 +136,7 @@ def test_no_match(self): class TestCheckErrorPattern: - def test_message_hash_match(self): + def test_message_hash_match(self) -> None: result = check_error_pattern( "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", "IndexOutOfBounds", @@ -144,7 +145,7 @@ def test_message_hash_match(self): assert result.duplicate is True assert result.confidence == "high" - def test_variant_only_does_not_match(self): + def test_variant_only_does_not_match(self) -> None: result = check_error_pattern( "nomatchhash", "ScalarMismatch", @@ -152,13 +153,13 @@ def test_variant_only_does_not_match(self): ) assert result.duplicate is False - def test_no_match(self): + def test_no_match(self) -> None: result = check_error_pattern("nomatch", "UnknownVariant", EXISTING_ISSUES) assert result.duplicate is False class TestCheckDuplicate: - def test_seed_hash_match_first(self, issues_file): + def test_seed_hash_match_first(self, issues_file) -> None: """Seed hash match should return immediately with check_order=1.""" crash_info = _make_crash_info( seed_hash="aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", @@ -168,7 +169,7 @@ def test_seed_hash_match_first(self, issues_file): assert result.check == "seed_hash" assert result.check_order == 1 - def test_panic_location_match_second(self, issues_file): + def test_panic_location_match_second(self, issues_file) -> None: """Panic location should be checked after seed hash.""" crash_info = _make_crash_info( panic_location="vortex-array/src/compute/slice.rs:142", @@ -178,7 +179,7 @@ def test_panic_location_match_second(self, issues_file): assert result.check == "panic_location" assert result.check_order == 2 - def test_stack_trace_match_third(self, issues_file): + def test_stack_trace_match_third(self, issues_file) -> None: """Stack trace should be checked after panic location.""" crash_info = _make_crash_info( stack_trace_hash="bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb", @@ -188,7 +189,7 @@ def test_stack_trace_match_third(self, issues_file): assert result.check == "stack_trace" assert result.check_order == 3 - def test_error_pattern_match_fourth(self, issues_file): + def test_error_pattern_match_fourth(self, issues_file) -> None: """Error pattern should be checked last.""" crash_info = _make_crash_info( message_hash="cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", @@ -198,13 +199,13 @@ def test_error_pattern_match_fourth(self, issues_file): assert result.check == "error_pattern" assert result.check_order == 4 - def test_no_match(self, issues_file): + def test_no_match(self, issues_file) -> None: """Should return no duplicate when nothing matches.""" crash_info = _make_crash_info() result = check_duplicate(crash_info, issues_file) assert result.duplicate is False - def test_empty_issues(self, temp_dir): + def test_empty_issues(self, temp_dir) -> None: """Empty issues file should return no duplicate.""" empty_file = temp_dir / "empty.json" empty_file.write_text("[]") @@ -212,7 +213,7 @@ def test_empty_issues(self, temp_dir): result = check_duplicate(crash_info, str(empty_file)) assert result.duplicate is False - def test_missing_issues_file(self, temp_dir): + def test_missing_issues_file(self, temp_dir) -> None: """Missing issues file should return no duplicate.""" crash_info = _make_crash_info() result = check_duplicate(crash_info, str(temp_dir / "nonexistent.json")) @@ -318,18 +319,18 @@ class TestEndToEndDedup: """ @pytest.fixture - def crash1_info(self, temp_dir): + def crash1_info(self, temp_dir) -> CrashInfo: log_path = temp_dir / "mask_crash.log" log_path.write_text(MASK_STRUCT_CAST_LOG) return extract_crash_info(str(log_path)) @pytest.fixture - def crash2_info(self, temp_dir): + def crash2_info(self, temp_dir) -> CrashInfo: log_path = temp_dir / "decimal_crash.log" log_path.write_text(DECIMAL_SUM_LOG) return extract_crash_info(str(log_path)) - def test_extraction_skips_boilerplate(self, crash1_info, crash2_info): + def test_extraction_skips_boilerplate(self, crash1_info, crash2_info) -> None: """Both crashes should extract real locations, not vortex-error.""" # Crash 1: mask/struct cast assert "vortex-error" not in crash1_info.panic_location @@ -343,16 +344,16 @@ def test_extraction_skips_boilerplate(self, crash1_info, crash2_info): assert "constructor.rs:61" in crash2_info.panic_location assert "decimal" in crash2_info.crash_location - def test_stack_frames_are_different(self, crash1_info, crash2_info): + def test_stack_frames_are_different(self, crash1_info, crash2_info) -> None: """The two crashes should produce entirely different stack frames.""" assert crash1_info.stack_frames != crash2_info.stack_frames assert crash1_info.stack_trace_hash != crash2_info.stack_trace_hash - def test_panic_locations_are_different(self, crash1_info, crash2_info): + def test_panic_locations_are_different(self, crash1_info, crash2_info) -> None: """The two crashes should have different panic locations.""" assert crash1_info.panic_location != crash2_info.panic_location - def test_no_high_confidence_match(self, crash1_info, crash2_info, temp_dir): + def test_no_high_confidence_match(self, crash1_info, crash2_info, temp_dir) -> None: """Crash 2 must NOT match an issue created from crash 1 at high or exact confidence. The old bug would match on 'lib.rs:310' (panic_location check, high confidence). @@ -382,7 +383,7 @@ def test_no_high_confidence_match(self, crash1_info, crash2_info, temp_dir): assert result.check != "stack_trace", f"False stack_trace match! debug={result.debug}" assert result.confidence != "exact", f"False exact match! debug={result.debug}" - def test_same_crash_does_match(self, crash1_info, temp_dir): + def test_same_crash_does_match(self, crash1_info, temp_dir) -> None: """A second occurrence of the SAME crash should still be detected.""" issue_body = _build_issue_body(crash1_info) issues_path = temp_dir / "issues.json" @@ -405,7 +406,7 @@ def test_same_crash_does_match(self, crash1_info, temp_dir): assert result.check in ("panic_location", "stack_trace") assert result.confidence == "high" - def test_debug_info_is_present(self, crash1_info, crash2_info, temp_dir): + def test_debug_info_is_present(self, crash1_info, crash2_info, temp_dir) -> None: """Dedup results should include debug details for diagnosis.""" issue_body = _build_issue_body(crash1_info) issues_path = temp_dir / "issues.json" @@ -534,18 +535,18 @@ class TestIssue6429: """ @pytest.fixture - def file_io_info(self, temp_dir): + def file_io_info(self, temp_dir) -> CrashInfo: p = temp_dir / "file_io.log" p.write_text(ISSUE_6429_FILE_IO_LOG) return extract_crash_info(str(p)) @pytest.fixture - def array_ops_info(self, temp_dir): + def array_ops_info(self, temp_dir) -> CrashInfo: p = temp_dir / "array_ops.log" p.write_text(ISSUE_6429_ARRAY_OPS_LOG) return extract_crash_info(str(p)) - def test_extraction_skips_vortex_error(self, file_io_info, array_ops_info): + def test_extraction_skips_vortex_error(self, file_io_info, array_ops_info) -> None: """Neither crash should reference vortex-error in extracted fields.""" for info in (file_io_info, array_ops_info): assert "vortex-error" not in info.panic_location @@ -554,14 +555,14 @@ def test_extraction_skips_vortex_error(self, file_io_info, array_ops_info): assert "constructor.rs:61" in info.panic_location assert "decimal" in info.crash_location - def test_no_noise_in_stack_frames(self, file_io_info, array_ops_info): + def test_no_noise_in_stack_frames(self, file_io_info, array_ops_info) -> None: """Stack frames should not contain any noise.""" for info in (file_io_info, array_ops_info): assert all("vortex_expect" not in f for f in info.stack_frames) assert all("{closure" not in f for f in info.stack_frames) assert "decimal" in info.stack_frames - def test_same_bug_matches_correctly(self, file_io_info, array_ops_info, temp_dir): + def test_same_bug_matches_correctly(self, file_io_info, array_ops_info, temp_dir) -> None: """array_ops crash should match the file_io issue — same bug.""" issue_body = _build_issue_body(file_io_info) issues_path = temp_dir / "issues.json" @@ -582,7 +583,9 @@ def test_same_bug_matches_correctly(self, file_io_info, array_ops_info, temp_dir assert result.duplicate is True assert result.confidence == "high" - def test_match_reason_references_real_site(self, file_io_info, array_ops_info, temp_dir): + def test_match_reason_references_real_site( + self, file_io_info, array_ops_info, temp_dir + ) -> None: """The match reason must reference the real crash, not boilerplate.""" issue_body = _build_issue_body(file_io_info) issues_path = temp_dir / "issues.json" diff --git a/.github/scripts/fuzz_report/tests/test_extract.py b/.github/scripts/fuzz_report/tests/test_extract.py index 990fe4861ec..d454dd1de69 100644 --- a/.github/scripts/fuzz_report/tests/test_extract.py +++ b/.github/scripts/fuzz_report/tests/test_extract.py @@ -121,26 +121,26 @@ class TestIsNoiseFrame: still error-handling boilerplate, driven by the NOISE_FRAME_PATHS deny list. """ - def test_deny_list_is_not_empty(self): + def test_deny_list_is_not_empty(self) -> None: assert len(NOISE_FRAME_PATHS) > 0 @pytest.mark.parametrize("path", NOISE_FRAME_PATHS) - def test_all_deny_list_entries_are_noise(self, path: str): + def test_all_deny_list_entries_are_noise(self, path: str) -> None: assert _is_noise_frame("some_func", f"{path}:1:1") - def test_closure_in_vortex_error_is_noise(self): + def test_closure_in_vortex_error_is_noise(self) -> None: assert _is_noise_frame( "{closure#1}", "vortex-error/src/lib.rs:457:9", ) - def test_bare_closure_is_noise(self): + def test_bare_closure_is_noise(self) -> None: assert _is_noise_frame("{closure#0}", "some/other/path.rs:1:1") - def test_real_frame_is_not_noise(self): + def test_real_frame_is_not_noise(self) -> None: assert not _is_noise_frame("decimal", "vortex-scalar/src/constructor.rs:61:10") - def test_real_frame_with_generics_is_not_noise(self): + def test_real_frame_with_generics_is_not_noise(self) -> None: assert not _is_noise_frame( "invoke", "vortex-array/src/vtable/compute.rs:120:9", @@ -152,64 +152,64 @@ class TestIsNoiseFunc: that lack file paths (libfuzzer ``#N 0x… in func``, dash ``N: 0x… - func``). """ - def test_prefix_list_is_not_empty(self): + def test_prefix_list_is_not_empty(self) -> None: assert len(NOISE_FUNC_PREFIXES) > 0 @pytest.mark.parametrize("prefix", NOISE_FUNC_PREFIXES) - def test_all_prefixes_are_noise(self, prefix: str): + def test_all_prefixes_are_noise(self, prefix: str) -> None: assert _is_noise_func(f"{prefix}some_function") - def test_std_panicking_is_noise(self): + def test_std_panicking_is_noise(self) -> None: assert _is_noise_func("std::panicking::begin_panic_handler") - def test_core_panicking_is_noise(self): + def test_core_panicking_is_noise(self) -> None: assert _is_noise_func("core::panicking::panic_fmt") - def test_dunder_sanitizer_is_noise(self): + def test_dunder_sanitizer_is_noise(self) -> None: assert _is_noise_func("__sanitizer_print_stack_trace") - def test_name_list_is_not_empty(self): + def test_name_list_is_not_empty(self) -> None: assert len(NOISE_FUNC_NAMES) > 0 @pytest.mark.parametrize("name", sorted(NOISE_FUNC_NAMES)) - def test_all_exact_names_are_noise(self, name: str): + def test_all_exact_names_are_noise(self, name: str) -> None: assert _is_noise_func(name) - def test_vortex_expect_is_noise(self): + def test_vortex_expect_is_noise(self) -> None: assert _is_noise_func("vortex_expect") - def test_vortex_expect_with_generics_is_noise(self): + def test_vortex_expect_with_generics_is_noise(self) -> None: assert _is_noise_func( "vortex_expect" ) - def test_vortex_unwrap_is_noise(self): + def test_vortex_unwrap_is_noise(self) -> None: assert _is_noise_func("vortex_unwrap") - def test_fuzzer_print_stack_trace_is_noise(self): + def test_fuzzer_print_stack_trace_is_noise(self) -> None: assert _is_noise_func("fuzzer::PrintStackTrace") - def test_fuzzer_prefix_is_noise(self): + def test_fuzzer_prefix_is_noise(self) -> None: assert _is_noise_func("fuzzer::Fuzzer::ExecuteCallback") - def test_vortex_func_is_not_noise(self): + def test_vortex_func_is_not_noise(self) -> None: assert not _is_noise_func("vortex_array::compute::slice::slice_primitive") - def test_fuzz_func_is_not_noise(self): + def test_fuzz_func_is_not_noise(self) -> None: assert not _is_noise_func("fuzz::array::run_fuzz_action") - def test_plain_func_is_not_noise(self): + def test_plain_func_is_not_noise(self) -> None: assert not _is_noise_func("decimal") class TestExtractPanicLocation: - def test_standard_format(self): + def test_standard_format(self) -> None: assert extract_panic_location(INDEX_BOUNDS_LOG) == "vortex-array/src/compute/slice.rs:142" - def test_unknown_when_missing(self): + def test_unknown_when_missing(self) -> None: assert extract_panic_location("no panic here") == "unknown" - def test_panicked_at_noise_path_is_skipped(self): + def test_panicked_at_noise_path_is_skipped(self) -> None: """vortex_expect panics report vortex-error/src/lib.rs as the `panicked at` location. This is the macro site, not the real crash. The extractor must skip it and find the real location from the @@ -230,7 +230,7 @@ def test_panicked_at_noise_path_is_skipped(self): assert "lib.rs" not in loc assert "constructor.rs:61" in loc - def test_fallback_skips_noise_paths(self): + def test_fallback_skips_noise_paths(self) -> None: """When the `panicked at` line is absent, the fallback regex scans for vortex paths in the log. It must skip NOISE_FRAME_PATHS like vortex-error/src/lib.rs and return the real crash site instead. @@ -249,11 +249,11 @@ def test_fallback_skips_noise_paths(self): class TestExtractCrashLocation: - def test_with_vortex_frame(self): + def test_with_vortex_frame(self) -> None: loc = extract_crash_location(LIBFUZZER_FRAME_LOG) assert "vortex" in loc - def test_fallback_to_panic_location(self): + def test_fallback_to_panic_location(self) -> None: # Log with panic but no stack frames in "#N 0x... in ..." format log = """thread 'main' panicked at vortex-array/src/compute/slice.rs:142:5: index out of bounds @@ -261,7 +261,7 @@ def test_fallback_to_panic_location(self): loc = extract_crash_location(log) assert "slice.rs:142" in loc - def test_skips_vortex_error_boilerplate(self): + def test_skips_vortex_error_boilerplate(self) -> None: """Two layers of noise filtering in the Rust backtrace format: Layer 1 (implicit via regex): Frames from /rustc/ stdlib paths like @@ -282,7 +282,7 @@ def test_skips_vortex_error_boilerplate(self): # Result: the real crash site assert "decimal" in loc - def test_skips_fuzzer_print_stack_trace(self): + def test_skips_fuzzer_print_stack_trace(self) -> None: """libfuzzer inserts its own C++ frames like fuzzer::PrintStackTrace early in the crash handler stack. These must be skipped. """ @@ -312,48 +312,48 @@ def test_skips_fuzzer_print_stack_trace(self): class TestExtractPanicMessage: - def test_index_bounds(self): + def test_index_bounds(self) -> None: msg = extract_panic_message(INDEX_BOUNDS_LOG) assert "index out of bounds" in msg - def test_scalar_mismatch(self): + def test_scalar_mismatch(self) -> None: msg = extract_panic_message(SCALAR_MISMATCH_LOG) assert "mismatch" in msg.lower() - def test_error_format(self): + def test_error_format(self) -> None: log = "==123== ERROR: libFuzzer: timeout after 120 seconds" msg = extract_panic_message(log) assert "timeout" in msg.lower() class TestExtractErrorVariant: - def test_index_out_of_bounds(self): + def test_index_out_of_bounds(self) -> None: assert extract_error_variant(INDEX_BOUNDS_LOG) == "IndexOutOfBounds" - def test_scalar_mismatch(self): + def test_scalar_mismatch(self) -> None: assert extract_error_variant(SCALAR_MISMATCH_LOG) == "ScalarMismatch" - def test_timeout(self): + def test_timeout(self) -> None: assert extract_error_variant(TIMEOUT_LOG) == "Timeout" - def test_unknown(self): + def test_unknown(self) -> None: assert extract_error_variant("some random log") == "unknown" class TestExtractStackFrames: - def test_dash_format(self): + def test_dash_format(self) -> None: frames = extract_stack_frames(INDEX_BOUNDS_LOG) assert len(frames) > 0 assert any("vortex" in f for f in frames) - def test_in_format(self): + def test_in_format(self) -> None: frames = extract_stack_frames(LIBFUZZER_FRAME_LOG) assert len(frames) > 0 assert any("vortex" in f for f in frames) # std:: frames should be filtered out assert all(not f.startswith("std::") for f in frames) - def test_in_format_non_vortex_crash(self): + def test_in_format_non_vortex_crash(self) -> None: """Crashes in non-vortex code (e.g. fuzz/) should still be captured.""" log = """\ stack backtrace: @@ -368,11 +368,11 @@ def test_in_format_non_vortex_crash(self): assert all(not f.startswith("std::") for f in frames) assert all(not f.startswith("__") for f in frames) - def test_no_frames(self): + def test_no_frames(self) -> None: frames = extract_stack_frames("no stack trace here") assert frames == ["unknown"] - def test_skips_vortex_error_boilerplate(self): + def test_skips_vortex_error_boilerplate(self) -> None: """Two layers of noise filtering in the Rust backtrace format: Layer 1 (implicit via regex): Frames from /rustc/ stdlib paths like @@ -398,23 +398,23 @@ def test_skips_vortex_error_boilerplate(self): class TestExtractStackTraceRaw: - def test_extracts_backtrace(self): + def test_extracts_backtrace(self) -> None: raw = extract_stack_trace_raw(INDEX_BOUNDS_LOG) assert "stack backtrace:" in raw assert "vortex_array" in raw - def test_empty_when_missing(self): + def test_empty_when_missing(self) -> None: raw = extract_stack_trace_raw("no stack trace") assert raw == "" class TestExtractDebugOutput: - def test_extracts_debug(self): + def test_extracts_debug(self) -> None: output = extract_debug_output(DEBUG_OUTPUT_LOG) assert "Array" in output assert "Int32" in output - def test_empty_when_missing(self): + def test_empty_when_missing(self) -> None: output = extract_debug_output("no debug output") assert output == "" @@ -431,20 +431,20 @@ class TestGetCrashType: ("", "unknown"), ], ) - def test_crash_types(self, filename: str, expected: str): + def test_crash_types(self, filename: str, expected: str) -> None: assert get_crash_type(filename) == expected class TestNormalizeMessage: - def test_replaces_numbers(self): + def test_replaces_numbers(self) -> None: assert normalize_message("index 15 of len 10") == "index N of len N" - def test_preserves_text(self): + def test_preserves_text(self) -> None: assert normalize_message("no numbers here") == "no numbers here" class TestExtractCrashInfo: - def test_full_extraction(self): + def test_full_extraction(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: log_path = Path(tmpdir) / "fuzz_output.log" crash_path = Path(tmpdir) / "crash-abc123" @@ -462,7 +462,7 @@ def test_full_extraction(self): assert info.stack_trace_raw != "" assert info.crash_location != "unknown" - def test_without_crash_file(self): + def test_without_crash_file(self) -> None: with tempfile.NamedTemporaryFile(mode="w", suffix=".log", delete=False) as f: f.write(INDEX_BOUNDS_LOG) f.flush() @@ -474,7 +474,7 @@ def test_without_crash_file(self): finally: Path(log_path).unlink() - def test_serialization_roundtrip(self): + def test_serialization_roundtrip(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: log_path = Path(tmpdir) / "fuzz_output.log" log_path.write_text(INDEX_BOUNDS_LOG) diff --git a/.github/scripts/fuzz_report/tests/test_template.py b/.github/scripts/fuzz_report/tests/test_template.py index 31e45c0c0ac..25233ba9ead 100644 --- a/.github/scripts/fuzz_report/tests/test_template.py +++ b/.github/scripts/fuzz_report/tests/test_template.py @@ -14,7 +14,7 @@ @pytest.fixture -def simple_template(tmp_path): +def simple_template(tmp_path) -> str: """Create a simple template file.""" content = "# {{TITLE}}\n\nTarget: {{TARGET}}\nValue: {{VALUE}}\nMissing: {{MISSING}}\n" path = tmp_path / "simple.md" @@ -23,7 +23,7 @@ def simple_template(tmp_path): @pytest.fixture -def conditional_template(tmp_path): +def conditional_template(tmp_path) -> str: """Create a template with Jinja2 conditionals.""" content = "# Title\n{% if ANALYSIS %}\n### Analysis\n{{ANALYSIS}}\n{% endif %}\nFooter\n" path = tmp_path / "conditional.md" @@ -32,19 +32,19 @@ def conditional_template(tmp_path): @pytest.fixture -def new_issue_template(): +def new_issue_template() -> str: """Path to the actual new_issue.md template.""" return str(Path(__file__).parent.parent / "templates" / "new_issue.md") @pytest.fixture -def related_comment_template(): +def related_comment_template() -> str: """Path to the actual related_comment.md template.""" return str(Path(__file__).parent.parent / "templates" / "related_comment.md") class TestRenderTemplate: - def test_with_variables(self, simple_template): + def test_with_variables(self, simple_template) -> None: variables = { "TITLE": "Test Title", "TARGET": "file_io", @@ -56,7 +56,7 @@ def test_with_variables(self, simple_template): assert "42" in result assert "(not set)" in result # MISSING - def test_with_env_variables(self, simple_template): + def test_with_env_variables(self, simple_template) -> None: os.environ["TITLE"] = "Env Title" os.environ["TARGET"] = "env_target" os.environ["VALUE"] = "99" @@ -70,7 +70,7 @@ def test_with_env_variables(self, simple_template): del os.environ["TARGET"] del os.environ["VALUE"] - def test_variables_override_env(self, simple_template): + def test_variables_override_env(self, simple_template) -> None: os.environ["TITLE"] = "Env Title" try: result = render_template( @@ -83,13 +83,13 @@ def test_variables_override_env(self, simple_template): finally: del os.environ["TITLE"] - def test_missing_variable(self, simple_template): + def test_missing_variable(self, simple_template) -> None: result = render_template(simple_template, {}, use_env=False) assert result.count("(not set)") == 4 class TestConditionals: - def test_if_shown_when_truthy(self, conditional_template): + def test_if_shown_when_truthy(self, conditional_template) -> None: result = render_template( conditional_template, {"ANALYSIS": "Root cause is X"}, @@ -98,7 +98,7 @@ def test_if_shown_when_truthy(self, conditional_template): assert "### Analysis" in result assert "Root cause is X" in result - def test_if_hidden_when_empty(self, conditional_template): + def test_if_hidden_when_empty(self, conditional_template) -> None: result = render_template( conditional_template, {"ANALYSIS": ""}, @@ -106,7 +106,7 @@ def test_if_hidden_when_empty(self, conditional_template): ) assert "### Analysis" not in result - def test_if_hidden_when_not_set(self, conditional_template): + def test_if_hidden_when_not_set(self, conditional_template) -> None: result = render_template( conditional_template, {}, @@ -114,13 +114,13 @@ def test_if_hidden_when_not_set(self, conditional_template): ) assert "### Analysis" not in result - def test_footer_always_present(self, conditional_template): + def test_footer_always_present(self, conditional_template) -> None: result = render_template(conditional_template, {}, use_env=False) assert "Footer" in result class TestRenderTemplateToFile: - def test_writes_to_file(self, simple_template, tmp_path): + def test_writes_to_file(self, simple_template, tmp_path) -> None: variables = {"TITLE": "Test", "TARGET": "test", "VALUE": "1"} output_path = tmp_path / "output.md" render_template_to_file(simple_template, str(output_path), variables, use_env=False) @@ -154,7 +154,7 @@ class TestDownstreamGrepCompatibility: "CLAUDE_ANALYSIS": "The crash is caused by an off-by-one error.", } - def test_target_grep_pattern(self, new_issue_template): + def test_target_grep_pattern(self, new_issue_template) -> None: """Verify **Target**: `value` matches downstream grep.""" rendered = render_template(new_issue_template, self.SAMPLE_VARS, use_env=False) # Simulate: grep -oP '(?<=\*\*Target\*\*: `)[^`]+' @@ -162,7 +162,7 @@ def test_target_grep_pattern(self, new_issue_template): assert match is not None, f"Target pattern not found in:\n{rendered}" assert match.group(1) == "file_io" - def test_crash_file_grep_pattern(self, new_issue_template): + def test_crash_file_grep_pattern(self, new_issue_template) -> None: """Verify **Crash File**: `value` matches downstream grep.""" rendered = render_template(new_issue_template, self.SAMPLE_VARS, use_env=False) # Simulate: grep -oP '(?<=\*\*Crash File\*\*: `)[^`]+' @@ -170,7 +170,7 @@ def test_crash_file_grep_pattern(self, new_issue_template): assert match is not None, f"Crash File pattern not found in:\n{rendered}" assert match.group(1) == "crash-abc123" - def test_artifact_url_grep_pattern(self, new_issue_template): + def test_artifact_url_grep_pattern(self, new_issue_template) -> None: """Verify artifact URL matches downstream grep.""" rendered = render_template(new_issue_template, self.SAMPLE_VARS, use_env=False) # Simulate: grep -oP 'https://[^\s]+/artifacts/[0-9]+' @@ -178,33 +178,33 @@ def test_artifact_url_grep_pattern(self, new_issue_template): assert match is not None, f"Artifact URL pattern not found in:\n{rendered}" assert "67890" in match.group(0) - def test_hidden_hashes_present(self, new_issue_template): + def test_hidden_hashes_present(self, new_issue_template) -> None: """Verify hidden HTML comment with hashes is present.""" rendered = render_template(new_issue_template, self.SAMPLE_VARS, use_env=False) assert " + + +# string-bench + +Compares Vortex string encoders on configured real-world Utf8 columns. + +## What it measures + +Three metrics per (column, encoder), from the default `vortex` suite: + +| Metric | Meaning | Table unit | `gh-json` unit | +| --- | --- | --- | --- | +| `size` | Serialized file bytes over canonical uncompressed bytes | % | % | +| `write` | Repartition + zone stats + compress (string scheme and children) + layout + serialize | MB/s | ms | +| `read` | Open + scan, decoding each row split to canonical form | MB/s | ms | + +Values are the median of `--iterations` runs. The `gh-json` output reports +durations in milliseconds so that lower is better for every emitted metric; the +human table reports MB/s over canonical uncompressed bytes. + +### The canonical baseline + +Every metric normalizes against canonical uncompressed bytes: one 16-byte view +per row plus the bytes of the strings too long to inline, so a string of 12 +bytes or fewer costs only its view. The codec suite uses the same baseline. + +### How faithful the timings are + +Write and read run the real Vortex writer and scan on a current-thread runtime +with no worker pool, so both timings are single-threaded CPU costs. + +`read` fuses the canonical decode into each row split's scan task — the shape +production uses, where `into_record_batch_stream` fuses the Arrow conversion — +and drops each decoded chunk before the next split runs, so steady-state memory +is one chunk, not the whole column. Splits are awaited one at a time rather than +through the scan's own stream, which spawns `concurrency * +available_parallelism()` of them at once; on one thread that read-ahead buys no +parallelism, it only holds more chunks in memory and ties the result to the +host's core count. + +Excluded by design: + +- **Physical I/O** — the file lives in a `Vec` and `open_buffer` slices it, + so there is no read-driver request coalescing, no segment cache, and none of + the copying a real file or object-store read does. +- **Parallel scan throughput** — production spreads splits across workers, so + these per-thread costs do not translate directly into query time. +- **Filters and projections** — the whole column is read with no predicate, so + `read` is the full-decode cost, not decode-with-mask or compute pushdown. + +`size` covers encoded children, metadata, padding, and file markers. + +### Codec diagnostic + +`--suite codec` runs a separate microbenchmark: one whole-column encoder call, +with no Vortex layout, child compression, serialization, or I/O. + +It is not part of the tracked set. Its size metric is also not interchangeable +with the file suite's: it trains one dictionary or symbol table for the entire +column and leaves the encoded array's children uncompressed, whereas Vortex +writes in chunks and compresses those children. + +## Inputs and local data + +The current input catalog is: + +| Output name | Source | +| --- | --- | +| `clickbench/URL/shard-N` | ClickBench `hits_N.parquet`, `URL` column | +| `tpch/l_comment` | TPC-H SF1 `lineitem`, `l_comment` column | + +On first use, the selected ClickBench shard is downloaded through +`vortex-bench`'s shared idempotent downloader. TPC-H SF1 is generated +deterministically through the shared TPC-H dataset helper. Both are stored +under `vortex-bench/data/`, which is gitignored, and reused on later runs. +Input preparation is outside benchmark timing. + +## Running locally + +```bash +# Tracked metrics: size, write, read for every configured column and encoder. +cargo run -p string-bench --profile release_debug + +# Focus on selected columns or encoders. +cargo run -p string-bench --profile release_debug -- \ + --columns URL --encoders onpair + +# Add the direct codec microbenchmark. +cargo run -p string-bench --profile release_debug -- \ + --suite both + +# Emit benchmark-comparator JSONL. +cargo run -p string-bench --profile release_debug -- \ + --display-format gh-json --output-path results.json +``` + +Run `cargo run -p string-bench -- --help` for all +filters and tuning options. + +Before timing, the benchmark checks that each requested encoding was produced +and, unless `--no-verify` is set, compares decoded output with the input. + +## CI + +The develop benchmark workflow runs the default suite after each merge to +`develop` and publishes the results to the shared benchmark history. + +Metric names are `//`, for example +`read/clickbench/URL/shard-0/onpair-12` and `size/tpch/l_comment/fsst`. The unit +is reported separately in the JSON output and the CI table, which groups rows by +unit — so metric-first names keep the rows you compare adjacent. diff --git a/benchmarks/string-bench/src/codec.rs b/benchmarks/string-bench/src/codec.rs new file mode 100644 index 00000000000..046cb230bb2 --- /dev/null +++ b/benchmarks/string-bench/src/codec.rs @@ -0,0 +1,276 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Direct array-level codec microbenchmarks. +//! +//! This module deliberately measures the selected encoder's train + compress +//! path over one canonical whole-column array. It does not exercise the Vortex +//! file layout, per-chunk dictionaries, child compression, or file I/O. + +use std::hint::black_box; +use std::time::Duration; +use std::time::Instant; + +use anyhow::Result; +use anyhow::bail; +use vortex::array::ArrayRef; +use vortex::array::ExecutionCtx; +use vortex::array::IntoArray; +use vortex::array::VortexSessionExecute; +use vortex::array::arrays::VarBinViewArray; +use vortex_bench::Format; +use vortex_bench::measurements::CustomUnitMeasurement; +use vortex_fsst::fsst_compress; +use vortex_fsst::fsst_train_compressor; +use vortex_onpair::Config; +use vortex_onpair::DEFAULT_CONFIG; +use vortex_onpair::MaxDictBits; +use vortex_onpair::onpair_compress; + +use crate::SESSION; +use crate::StringColumn; +use crate::StringEncoder; +use crate::duration_ms; +use crate::median; +use crate::onpair_label; +use crate::prepare_column; +use crate::throughput; +use crate::verify_canonicalized; + +/// A fully configured direct array-compression candidate: an encoder family +/// plus the fixed configuration that path needs. +pub enum DirectCandidate { + /// OnPair with a deterministic config, whose dictionary budget is the only + /// thing the benchmark varies. + OnPair(Config), + /// FSST (no configuration). + Fsst, +} + +impl DirectCandidate { + /// OnPair with a dictionary budget of up to `2^max_dict_bits` tokens, + /// reusing every other default. `max_dict_bits` must be in `9..=16`; OnPair + /// stores codes as `u16`, so all of those fit. + pub fn on_pair(max_dict_bits: u8) -> Result { + let max_dict_bits = MaxDictBits::new(max_dict_bits).map_err(|e| { + anyhow::anyhow!( + "invalid maximum dictionary bit width {max_dict_bits} (want 9..=16): {e}" + ) + })?; + Ok(Self::OnPair(Config { + max_dict_bits, + ..DEFAULT_CONFIG + })) + } + + /// The encoder family, for the post-compression type check. + fn family(&self) -> StringEncoder { + match self { + Self::OnPair(_) => StringEncoder::OnPair, + Self::Fsst => StringEncoder::Fsst, + } + } + + /// Stable label used in benchmark output, e.g. `onpair-12` or `fsst`. + fn label(&self) -> String { + match self { + Self::OnPair(config) => onpair_label(config), + Self::Fsst => StringEncoder::Fsst.label().to_string(), + } + } + + /// Run the encoder's direct array-level train + compress path. + pub(crate) fn compress(&self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> Result { + match self { + Self::OnPair(config) => Ok(onpair_compress(array, *config, ctx)?), + Self::Fsst => { + let compressor = fsst_train_compressor(array, ctx)?; + Ok(fsst_compress(array, &compressor, ctx)?.into_array()) + } + } + } +} + +/// Measured in-memory statistics for a single string column under one encoder +/// configuration (an OnPair dictionary size, or FSST). +pub struct ColumnResult { + /// Column identifier. + pub name: String, + /// Encoder configuration label, e.g. `onpair-12` or `fsst`. + pub encoder: String, + /// Number of rows. + pub rows: usize, + /// Canonical uncompressed array bytes used to normalize size and + /// throughput: one 16-byte view per row plus the bytes of the strings too + /// long to inline. + pub uncompressed_bytes: u64, + /// Buffer bytes referenced by the encoded array. + pub encoded_bytes: u64, + /// Direct array-level train + compress times, one per iteration. + pub compression_runs: Vec, +} + +impl ColumnResult { + /// Median direct array compression time across iterations. + fn compression_median(&self) -> Duration { + median(&self.compression_runs) + } + + /// Encoded array buffer bytes as a percentage of canonical uncompressed + /// array bytes. Lower is better. + /// + /// This direct representation has not had its children compressed by the + /// file writer. + pub fn encoded_size_pct(&self) -> f64 { + self.encoded_bytes as f64 / self.uncompressed_bytes as f64 * 100.0 + } + + /// Direct compression throughput in MB/s of canonical uncompressed array + /// bytes, from the median run. + pub fn compression_mbps(&self) -> f64 { + throughput(self.uncompressed_bytes, self.compression_median()) / 1e6 + } + + /// Emit lower-is-better timings and size percentages as Vortex custom-unit + /// metrics, named `codec///`. + /// + /// The `codec/` prefix keeps these distinct from the tracked file metrics, + /// which CI reports: this suite is a local diagnostic, so its measurements + /// only reach `gh-json` when it is explicitly selected. + /// + /// `Format` has no in-memory Vortex variant, so these measurements use + /// `Format::OnDiskVortex` as their reporting target. + pub fn measurements(&self) -> Vec { + let suffix = format!("{}/{}", self.name, self.encoder); + vec![ + CustomUnitMeasurement { + name: format!("codec/size/{suffix}"), + format: Format::OnDiskVortex, + unit: "%".into(), + value: self.encoded_size_pct(), + }, + CustomUnitMeasurement { + name: format!("codec/compress/{suffix}"), + format: Format::OnDiskVortex, + unit: "ms".into(), + value: duration_ms(self.compression_median()), + }, + ] + } +} + +/// Time `candidate`'s direct array-level train + compress path on `column` across +/// `iterations` runs, recording every run so the caller can report the median. +/// +/// Fails if the column contains nulls or does not compress to `candidate`'s +/// encoding, rather than silently changing the measured workload. +pub fn bench_column( + column: &StringColumn, + iterations: usize, + warmup: usize, + candidate: &DirectCandidate, + verify: bool, +) -> Result { + crate::validate_iterations(iterations)?; + let mut ctx = SESSION.create_execution_ctx(); + + let (canonical, uncompressed_bytes) = prepare_column(column, &mut ctx)?; + let input = canonical.clone().into_array(); + let rows = canonical.len(); + + // At least one warm-up produces the reference array for the encoding check, + // size metric, and canonicalization verification. Extra runs stabilize timings. + let mut compressed = candidate.compress(&input, &mut ctx)?; + for _ in 1..warmup.max(1) { + let mut warm_ctx = SESSION.create_execution_ctx(); + compressed = candidate.compress(&input, &mut warm_ctx)?; + } + + if !candidate.family().matches(&compressed) { + bail!( + "column {} did not compress to {} (got {}); its data may be unsupported \ + by the encoder", + column.name, + candidate.label(), + compressed.encoding_id(), + ); + } + let encoded_bytes = compressed.nbytes(); + + // One-time correctness check before timing: a byte-wrong canonicalization + // must not be reported as a fast one. + if verify { + let canonicalized = compressed.execute::(&mut ctx)?; + verify_canonicalized( + &format!("{} [{}]", column.name, candidate.label()), + &canonical, + &canonicalized, + &mut ctx, + )?; + } + + // Each timed run re-compresses the source array from scratch; the fresh + // context is cheap per-run isolation, not a cache reset (`ExecutionCtx` + // holds no cross-run result cache). `black_box` keeps the result live. + let mut compression_runs = Vec::with_capacity(iterations); + for _ in 0..iterations { + let mut ctx = SESSION.create_execution_ctx(); + let start = Instant::now(); + let compressed = candidate.compress(&input, &mut ctx)?; + compression_runs.push(start.elapsed()); + black_box(compressed); + } + + Ok(ColumnResult { + name: column.name.clone(), + encoder: candidate.label(), + rows, + uncompressed_bytes, + encoded_bytes, + compression_runs, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn direct_benchmark_smoke_tests_each_encoder() -> Result<()> { + let (column, expected_uncompressed_bytes) = crate::repeated_fixture(); + + for (candidate, expected_label) in [ + (DirectCandidate::on_pair(12)?, "onpair-12"), + (DirectCandidate::Fsst, "fsst"), + ] { + let result = bench_column(&column, 1, 0, &candidate, true)?; + + assert_eq!(result.encoder, expected_label); + assert_eq!(result.rows, 128); + assert_eq!(result.uncompressed_bytes, expected_uncompressed_bytes); + assert!(result.encoded_bytes > 0); + assert_eq!(result.compression_runs.len(), 1); + } + Ok(()) + } + + #[test] + fn machine_readable_metric_schema_is_stable() { + let result = ColumnResult { + name: "fixture".to_string(), + encoder: "fsst".to_string(), + rows: 1, + uncompressed_bytes: 2, + encoded_bytes: 1, + compression_runs: vec![Duration::from_millis(12)], + }; + + assert_eq!( + crate::measurement_rows(&result.measurements()), + [ + "codec/size/fixture/fsst % 50", + "codec/compress/fixture/fsst ms 12", + ] + ); + } +} diff --git a/benchmarks/string-bench/src/lib.rs b/benchmarks/string-bench/src/lib.rs new file mode 100644 index 00000000000..40e73cf4b15 --- /dev/null +++ b/benchmarks/string-bench/src/lib.rs @@ -0,0 +1,499 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! String-compression benchmarks for Vortex. +//! +//! Two intentionally separate suites: +//! +//! * `serialized` is the default and the only one CI tracks. It writes a column +//! to an in-memory Vortex file with one encoder forced, then reads it back, +//! reporting `size`, `write`, and `read`. +//! * `codec` is a local diagnostic (`--suite codec`) that times one encoder's +//! array-level train + compress call and measures the encoded array's buffer +//! bytes, with no layout, child compression, or I/O in the way. Use it to tell +//! an encoder-level change apart from a change in the file stack. +//! +//! Both suites force the requested encoder rather than letting the btrblocks +//! selector pick, so neither can silently measure a different encoding. Their +//! size metrics are still not interchangeable: the codec path trains one state +//! over the whole column and leaves children uncompressed, while the file path +//! trains one per compression chunk and compresses children. See the README +//! ("Codec diagnostic"). + +use std::path::PathBuf; +use std::sync::LazyLock; +use std::time::Duration; + +use anyhow::Result; +use anyhow::bail; +use clap::ValueEnum; +use futures::StreamExt; +use futures::TryStreamExt; +use parquet::arrow::ParquetRecordBatchStreamBuilder; +use parquet::arrow::ProjectionMask; +use tokio::fs::File; +use vortex::VortexSessionDefault; +use vortex::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; +use vortex::array::ArrayRef; +use vortex::array::ExecutionCtx; +use vortex::array::IntoArray; +use vortex::array::arrays::ChunkedArray; +use vortex::array::arrays::StructArray; +use vortex::array::arrays::VarBinViewArray; +use vortex::array::arrays::struct_::StructArrayExt; +use vortex::io::session::RuntimeSessionExt; +use vortex::session::VortexSession; +use vortex_arrow::ArrowSessionExt; +use vortex_bench::Format; +use vortex_bench::IdempotentPath; +use vortex_bench::datasets::Dataset; +use vortex_bench::datasets::data_downloads::download_data; +use vortex_bench::datasets::tpch_l_comment::TPCHLCommentCanonical; +use vortex_fsst::FSST; +use vortex_onpair::Config; +use vortex_onpair::OnPair; + +mod codec; +pub use codec::ColumnResult; +pub use codec::DirectCandidate; +pub use codec::bench_column; + +const CLICKBENCH_SHARD_COUNT: u32 = 100; +const CLICKBENCH_URL_PREFIX: &str = + "https://pub-3ba949c0f0354ac18db1f0f14f0a2c52.r2.dev/clickbench/parquet_many/"; + +/// Serialized write and read benchmark: write → open → scan, with each row +/// split decoded to canonical form inside its own scan task. +mod serialized; +pub use serialized::*; + +/// Session with the available string encodings and their canonicalize kernels +/// registered for benchmarking. +pub static SESSION: LazyLock = + LazyLock::new(|| VortexSession::default().with_tokio()); + +/// A named string column, canonicalized to a `VarBinViewArray`-backed array, +/// ready for either benchmark path. +pub struct StringColumn { + /// Human-readable column identifier used in measurement names. + pub name: String, + /// Canonical (`VarBinViewArray`) Utf8 data. + pub array: ArrayRef, +} + +/// String encoder selected by the benchmark. +#[derive(Clone, Copy, Debug, Eq, PartialEq, ValueEnum)] +pub enum StringEncoder { + /// The OnPair encoder. + #[value(name = "onpair")] + OnPair, + /// The FSST encoder. + #[value(name = "fsst")] + Fsst, +} + +impl StringEncoder { + /// Stable short label used in benchmark output. + pub fn label(self) -> &'static str { + match self { + Self::OnPair => "onpair", + Self::Fsst => "fsst", + } + } + + /// Whether `array` is this encoder's Vortex encoding. Used to confirm a + /// column was compressed (or serialized) with the requested scheme. + pub fn matches(self, array: &ArrayRef) -> bool { + match self { + Self::OnPair => array.as_opt::().is_some(), + Self::Fsst => array.as_opt::().is_some(), + } + } +} + +impl std::fmt::Display for StringEncoder { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.label()) + } +} + +/// Label for an OnPair configuration, e.g. `onpair-12`. The dictionary budget is +/// read from the config that was actually used, so a label can never claim a +/// budget the measured data was not encoded with. +pub(crate) fn onpair_label(config: &Config) -> String { + format!("onpair-{}", config.max_dict_bits.value()) +} + +/// Bytes per second for `bytes` processed in `elapsed`. +pub(crate) fn throughput(bytes: u64, elapsed: Duration) -> f64 { + let secs = elapsed.as_secs_f64(); + if secs <= 0.0 { + 0.0 + } else { + bytes as f64 / secs + } +} + +/// Duration in milliseconds for lower-is-better machine-readable output. +pub(crate) fn duration_ms(duration: Duration) -> f64 { + duration.as_secs_f64() * 1e3 +} + +fn validate_iterations(iterations: usize) -> Result<()> { + if iterations == 0 { + bail!("iterations must be greater than zero"); + } + Ok(()) +} + +/// Median of `runs` (the two middle values are averaged for an even count). +/// Empty → zero. +pub(crate) fn median(runs: &[Duration]) -> Duration { + if runs.is_empty() { + return Duration::ZERO; + } + let mut sorted = runs.to_vec(); + sorted.sort_unstable(); + let n = sorted.len(); + if n % 2 == 1 { + sorted[n / 2] + } else { + (sorted[n / 2 - 1] + sorted[n / 2]) / 2 + } +} + +/// Return the non-empty, all-valid Utf8 input in compact canonical form, plus +/// the canonical baseline every metric normalizes against: one 16-byte view per +/// row plus the bytes of the strings too long to inline. Preparation and size +/// accounting are outside all timed regions. +pub(crate) fn prepare_column( + column: &StringColumn, + ctx: &mut ExecutionCtx, +) -> Result<(VarBinViewArray, u64)> { + let canonical = column.array.clone().execute::(ctx)?; + if !canonical.dtype().is_utf8() { + bail!( + "column {} has dtype {}; string-bench requires Utf8 input", + column.name, + canonical.dtype(), + ); + } + let valid_count = canonical + .validity()? + .execute_mask(canonical.len(), ctx)? + .true_count(); + if valid_count != canonical.len() { + bail!( + "column {} contains {} null strings; string-bench requires all rows to be valid", + column.name, + canonical.len() - valid_count, + ); + } + + // Use an aggressively compacted array so the baseline cannot include + // retained backing-buffer regions that the compressors never see. + let canonical = canonical.compact_with_threshold(1.0, ctx)?; + let uncompressed_bytes = u64::try_from(uncompressed_size_in_bytes(canonical.as_ref(), ctx)?)?; + if uncompressed_bytes == 0 { + bail!("column {} has zero uncompressed bytes", column.name); + } + + Ok((canonical, uncompressed_bytes)) +} + +/// Assert the canonicalized column matches the original row-for-row (dtype, +/// length, validity, and bytes). This runs once before timing. +pub(crate) fn verify_canonicalized( + label: &str, + expected: &VarBinViewArray, + canonicalized: &VarBinViewArray, + ctx: &mut ExecutionCtx, +) -> Result<()> { + if canonicalized.dtype() != expected.dtype() { + bail!( + "{label}: canonicalized dtype {} != original {}", + canonicalized.dtype(), + expected.dtype(), + ); + } + + let len = expected.len(); + if canonicalized.len() != len { + bail!( + "{label}: canonicalized row count {} != original {len}", + canonicalized.len(), + ); + } + // Materialize both validity masks up front so the per-row `value(i)` and + // `bytes_at(i)` below are O(1); this one-shot check stays off the timed path. + let expected_valid = expected.validity()?.execute_mask(len, ctx)?; + let canonicalized_valid = canonicalized.validity()?.execute_mask(len, ctx)?; + for i in 0..len { + let valid = expected_valid.value(i); + if valid != canonicalized_valid.value(i) { + bail!("{label}: validity mismatch at row {i}"); + } + if valid && expected.bytes_at(i).as_slice() != canonicalized.bytes_at(i).as_slice() { + bail!("{label}: canonicalized value differs from input at row {i}"); + } + } + tracing::debug!("{label}: canonicalized output verified ({len} rows)"); + Ok(()) +} + +/// Canonicalize `array` (a struct, possibly chunked) and pull out `field` as a +/// canonical `VarBinViewArray`-backed Utf8 column. +fn to_utf8_column(array: ArrayRef, field: &str, ctx: &mut ExecutionCtx) -> Result { + let structs = array.execute::(ctx)?; + let column = structs.unmasked_field_by_name(field)?.clone(); + Ok(column.execute::(ctx)?.into_array()) +} + +/// Load the ClickBench `hits` `URL` column from shard `shard`, downloading the +/// shard first if it is not already present locally. +pub async fn load_clickbench_url(shard: u32, ctx: &mut ExecutionCtx) -> Result { + let path = download_clickbench_shard(shard).await?; + let struct_array = read_parquet_projected(path, "URL").await?; + Ok(StringColumn { + name: format!("clickbench/URL/shard-{shard}"), + array: to_utf8_column(struct_array, "URL", ctx)?, + }) +} + +async fn download_clickbench_shard(shard: u32) -> Result { + if shard >= CLICKBENCH_SHARD_COUNT { + bail!( + "invalid ClickBench shard {shard} (want 0..{})", + CLICKBENCH_SHARD_COUNT - 1 + ); + } + let filename = format!("hits_{shard}.parquet"); + let path = "clickbench_partitioned" + .to_data_path() + .join(Format::Parquet.name()) + .join(&filename); + download_data(path, format!("{CLICKBENCH_URL_PREFIX}{filename}")).await +} + +/// Load the TPC-H `l_comment` column (from the first `lineitem` parquet shard), +/// generating the TPC-H data if needed. +pub async fn load_tpch_l_comment(ctx: &mut ExecutionCtx) -> Result { + let path = TPCHLCommentCanonical.to_parquet_path().await?; + let struct_array = read_parquet_projected(path, "l_comment").await?; + Ok(StringColumn { + name: "tpch/l_comment".to_string(), + array: to_utf8_column(struct_array, "l_comment", ctx)?, + }) +} + +/// Read a single column from a parquet file (projected at the parquet level to +/// avoid decoding the other columns) into a chunked struct array. +async fn read_parquet_projected(path: PathBuf, column: &str) -> Result { + let file = File::open(&path).await?; + let builder = ParquetRecordBatchStreamBuilder::new(file).await?; + let col_idx = builder.schema().index_of(column)?; + let mask = ProjectionMask::roots(builder.parquet_schema(), [col_idx]); + let reader = builder.with_projection(mask).build()?; + + let chunks: Vec = reader + .map(|batch| { + batch.map_err(anyhow::Error::from).and_then(|rb| { + let schema = rb.schema(); + SESSION + .arrow() + .from_arrow_record_batch(rb, &schema) + .map_err(anyhow::Error::from) + }) + }) + .try_collect() + .await?; + + Ok(ChunkedArray::from_iter(chunks).into_array()) +} + +/// 128 rows cycling over 8 distinct outlined values, plus the canonical +/// uncompressed size they must account for (one 16-byte view per row). +#[cfg(test)] +fn repeated_fixture() -> (StringColumn, u64) { + let values = (0..128) + .map(|i| { + format!( + "https://example.com/path/to/a/repeated/string-value-{}", + i % 8 + ) + }) + .collect::>(); + let outlined_bytes = values.iter().map(|value| value.len() as u64).sum::(); + ( + StringColumn { + name: "fixture".to_string(), + array: VarBinViewArray::from_iter_str(values.iter()).into_array(), + }, + 128 * 16 + outlined_bytes, + ) +} + +/// Compact `(name, unit, value)` rendering of emitted measurements, for the +/// metric-schema tests that guard the tracked CI names. +#[cfg(test)] +pub(crate) fn measurement_rows( + measurements: &[vortex_bench::measurements::CustomUnitMeasurement], +) -> Vec { + measurements + .iter() + .map(|m| format!("{} {} {}", m.name, m.unit, m.value)) + .collect() +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use vortex::array::VortexSessionExecute; + + use super::*; + + /// The size baseline every metric normalizes against: one 16-byte view per + /// row plus the bytes of the strings too long to inline, counting neither + /// buffer regions a slice merely retains nor anything but the live rows. + #[test] + fn canonical_baseline_counts_views_and_live_outlined_bytes() -> Result<()> { + let kept = "second outlined string kept by the slice"; + let chunk_tail = "an outlined string in the second chunk"; + let cases = [ + ( + "inline and outlined", + VarBinViewArray::from_iter_str(["cat", "hello", "abcdefghijklmnop"]).into_array(), + 3, + 3 * 16 + 16, + ), + ( + "sliced", + VarBinViewArray::from_iter_str([ + "first outlined string dropped by the slice", + kept, + ]) + .into_array() + .slice(1..2)?, + 1, + 16 + kept.len() as u64, + ), + ( + "chunked", + ChunkedArray::from_iter([ + VarBinViewArray::from_iter_str(["alpha", ""]).into_array(), + VarBinViewArray::from_iter_str([chunk_tail]).into_array(), + ]) + .into_array(), + 3, + 3 * 16 + chunk_tail.len() as u64, + ), + ]; + let mut ctx = SESSION.create_execution_ctx(); + + for (name, array, rows, expected_bytes) in cases { + let column = StringColumn { + name: name.to_string(), + array, + }; + let (canonical, uncompressed_bytes) = prepare_column(&column, &mut ctx)?; + + assert_eq!(canonical.len(), rows, "{name} row count"); + assert_eq!(uncompressed_bytes, expected_bytes, "{name} size"); + } + Ok(()) + } + + #[test] + fn verify_canonicalized_rejects_mismatches() { + let cases = [ + ( + "dtype", + VarBinViewArray::from_iter_str(["alpha"]), + VarBinViewArray::from_iter_bin(["alpha".as_bytes()]), + ), + ( + "length", + VarBinViewArray::from_iter_str(["alpha", "beta"]), + VarBinViewArray::from_iter_str(["alpha"]), + ), + ( + "validity", + VarBinViewArray::from_iter_nullable_str([Some("alpha"), Some("beta")]), + VarBinViewArray::from_iter_nullable_str([Some("alpha"), None]), + ), + ( + "value", + VarBinViewArray::from_iter_str(["alpha", "beta"]), + VarBinViewArray::from_iter_str(["alpha", "BETA"]), + ), + ]; + let mut ctx = SESSION.create_execution_ctx(); + + for (kind, expected, actual) in cases { + assert!( + verify_canonicalized("fixture", &expected, &actual, &mut ctx).is_err(), + "{kind} mismatch must be rejected" + ); + } + } + + #[test] + fn median_is_stable_for_odd_and_even_runs() { + let runs = [ + Duration::from_nanos(30), + Duration::from_nanos(10), + Duration::from_nanos(20), + ]; + assert_eq!(median(&runs), Duration::from_nanos(20)); + + let runs = [Duration::from_nanos(10), Duration::from_nanos(30)]; + assert_eq!(median(&runs), Duration::from_nanos(20)); + assert_eq!(median(&[]), Duration::ZERO); + } + + #[test] + fn iterations_must_be_positive() -> Result<()> { + assert!(validate_iterations(0).is_err()); + validate_iterations(1) + } + + #[tokio::test] + async fn clickbench_shard_rejects_out_of_range() { + assert!( + download_clickbench_shard(CLICKBENCH_SHARD_COUNT) + .await + .is_err() + ); + } + + #[test] + fn prepare_column_rejects_invalid_input() { + let cases = [ + ( + "null", + VarBinViewArray::from_iter_nullable_str([Some("alpha"), None]).into_array(), + ), + ( + "empty", + VarBinViewArray::from_iter_str(std::iter::empty::<&str>()).into_array(), + ), + ( + "binary", + VarBinViewArray::from_iter_bin([b"alpha".as_slice()]).into_array(), + ), + ]; + let mut ctx = SESSION.create_execution_ctx(); + + for (name, array) in cases { + let column = StringColumn { + name: name.to_string(), + array, + }; + assert!( + prepare_column(&column, &mut ctx).is_err(), + "{name} input must be rejected" + ); + } + } +} diff --git a/benchmarks/string-bench/src/main.rs b/benchmarks/string-bench/src/main.rs new file mode 100644 index 00000000000..d03304d14c7 --- /dev/null +++ b/benchmarks/string-bench/src/main.rs @@ -0,0 +1,395 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::io::Write; +use std::num::NonZeroUsize; +use std::path::PathBuf; + +use anyhow::Result; +use anyhow::bail; +use clap::Parser; +use clap::ValueEnum; +use indicatif::ProgressBar; +use regex::Regex; +use string_bench::ColumnResult; +use string_bench::DirectCandidate; +use string_bench::SESSION; +use string_bench::SerializedResult; +use string_bench::StringColumn; +use string_bench::StringEncoder; +use string_bench::bench_column; +use string_bench::bench_serialized_with_session; +use string_bench::load_clickbench_url; +use string_bench::load_tpch_l_comment; +use tabled::builder::Builder; +use tabled::settings::Style; +use vortex::VortexSessionDefault; +use vortex::array::ExecutionCtx; +use vortex::array::VortexSessionExecute; +use vortex::io::runtime::BlockingRuntime; +use vortex::io::runtime::current::CurrentThreadRuntime; +use vortex::io::session::RuntimeSessionExt; +use vortex::session::VortexSession; +use vortex_bench::LogFormat; +use vortex_bench::create_output_writer; +use vortex_bench::display::DisplayFormat; +use vortex_bench::display::print_measurements_json; +use vortex_bench::measurements::CustomUnitMeasurement; +use vortex_bench::setup_logging_and_tracing_with_format; + +/// The benchmark ID used for the output path. +const BENCHMARK_ID: &str = "string"; + +/// Repo-relative path of the suite explainer linked from CI benchmark PR comments. +const DOC_PATH: &str = "benchmarks/string-bench/README.md"; + +const CODEC_ONPAIR_DICT_BITS: [u8; 2] = [12, 16]; + +/// The configured benchmark inputs. Add a variant to extend the catalog. +enum Input { + /// The `URL` column of one ClickBench `hits` shard. + ClickBenchUrl(u32), + /// The TPC-H SF1 `lineitem.l_comment` column. + TpchLComment, +} + +impl Input { + /// Name the `--columns` regex is matched against. + fn name(&self) -> &'static str { + match self { + Self::ClickBenchUrl(_) => "URL", + Self::TpchLComment => "l_comment", + } + } + + async fn load(&self, ctx: &mut ExecutionCtx) -> Result { + match *self { + Self::ClickBenchUrl(shard) => load_clickbench_url(shard, ctx).await, + Self::TpchLComment => load_tpch_l_comment(ctx).await, + } + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq, ValueEnum)] +enum BenchmarkSuite { + /// Run only the Vortex file write/read benchmark: the tracked suite. + #[value(name = "vortex")] + Vortex, + /// Run only the direct whole-column codec microbenchmark: a local + /// diagnostic for separating encoder cost from the Vortex file stack. + #[value(name = "codec")] + Codec, + /// Run both benchmark paths. + #[value(name = "both")] + Both, +} + +#[derive(Parser, Debug)] +#[command(version, about, long_about = None)] +struct Args { + /// Benchmark path to run: the tracked Vortex file suite, the direct codec + /// diagnostic, or both. + #[arg(long, value_enum, default_value_t = BenchmarkSuite::Vortex)] + suite: BenchmarkSuite, + /// Timed runs per (column, encoder, config); the median is reported. + #[arg(short, long, default_value = "10")] + iterations: NonZeroUsize, + /// Untimed warm-up runs before timing, per (column, encoder, config). At + /// least one warm-up always runs; correctness validation is also untimed. + #[arg(long, default_value_t = 3)] + warmup: usize, + /// Regex filter matched against configured input-column identifiers. + #[arg(long)] + columns: Option, + /// ClickBench shard index to read the `URL` column from. + #[arg(long, default_value_t = 0)] + clickbench_shard: u32, + /// Skip the one-time canonicalized-output correctness check before timing. + #[arg(long)] + no_verify: bool, + /// Encoders to benchmark (comma-separated). + #[arg(long, value_delimiter = ',', default_values_t = vec![ + StringEncoder::OnPair, + StringEncoder::Fsst, + ])] + encoders: Vec, + /// Output format: `table` for humans, `gh-json` for machine-readable JSONL. + #[arg(short, long, default_value_t, value_enum)] + display_format: DisplayFormat, + /// Write output to this file instead of stdout. + #[arg(short, long)] + output_path: Option, + /// Enable verbose (debug-level) logging. + #[arg(short, long)] + verbose: bool, + /// Enable span tracing output. + #[arg(long)] + tracing: bool, + /// Format for the primary stderr log sink. + #[arg(long, value_enum, default_value_t = LogFormat::Text)] + log_format: LogFormat, +} + +fn contains_duplicates(values: &[T]) -> bool { + values + .iter() + .enumerate() + .any(|(index, value)| values[..index].contains(value)) +} + +fn main() -> Result<()> { + let args = Args::parse(); + setup_logging_and_tracing_with_format(args.verbose, args.tracing, args.log_format)?; + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + runtime.block_on(run(args)) +} + +async fn run(args: Args) -> Result<()> { + let iterations = args.iterations.get(); + let run_codec = matches!(args.suite, BenchmarkSuite::Codec | BenchmarkSuite::Both); + let run_vortex = matches!(args.suite, BenchmarkSuite::Vortex | BenchmarkSuite::Both); + + if args.encoders.is_empty() { + bail!("--encoders must select at least one encoder"); + } + if contains_duplicates(&args.encoders) { + bail!("--encoders must not contain duplicates"); + } + // The codec suite sweeps OnPair dictionary budgets; the file suite cannot, + // since btrblocks picks the budget it compresses with. + let mut candidates: Vec = Vec::new(); + if run_codec { + if args.encoders.contains(&StringEncoder::OnPair) { + for bits in CODEC_ONPAIR_DICT_BITS { + candidates.push(DirectCandidate::on_pair(bits)?); + } + } + if args.encoders.contains(&StringEncoder::Fsst) { + candidates.push(DirectCandidate::Fsst); + } + } + + let filter = args.columns.as_deref().map(Regex::new).transpose()?; + let mut ctx = SESSION.create_execution_ctx(); + let mut columns = Vec::new(); + for input in [ + Input::ClickBenchUrl(args.clickbench_shard), + Input::TpchLComment, + ] + .into_iter() + .filter(|input| filter.as_ref().is_none_or(|f| f.is_match(input.name()))) + { + columns.push(input.load(&mut ctx).await?); + } + if columns.is_empty() { + bail!("no input columns matched the --columns filter"); + } + + let verify = !args.no_verify; + + // `candidates` is empty unless the codec suite runs. + let steps_per_column = candidates.len() + if run_vortex { args.encoders.len() } else { 0 }; + let progress = ProgressBar::new((columns.len() * steps_per_column) as u64); + + let mut column_results: Vec = Vec::new(); + let mut serialized_results: Vec = Vec::new(); + + for column in &columns { + if run_codec { + column_results.extend(run_codec_column( + column, + iterations, + args.warmup, + verify, + &candidates, + &progress, + )?); + } + + if run_vortex { + serialized_results.extend(run_vortex_column( + column, + iterations, + args.warmup, + verify, + &args.encoders, + &progress, + )?); + } + } + progress.finish(); + + let mut writer = create_output_writer(&args.display_format, args.output_path, BENCHMARK_ID)?; + match args.display_format { + DisplayFormat::Table => { + if !serialized_results.is_empty() { + render_serialized_table(&mut writer, &serialized_results)?; + } + if !column_results.is_empty() { + render_codec_table(&mut writer, &column_results)?; + } + } + DisplayFormat::GhJson => { + let mut measurements: Vec = Vec::new(); + for result in &serialized_results { + measurements.extend(result.measurements()); + } + for result in &column_results { + measurements.extend(result.measurements()); + } + print_measurements_json(&mut writer, measurements, DOC_PATH)?; + } + } + + Ok(()) +} + +fn run_codec_column( + column: &StringColumn, + iterations: usize, + warmup: usize, + verify: bool, + candidates: &[DirectCandidate], + progress: &ProgressBar, +) -> Result> { + let mut results = Vec::with_capacity(candidates.len()); + for candidate in candidates { + let result = bench_column(column, iterations, warmup, candidate, verify)?; + tracing::info!( + "{} [{}]: {} rows, {:.3} GB canonical, size {:.2}%, compress {:.2} MB/s", + result.name, + result.encoder, + result.rows, + result.uncompressed_bytes as f64 / 1e9, + result.encoded_size_pct(), + result.compression_mbps(), + ); + results.push(result); + progress.inc(1); + } + Ok(results) +} + +fn run_vortex_column( + column: &StringColumn, + iterations: usize, + warmup: usize, + verify: bool, + encoders: &[StringEncoder], + progress: &ProgressBar, +) -> Result> { + let runtime = CurrentThreadRuntime::new(); + let session = VortexSession::default().with_handle(runtime.handle()); + let mut results = Vec::with_capacity(encoders.len()); + for &encoder in encoders { + let result = runtime.block_on(bench_serialized_with_session( + &session, column, iterations, warmup, verify, encoder, + ))?; + tracing::info!( + "{} [{}]: size {:.2}% | write {:.2} MB/s | read {:.2} MB/s", + result.name, + result.encoder, + result.size_pct(), + result.write_mbps(), + result.read_mbps(), + ); + results.push(result); + progress.inc(1); + } + Ok(results) +} + +/// Render one titled table, prefixed by a one-line note on how to read its +/// metrics. +fn render_table( + writer: &mut dyn Write, + title: &str, + note: &str, + header: &[&str], + rows: Vec>, +) -> Result<()> { + let mut builder = Builder::default(); + builder.push_record(header.iter().copied()); + for row in rows { + builder.push_record(row); + } + let mut table = builder.build(); + table.with(Style::modern()); + writeln!(writer, "{title}\n {note}")?; + writeln!(writer, "{table}")?; + Ok(()) +} + +/// Render the three tracked metrics, one row per (column, encoder). +fn render_serialized_table(writer: &mut dyn Write, results: &[SerializedResult]) -> Result<()> { + render_table( + writer, + "Vortex file: size, write, read", + "Size = % of canonical uncompressed bytes; write = repartition + zone stats + compress \ + (string scheme and children) + layout + serialize; read = open + scan, decoding each row \ + split to canonical form in its own task. Single-threaded, in-memory; MB/s over canonical \ + uncompressed bytes.", + &[ + "Column", + "Encoder", + "Size (%)", + "Write (MB/s)", + "Read (MB/s)", + ], + results + .iter() + .map(|r| { + vec![ + r.name.clone(), + r.encoder.clone(), + format!("{:.2}", r.size_pct()), + format!("{:.2}", r.write_mbps()), + format!("{:.2}", r.read_mbps()), + ] + }) + .collect(), + ) +} + +/// Render the direct codec microbenchmark metrics. +fn render_codec_table(writer: &mut dyn Write, results: &[ColumnResult]) -> Result<()> { + render_table( + writer, + "\nDirect codec microbenchmark (diagnostic)", + "One whole-column codec state, no Vortex layout, child compression, or file I/O. \ + Not comparable with the file suite's size.", + &["Column", "Encoder", "Size (%)", "Compress (MB/s)"], + results + .iter() + .map(|r| { + vec![ + r.name.clone(), + r.encoder.clone(), + format!("{:.2}", r.encoded_size_pct()), + format!("{:.2}", r.compression_mbps()), + ] + }) + .collect(), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rejects_zero_iterations() { + assert!(Args::try_parse_from(["string-bench", "--iterations", "0"]).is_err()); + } + + /// CI passes no `--suite`, so the default is what gets tracked. + #[test] + fn suite_defaults_to_vortex() -> Result<()> { + let args = Args::try_parse_from(["string-bench"])?; + + assert_eq!(args.suite, BenchmarkSuite::Vortex); + Ok(()) + } +} diff --git a/benchmarks/string-bench/src/serialized.rs b/benchmarks/string-bench/src/serialized.rs new file mode 100644 index 00000000000..62de3dd9a51 --- /dev/null +++ b/benchmarks/string-bench/src/serialized.rs @@ -0,0 +1,467 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Serialized Vortex benchmarks. Each timed iteration writes a fresh in-memory +//! file with one string encoder forced and then reads it back: +//! +//! * **write** runs the full default write pipeline — repartition into row +//! blocks, zoned statistics, dictionary probe, coalesce, compress with the +//! forced string scheme and its children, layout, serialize into a buffer; +//! * **read** opens that buffer and runs the scan, decoding each row split to +//! canonical `VarBinViewArray` form inside its own scan task and dropping it +//! before the next split runs — the shape production uses, where +//! `into_record_batch_stream` fuses the Arrow conversion into the split task. +//! +//! Both run on a current-thread runtime, so these are single-threaded CPU costs +//! and exclude physical I/O. + +use std::hint::black_box; +use std::io::Cursor; +use std::sync::Arc; +use std::time::Duration; +use std::time::Instant; + +use anyhow::Result; +use anyhow::bail; +use bytes::Bytes; +use futures::TryStreamExt; +use vortex::array::ArrayRef; +use vortex::array::ExecutionCtx; +use vortex::array::IntoArray; +use vortex::array::VortexSessionExecute; +use vortex::array::arrays::ChunkedArray; +use vortex::array::arrays::VarBinViewArray; +use vortex::compressor::BtrBlocksCompressorBuilder; +use vortex::file::OpenOptionsSessionExt; +use vortex::file::WriteOptionsSessionExt; +use vortex::file::WriteStrategyBuilder; +use vortex::layout::LayoutStrategy; +use vortex::session::VortexSession; +use vortex_bench::Format; +use vortex_bench::measurements::CustomUnitMeasurement; +use vortex_btrblocks::SchemeExt; +use vortex_btrblocks::SchemeId; +use vortex_btrblocks::schemes::integer::DeltaScheme; +use vortex_btrblocks::schemes::string::FSSTScheme; +use vortex_btrblocks::schemes::string::NullDominatedSparseScheme; +use vortex_btrblocks::schemes::string::OnPairScheme; +use vortex_btrblocks::schemes::string::StringDictScheme; +use vortex_onpair::DEFAULT_CONFIG; + +use crate::StringColumn; +use crate::StringEncoder; +use crate::duration_ms; +use crate::median; +use crate::onpair_label; +use crate::prepare_column; +use crate::throughput; +use crate::verify_canonicalized; + +/// The btrblocks string schemes that `BtrBlocksCompressorBuilder::default()` can +/// choose between. Forcing one encoder excludes every entry except its own +/// scheme, so this list must track the default scheme set: add a row whenever a +/// new string encoder becomes selectable by default (e.g. Zstd). +fn default_string_scheme_ids() -> Vec { + vec![ + StringDictScheme.id(), + FSSTScheme.id(), + OnPairScheme.id(), + NullDominatedSparseScheme.id(), + ] +} + +impl StringEncoder { + /// The btrblocks string scheme that produces this encoder's on-disk arrays. + fn scheme_id(self) -> SchemeId { + match self { + Self::OnPair => OnPairScheme.id(), + Self::Fsst => FSSTScheme.id(), + } + } +} + +/// Label for one encoder on the serialized path, carrying the configuration the +/// file was written with, e.g. `onpair-12` or `fsst`. Changing that config +/// renames the metric and so restarts its benchmark history, which is intended: +/// neither size nor read time is comparable across dictionary budgets. +fn serialized_encoder_label(encoder: StringEncoder) -> String { + match encoder { + // The config `OnPairScheme` compresses with. When btrblocks gains a + // configurable budget, pass the benchmark's own config here. + StringEncoder::OnPair => onpair_label(&DEFAULT_CONFIG), + StringEncoder::Fsst => encoder.label().to_string(), + } +} + +/// An in-memory Vortex file produced for one column and encoder. +struct SerializedFile { + file_bytes: u64, + /// Number of row splits the scan will produce, i.e. the number of encoded + /// arrays the read decodes one at a time. + chunk_count: usize, + strategy: Arc, +} + +/// The three core metrics for one serialized write and read of a string column +/// under one encoder: `size`, `write`, and `read`. +pub struct SerializedResult { + /// Column identifier. + pub name: String, + /// Encoder forced for this file, with its configuration, e.g. `onpair-12` + /// or `fsst`. + pub encoder: String, + /// Number of rows. + pub rows: usize, + /// Canonical uncompressed array bytes used to normalize size and + /// throughput: one 16-byte view per row plus the bytes of the strings too + /// long to inline. + pub uncompressed_bytes: u64, + /// Serialized Vortex file size (bytes). + pub file_bytes: u64, + /// Per-iteration write times: the whole default write pipeline, from + /// repartitioning and zoned statistics through compression, layout, and + /// serialization into an in-memory buffer. + pub write_runs: Vec, + /// Per-iteration read times: open the file, then run the scan with each row + /// split decoded to canonical form inside its own task. + pub read_runs: Vec, +} + +impl SerializedResult { + /// Median serialized-write throughput in MB/s of canonical uncompressed + /// array bytes. + pub fn write_mbps(&self) -> f64 { + throughput(self.uncompressed_bytes, median(&self.write_runs)) / 1e6 + } + + /// Median read throughput in MB/s of canonical uncompressed bytes. + pub fn read_mbps(&self) -> f64 { + throughput(self.uncompressed_bytes, median(&self.read_runs)) / 1e6 + } + + /// Complete serialized file bytes as a percentage of canonical + /// uncompressed array bytes. Lower is better. + pub fn size_pct(&self) -> f64 { + self.file_bytes as f64 / self.uncompressed_bytes as f64 * 100.0 + } + + /// Emit the three core metrics as Vortex-format custom-unit metrics, named + /// `//`. + pub fn measurements(&self) -> Vec { + let suffix = format!("{}/{}", self.name, self.encoder); + let ms = |name: String, runs: &[Duration]| CustomUnitMeasurement { + name, + format: Format::OnDiskVortex, + unit: "ms".into(), + value: duration_ms(median(runs)), + }; + vec![ + CustomUnitMeasurement { + name: format!("size/{suffix}"), + format: Format::OnDiskVortex, + unit: "%".into(), + value: self.size_pct(), + }, + ms(format!("write/{suffix}"), &self.write_runs), + ms(format!("read/{suffix}"), &self.read_runs), + ] + } +} + +/// Build the file writer strategy that forces one selected string scheme while +/// leaving editioned non-string child compression enabled. +fn serialized_write_strategy(encoder: StringEncoder) -> Arc { + let forced = encoder.scheme_id(); + let compressor = BtrBlocksCompressorBuilder::default().exclude_schemes( + default_string_scheme_ids() + .into_iter() + .filter(|&id| id != forced) + .chain([DeltaScheme::default().id()]), + ); + WriteStrategyBuilder::default() + .with_btrblocks_builder(compressor) + .build() +} + +/// Write one canonical string column to an in-memory Vortex file, forcing the +/// requested string encoder. +async fn write_serialized_file( + session: &VortexSession, + input: &ArrayRef, + strategy: &Arc, +) -> Result { + let mut buf = Vec::new(); + { + let mut cursor = Cursor::new(&mut buf); + session + .write_options() + .with_strategy(Arc::clone(strategy)) + .write(&mut cursor, input.to_array_stream()) + .await?; + } + Ok(Bytes::from(buf)) +} + +/// Time one complete read of a serialized Vortex buffer: open the file, then run +/// the scan with the canonical decode fused into each row split's task, dropping +/// each decoded chunk before the next split runs. +/// +/// The splits are awaited one at a time rather than through +/// `ScanBuilder::into_array_stream`, which spawns +/// `concurrency * available_parallelism()` of them at once. On a current-thread +/// runtime that read-ahead buys no parallelism; it only holds that many chunks in +/// memory and makes the result depend on the host's core count. Awaiting one at a +/// time keeps the per-split work identical to production while making the +/// measurement machine-independent. +async fn read_serialized_buffer(session: &VortexSession, data: Bytes) -> Result { + let decode_session = session.clone(); + + let start = Instant::now(); + let file = session.open_options().open_buffer(data)?; + let splits = file + .scan()? + .map(move |chunk: ArrayRef| { + let mut ctx = decode_session.create_execution_ctx(); + chunk.execute::(&mut ctx) + }) + .build()?; + + let mut rows = 0usize; + for split in splits { + if let Some(canonical) = split.await? { + rows += canonical.len(); + drop(black_box(canonical)); + } + } + + black_box(rows); + + Ok(start.elapsed()) +} + +/// Write one reference file, then check it is uniformly `encoder`-encoded and, +/// when `verify` is set, that it decodes back to the input. Runs once, outside +/// every timed region: its size is the reported `size` metric. +async fn prepare_serialized_file( + session: &VortexSession, + column: &StringColumn, + input: &ArrayRef, + canonical: &VarBinViewArray, + encoder: StringEncoder, + verify: bool, + ctx: &mut ExecutionCtx, +) -> Result { + let strategy = serialized_write_strategy(encoder); + let data = write_serialized_file(session, input, &strategy).await?; + let file_bytes = data.len() as u64; + + let file = session.open_options().open_buffer(data)?; + let chunks: Vec = file.scan()?.into_array_stream()?.try_collect().await?; + if chunks.is_empty() { + bail!("empty scan for column {}", column.name); + } + if let Some(unexpected) = chunks.iter().find(|array| !encoder.matches(array)) { + bail!( + "serialized column {} contains {} instead of {encoder} — the file was not \ + uniformly {encoder}-encoded (is {encoder} the smallest scheme?)", + column.name, + unexpected.encoding_id(), + ); + } + let chunk_count = chunks.len(); + if verify { + let canonicalized = ChunkedArray::from_iter(chunks) + .into_array() + .execute::(ctx)?; + verify_canonicalized( + &format!("{} [serialized read {encoder}]", column.name), + canonical, + &canonicalized, + ctx, + )?; + } + + Ok(SerializedFile { + file_bytes, + chunk_count, + strategy, + }) +} + +/// Time fresh in-memory serialized writes and reads for one string column and +/// encoder, on the given session's runtime. +pub async fn bench_serialized_with_session( + session: &VortexSession, + column: &StringColumn, + iterations: usize, + warmup: usize, + verify: bool, + encoder: StringEncoder, +) -> Result { + crate::validate_iterations(iterations)?; + let mut ctx = session.create_execution_ctx(); + + let (canonical, uncompressed_bytes) = prepare_column(column, &mut ctx)?; + let input = canonical.clone().into_array(); + let rows = canonical.len(); + let serialized = prepare_serialized_file( + session, column, &input, &canonical, encoder, verify, &mut ctx, + ) + .await?; + + let label = serialized_encoder_label(encoder); + tracing::debug!( + "{} [{label}]: {rows} rows, {} file bytes in {} row splits", + column.name, + serialized.file_bytes, + serialized.chunk_count, + ); + + // Warm up the same complete path that will be timed. The reference file + // above is used only for validation and size. + for _ in 0..warmup.max(1) { + let data = write_serialized_file(session, &input, &serialized.strategy).await?; + let _ = read_serialized_buffer(session, data).await?; + } + + let mut write_runs = Vec::with_capacity(iterations); + let mut read_runs = Vec::with_capacity(iterations); + + for _ in 0..iterations { + let start = Instant::now(); + let data = write_serialized_file(session, &input, &serialized.strategy).await?; + write_runs.push(start.elapsed()); + read_runs.push(read_serialized_buffer(session, data).await?); + } + + Ok(SerializedResult { + name: column.name.clone(), + encoder: label, + rows, + uncompressed_bytes, + file_bytes: serialized.file_bytes, + write_runs, + read_runs, + }) +} + +#[cfg(test)] +mod tests { + use vortex::VortexSessionDefault; + use vortex::array::Canonical; + use vortex::array::VortexSessionExecute; + use vortex::editions::CORE_2026_08_3; + use vortex::editions::EditionSessionExt; + use vortex::io::runtime::BlockingRuntime; + use vortex::io::runtime::current::CurrentThreadRuntime; + use vortex::io::session::RuntimeSessionExt; + use vortex_btrblocks::ALL_SCHEMES; + use vortex_btrblocks::SchemeExt; + + use super::*; + + #[test] + fn forced_scheme_inventory_matches_default_utf8_schemes() { + // Every default scheme whose dtype gate accepts canonical Utf8 must be + // excluded when another root string encoding is forced. + let canonical = Canonical::VarBinView(VarBinViewArray::from_iter_str(["value"])); + let mut actual = ALL_SCHEMES + .iter() + .filter(|scheme| scheme.matches(&canonical)) + .map(|scheme| scheme.id()) + .collect::>(); + let mut expected = default_string_scheme_ids(); + actual.sort_unstable_by_key(|id| id.to_string()); + expected.sort_unstable_by_key(|id| id.to_string()); + + assert_eq!(actual, expected); + } + + /// The three core metric names are the tracked CI schema: renaming one + /// silently restarts its benchmark history. + #[test] + fn machine_readable_metric_schema_is_stable() { + let result = SerializedResult { + name: "fixture".to_string(), + encoder: "fsst".to_string(), + rows: 1, + uncompressed_bytes: 2, + file_bytes: 1, + write_runs: vec![Duration::from_millis(1)], + read_runs: vec![Duration::from_millis(6)], + }; + + assert_eq!( + crate::measurement_rows(&result.measurements()), + [ + "size/fixture/fsst % 50", + "write/fixture/fsst ms 1", + "read/fixture/fsst ms 6", + ] + ); + } + + /// Also pins the encoder labels: a change to the btrblocks OnPair default + /// renames the tracked metrics and restarts their history. + #[test] + fn serialized_benchmark_smoke_tests_each_encoder() -> Result<()> { + let (column, expected_uncompressed_bytes) = crate::repeated_fixture(); + let runtime = CurrentThreadRuntime::new(); + let session = VortexSession::default().with_handle(runtime.handle()); + session.enable_edition(CORE_2026_08_3)?; + + for (encoder, expected_label) in [ + (StringEncoder::OnPair, "onpair-12"), + (StringEncoder::Fsst, "fsst"), + ] { + let result = runtime.block_on(bench_serialized_with_session( + &session, &column, 1, 0, true, encoder, + ))?; + + assert_eq!(result.encoder, expected_label); + assert_eq!(result.rows, 128); + assert_eq!(result.uncompressed_bytes, expected_uncompressed_bytes); + assert!(result.file_bytes > 0); + assert_eq!(result.write_runs.len(), 1); + assert_eq!(result.read_runs.len(), 1); + } + Ok(()) + } + + /// Enough poorly-compressible rows to span several row splits, so the + /// pre-timing check covers a file the read decodes one split at a time. + #[test] + fn serialized_verification_handles_multiple_chunks() -> Result<()> { + let values = (0..50_000_u64) + .map(|i| { + let mixed = i.wrapping_mul(0x9e37_79b9_7f4a_7c15); + format!( + "https://example.com/users/{i:016x}/events/{mixed:016x}/\ + common/repeated/string/payload" + ) + }) + .collect::>(); + let column = StringColumn { + name: "multi-chunk".to_string(), + array: VarBinViewArray::from_iter_str(values.iter()).into_array(), + }; + let runtime = CurrentThreadRuntime::new(); + let session = VortexSession::default().with_handle(runtime.handle()); + let mut ctx = session.create_execution_ctx(); + let (canonical, _) = prepare_column(&column, &mut ctx)?; + let input = canonical.clone().into_array(); + + let serialized = runtime.block_on(prepare_serialized_file( + &session, + &column, + &input, + &canonical, + StringEncoder::Fsst, + true, + &mut ctx, + ))?; + + assert!(serialized.chunk_count > 1); + Ok(()) + } +} diff --git a/docs/AGENTS.md b/docs/AGENTS.md new file mode 100644 index 00000000000..b45df17793c --- /dev/null +++ b/docs/AGENTS.md @@ -0,0 +1,12 @@ +# Documentation Guidance + +Applies to files under `docs/`. + +- Documentation should focus on behavior. Include implementation details only for complex behavior + or private functions. +- Keep prose terse and accurate, and link to relevant external resources, project documentation, + or API types when useful. +- Follow the [documentation development workflow](../CONTRIBUTING.md#documentation) in the + contributor guide. + Keep that contributor guide as the source of truth for Sphinx commands shared by humans and + agents, and do not invoke `sphinx-build` directly. diff --git a/docs/Doxyfile.cpp b/docs/Doxyfile.cpp deleted file mode 100644 index 23072f7479a..00000000000 --- a/docs/Doxyfile.cpp +++ /dev/null @@ -1,28 +0,0 @@ -# Doxygen configuration for Vortex C++ API documentation. -# XML output is consumed by Sphinx via the Breathe extension. - -PROJECT_NAME = "Vortex C++" -OUTPUT_DIRECTORY = _build/doxygen-cpp - -# Input sources -INPUT = ../lang/cpp/include/vortex -FILE_PATTERNS = *.hpp -RECURSIVE = NO - -# We only care about XML output for Breathe -GENERATE_XML = YES -GENERATE_HTML = NO -GENERATE_LATEX = NO -XML_PROGRAMLISTING = YES - -# Extract everything, even if not fully documented yet -EXTRACT_ALL = YES -EXTRACT_PRIVATE = NO -EXTRACT_STATIC = YES - -# Preprocessing — resolve includes but don't expand macros -ENABLE_PREPROCESSING = YES -MACRO_EXPANSION = NO - -# Suppress warnings about undocumented members (WIP API) -WARN_IF_UNDOCUMENTED = NO diff --git a/docs/Makefile b/docs/Makefile index 3daf92fd5ae..115067d75e8 100644 --- a/docs/Makefile +++ b/docs/Makefile @@ -1,24 +1,21 @@ -# Minimal makefile for Sphinx documentation -# - -# You can set these variables from the command line, and also -# from the environment for the first two. SPHINXOPTS ?= --fail-on-warning --keep-going SPHINXBUILD ?= sphinx-build SOURCEDIR = . BUILDDIR = _build -.PHONY: Makefile help +.PHONY: Makefile check help serve + +check: + @$(MAKE) --no-print-directory clean + @$(MAKE) --no-print-directory html + @$(MAKE) --no-print-directory doctest -# Put it first so that "make" without argument is like "make help". help: @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) -# Catch-all target: route all unknown targets to Sphinx using the new -# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). %: Makefile @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) serve: @rm -rf "$(BUILDDIR)" - sphinx-autobuild "$(SOURCEDIR)" "$(BUILDDIR)/html" --re-ignore '$(BUILDDIR)/.*' + uv run --all-packages --with sphinx-autobuild sphinx-autobuild "$(SOURCEDIR)" "$(BUILDDIR)/html" --re-ignore '$(BUILDDIR)/.*' --watch _theme --watch _ext diff --git a/docs/README.md b/docs/README.md index cbe65659a20..849c4fc25dc 100644 --- a/docs/README.md +++ b/docs/README.md @@ -1,36 +1,5 @@ # Vortex Documentation -## Building - -First, you must compile the vortex-python Rust code into a native library because the Python package -inherits some of its doc strings from Rust docstrings: - -``` -cd ../vortex-python && uv run maturin develop -``` - -The docs also require the [`doxygen`](https://www.doxygen.nl/) tool, which can be installed with: - -``` -brew install doxygen -``` - -Build the Vortex docs: - -``` -uv run make html -``` - -## Development - -Live-reloading (ish) build of the docs: - -``` -uv run make serve -``` - -## Python Doctests - -``` -uv run make doctest -``` +See the contributor guide's [documentation development +workflow](../CONTRIBUTING.md#documentation) for environment setup, validation, and live-reloading +instructions. Use `docs/Makefile` as the canonical interface for all Sphinx tasks. diff --git a/docs/_ext/vortex_theme.py b/docs/_ext/vortex_theme.py new file mode 100644 index 00000000000..59069cb3b13 --- /dev/null +++ b/docs/_ext/vortex_theme.py @@ -0,0 +1,87 @@ +""" +Build hooks for our theme. +1. Add pygment styles to vortex.css +2. Load only one CSS and no JS per page except search page +3. Rewrite TOC into detail/summary with regex so it loads without JS +4. Minify CSS +""" + +from __future__ import annotations + +import re +from pathlib import Path +from typing import TYPE_CHECKING + +import rcssmin +from pygments.formatters.html import HtmlFormatter + +if TYPE_CHECKING: + from typing import Any + + from sphinx.application import Sphinx + +_LIGHT_STYLE = "a11y-high-contrast-light" +_DARK_STYLE = "a11y-high-contrast-dark" + +_DOCS_DIR = Path(__file__).parent.parent +_CSS_SOURCE = _DOCS_DIR / "_theme" / "vortex" / "vortex.css" +_GEN_STATIC = _DOCS_DIR / "_build" / "_gen_static" + +_PAGE_CSS = {"_static/vortex.css"} +_PAGE_JS: set[str] = set() +_SEARCH_JS = { + "_static/documentation_options.js", + "_static/doctools.js", + "_static/sphinx_highlight.js", +} + + +def _css(app: Sphinx) -> None: + if app.builder.format != "html": + return + light = HtmlFormatter(style=_LIGHT_STYLE).get_style_defs(".highlight") + dark = HtmlFormatter(style=_DARK_STYLE).get_style_defs(".highlight") + css = rcssmin.cssmin(f"{_CSS_SOURCE.read_text()}\n{light}\n@media (prefers-color-scheme: dark) {{\n{dark}\n}}\n") + out = _GEN_STATIC / "vortex.css" + if not out.exists() or out.read_text() != css: + out.write_text(css) + + +def _filter_assets(app: Sphinx, pagename: str, templatename: str, context: dict[str, Any], doctree: object) -> None: + allowed_js = _SEARCH_JS if pagename == "search" else _PAGE_JS + context["script_files"] = [ + js for js in context.get("script_files", ()) if str(getattr(js, "filename", js)) in allowed_js + ] + context["css_files"] = [ + css for css in context.get("css_files", ()) if str(getattr(css, "filename", css)) in _PAGE_CSS + ] + + +_SUMMARY = '' +_TOC_BRANCH = re.compile(r"\s*]*)>") +_TOC_BRANCH_END = re.compile(r"\s*") +_TOC_CURRENT_BRANCH = re.compile(rf"
({re.escape(_SUMMARY)}]*\bcurrent\b)") +_TOC_CURRENT_PARENT = re.compile(r'().)*)
', re.S) +_TOC_CURRENT_LINK = re.compile(r' None: + toctree = context.get("toctree") + if toctree is None: + context["vortex_globaltoc"] = "" + return + html = toctree(maxdepth=3, collapse=False, includehidden=True, titles_only=True) or "" + html = _TOC_CURRENT_LINK.sub(r'
{_SUMMARY}", html) + html = _TOC_CURRENT_BRANCH.sub(r"
\1", html) + html = _TOC_CURRENT_PARENT.sub(r"\1
", html) + html = _TOC_BRANCH_END.sub("
", html) + context["vortex_globaltoc"] = html + + +def setup(app: Sphinx) -> dict[str, Any]: + _GEN_STATIC.mkdir(parents=True, exist_ok=True) + app.connect("builder-inited", _css) + app.connect("html-page-context", _filter_assets) + app.connect("html-page-context", _toc) + return {"parallel_read_safe": True, "parallel_write_safe": True} diff --git a/docs/_headers b/docs/_headers new file mode 100644 index 00000000000..dfef48d6714 --- /dev/null +++ b/docs/_headers @@ -0,0 +1,13 @@ +# Cloudflare Pages headers + +/* + Content-Security-Policy: default-src 'none'; script-src 'self'; style-src 'self'; style-src-attr 'unsafe-inline'; img-src 'self' data:; connect-src 'self'; form-action 'self'; base-uri 'none'; frame-ancestors 'none' + X-Content-Type-Options: nosniff + Referrer-Policy: strict-origin-when-cross-origin + Cross-Origin-Opener-Policy: same-origin + +/_static/vortex-jni/* + Content-Security-Policy: default-src 'self'; script-src 'self' 'unsafe-inline'; style-src 'self' 'unsafe-inline'; img-src 'self' data: + +/_static/vortex-spark/* + Content-Security-Policy: default-src 'self'; script-src 'self' 'unsafe-inline'; style-src 'self' 'unsafe-inline'; img-src 'self' data: diff --git a/docs/_static/style.css b/docs/_static/style.css deleted file mode 100644 index 603ebc5705d..00000000000 --- a/docs/_static/style.css +++ /dev/null @@ -1,46 +0,0 @@ -/* -- Landing page wordmark ------------------------------------------------ */ - -.vortex-wordmark { - display: block; - max-width: min(400px, 100%); - margin: 1.5rem auto 2.5rem; -} - -/* -- Headings ------------------------------------------------------------- */ - -h2 { - font-size: 1.75rem; -} - -h3 { - font-size: 1.5rem; -} - -h4 { - font-size: 1.25rem; -} - -/* -- Landing page cards --------------------------------------------------- */ - -.sd-card .sd-card-title { - font-size: 1.15rem; - font-weight: 600; - color: var(--pst-color-primary); - border-bottom: 1px solid var(--pst-color-border); - padding-bottom: 0.4rem; - margin-bottom: 0.6rem; -} - -.sd-card .sd-card-body { - font-size: 0.92rem; - line-height: 1.55; -} - -.sd-card { - transition: box-shadow 0.15s ease-in-out, border-color 0.15s ease-in-out; -} - -.sd-card:hover { - border-color: var(--pst-color-primary); - box-shadow: 0 2px 8px rgba(0, 0, 0, 0.08); -} diff --git a/docs/_theme/vortex/domainindex.html b/docs/_theme/vortex/domainindex.html new file mode 100644 index 00000000000..74777ba3251 --- /dev/null +++ b/docs/_theme/vortex/domainindex.html @@ -0,0 +1,40 @@ +{# + sphinx/themes/basic/domainindex.html but without + DOCUMENTATION_OPTIONS.COLLAPSE_INDEX script which we don't need +#} +{%- extends "layout.html" %} +{% set title = indextitle %} +{% block body %} + + {%- set groupid = idgen() %} + +

{{ indextitle }}

+ +
+ {%- for (letter, entries) in content %} + {{ letter }} + {%- if not loop.last %} | {% endif %} + {%- endfor %} +
+ + + {%- for letter, entries in content %} + + + {%- for (name, grouptype, page, anchor, extra, qualifier, description) + in entries %} + + + + {%- endfor %} + {%- endfor %} +
 
+ {{ letter }}
{% if grouptype == 2 %}   {% endif %} + {% if page %}{% endif -%} + {{ name|e }} + {%- if page %}{% endif %} + {%- if extra %} ({{ extra|e }}){% endif -%} + {% if qualifier %}{{ qualifier|e }}:{% endif %} + {{ description|e }}
+ +{% endblock %} diff --git a/docs/_theme/vortex/layout.html b/docs/_theme/vortex/layout.html new file mode 100644 index 00000000000..30be1fd64b6 --- /dev/null +++ b/docs/_theme/vortex/layout.html @@ -0,0 +1,89 @@ + +{%- if not embedded and docstitle %} + {%- set titlesuffix = " — "|safe + docstitle|e %} +{%- else %} + {%- set titlesuffix = "" %} +{%- endif %} + + + + + {{- metatags }} + {%- block htmltitle %} + {{ title|striptags|e }}{{ titlesuffix }} + {%- endblock %} + + + + {%- if pageurl %} + + {%- endif %} + + {%- for css in css_files %} + {{ css_tag(css) }} + {%- endfor %} + {%- block scripts %} + {%- for js in script_files %} + {{ js_tag(js) }} + {%- endfor %} + {%- endblock %} + {%- if favicon_url %} + + {%- endif %} + {%- if pageurl %} + + {%- endif %} + {%- if hasdoc('genindex') %} + + {%- endif %} + {%- if hasdoc('search') %} + + {%- endif %} + {%- block extrahead %}{% endblock %} + + + + +
+ + + + +
+
+ +
+ {%- block body %}{% endblock %} +
+ {%- if page_source_suffix %} + {{ _('Edit this page') }} · + {%- endif %} + © {{ copyright|e }}. + {%- if show_sphinx %} {% trans sphinx_version=sphinx_version|e %}Created using Sphinx {{ sphinx_version }}.{% endtrans %}{% endif %} +
+
+ {%- if display_toc %} + + {%- endif %} +
+ + diff --git a/docs/_theme/vortex/theme.toml b/docs/_theme/vortex/theme.toml new file mode 100644 index 00000000000..e82120357db --- /dev/null +++ b/docs/_theme/vortex/theme.toml @@ -0,0 +1,7 @@ +[theme] +inherit = "basic" +stylesheets = [ + "vortex.css", +] + +[options] diff --git a/docs/_theme/vortex/vortex.css b/docs/_theme/vortex/vortex.css new file mode 100644 index 00000000000..a70c4dbc268 --- /dev/null +++ b/docs/_theme/vortex/vortex.css @@ -0,0 +1,304 @@ +:root { + --bg: #ffffff; + --bg-alt: #f5f6f7; + --fg: #101828; + --fg-muted: #4a5565; + --link: #0e7490; + --border: #e5e7eb; + --accent: #0e7490; + --warn: #9a6700; + --danger: #a40e26; + --mark: #fff8c5; + color-scheme: light dark; +} +@media (prefers-color-scheme: dark) { + :root { + --bg: #101010; + --bg-alt: #1a1a1a; + --fg: #ececec; + --fg-muted: #99a1af; + --link: #2cb9d1; + --border: #272727; + --accent: #2cb9d1; + --warn: #d29922; + --danger: #ff8182; + --mark: #693e00; + } + .only-light { display: none !important; } +} +@media (prefers-color-scheme: light) { + .only-dark { display: none !important; } +} +* { box-sizing: border-box; } +body { + margin: 0; + background: var(--bg); + color: var(--fg); + font: 1rem/1.6 system-ui, -apple-system, "Segoe UI", Roboto, sans-serif; +} +img, svg { max-width: 100%; height: auto; } +a { color: var(--link); text-decoration: none; } +a:hover { text-decoration: underline; } +:focus-visible { outline: 2px solid var(--accent); outline-offset: 2px; } + +h1, h2, h3, h4, h5, h6 { line-height: 1.25; margin: 1.5em 0 0.5em; } +h1 { font-size: 1.75rem; margin-top: 0.5em; } +h2 { font-size: 1.35rem; } +h3 { font-size: 1.15rem; } + +.visually-hidden { + position: absolute; + width: 1px; height: 1px; + margin: -1px; padding: 0; + overflow: hidden; + clip-path: inset(50%); + white-space: nowrap; +} +.skip-link { + position: absolute; + left: 0.5rem; top: -10rem; + z-index: 3; + padding: 0.5rem 1rem; + background: var(--bg-alt); +} +.skip-link:focus { top: 0.5rem; } + +.header { + display: flex; + align-items: center; + gap: 1rem; + padding: 0.5rem 1rem; + border-bottom: 1px solid var(--border); +} +.logo { + display: flex; + align-items: center; + gap: 0.5rem; + color: var(--fg); + font-weight: 600; + font-size: 1.1rem; +} +.search { flex: 1; max-width: 24rem; } +.search input { + width: 100%; + padding: 0.3rem 0.6rem; + border: 1px solid var(--border); + border-radius: 0.4rem; + background: var(--bg); + color: var(--fg); + font: inherit; +} +.extlinks { + display: flex; + align-items: center; + gap: 0.75rem; + margin-left: auto; +} +.extlinks a { color: var(--fg-muted); } +.extlinks svg { display: block; } +.nav-toggle { display: none; color: var(--fg-muted); cursor: pointer; } +#nav-toggle:focus-visible ~ .header .nav-toggle { + outline: 2px solid var(--accent); +} + +.layout { + display: flex; + gap: 2rem; + max-width: 90rem; + margin: 0 auto; + padding: 0 1rem; +} +.sidebar { + flex: 0 0 16rem; + position: sticky; + top: 0; + max-height: 100vh; + overflow-y: auto; + padding: 1rem 0.5rem 2rem 0; + border-right: 1px solid var(--border); + font-size: 0.9rem; +} +.content { + flex: 1; + min-width: 0; + max-width: 48rem; + padding: 1rem 0 3rem; +} +.pagetoc { + flex: 0 0 13rem; + position: sticky; + top: 0; + max-height: 100vh; + overflow-y: auto; + padding: 1rem 0; + font-size: 0.85rem; +} +.pagetoc > nav > p { margin: 0 0 0.25rem; font-weight: 600; } +.pagetoc ul { list-style: none; margin: 0; padding-left: 0; } +.pagetoc ul ul { padding-left: 0.9rem; } +.pagetoc a { color: var(--fg-muted); display: inline-block; padding: 0.1rem 0; } +.pagetoc a:hover { color: var(--link); } + +.sidebar ul { list-style: none; margin: 0; padding-left: 0; } +.sidebar ul ul { padding-left: 0.9rem; } +.sidebar li { position: relative; } +.sidebar a { + display: block; + padding: 0.15rem 1.6rem 0.15rem 0.4rem; + border-radius: 0.3rem; + color: var(--fg-muted); +} +.sidebar a:hover { color: var(--link); text-decoration: none; } +.sidebar a.current { + color: var(--accent); + font-weight: 600; + background: var(--bg-alt); +} +.sidebar summary { + position: absolute; + top: 0.15rem; right: 0.2rem; + width: 1.3rem; + list-style: none; + cursor: pointer; + text-align: center; + color: var(--fg-muted); + user-select: none; +} +.sidebar summary::-webkit-details-marker { display: none; } +.sidebar summary::after { content: "➤"; display: inline-block; } +.sidebar details[open] > summary::after { transform: rotate(90deg); } + +.footer { + margin-top: 3rem; + padding-top: 1rem; + border-top: 1px solid var(--border); + font-size: 0.85rem; + color: var(--fg-muted); +} + +pre, code, .sig { + font-family: monospace; + font-size: 0.875em; +} +pre { + padding: 0.75rem 1rem; + overflow-x: auto; + line-height: 1.45; +} +[class*="highlight-"] { margin: 1rem 0; } +.highlight pre { margin: 0; border: 1px solid var(--border); border-radius: 0.4rem; } +.literal { + background: var(--bg-alt); + padding: 0.1em 0.3em; + border-radius: 0.3em; + overflow-wrap: break-word; +} +dl.py, dl.c { margin: 1.5rem 0; } +.sig { + padding: 0.4rem 0.6rem; + background: var(--bg-alt); + border-left: 3px solid var(--accent); + border-radius: 0 0.3rem 0.3rem 0; + overflow-x: auto; +} +.sig .sig-name { font-weight: 700; } +.sig + dd, .field-list + * { margin-top: 0.5rem; } +dd { margin-left: 1.5rem; } +.field-list { display: grid; grid-template-columns: max-content auto; gap: 0.25rem 1rem; } +.field-list > dt { font-weight: 600; } +.field-list > dd { margin: 0; } +.property { color: var(--fg-muted); font-style: normal; } + +a.headerlink { + visibility: hidden; + margin-left: 0.3rem; + color: var(--fg-muted); + font-size: 0.8em; +} +h1:hover > a.headerlink, h2:hover > a.headerlink, h3:hover > a.headerlink, +h4:hover > a.headerlink, dt:hover > a.headerlink { + visibility: visible; +} +a.headerlink:focus-visible { + visibility: visible; +} + +.admonition { + margin: 1rem 0; + padding: 0.5rem 1rem; + border: 1px solid var(--border); + border-left: 3px solid var(--accent); + border-radius: 0.3rem; + background: var(--bg-alt); + font-size: 0.95rem; +} +.admonition > .admonition-title { font-weight: 600; margin: 0 0 0.25rem; } +.admonition > :last-child { margin-bottom: 0; } +.warning { border-left-color: var(--warn); } +.warning > .admonition-title { color: var(--warn); } +.danger, .error { border-left-color: var(--danger); } +.danger > .admonition-title, .error > .admonition-title { color: var(--danger); } + +table.docutils { + border-collapse: collapse; + display: block; + overflow-x: auto; + margin: 1rem 0; +} +table.docutils th, table.docutils td { + border: 1px solid var(--border); + padding: 0.3rem 0.75rem; + text-align: left; +} +table.docutils th { background: var(--bg-alt); } + +.page-index .content > section > h1 { + position: absolute; + width: 1px; height: 1px; + margin: -1px; + overflow: hidden; + clip-path: inset(50%); +} +.vortex-wordmark { display: block; max-width: min(400px, 100%); margin: 1.5rem auto 2.5rem; } +.cards { + display: grid; + grid-template-columns: repeat(auto-fit, minmax(15rem, 1fr)); + gap: 1rem; +} +.card { + position: relative; + border: 1px solid var(--border); + border-radius: 0.4rem; + padding: 0.75rem 1rem; + font-size: 0.92rem; +} +.card > p:first-child { margin-top: 0; } +.card > p:first-child a::after { content: ""; position: absolute; inset: 0; } +.card > :last-child { margin-bottom: 0; } + +#search-results .search { list-style: none; padding-left: 0; } +#search-results .search li { margin: 1rem 0; } +span.highlighted { background: var(--mark); } + +.indextable { width: 100%; } +.indextable td { vertical-align: top; padding: 0.1rem 0.5rem; } + +blockquote { margin: 1rem 0; padding-left: 1rem; border-left: 3px solid var(--border); color: var(--fg-muted); } +hr { border: 0; border-top: 1px solid var(--border); margin: 2rem 0; } + +@media (max-width: 70em) { + .pagetoc { display: none; } +} +@media (max-width: 50em) { + .layout { display: block; } + .sidebar { + display: none; + position: static; + max-height: none; + border-right: 0; + border-bottom: 1px solid var(--border); + } + #nav-toggle:checked ~ .layout .sidebar { display: block; } + .nav-toggle { display: block; } + .search { max-width: none; } +} diff --git a/docs/api/index.md b/docs/api/index.md index c4fcf3fe182..343661f6717 100644 --- a/docs/api/index.md +++ b/docs/api/index.md @@ -11,6 +11,5 @@ maxdepth: 2 python/index Rust API c/index -cpp/index java/index ``` diff --git a/docs/api/java/index.rst b/docs/api/java/index.rst index 9cb815a0aba..a687b8aac14 100644 --- a/docs/api/java/index.rst +++ b/docs/api/java/index.rst @@ -32,6 +32,7 @@ The Java bindings are supported on the following architectures: * x86_64 Linux * ARM64 Linux * Apple Silicon macOS +* x86_64 Windows They support any Linux distribution with a GLIBC version >= 2.31. This includes @@ -46,14 +47,31 @@ Here's a basic example of using the Vortex Java API to read a Vortex file: .. code-block:: java - import dev.vortex.api.File; - import dev.vortex.api.Array; - - // Open a Vortex file - File vortexFile = File.open("path/to/file.vortex"); - - // Read arrays from the file - Array array = vortexFile.readArray(); - - // Work with the array data - System.out.println("Array length: " + array.getLength()); + import dev.vortex.api.DataSource; + import dev.vortex.api.Partition; + import dev.vortex.api.Scan; + import dev.vortex.api.ScanOptions; + import dev.vortex.api.Session; + import dev.vortex.arrow.ArrowAllocation; + import org.apache.arrow.memory.BufferAllocator; + import org.apache.arrow.vector.VectorSchemaRoot; + import org.apache.arrow.vector.ipc.ArrowReader; + + BufferAllocator allocator = ArrowAllocation.rootAllocator(); + Session session = Session.create(); + DataSource source = DataSource.open(session, "path/to/file.vortex"); + + // A scan yields one partition per chunk of the file. + Scan scan = source.scan(ScanOptions.of()); + while (scan.hasNext()) { + Partition partition = scan.next(); + try (ArrowReader reader = partition.scanArrow(allocator)) { + while (reader.loadNextBatch()) { + VectorSchemaRoot batch = reader.getVectorSchemaRoot(); + System.out.println("read " + batch.getRowCount() + " rows"); + } + } + } + +Data crosses the JNI boundary as Arrow record batches, so the buffers stay in native memory and +are read from Java through the Arrow C Data Interface. diff --git a/docs/api/python/datasets.rst b/docs/api/python/datasets.rst new file mode 100644 index 00000000000..24e27e5fba4 --- /dev/null +++ b/docs/api/python/datasets.rst @@ -0,0 +1,53 @@ +Hugging Face Datasets +===================== + +Vortex files can be loaded directly as Hugging Face ``datasets`` objects. Install the optional +dependencies with: + +.. code-block:: bash + + pip install vortex-data[hf] + +``vortex.datasets.load_dataset`` accepts a local path or glob, a URL readable by Vortex +(``https://``, ``s3://``, ``gs://``, ``az://``), a Hugging Face Hub dataset repository id, or an +``hf://`` URI of the form ``hf://datasets//[@][/]`` (revisions +containing ``/`` must be percent-encoded, e.g. ``@refs%2Fconvert%2Fparquet``). +Unlike ``datasets.load_dataset``, it defaults to ``streaming=True`` and returns a streaming +dataset that keeps Vortex in charge of reading. Column selection, Vortex filter expressions, and +row limits are pushed down into each Vortex scan before examples are yielded to Hugging Face +transforms. + +Hub repositories are streamed in place: files are read with HTTP range requests, so only the +projected columns and matching rows are ever transferred. Private and gated repositories +authenticate with the ``token`` argument, ``HF_TOKEN``, or the locally saved login — see +:doc:`store/huggingface` for the full precedence. Files are downloaded (with the usual Hub +caching) only when ``streaming=False`` or ``local_files_only=True``. + +.. code-block:: python + + import vortex as vx + + # Stream a local directory of Vortex files as a Hugging Face IterableDataset. + ds = vx.datasets.load_dataset("path/to/dir", split="train") + + # Stream a Hub dataset directly, without downloading it. + ds = vx.datasets.load_dataset("username/dataset", split="train") + + # hf:// URIs select a revision and a path or glob within the repository. + ds = vx.datasets.load_dataset("hf://datasets/username/dataset@main/data/*.vortex", split="train") + + # Column projection and Vortex filter expressions are pushed into the scan. + ds = ds.select_columns(["text", "label"]).filter(vx.expr.column("label") == 1) + + for example in ds.take(100): + ... + + # Materialize an in-memory dataset instead, optionally selecting splits via data_files. + splits = vx.datasets.load_dataset( + "username/dataset", + data_files={"train": "train-*.vortex", "validation": "validation.vortex"}, + streaming=False, + ) + +When loading multiple splits, pass a ``data_files`` mapping from split name to file patterns; a +list of splits without such a mapping cannot be resolved and raises an error. diff --git a/docs/api/python/expr.rst b/docs/api/python/expr.rst index 6361b558347..9ee4a7d4450 100644 --- a/docs/api/python/expr.rst +++ b/docs/api/python/expr.rst @@ -11,25 +11,174 @@ the following expression represents the set of rows for which the `age` column l >>> age = vortex.expr.column("age") >>> (23 > age) & (age < 55) # doctest: +SKIP +Expressions are picklable, so a filter built in one process can be sent to another (for example to a +``multiprocessing`` worker or a Ray task). Pickling uses the same protobuf wire format exposed by +:meth:`vortex.expr.Expr.serialize` and :func:`vortex.expr.deserialize`. + .. autosummary:: :nosignatures: - ~vortex.expr.column ~vortex.expr.Expr + ~vortex.expr.root + ~vortex.expr.column + ~vortex.expr.literal + ~vortex.expr.get_item + ~vortex.expr.not_ + ~vortex.expr.and_ + ~vortex.expr.or_ + ~vortex.expr.and_collect + ~vortex.expr.or_collect + ~vortex.expr.eq + ~vortex.expr.not_eq + ~vortex.expr.gt + ~vortex.expr.gt_eq + ~vortex.expr.lt + ~vortex.expr.lt_eq + ~vortex.expr.add + ~vortex.expr.sub + ~vortex.expr.mul + ~vortex.expr.div + ~vortex.expr.between + ~vortex.expr.is_null + ~vortex.expr.is_not_null + ~vortex.expr.fill_null + ~vortex.expr.like + ~vortex.expr.ilike + ~vortex.expr.not_like + ~vortex.expr.not_ilike + ~vortex.expr.byte_length + ~vortex.expr.select + ~vortex.expr.select_exclude + ~vortex.expr.pack + ~vortex.expr.merge + ~vortex.expr.list_contains + ~vortex.expr.list_length + ~vortex.expr.list_sum + ~vortex.expr.case_when + ~vortex.expr.zip_ + ~vortex.expr.mask + ~vortex.expr.cast + ~vortex.expr.ext_storage + ~vortex.expr.variant_get + ~vortex.expr.deserialize .. raw:: html
+Leaves and scope +---------------- + +.. autofunction:: vortex.expr.root + .. autofunction:: vortex.expr.column +.. autofunction:: vortex.expr.literal + +.. autofunction:: vortex.expr.get_item + +Boolean logic +------------- + .. autofunction:: vortex.expr.not_ .. autofunction:: vortex.expr.and_ -.. autofunction:: vortex.expr.root +.. autofunction:: vortex.expr.or_ -.. autofunction:: vortex.expr.literal +.. autofunction:: vortex.expr.and_collect + +.. autofunction:: vortex.expr.or_collect + +Comparisons and arithmetic +-------------------------- + +.. autofunction:: vortex.expr.eq + +.. autofunction:: vortex.expr.not_eq + +.. autofunction:: vortex.expr.gt + +.. autofunction:: vortex.expr.gt_eq + +.. autofunction:: vortex.expr.lt + +.. autofunction:: vortex.expr.lt_eq + +.. autofunction:: vortex.expr.add + +.. autofunction:: vortex.expr.sub + +.. autofunction:: vortex.expr.mul + +.. autofunction:: vortex.expr.div + +.. autofunction:: vortex.expr.between + +Nullability +----------- + +.. autofunction:: vortex.expr.is_null + +.. autofunction:: vortex.expr.is_not_null + +.. autofunction:: vortex.expr.fill_null + +Strings +------- + +.. autofunction:: vortex.expr.like + +.. autofunction:: vortex.expr.ilike + +.. autofunction:: vortex.expr.not_like + +.. autofunction:: vortex.expr.not_ilike + +.. autofunction:: vortex.expr.byte_length + +Structs +------- + +.. autofunction:: vortex.expr.select + +.. autofunction:: vortex.expr.select_exclude + +.. autofunction:: vortex.expr.pack + +.. autofunction:: vortex.expr.merge + +Lists +----- + +.. autofunction:: vortex.expr.list_contains + +.. autofunction:: vortex.expr.list_length + +.. autofunction:: vortex.expr.list_sum + +Conditionals and conversions +---------------------------- + +.. autofunction:: vortex.expr.case_when + +.. autofunction:: vortex.expr.zip_ + +.. autofunction:: vortex.expr.mask + +.. autofunction:: vortex.expr.cast + +.. autofunction:: vortex.expr.ext_storage + +.. autofunction:: vortex.expr.variant_get + +Serialization +------------- + +.. autofunction:: vortex.expr.deserialize + +The expression class +-------------------- .. autoclass:: vortex.expr.Expr :members: diff --git a/docs/api/python/index.rst b/docs/api/python/index.rst index 1ef85a4acb6..c50c45efe8c 100644 --- a/docs/api/python/index.rst +++ b/docs/api/python/index.rst @@ -16,7 +16,7 @@ Optional integrations can be installed as extras: .. code-block:: bash - pip install vortex-data[polars,pandas,numpy,duckdb,ray] + pip install vortex-data[polars,pandas,numpy,duckdb,ray,hf] Compatibility @@ -27,6 +27,8 @@ The Python bindings require Python 3.11 or newer. Pre-built wheels are available * x86_64 Linux * ARM64 Linux * Apple Silicon macOS +* x86_64 macOS +* x86_64 Windows They support any Linux distribution with a GLIBC version >= 2.17. This includes @@ -66,5 +68,6 @@ API Reference io store dataset + datasets runtime type_aliases diff --git a/docs/api/python/store.rst b/docs/api/python/store.rst index 305657c65f1..59bd6661695 100644 --- a/docs/api/python/store.rst +++ b/docs/api/python/store.rst @@ -11,8 +11,10 @@ Vortex arrays support reading and writing to many object storage systems: store/gcs store/azure store/http + store/huggingface store/local store/memory + store/opendal store/config .. autofunction:: vortex.store.from_url diff --git a/docs/api/python/store/huggingface.rst b/docs/api/python/store/huggingface.rst new file mode 100644 index 00000000000..fd6831ceeb4 --- /dev/null +++ b/docs/api/python/store/huggingface.rst @@ -0,0 +1,97 @@ +================== +Hugging Face Hub +================== + +Vortex reads Hugging Face Hub repositories over ``hf://`` URLs. A Hub repository is a set of files +behind an HTTP endpoint that honours range requests, so no cloud SDK is involved and no extra build +feature is needed. + +.. list-table:: + :header-rows: 1 + + * - URL + - Repository kind + * - ``hf://datasets//[@][/]`` + - Dataset + * - ``hf://spaces//[@][/]`` + - Space + * - ``hf:///[@][/]`` + - Model + +```` is a branch, tag or commit, defaulting to ``main``. A revision containing ``/`` must +be percent-encoded, e.g. ``hf://datasets/org/name@refs%2Fconvert%2Fparquet/data/train.vortex``. + +Configuration comes from the same environment variables ``huggingface_hub`` reads: + +.. list-table:: + :header-rows: 1 + + * - Variable + - Meaning + * - ``HF_TOKEN`` + - API token for private and gated repositories. Falls back to the token file at + ``HF_TOKEN_PATH``, then ``$HF_HOME/token``, then ``$HOME/.cache/huggingface/token``. + * - ``HF_ENDPOINT`` + - Hub endpoint, defaulting to ``https://huggingface.co``. + +Reading from the Hub +==================== + +Pass an ``hf://`` URL directly. Public repositories need no credentials; private and gated ones +authenticate from ``HF_TOKEN`` or the saved login: + +.. code-block:: python + + import vortex as vx + + vxf = vx.open("hf://datasets/org/name/data/train.vortex") + for batch in vxf.to_arrow(): + ... + +:class:`vortex.store.HfStore` +============================= + +.. py:class:: vortex.store.HfStore(repo_id, *, repo_type="dataset", revision=None, token=None, endpoint=None) + + A Hugging Face Hub object store, rooted at one repository and revision. + + A URL is enough for most reads, so reach for this class only for the two things a URL cannot + express: a token held in a variable rather than the environment, and a read that must stay + anonymous even though the environment offers credentials. + + Because the store is rooted at the repository and revision, the path passed alongside it is a + path *within* the repository. + + :param repo_id: The repository, as ``"/"``. + :param repo_type: ``"dataset"``, ``"model"`` or ``"space"``. Defaults to ``"dataset"``. + :param revision: A branch, tag or commit. Defaults to ``main``. Unlike in a URL, a revision + containing ``/`` is passed literally — the store percent-encodes it. + :param token: ``None`` (the default) or ``True`` authenticates from ``HF_TOKEN`` or the saved + login; ``False`` forces an anonymous read even when credentials are available; a string is + used as the token directly. + :param endpoint: Hub endpoint. Defaults to ``HF_ENDPOINT``, then ``https://huggingface.co``. + +.. code-block:: python + + import vortex as vx + from vortex.store import HfStore + + store = HfStore("org/name", revision="refs/convert/parquet", token="hf_...") + + # With `store=`, the path is a path within the repository. + vxf = vx.open("data/train.vortex", store=store) + +Listing +======= + +The Hub does not implement WebDAV ``PROPFIND``, which is how object-store HTTP listing works, so a +Hub store cannot list a prefix. Opening a known path works, since that is a ``HEAD`` plus ranged +``GET``. To expand a glob, list the repository through the Hub's own API first — which is what +``vortex.datasets.load_dataset`` does — and then open each path it returns. + +Hugging Face Datasets +===================== + +``vortex.datasets.load_dataset`` builds on this to load Vortex files from the Hub as Hugging Face +``Datasets`` objects, expanding globs and pushing projections, filters and row limits into each +scan. See :doc:`../datasets`. diff --git a/docs/api/python/store/opendal.rst b/docs/api/python/store/opendal.rst new file mode 100644 index 00000000000..094460c5198 --- /dev/null +++ b/docs/api/python/store/opendal.rst @@ -0,0 +1,143 @@ +=========================== +OpenDAL (COS, OSS, GooseFS) +=========================== + +Vortex can read from and write to Tencent Cloud COS, Alibaba Cloud OSS, and Tencent Cloud +GooseFS through `OpenDAL `_, which provides native service +support. + +The Python bindings always include these stores. Other consumers opt in with the ``opendal`` +Cargo feature (e.g. ``cargo build -p vortex-jni --features opendal``). + +.. list-table:: + :header-rows: 1 + + * - Scheme + - Service + - Endpoint / master variable + - Credential variables + * - ``cos://`` + - Tencent Cloud COS + - ``COS_ENDPOINT`` + - ``TENCENTCLOUD_SECRET_ID``, ``TENCENTCLOUD_SECRET_KEY`` + * - ``oss://`` + - Alibaba Cloud OSS + - ``OSS_ENDPOINT`` + - ``ALIBABA_CLOUD_ACCESS_KEY_ID``, ``ALIBABA_CLOUD_ACCESS_KEY_SECRET`` + * - ``goosefs://`` + - Tencent Cloud GooseFS + - ``GOOSEFS_MASTER_ADDR`` + - (optional) ``auth_type`` / ``auth_username`` properties + +:class:`vortex.store.CosStore` +============================== + +.. py:class:: vortex.store.CosStore(bucket, endpoint, *, secret_id=None, secret_key=None, root=None, disable_config_load=False) + + A Tencent Cloud COS object store, backed by OpenDAL. Construct it with explicit + configuration and pass it to + :func:`vortex.io.read_url` / :func:`vortex.io.write` via the ``store=`` argument, + exactly like the built-in S3/Azure/GCS stores. + + :param bucket: COS bucket name (e.g. ``"my-bucket"``). + :param endpoint: COS endpoint (e.g. ``"https://cos.ap-guangzhou.myqcloud.com"``). + :param secret_id: Optional Tencent Cloud secret id. Maps to the ``TENCENTCLOUD_SECRET_ID`` + environment variable when unset. + :param secret_key: Optional Tencent Cloud secret key. Maps to the + ``TENCENTCLOUD_SECRET_KEY`` environment variable when unset. + :param root: Optional key prefix applied to every operation. + :param disable_config_load: When ``True``, disable OpenDAL's automatic config loading + and rely only on the explicit configuration. Defaults to ``False``. + +:class:`vortex.store.GoosefsStore` +================================== + +.. py:class:: vortex.store.GoosefsStore(master_addr, *, root=None, block_size=None, chunk_size=None, write_type=None, auth_type=None, auth_username=None) + + A Tencent Cloud GooseFS object store, backed by OpenDAL. Construct it with explicit + configuration and pass it to + :func:`vortex.io.read_url` / :func:`vortex.io.write` via the ``store=`` argument, + exactly like the built-in S3/Azure/GCS stores. + + :param master_addr: GooseFS master address(es). Single master: + ``"10.0.0.1:9200"``. HA (comma-separated): + ``"10.0.0.1:9200,10.0.0.2:9200,10.0.0.3:9200"``. + :param root: Optional key prefix applied to every operation. + :param block_size: Block size in bytes for new files (default: 64 MiB). + :param chunk_size: Chunk size in bytes for streaming RPCs (default: 1 MiB). + :param write_type: Default write type: ``"must_cache"``, ``"cache_through"``, + ``"through"``, or ``"async_through"``. + :param auth_type: Authentication type: ``"nosasl"`` or ``"simple"`` (default: + ``"simple"``). + :param auth_username: Authentication username (default: current OS user). + +Reading from COS +================ + +Pass a ``cos://`` URL directly. Credentials and the endpoint are picked up from the environment +variables OpenDAL's COS builder reads (``TENCENTCLOUD_SECRET_ID``, ``TENCENTCLOUD_SECRET_KEY`` and +``COS_ENDPOINT``): + +.. code-block:: python + + import vortex as vx + + a = vx.io.read_url("cos://my-bucket/path/to/dataset.vortex") + +Or configure explicitly with :class:`~vortex.store.CosStore` and pass it to +:func:`vortex.io.read_url` via ``store=``: + +.. code-block:: python + + from vortex.io import read_url + from vortex.store import CosStore + + store = CosStore( + bucket="my-bucket", + endpoint="https://cos.ap-guangzhou.myqcloud.com", + secret_id="AKID...", + secret_key="...", + ) + + # When `store=` is supplied, the path is a key within the store, so the scheme and + # bucket are not part of the path passed to read_url. + a = read_url("path/to/dataset.vortex", store=store) + +Reading from OSS +================ + +Alibaba Cloud OSS is reachable by URL. There is no standalone ``OssStore`` class yet, so +configuration comes from the environment variables OpenDAL's OSS builder reads +(``ALIBABA_CLOUD_ACCESS_KEY_ID``, ``ALIBABA_CLOUD_ACCESS_KEY_SECRET`` and ``OSS_ENDPOINT``): + +.. code-block:: python + + import vortex as vx + + a = vx.io.read_url("oss://my-bucket/path/to/dataset.vortex") + +Reading from GooseFS +==================== + +Pass a ``goosefs://`` URL directly. The master address is taken from the URL authority, or +from the ``GOOSEFS_MASTER_ADDR`` environment variable when the authority is empty: + +.. code-block:: python + + import vortex as vx + + a = vx.io.read_url("goosefs://10.0.0.1:9200/path/to/dataset.vortex") + +Or configure explicitly with :class:`~vortex.store.GoosefsStore` and pass it to +:func:`vortex.io.read_url` via ``store=``: + +.. code-block:: python + + from vortex.io import read_url + from vortex.store import GoosefsStore + + store = GoosefsStore(master_addr="10.0.0.1:9200") + + # When `store=` is supplied, the path is a key within the store, so the scheme and + # master address are not part of the path passed to read_url. + a = read_url("path/to/dataset.vortex", store=store) diff --git a/docs/concepts/expressions.md b/docs/concepts/expressions.md index 3c4704ca71a..81317a2b888 100644 --- a/docs/concepts/expressions.md +++ b/docs/concepts/expressions.md @@ -7,12 +7,12 @@ to the returned data. ## Scalar Functions Expressions are defined as abstract scalar functions. These vtables define the signature of the function, properties -such as whether the function is null-sensitive, and the actual logic for executing the function over input arrays. +such as whether the function is strict, and the actual logic for executing the function over input arrays. The built-in scalar functions can be found in the `vortex-array::expr` module, with additional use-case specific functions provided by integration and plugin crates. -In the future, we plan to add full support for geospatial functions, date-time functions, and more. +In the future, we plan to add full support for spatial functions, date-time functions, and more. ## Scalar Function Arrays @@ -50,14 +50,14 @@ scalar_fn(struct.pack): Delaying the computation in this way allows for many more optimizations to be performed later, including fused computation of several expressions at once, as is commonly done by the Vortex GPU compute context. -## Type Coercion +## Type Checking Vortex expressions are strictly typed. This means that the input data types to an expression must exactly match the expected data types of the expression's signature. -For example, all binary functions require the same data type for both inputs. Any coercion of types should be performed -by the caller before constructing the expression. This allows Vortex to be agnostic to the type coercion rules of -different compute engines. +For example, all binary functions require the same data type for both inputs. Any casts should be performed by the caller +before constructing the expression. This allows Vortex to be agnostic to the type conversion rules of different compute +engines. -The notable exception to this rule, is that many expressions permit null-coercion. For example, the equality comparison -function allows comparing a `u32` and `u32?` array, but not a `u32` and `i32` array. +The notable exception to this rule is nullability. For example, the equality comparison function allows comparing a +`u32` and `u32?` array, but not a `u32` and `i32` array. diff --git a/docs/conf.py b/docs/conf.py index a2cadbb255f..049da55b104 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -1,11 +1,14 @@ import doctest import os import re +import sys from pathlib import Path import hawkmoth.docstring from sphinx.util import logging +sys.path.insert(0, str(Path(__file__).parent / "_ext")) + log = logging.getLogger("vortex.docs.conf") # Configuration file for the Sphinx documentation builder. @@ -23,7 +26,6 @@ # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration extensions = [ - "breathe", # C++ API (Doxygen -> Sphinx bridge) "hawkmoth", # C API "myst_parser", # Markdown support "sphinx.ext.autodoc", @@ -31,17 +33,12 @@ "sphinx.ext.doctest", "sphinx.ext.intersphinx", "sphinx.ext.napoleon", - "sphinx_copybutton", - "sphinx_design", - "sphinx_inline_tabs", - "sphinxcontrib.bibtex", - "sphinxcontrib.mermaid", - "sphinxext.opengraph", + "vortex_theme", ] templates_path = ["_templates"] html_show_sourcelink = False -exclude_patterns = ["_build", "Thumbs.db", ".DS_Store", "README.md"] +exclude_patterns = ["_build", "Thumbs.db", ".DS_Store", "README.md", "AGENTS.md"] intersphinx_mapping = { "python": ("https://docs.python.org/3", None), @@ -54,7 +51,20 @@ git_root = Path(__file__).parent.parent nitpicky = True # ensures all :class:, :obj:, etc. links are valid -nitpick_ignore = [] +nitpick_ignore = [ + # `vortex.store.CosStore` / `GoosefsStore` are re-exported through private modules, + # and the `ObjectStore` type alias resolves to those private paths. The public + # classes are fully documented in `opendal.rst`; the private paths are intentionally not. + ("py:class", "vortex.store._cos.CosStore"), + ("py:class", "vortex.store._goosefs.GoosefsStore"), + # `vortex.store.CosStore` / `GoosefsStore` / `HfStore` are the native classes re-exported + # through private modules, so annotations resolve to their `vortex._lib` module paths. The + # public classes are fully documented in `opendal.rst` / `huggingface.rst`; the native paths + # are intentionally not. + ("py:class", "vortex._lib.CosStore"), + ("py:class", "vortex._lib.GoosefsStore"), + ("py:class", "vortex._lib.HfStore"), +] doctest_global_setup = "import pyarrow; import vortex; import vortex as vx; import random; random.seed(a=0)" doctest_default_flags = ( @@ -71,59 +81,14 @@ # -- Options for HTML output ------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output -html_theme = "pydata_sphinx_theme" -html_static_path = ["_static"] -html_css_files = ["style.css"] +html_theme = "vortex" +html_theme_path = ["_theme"] +html_baseurl = "https://docs.vortex.dev/" +html_static_path = ["_static", "_build/_gen_static"] +html_extra_path = ["_headers"] # Cloudflare Pages CSP headers html_favicon = "_static/vortex_logo.svg" # relative to _static/ -# -- Options for PyData Sphinx Theme ---------------------------------------- - -html_theme_options = { - "logo": { - "image_light": "_static/vortex_logo.svg", - "image_dark": "_static/vortex_logo_dark_theme.svg", - }, - "github_url": "https://github.com/vortex-data/vortex", - "icon_links": [ - { - "name": "PyPI", - "url": "https://pypi.org/project/vortex-data", - "icon": "fa-brands fa-python", - }, - { - "name": "Crates.io", - "url": "https://crates.io/crates/vortex", - "icon": "fa-brands fa-rust", - }, - ], - "header_links_before_dropdown": 7, - "navbar_align": "left", - "show_nav_level": 2, - "navigation_depth": 3, - "show_toc_level": 2, -} - -# -- Options for OpenGraph --------------------------------------------------- - -ogp_site_url = "https://docs.vortex.dev" -ogp_image = "https://docs.vortex.dev/_static/vortex_logo.svg" - -# -- Options for Sphinx BibTEX ------------------------------------------- - -bibtex_bibfiles = ["references.bib"] - -# -- Options for Breathe C++ API gen ------------------------------------ - -_doxygen_xml_dir = str(Path(__file__).parent / "_build" / "doxygen-cpp" / "xml") - -os.makedirs(os.path.dirname(_doxygen_xml_dir), exist_ok=True) - -# if not shutil.which("doxygen"): -# raise RuntimeError("doxygen is required to build the docs but was not found on PATH") -# subprocess.run(["doxygen", "Doxyfile.cpp"], cwd=Path(__file__).parent, check=True) - -breathe_projects = {"vortex-cpp": _doxygen_xml_dir} -breathe_default_project = "vortex-cpp" +os.makedirs(Path(__file__).parent / "_build" / "_gen_static", exist_ok=True) # -- Options for hawkmoth C API gen ---------------------------- @@ -151,7 +116,7 @@ C_DOCS: set[str] | None = None -def _replace_rust_references(app, lines, transform, options): +def _replace_rust_references(app, lines, transform, options) -> None: """Replace Rust references with C equivalents in hawkmoth docstrings. See: https://hawkmoth.readthedocs.io/en/stable/extending.html#event-hawkmoth-process-docstring @@ -189,7 +154,7 @@ def _replace_rust_references(app, lines, transform, options): # Pattern to match [`crate::path::to::function`] pattern = r"\[`([^:]+::)*?(vx_[^`]+)`\]" - def replace_match(match): + def replace_match(match) -> str: # Extract the function name (already starts with vx_) # TODO(ngates): detect if the reference is a function or a type func_name = match.group(2) @@ -218,7 +183,7 @@ def replace_match(match): lines[i] = re.sub(pattern, replace_match, line) -def _post_process(app, builder): +def _post_process(app, builder) -> None: """Post-process the documentation after writing.""" global C_DOCS if C_DOCS: @@ -235,7 +200,9 @@ def _post_process(app, builder): os.environ["POLARS_TABLE_WIDTH"] = "80" -def _convert_python_fenced_blocks_from_rust_to_valid_reST_blocks(app, what, name, obj, options, lines: list[str]): +def _convert_python_fenced_blocks_from_rust_to_valid_reST_blocks( + app, what, name, obj, options, lines: list[str] +) -> None: """Remove Markdown-style code fences from Python docs written in Rust. We would like `cargo test` to Just Work (TM). Unfortunately, by default, it executes any @@ -287,7 +254,7 @@ def _convert_python_fenced_blocks_from_rust_to_valid_reST_blocks(app, what, name in_block = False -def _resolve_breathe_cpp_references(app, env, node, contnode): +def _resolve_breathe_cpp_references(app, env, node, contnode) -> object | None: """Resolve relative C++ references emitted by Breathe. Breathe emits cross-namespace parameter types with relative qualifiers (e.g. ``scalar::Scalar`` @@ -310,7 +277,7 @@ def _resolve_breathe_cpp_references(app, env, node, contnode): ) -def setup(app): +def setup(app) -> None: app.connect("hawkmoth-process-docstring", _replace_rust_references) app.connect("write-started", _post_process) app.connect("autodoc-process-docstring", _convert_python_fenced_blocks_from_rust_to_valid_reST_blocks) diff --git a/docs/developer-guide/benchmarking.md b/docs/developer-guide/benchmarking.md index dfa0fa1cd45..48df9a2aa21 100644 --- a/docs/developer-guide/benchmarking.md +++ b/docs/developer-guide/benchmarking.md @@ -119,10 +119,20 @@ const VECTOR_SIZE: &[usize] = &[16, 256, 2048, 8192]; fn my_bench(bencher: Bencher, num_indices: usize) { ... } ``` -### Keep per-iteration execution time under ~1 ms +### Keep per-iteration execution time under 1 ms -Each individual iteration of the benchmarked closure should complete in -**less than 1ms**. This is to keep benchmarks snappy, locally and on CI. +**1 ms is the maximum, not a soft target.** Each individual iteration of the benchmarked +closure must complete in less than 1 ms. This is to keep benchmarks snappy, locally and on +CI. + +A benchmark that needs longer than that is measuring too much work at once. Shrink the +input size until a single iteration fits, split it into smaller parameterized cases, or +gate it with `#[cfg(not(codspeed))]` if it genuinely cannot be made to fit. + +The number to check against the budget is the per-iteration time, not the time the whole +benchmark binary takes. CodSpeed reports exactly that: its performance report on a pull +request lists the per-iteration time under `HEAD` for every benchmark the pull request adds +or changes, so check any new benchmark there before merging. ### Gate CodSpeed-incompatible benchmarks @@ -208,14 +218,20 @@ uv run --project bench-orchestrator vx-bench run tpch \ Benchmarks run automatically on all commits to `develop` and can be run on-demand for PRs: -- **Post-commit** -- compression, random access, and SQL benchmarks run on every commit to - `develop`, with results uploaded for historical tracking. -- **PR benchmarks** -- triggered by the `action/benchmark` label. Results are compared against - the latest `develop` run and posted as a PR comment. -- **SQL benchmarks** -- triggered by the `action/benchmark-sql` label. Runs the base SQL matrix, - which excludes Appian, TPC-H SF=10 on S3, `vortex-compact`, and `duckdb:duckdb`. -- **Full SQL benchmarks** -- triggered by the `action/benchmark-sql-full` label. Runs the full - SQL matrix of suites, engines, formats, and storage backends (NVMe, S3). +- **Post-commit** -- compression, string encoding, random access, and SQL benchmarks run on every + commit to `develop`, with results uploaded for historical tracking. +- **Random access** -- `action/bench-random-access` runs only the random-access benchmark. +- **Compression** -- `action/bench-compress` runs only the compression benchmark. +- **String encoding** -- `action/bench-string` runs only the string encoding benchmark. +- **GPU compression** -- `action/bench-gpu-compress` runs the allow-listed Vortex decompression + cases on a GPU runner. +- **SQL** -- `action/bench-sql` runs the `pr` preset, which excludes `vortex-compact`. +- **SQL Compact** -- `action/bench-sql-compact` runs the `pr-compact` preset, which benchmarks + `vortex-compact` plus Parquet control rows used to distinguish code changes from runner drift. +- **All CPU benchmarks** -- `action/bench-all` runs random access, compression, string encoding, + and the `pr-all` SQL preset, which combines the `pr` and `pr-compact` coverage without + repeating shared jobs. Do not combine it with other benchmark labels; GPU compression is the + only exception. All CI benchmarks run on dedicated instances with the `release_debug` profile and `-C target-cpu=native` to produce representative numbers. diff --git a/docs/developer-guide/embedding/cxx.md b/docs/developer-guide/embedding/cxx.md deleted file mode 100644 index 40238d36282..00000000000 --- a/docs/developer-guide/embedding/cxx.md +++ /dev/null @@ -1,12 +0,0 @@ -# C++ Integration - -:::{warning} -This page is under construction. -::: - -Planned content: - -- The C++ wrapper around the C FFI -- Building with CMake -- RAII wrappers for Vortex objects -- Working with arrays and dtypes from C++ diff --git a/docs/developer-guide/embedding/ffi.md b/docs/developer-guide/embedding/ffi.md deleted file mode 100644 index 09517fe87e0..00000000000 --- a/docs/developer-guide/embedding/ffi.md +++ /dev/null @@ -1,14 +0,0 @@ -# C FFI Integration - -:::{warning} -This page is under construction. -::: - -Planned content: - -- Overview of the Vortex C API -- Building and linking the shared library -- Session management from C -- Working with arrays and dtypes from C -- Error handling -- Memory management and ownership rules diff --git a/docs/developer-guide/embedding/gpu.md b/docs/developer-guide/embedding/gpu.md deleted file mode 100644 index d92a9ef0370..00000000000 --- a/docs/developer-guide/embedding/gpu.md +++ /dev/null @@ -1,14 +0,0 @@ -# GPU Acceleration - -:::{warning} -This page is under construction. -::: - -Planned content: - -- CUDA requirements and supported hardware -- Building with GPU support (Linux only) -- GPU-accelerated decompression -- GPU compute kernels -- Memory management between host and device -- Current limitations and roadmap diff --git a/docs/developer-guide/embedding/index.md b/docs/developer-guide/embedding/index.md deleted file mode 100644 index 8c8154718ca..00000000000 --- a/docs/developer-guide/embedding/index.md +++ /dev/null @@ -1,29 +0,0 @@ -# Embedding Vortex - -:::{warning} -This section is under construction. For guidance on embedding Vortex, please join the -[Vortex Slack channel](https://vortex.dev/slack) -or open a [GitHub Issue](https://github.com/vortex-data/vortex/issues/new/choose). -::: - -Vortex can be embedded into applications and services via its C FFI, C++ wrapper, or the Scan API. -The following topics are planned for this section: - -- **C FFI** -- the Vortex C API, building and linking, session management, arrays, dtypes, - error handling, and memory ownership. -- **C++** -- the C++ wrapper around the C FFI, CMake integration, and RAII wrappers. -- **Scan API** -- serving Vortex data to query engines, wire format serialization, filter and - projection pushdown, and custom scan providers. -- **GPU Acceleration** -- CUDA requirements, GPU-accelerated decompression and compute, - host/device memory management, and current limitations. - -```{toctree} ---- -hidden: true ---- - -ffi -cxx -scan-api -gpu -``` diff --git a/docs/developer-guide/embedding/scan-api.md b/docs/developer-guide/embedding/scan-api.md deleted file mode 100644 index b2f8db81122..00000000000 --- a/docs/developer-guide/embedding/scan-api.md +++ /dev/null @@ -1,13 +0,0 @@ -# Embedding the Scan API - -:::{warning} -This page is under construction. -::: - -Planned content: - -- Using the Scan API to serve Vortex data to query engines -- Wire format serialization and deserialization -- Filter and projection pushdown -- Implementing a custom scan provider -- Integration patterns for storage servers diff --git a/docs/developer-guide/index.md b/docs/developer-guide/index.md index e6eec4bfb70..9afff8877aa 100644 --- a/docs/developer-guide/index.md +++ b/docs/developer-guide/index.md @@ -8,7 +8,6 @@ maxdepth: 2 --- extending/index -embedding/index language-bindings benchmarking ``` @@ -24,6 +23,7 @@ internals/session internals/async-runtime internals/vtables internals/execution +internals/scan-planning internals/stats-pruning internals/io internals/serialization @@ -37,6 +37,5 @@ caption: Integrations --- integrations/datafusion -integrations/duckdb integrations/spark ``` diff --git a/docs/developer-guide/integrations/duckdb.md b/docs/developer-guide/integrations/duckdb.md deleted file mode 100644 index 060f5fca973..00000000000 --- a/docs/developer-guide/integrations/duckdb.md +++ /dev/null @@ -1,69 +0,0 @@ -# DuckDB - -The `vortex-duckdb` crate integrates Vortex as a DuckDB extension, registering table functions -(`vortex_scan` and `read_vortex`) that allow DuckDB to scan Vortex files directly. - -## Registration - -Vortex registers itself as a DuckDB table function. The function accepts a file path or glob -pattern as its argument. During the bind phase, the first matching file is opened to extract the -schema and register result columns with DuckDB's planner. The remaining files are opened lazily -during execution. - -## Multiple Files - -Glob patterns are expanded at bind time to produce a list of file paths. Files are opened -concurrently during the global initialization phase, with a concurrency limit proportional to -the number of DuckDB worker threads to keep the I/O pipeline saturated without overwhelming -the system. - -Active file scans are driven concurrently via `try_flatten_unordered`, ensuring that DuckDB's -execution threads always have data to consume while background I/O proceeds. File footers are -cached to avoid redundant parsing when the same file appears in multiple queries. - -## Threading Model - -DuckDB uses a thread-per-core execution model where each worker thread pulls data from a shared -scan state. The Vortex integration bridges this with a single `CurrentThreadRuntime` instance -that all async work is dispatched through. - -The global initialization phase creates an async stream of scan results and wraps it in a -`ThreadSafeIterator` backed by a multi-producer multi-consumer channel. Each DuckDB worker -thread calls into the iterator to pull the next batch, and the CRT's smol executor drives I/O -progress from whichever thread happens to be polling. This allows DuckDB to retain its -thread-per-core scheduling while Vortex handles async I/O transparently. - -Because DuckDB threads block while waiting for data, the CRT may need background worker threads -to ensure I/O continues to make progress when all DuckDB threads are busy exporting data. See -the [runtime documentation](../internals/async-runtime.md) for more on this trade-off. - -## Filter and Projection Pushdown - -DuckDB's planner pushes filter predicates into the scan via the `pushdown_complex_filter` -callback. These are converted from DuckDB's bound expression representation into Vortex -expressions and stored alongside any table filter expressions. During scanning, the combined -filter is applied to the `ScanBuilder` for each file. - -Files can be pruned entirely before opening if their statistics prove that no rows can match -the filter. - -Projection pushdown maps DuckDB's requested column indices to Vortex field names and passes -them as a projection expression to the scan. - -## Data Export - -Scan results are exported to DuckDB's native `DataChunk` and `Vector` format. The exporter -walks the Vortex array tree and uses encoding-aware exporters where possible -- for example, -constant arrays, run-end arrays, and dictionary arrays can be exported without first -decompressing to a canonical form. Arrays that lack a specialized exporter fall back to -canonical (Arrow-compatible) conversion before export. - -Results are exported in chunks matching DuckDB's standard vector size to align with its -vectorized execution model. - -## Future Work - -The current integration builds directly on the `ScanBuilder`, layout reader, and file APIs. -Future work will migrate it to use the [Scan API](/concepts/scanning) `Source` trait, unifying -file discovery, multi-file coordination, and pushdown behind a single interface shared across -all engine integrations. diff --git a/docs/developer-guide/integrations/spark.md b/docs/developer-guide/integrations/spark.md index 2342f95063f..6758645c0dc 100644 --- a/docs/developer-guide/integrations/spark.md +++ b/docs/developer-guide/integrations/spark.md @@ -39,8 +39,9 @@ Projection pushdown is supported through Spark's `SupportsPushDownRequiredColumn The scan builder prunes the column list to only those referenced by the query, and the pruned column set is passed to the native scan via `ScanOptions`. -Filter pushdown infrastructure exists in the `ScanOptions` type but is not yet connected to -Spark's `SupportsPushDownFilters` interface. This is planned future work. +Filter pushdown is supported through `SupportsPushDownV2Filters`. The scan builder converts each +predicate it recognizes into a Vortex expression and keeps the rest for Spark to evaluate after the +scan; the converted expression reaches native code as the `ScanOptions` filter. ## Data Export @@ -53,5 +54,4 @@ through direct byte buffers. The current integration builds directly on the native file and scan APIs via JNI. Future work will migrate it to use the [Scan API](/concepts/scanning) `Source` trait, which will provide a -standard interface for file discovery, partitioning, and pushdown. This will also enable -connecting Spark's filter pushdown to Vortex's expression-based filtering. +standard interface for file discovery, partitioning, and pushdown. diff --git a/docs/developer-guide/internals/scan-planning.md b/docs/developer-guide/internals/scan-planning.md new file mode 100644 index 00000000000..e03131924ac --- /dev/null +++ b/docs/developer-guide/internals/scan-planning.md @@ -0,0 +1,63 @@ +# Scan Plans + +A scan plan is the physical plan for satisfying one scan query. It is a tree of physical operators +over a row domain, describing the reads and derived work needed to produce that query's result. + +## Operators, not layout mirrors + +Plan operators describe *what work happens*, not *which layout produced it*. Their identity and +operator-specific state are independent of the source layout kind. The complete plan node is not: +its common lazy-child container can own hidden source state used to materialize individual children +on demand. + +| Operator | Work | +| --- | --- | +| `SegmentScan` | read one segment and decode it to an array | +| `Concat` | concatenate its children row-wise | +| `Pack` | assemble a struct from one child per field, plus optional validity | +| `Take` | index `values` by `codes` | +| `ListPack` | assemble a list from elements and offsets, plus optional validity | +| `Eval` | apply an expression to its child | +| `RowIdx` | offset row numbers into the file's row domain | + +Naming operators for what they compute is what lets one rule cover every case. `Concat` of +`Concat` flattens on shape alone, and `Take` over `SegmentScan` is the dictionary pushdown, +regardless of the source layout. + +The stored layout tree describes all physical data in a file. A plan is query-specific: it is built +from that tree for one projection, filter, and row domain. Different queries over the same file can +therefore produce different plans. + +## Optimization + +Child replacement is implemented by the common plan container rather than by every operator. It +replaces the external child container, clones `PlanData`, then invokes the operator's +`PlanVTable::with_children` callback to validate the new children and refresh derived caches such +as `Concat` row offsets. Rules therefore rewrite the generic tree without reconstructing common +plan fields inside each operator. + +Optimization rewrites the initial tree so that each expression is evaluated as close as possible to +the physical data that can satisfy it. Every rewrite must preserve the query result, including its +dtype, row domain, row order, row identity, null behavior, and observable errors. + +Planning does not read segment data. It constructs and optimizes a description of the work that a +later execution stage will perform. + +## Vtables + +Each operator is a small vtable type implementing `PlanVTable`, paired with a `Plan` container +over a shared `PlanRef`. `PlanRef` points to one allocation whose ordinary fields hold the operator +ID, dtype, row count, and lazy children. Only the unsized tail containing the vtable and +`V::PlanData` is erased behind `dyn DynPlan`, so common-field reads do not use dynamic dispatch. +`Plan` provides typed access to that operator data through `Deref`. + +`PlanVTable` also carries `id` and a `Metadata` codec. Operators with no unrecoverable state +already serialize their metadata; the ones holding a read context or a bound expression return +`None` until those codecs exist. + +## Future work + +Plans currently stop at construction and optimization. Still to come: a plan registry and foreign +operator placeholder so third-party operators survive a round trip, a serialization envelope, and +an execution stage that walks an optimized plan, reads the referenced segments, and returns the +query result. diff --git a/docs/developer-guide/internals/session.md b/docs/developer-guide/internals/session.md index 71158eaae20..e0dddfb4b74 100644 --- a/docs/developer-guide/internals/session.md +++ b/docs/developer-guide/internals/session.md @@ -15,10 +15,11 @@ The session is built on two primitives from the `vortex-session` crate: - **`VortexSession`** -- a cloneable, thread-safe map from Rust `TypeId` to a shared (`Arc`-wrapped) value. Any type that is `Clone + Send + Sync + Debug + 'static` can be stored as a session variable. -- **`Registry`** -- a concurrent map from string IDs to values of type `T`, used by each - component to look up registered plugins at runtime. +- **`ArcSwapMap`** -- a concurrent copy-on-write map, used both as the backing store of + `VortexSession` itself and (keyed by interned string `Id`s) as the plugin registries each + component uses to look up registered plugins at runtime. -Because `VortexSession` is backed by an `ArcSwap`, cloning is cheap and all clones share the +Because both primitives are backed by an `ArcSwap`, cloning is cheap and all clones share the same state, with lock-free reads and copy-on-write writes. This makes it safe to hand the session to multiple threads, tasks, or I/O operations without coordination. @@ -77,8 +78,11 @@ session.write_options() .await?; // Scanning a layout +let filter = expr + .optimize_recursive(layout_reader.dtype())? + .bind(layout_reader.dtype())?; ScanBuilder::new(session.clone(), layout_reader) - .with_filter(expr) + .with_filter(filter) .into_array_stream()?; ``` diff --git a/docs/generate_changelog.py b/docs/generate_changelog.py index f2bf5f74c14..e33f2b36df1 100644 --- a/docs/generate_changelog.py +++ b/docs/generate_changelog.py @@ -5,15 +5,25 @@ import re import subprocess import sys +from collections.abc import Iterable from datetime import datetime from pathlib import Path +from typing import TypedDict REPO = "vortex-data/vortex" NUM_RELEASES = 5 CHANGELOG_DIR = Path(__file__).parent / "project" / "changelog" -def fetch_releases(): +class Release(TypedDict): + tagName: str + isDraft: bool + isPrerelease: bool + publishedAt: str + body: str + + +def fetch_releases() -> list[Release]: # List releases (body is not available in list output). result = subprocess.run( [ @@ -50,7 +60,7 @@ def fetch_releases(): return releases -def linkify(body): +def linkify(body: str) -> str: """Convert PR/issue numbers and GitHub usernames to markdown links.""" # (#1234) -> ([#1234](https://github.com/REPO/pull/1234)) body = re.sub( @@ -68,7 +78,7 @@ def linkify(body): return body -def write_release_page(release): +def write_release_page(release: Release) -> str: tag = release["tagName"] body = linkify(release["body"].replace("\r\n", "\n").strip()) @@ -92,7 +102,7 @@ def write_release_page(release): return filename -def write_index(filenames): +def write_index(filenames: Iterable[str]) -> None: toctree_entries = "\n".join(filenames) content = f"""# Changelog @@ -111,7 +121,7 @@ def write_index(filenames): (CHANGELOG_DIR / "index.md").write_text(content) -def main(): +def main() -> None: releases = fetch_releases() if not releases: print("warning: no releases found, skipping changelog generation", file=sys.stderr) diff --git a/docs/api/cpp/index.rst b/docs/getting-started/cpp.rst similarity index 66% rename from docs/api/cpp/index.rst rename to docs/getting-started/cpp.rst index 642b3d18ce9..47a9a19cbc4 100644 --- a/docs/api/cpp/index.rst +++ b/docs/getting-started/cpp.rst @@ -1,81 +1,69 @@ -C++ API -======= +C++ Quickstart +============== -The Vortex C++ API allows you to read and write ``.vortex`` files directly or -via an Arrow compatibility layer like `nanoarrow -`_. The only dependency apart from Vortex -is ``nanoarrow``. +The C++ API reads and writes ``.vortex`` files, with optional Arrow interoperability +through `nanoarrow `_. .. note:: - C++ API is a work in progress. Please reach out to us if you are interested - in using Vortex from C++ or you want a feature not covered yet e.g. - extension support. + The C++ API is still evolving. Please reach out if you need a missing feature. -Installation ------------- +Building from source +-------------------- -We don't provide prebuilt library files (yet) so you will need to build Vortex -from source, and for that you will need: - -- C++20, -- Rust toolchain, -- and CMake 3.10. +Vortex currently provides source builds only, for GNU/Linux x86_64/aarch64 and +macOS arm64. You need a C compiler, a C++20 compiler, CMake 3.25+ with a single-config +generator, and Cargo/rustc 1.95+ on ``PATH``. Repository checkout: .. code-block:: bash git clone --depth 1 https://github.com/vortex-data/vortex cd vortex - cargo build --release -p vortex-ffi - cd vortex-cxx - mkdir build - cmake -Bbuild -DCMAKE_BUILD_TYPE=Release - # To build the examples, pass -DBUILD_EXAMPLES=1 - # cmake -Bbuild -DBUILD_EXAMPLES=1 + cmake -S lang/cpp -B build/cpp \ + -DCMAKE_BUILD_TYPE=Release \ + -DVORTEX_BUILD_EXAMPLES=ON + cmake --build build/cpp --parallel - cmake --build build -j +CMake builds the Rust FFI through Cargo. No separate ``cargo build`` is needed. +The first build may download dependencies. The commands above enable the examples +used below. Tests and examples are otherwise disabled by default. -This produces a shared and a static library which you can use directly or via +To embed Vortex in a CMake project, vendor or fetch a pinned checkout and link its +target: .. code-block:: cmake - # static library - target_link_libraries(target PRIVATE vortex_cxx) - # shared library - target_link_libraries(target PRIVATE vortex_cxx_shared) + add_subdirectory(path/to/vortex vortex) + target_link_libraries(target PRIVATE Vortex::cpp_static) + +The target includes the required headers and link dependencies. Vortex builds +static, position-independent libraries, with no installation rules or +``find_package(Vortex)`` package. If embedding it in a shared library, keep Vortex +symbols private using the parent's export policy. -Have a look at the `examples -`_ -directory as well. +See the `C++ README +`_ for build +options, platform limits, sanitizers, and CUDA deployment requirements, and the +`examples directory +`_ for complete +programs. Reading files ------------- Full source code for this example is `reader.cpp -`_. -For brevity we omit ``main()``, system includes and ``using`` directives. +`_ +. For brevity we omit ``main()``, system includes and ``using`` directives. Assuming you have Vortex files ``people0``, ``people1``, and ``me`` in a local folder, each containing U8 column "age" and U16 column "height", this is how you print all ages for specific heights: -.. code-block:: cpp - - Session session; - DataSource ds = DataSource::open(session, {"people*.vortex", "me.vortex"}); - Scan scan = ds.scan({.filter = col("height") >= lit(50)}); - - for (Partition &partition : scan.partitions()) { - for (Array &array : partition.batches()) { - Array age = array.field("age"); - PrimitiveView age_view = age.values(session); - std::span age_values = age_view.values(); - for (uint8_t value : age_values) { - std::cout << int(value) << " "; - } - } - } - std::cout << "\n"; +.. literalinclude:: ../../lang/cpp/examples/reader.cpp + :language: cpp + :start-after: docs:begin:example + :end-before: docs:end:example + :dedent: 4 DataSource and Scan ^^^^^^^^^^^^^^^^^^^ @@ -160,42 +148,13 @@ Writing files Now let's write the first files to be read by our previous example. Source code for this example is `writer.cpp -`_. - -.. code-block:: cpp - - Session session; - DataType dtype = dtype::struct_({ - {"age", dtype::uint8()}, - {"height", dtype::uint16(Nullable)}, - }); - - constexpr size_t SAMPLE_ROWS = 100; - std::vector age_buffer(SAMPLE_ROWS); - std::vector height_buffer(SAMPLE_ROWS); - for (size_t i = 0; i < SAMPLE_ROWS; ++i) { - age_buffer[i] = static_cast(i); - height_buffer[i] = static_cast((i + 1) % 200); - } - - Array age = Array::primitive(age_buffer); - Array array = make_struct({ - {"age", age}, - {"height", Array::primitive(height_buffer, AllValid)}, - }); - - Expression age_gt_10 = expr::gt(expr::col("age"), expr::lit(10)); - Array validity_array = array.apply(age_gt_10); - - Validity validity = Validity::from_array(validity_array); - Array array2 = make_struct({ - {"age", age}, - {"height", Array::primitive(height_buffer, validity)}, - }); - - Writer writer = Writer::open(session, argv[1], dtype); - writer.push({array, array2}); - writer.finish(); +`_. + +.. literalinclude:: ../../lang/cpp/examples/writer.cpp + :language: cpp + :start-after: docs:begin:example + :end-before: docs:end:example + :dedent: 4 DataType ^^^^^^^^ @@ -249,7 +208,7 @@ Now you can build the example and read back the generated files: .. code-block:: - ./build/examples/writer people0.vortex - ./build/examples/writer people1.vortex - ./build/examples/writer me.vortex - ./build/examples/reader + ./build/cpp/examples/writer people0.vortex + ./build/cpp/examples/writer people1.vortex + ./build/cpp/examples/writer me.vortex + ./build/cpp/examples/reader diff --git a/docs/getting-started/index.md b/docs/getting-started/index.md index 72b88db57d4..fb7ce0afbe9 100644 --- a/docs/getting-started/index.md +++ b/docs/getting-started/index.md @@ -24,5 +24,6 @@ caption: Language Libraries Python Quickstart Rust Quickstart +C++ Quickstart Java Quickstart ``` diff --git a/docs/getting-started/install.md b/docs/getting-started/install.md index 9ff30a30fcd..4192bec33f8 100644 --- a/docs/getting-started/install.md +++ b/docs/getting-started/install.md @@ -3,39 +3,32 @@ ## `vx` CLI The `vx` command-line tool lets you convert, inspect, browse, and query Vortex files from -the terminal. +the terminal. You have the following options: -::::{tab-set} +1. Binstall (recommended) downloads a pre-built binary. Requires +[cargo-binstall](https://github.com/cargo-bins/cargo-binstall): -:::{tab-item} Binstall (recommended) ```bash cargo binstall vortex-tui ``` -Downloads a pre-built binary. Requires [cargo-binstall](https://github.com/cargo-bins/cargo-binstall). -::: -:::{tab-item} pip +2. pip also installs the Python library; see the [Python quickstart](python.rst) for library usage: + ```bash pip install vortex-data ``` -This also installs the Python library. See the [Python quickstart](python.rst) for library usage. -::: -:::{tab-item} uvx +3. uvx runs the CLI without installing. Requires [uv](https://docs.astral.sh/uv/): + ```bash uvx --from vortex-data vx --help ``` -Runs the CLI without installing. Requires [uv](https://docs.astral.sh/uv/). -::: -:::{tab-item} Cargo +4. Cargo builds from source, which can be slow due to the large dependency tree: + ```bash cargo install vortex-tui ``` -Builds from source. This can be slow due to the large dependency tree. -::: - -:::: Verify the installation: diff --git a/docs/getting-started/query.md b/docs/getting-started/query.md index 5ff33e49d20..234a270d221 100644 --- a/docs/getting-started/query.md +++ b/docs/getting-started/query.md @@ -51,7 +51,8 @@ vx segments yellow_tripdata_2024-01.vortex vx inspect yellow_tripdata_2024-01.vortex ``` -All inspection commands support `--json` for machine-readable output. +`vx inspect` and `vx tree layout` take `--json` for machine-readable output; `vx segments` always +emits JSON. ## Standalone SQL diff --git a/docs/index.md b/docs/index.md index b75a438f913..9de33b387b0 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,7 +1,3 @@ ---- -sd_hide_title: true ---- - # Vortex :::{image} _static/vortex_wordmark.svg @@ -22,52 +18,45 @@ around the latest research from the database community. ## Where to start -::::{grid} 1 2 2 3 -:gutter: 3 +::::{container} cards -:::{grid-item-card} Read & write Vortex files -:link: getting-started/index -:link-type: doc +:::{container} card +[**Read & write Vortex files**](getting-started/index.md) -Get started with Vortex in **Python**, **Rust**, or **Java**. Convert from Parquet, compress -your data, and query it. +Get started with Vortex in **Python**, **Rust**, **C++**, or **Java**. Convert +from Parquet, compress your data, and query it. ::: -:::{grid-item-card} Use with a query engine -:link: user-guide/index -:link-type: doc +:::{container} card +[**Use with a query engine**](user-guide/index.md) Integrate Vortex with **DataFusion**, **DuckDB**, **Spark**, **Trino**, or **Ray** for accelerated queries over compressed data. ::: -:::{grid-item-card} Understand the architecture -:link: concepts/index -:link-type: doc +:::{container} card +[**Understand the architecture**](concepts/index.md) Learn how **DTypes**, **Arrays**, **Encodings**, **Layouts**, and the **Scan API** fit together as building blocks. ::: -:::{grid-item-card} Extend Vortex -:link: developer-guide/index -:link-type: doc +:::{container} card +[**Extend Vortex**](developer-guide/index.md) Write your own **encodings**, **layouts**, **compute functions**, or **extension types** from Rust or Python. ::: -:::{grid-item-card} Create an engine integration -:link: developer-guide/index -:link-type: doc +:::{container} card +[**Create an engine integration**](developer-guide/index.md) Build a **query engine connector** or **data source** using the **Scan API**, **C FFI**, or **C++ wrapper**. ::: -:::{grid-item-card} Internals -:link: developer-guide/index -:link-type: doc +:::{container} card +[**Internals**](developer-guide/index.md) Explore the **crate architecture**, **async runtime**, **session system**, and integration internals. Build and benchmark locally. diff --git a/docs/make.bat b/docs/make.bat deleted file mode 100644 index 2c8600c978e..00000000000 --- a/docs/make.bat +++ /dev/null @@ -1,35 +0,0 @@ -@ECHO OFF - -pushd %~dp0 - -REM Command file for Sphinx documentation - -if "%SPHINXBUILD%" == "" ( - set SPHINXBUILD=uv run sphinx-build -) -set SOURCEDIR=. -set BUILDDIR=_build - -%SPHINXBUILD% >NUL 2>NUL -if errorlevel 9009 ( - echo. - echo.The 'sphinx-build' command was not found. Make sure you have Sphinx - echo.installed, then set the SPHINXBUILD environment variable to point - echo.to the full path of the 'sphinx-build' executable. Alternatively you - echo.may add the Sphinx directory to PATH. - echo. - echo.If you don't have Sphinx installed, grab it from - echo.https://www.sphinx-doc.org/ - exit /b 1 -) - -if "%1" == "" goto help - -%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% -goto end - -:help -%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% - -:end -popd diff --git a/docs/project/bindings.md b/docs/project/bindings.md deleted file mode 100644 index 3311a7a3663..00000000000 --- a/docs/project/bindings.md +++ /dev/null @@ -1,187 +0,0 @@ -# Language Bindings Evolution - -This document describes the evolution strategy for Vortex's language bindings. It covers the -current state, architectural principles, open questions, per-language migration plans, and a phased -roadmap. For the tier model and API surface definitions, see the -[Language Bindings](../developer-guide/language-bindings.md) developer guide page. - -## Current State - -| Language | Technology | API Surface | -|---------------|------------|------------------------------------------------------------------------------------------| -| Rust | Native | Full access (Tier 3). All array plugins, compute plugins, layouts, and extension DTypes. | -| Python | PyO3 | Near Tier 2. Native expression construction, array access, file I/O. | -| C | cbindgen | Tier ~1. File I/O, basic scan, Arrow stream export. | -| C++ | cxx | Tier ~1. Wraps Rust via cxx bridge. File I/O and scan. | -| Java (JNI) | JNI | Tier ~1. Arrow-based file I/O and basic scan via JNI. | -| Java (Panama) | Panama FFI | Not yet implemented. Target Tier 2 via direct C ABI access. Requires JDK 22+. | - -## Architecture Principle - -**The C API is the foundation ABI for non-Rust bindings.** The C API is not yet ABI-stable — it -evolves as fast as the rest of the project and should be **statically linked** to avoid breakage. -In the future, a subset of the API will be flagged as stable for use via dynamically linked -libraries. - -Each language chooses the binding strategy that makes the most sense: - -- **Python** uses PyO3 to call Rust directly and will continue to do so. PyO3 provides the best - ergonomics and performance for Python. -- **Other languages** use Rust-specific bindings where it makes sense (e.g. cxx for C++), or wrap - the C API otherwise (e.g. Java Panama, future C++ migration). -- New capabilities should be exposed through the C API so that any language can consume them, but - languages are not required to go through C if a more direct path exists. -- The C API's error handling, memory ownership, and lifetime conventions must be documented - regardless of stability guarantees, since all consumers depend on them. - -## Open Questions - -### Array Execution Model - -The array execution model (lazy vs eager evaluation) is still being designed. This directly affects -what bindings can expose: - -- **Eager execution** is simpler for bindings — arrays are materialized values that can be passed - across the FFI boundary. -- **Lazy execution** requires bindings to handle expression graphs, materialization triggers, and - potentially device placement (CPU vs GPU). - -Until the execution model stabilizes, bindings should focus on the Scan API (Tier 1) where results -are always materialized Arrow streams, and avoid exposing internal array operations that may change. - -### Stable Rust Plugin API - -Should Rust plugins go through the C ABI for dynamic loading, or should they use a Rust-native -trait-based plugin interface? - -- **C ABI approach:** Maximum portability, works across Rust compiler versions. Higher development - cost per plugin. -- **Rust-native approach:** Ergonomic, zero-cost, but ties plugins to a specific Rust compiler - version and Vortex ABI. - -This is a future concern — the current priority is stabilizing the C API for non-Rust consumers. - -### Panama Timeline for Java - -Java's Foreign Function & Memory API (Panama) graduated to a stable API in JDK 22. Key -considerations: - -- Trino already supports JDK 22 and can adopt Panama immediately. -- Spark targets older LTS releases and will not support Panama for some time. -- JNI must remain supported as the primary path for Spark and any other engine on older JDKs. -- Panama provides direct C ABI access, eliminating JNI overhead and enabling Tier 2 capabilities. - -## Per-Language Migration Plan - -### C FFI - -Expand from Tier ~1 to Tier 2: - -- **Array stream export:** Expose Vortex array streams (not just Arrow) through the C API. -- **Array tree inspection:** Provide functions to walk array children, read encoding IDs, and - access array metadata. -- **Error handling:** Stabilize the error reporting convention (error codes, error message - retrieval, thread-local error state). -- **Memory ownership:** Document and enforce ownership conventions (who allocates, who frees, - reference counting semantics). - -### C++ - -Migrate from cxx to wrapping the C API: - -- **RAII wrappers:** Provide C++ classes that manage lifetime of C API objects (files, scanners, - arrays, streams). -- **CMake integration:** Ship a CMake config so downstream projects can `find_package(Vortex)`. -- **Header generation:** Auto-generate C++ headers from the C API headers, adding type safety and - namespace scoping. -- **Target:** Tier 2 (native array access through the C API). - -### Python - -Already near Tier 2-3. Python will continue to use PyO3 for direct Rust interop. The migration -plan focuses on formalization: - -- **API stability:** Classify existing Python APIs as stable, experimental, or internal. Publish - this classification in the Python API docs. -- **Plugin API:** Expose plugin registration (custom array plugins, compute plugins) to Python via - PyO3. The registration model should mirror the C API's capabilities so that the same plugin - concepts are available in both languages. - -### Java (JNI) - -Maintain current JNI bindings for Arrow I/O and basic scan. This track targets broad JDK -compatibility and remains the primary integration point for Spark, which targets older LTS -releases. JNI stays at Tier 1. - -### Java (Panama) - -Build a new binding layer using Panama's Foreign Function & Memory API to call the C API directly. -This enables native array access, lower overhead, and a path to Tier 2 capabilities. Panama -bindings are opt-in for environments running JDK 22+. Trino already supports JDK 22 and is the -likely first adopter. Spark will not support Panama until it moves to a compatible JDK. Connectors -should abstract over the binding layer so they can use JNI or Panama transparently. - -### Rust - -Remains native Tier 3. Future considerations: - -- **Stable plugin ABI:** Investigate a C ABI-based plugin interface for dynamically loading - encoding crates. This would allow third-party encodings to be distributed as shared libraries - without requiring recompilation of Vortex. -- **Version compatibility:** A stable plugin ABI would need versioning and capability negotiation - to handle evolution of the array and compute plugin interfaces. - -## Phased Roadmap - -### Phase 1: Complete C API at Tier 1 - -- Complete the C API for file I/O, scan with serialized expressions, and Arrow stream output. -- Document the API: function signatures, error handling conventions, memory ownership rules. -- Add integration tests that exercise the C API from C and C++. -- Ensure Python and Java (JNI) bindings are fully functional at Tier 1. - -### Phase 2: Extend to Tier 2 - -- Add array stream export to the C API (Vortex arrays, not just Arrow). -- Add array tree inspection functions (walk children, read encoding IDs and metadata). -- Migrate C++ bindings from cxx to wrapping the C API with RAII wrappers. -- Prototype Java Panama bindings targeting Tier 2. -- Extend Python bindings to formalize Tier 2 APIs as stable. - -### Phase 3: Plugin API Exploration - -- Design the plugin registration interface in the C API (Tier 3 capabilities). -- Formalize the Python plugin API (array plugins, compute plugins, extension DTypes). -- Investigate Rust stable plugin ABI for dynamic loading of encoding crates. -- Evaluate Panama adoption timeline based on Spark/Trino JDK requirements. - -## Expression Strategy - -Expressions are the primary mechanism for pushing computation into Vortex (filters, projections, -computed columns). The bindings must support expression construction across all languages. - -**Serialized expressions (protobuf)** are the baseline for all languages. Any language that can -produce protobuf bytes can construct expressions — this is the lowest-common-denominator approach -and the one the C API should always accept. - -**Native expression construction** is a convenience layer built per-language. Python already has -this via PyO3. C++ and Java should have builder APIs that produce the same protobuf under the hood. - -However, serialized expressions have limitations that mean a mixed approach will likely remain -necessary: - -- **Dynamic expressions** (e.g. UDFs, closures, expressions that reference runtime state) cannot - be represented in serialized form. These require native expression handles passed through the - binding layer. -- **Large literal values** embedded in expressions (e.g. large `IN` lists) can be a performance - constraint when serialized, since the entire value must be copied through protobuf encoding and - decoding. - -The C API should therefore support both modes: - -- Accept serialized expression bytes (`vortex_scan_set_filter_bytes`). -- Provide builder functions for common expressions (`vortex_expr_column`, `vortex_expr_eq`, - `vortex_expr_and`, etc.) that return opaque expression handles. - -This dual approach lets simple integrations pass pre-built protobuf while giving interactive users -and query engines with dynamic expressions an ergonomic builder API. diff --git a/docs/project/contributing.md b/docs/project/contributing.md index 4721f335dfc..33699d12789 100644 --- a/docs/project/contributing.md +++ b/docs/project/contributing.md @@ -45,7 +45,9 @@ This section describes our expectations around disclosure, review, and accountab The repository has [Claude Code](https://docs.anthropic.com/en/docs/claude-code) configured as a GitHub Action. Users with write access to the repository can mention `@claude` in PR comments or issue comments to trigger AI-powered code reviews, request changes, or generate PRs. Configuration -for Claude's behavior lives in `CLAUDE.md` at the repository root. +for coding agents lives in the repository's +[`AGENTS.md`](https://github.com/vortex-data/vortex/blob/develop/AGENTS.md), with reusable workflows +under [`.agents/skills`](https://github.com/vortex-data/vortex/tree/develop/.agents/skills). ### Disclosure diff --git a/docs/project/index.md b/docs/project/index.md index 51d8b29d904..b5622159a52 100644 --- a/docs/project/index.md +++ b/docs/project/index.md @@ -7,7 +7,6 @@ maxdepth: 2 community roadmap -bindings contributing changelog/index ../references diff --git a/docs/project/roadmap.md b/docs/project/roadmap.md index f2a6d8ec7b7..5c081411710 100644 --- a/docs/project/roadmap.md +++ b/docs/project/roadmap.md @@ -40,9 +40,6 @@ ecosystem: - **Other languages** (e.g. Java) will initially focus on reading and writing Arrow data to and from Vortex files. -See the [Bindings Evolution](bindings.md) design doc for the full tier model, per-language migration -plans, and phased roadmap. - ### Tensor Extension DType A built-in extension DType for multi-dimensional tensor data, enabling native support for diff --git a/docs/pyproject.toml b/docs/pyproject.toml index 809bd55d67f..9501bdf3d42 100644 --- a/docs/pyproject.toml +++ b/docs/pyproject.toml @@ -4,23 +4,13 @@ version = "0.1.0" description = "Vortex documentation." authors = [] dependencies = [ - "pydata-sphinx-theme>=0.16.0", + "accessible-pygments>=0.0.5", "hawkmoth>=0.20.0", "myst-parser>=4.0.0", - "setuptools>=75.8.0", # Required by sphinxcontrib-bibtex - "sphinx-autobuild>=2024.10.3", - "sphinx-copybutton>=0.5.2", - "sphinx-design>=0.6.0", - "sphinx-inline-tabs>=2023.4.21", + "rcssmin>=1.1.2", "sphinx>=9.0.0", - "sphinxcontrib-bibtex>=2.6.3", - "sphinxcontrib-mermaid>=1.0.0", - "sphinxext-opengraph>=0.9.1", - "breathe>=4.35.0", "vortex-data", "libclang>=18.1.1", - # forced transitive bumps - "starlette>=0.49.1", ] requires-python = ">= 3.11" diff --git a/docs/references.bib b/docs/references.bib deleted file mode 100644 index e6ef8ca5c59..00000000000 --- a/docs/references.bib +++ /dev/null @@ -1,174 +0,0 @@ -@article{2023/BtrBlocks, - author = {Kuschewski, Maximilian and Sauerwein, David and Alhomssi, Adnan and Leis, Viktor}, - title = {BtrBlocks: Efficient Columnar Compression for Data Lakes}, - year = {2023}, - issue_date = {June 2023}, - publisher = {Association for Computing Machinery}, - address = {New York, NY, USA}, - volume = {1}, - number = {2}, - url = {https://doi.org/10.1145/3589263}, - doi = {10.1145/3589263}, - abstract = {Analytics is moving to the cloud and data is moving into data lakes. These reside on object storage services like S3 and enable seamless data sharing and system interoperability. To support this, many systems build on open storage formats like Apache Parquet. However, these formats are not optimized for remotely-accessed data lakes and today's high-throughput networks. Inefficient decompression makes scans CPU-bound and thus increases query time and cost. With this work we present BtrBlocks, an open columnar storage format designed for data lakes. BtrBlocks uses a set of lightweight encoding schemes, achieving fast and efficient decompression and high compression ratios.}, - journal = {Proc. ACM Manag. Data}, - month = jun, - articleno = {118}, - numpages = {26}, - keywords = {columnar storage, compression, data lake, query processing} -} - -@article{2023/FastLanes, - author = {Afroozeh, Azim and Boncz, Peter}, - title = {The FastLanes Compression Layout: Decoding > 100 Billion Integers per Second with Scalar Code}, - year = {2023}, - issue_date = {May 2023}, - publisher = {VLDB Endowment}, - volume = {16}, - number = {9}, - issn = {2150-8097}, - url = {https://doi.org/10.14778/3598581.3598587}, - doi = {10.14778/3598581.3598587}, - abstract = {The open-source FastLanes project aims to improve big data formats, such as Parquet, ORC and columnar database formats, in multiple ways. In this paper, we significantly accelerate decoding of all common Light-Weight Compression (LWC) schemes: DICT, FOR, DELTA and RLE through better data-parallelism. We do so by re-designing the compression layout using two main ideas: (i) generalizing the value interleaving technique in the basic operation of bit-(un)packing by targeting a virtual 1024-bits SIMD register, (ii) reordering the tuples in all columns of a table in the same Unified Transposed Layout that puts tuple chunks in a common "04261537" order (explained in the paper); allowing for maximum independent work for all possible basic SIMD lane widths: 8, 16, 32, and 64 bits.We address the software development, maintenance and future-proofness challenges of increasing hardware diversity, by defining a virtual 1024-bits instruction set that consists of simple operators supported by all SIMD dialects; and also, importantly, by scalar code. The interleaved and tuple-reordered layout actually makes scalar decoding faster, extracting more data-parallelism from today's wide-issue CPUs. Importantly, the scalar version can be fully auto-vectorized by modern compilers, eliminating technical debt in software caused by platform-specific SIMD intrinsics.Micro-benchmarks on Intel, AMD, Apple and AWS CPUs show that FastLanes accelerates decoding by factors (decoding >40 values per CPU cycle). FastLanes can make queries faster, as compressing the data reduces bandwidth needs, while decoding is almost free.}, - journal = {Proc. VLDB Endow.}, - month = may, - pages = {2132–2144}, - numpages = {13} -} - -@article{2023/ALP, - author = {Afroozeh, Azim and Kuffo, Leonardo X. and Boncz, Peter}, - title = {ALP: Adaptive Lossless floating-Point Compression}, - year = {2023}, - issue_date = {December 2023}, - publisher = {Association for Computing Machinery}, - address = {New York, NY, USA}, - volume = {1}, - number = {4}, - url = {https://doi.org/10.1145/3626717}, - doi = {10.1145/3626717}, - abstract = {IEEE 754 doubles do not exactly represent most real values, introducing rounding errors in computations and [de]serialization to text. These rounding errors inhibit the use of existing lightweight compression schemes such as Delta and Frame Of Reference (FOR), but recently new schemes were proposed: Gorilla, Chimp128, PseudoDecimals (PDE), Elf and Patas. However, their compression ratios are not better than those of general-purpose compressors such as Zstd; while [de]compression is much slower than Delta and FOR.We propose and evaluate ALP, that significantly improves these previous schemes in both speed and compression ratio (Figure 1). We created ALP after carefully studying the datasets used to evaluate the previous schemes. To obtain speed, ALP is designed to fit vectorized execution. This turned out to be key for also improving the compression ratio, as we found in-vector commonalities to create compression opportunities. ALP is an adaptive scheme that uses a strongly enhanced version of PseudoDecimals [31] to losslessly encode doubles as integers if they originated as decimals, and otherwise uses vectorized compression of the doubles' front bits. Its high speeds stem from our implementation in scalar code that auto-vectorizes, using building blocks provided by our FastLanes library [6], and an efficient two-stage compression algorithm that first samples row-groups and then vectors.}, - journal = {Proc. ACM Manag. Data}, - month = dec, - articleno = {230}, - numpages = {26}, - keywords = {big data formats, columnar storage, floating point compression, lightweight compression, lossless compression, vectorized execution} -} - -@inproceedings{2024/FastLanesGPU, - author = {Afroozeh, Azim and Felius, Lotte and Boncz, Peter}, - title = {Accelerating GPU Data Processing using FastLanes Compression}, - year = {2024}, - isbn = {9798400706677}, - publisher = {Association for Computing Machinery}, - address = {New York, NY, USA}, - url = {https://doi.org/10.1145/3662010.3663450}, - doi = {10.1145/3662010.3663450}, - abstract = {We show that compression can be a win-win for GPU data processing: it not only allows to store more data in GPU global memory, but can also accelerate data processing. We show that the complete redesign of compressed columnar storage in FastLanes, with its fully data-parallel bit-packing and encodings, also benefits GPU hardware. We micro-benchmark the performance of FastLanes on two GPU architectures (Nvidia T4 and V100) and integrate FastLanes in the Crystal GPU query processing prototype. Our experiments show that FastLanes decompression significantly outperforms previous decompression methods in micro-benchmarks, and can make end-to-end SSB queries up to twice faster compared to uncompressed query processing - in contrast to previous work where GPU decompression caused execution to slow down. We further discovered that an access granularity of decoding vectors of 1024 values is too large for a single GPU warp due to register pressure. We mitigate this here using mini-vectors - a future work question is how to further reduce this granularity with minimal impact on efficiency.}, - booktitle = {Proceedings of the 20th International Workshop on Data Management on New Hardware}, - articleno = {8}, - numpages = {11}, - location = {Santiago, AA, Chile}, - series = {DaMoN '24} -} - -@article{2020/FSST, - author = {Boncz, Peter and Neumann, Thomas and Leis, Viktor}, - title = {FSST: fast random access string compression}, - year = {2020}, - issue_date = {August 2020}, - publisher = {VLDB Endowment}, - volume = {13}, - number = {12}, - issn = {2150-8097}, - url = {https://doi.org/10.14778/3407790.3407851}, - doi = {10.14778/3407790.3407851}, - abstract = {Strings are prevalent in real-world data sets. They often occupy a large fraction of the data and are slow to process. In this work, we present Fast Static Symbol Table (FSST), a lightweight compression scheme for strings. On text data, FSST offers decompression and compression speed similar to or better than the best speed-optimized compression methods, such as LZ4, yet offers significantly better compression factors. Moreover, its use of a static symbol table allows random access to individual, compressed strings, enabling lazy decompression and query processing on compressed data. We believe these features will make FSST a valuable piece in the standard compression toolbox.}, - journal = {Proc. VLDB Endow.}, - month = jul, - pages = {2649–2661}, - numpages = {13} -} - -@article{2019/Procella, - title = {Procella: Unifying serving and analytical data at YouTube}, - author = {Biswapesh Chattopadhyay and Priyam Dutta and Weiran Liu and Ott Tinn and Andrew McCormick and Aniket Mokashi and Paul Harvey and Hector Gonzalez and David Lomax and Sagar Mittal and Roee Aharon Ebenstein and Nikita Mikhaylin and Hung-ching Lee and Xiaoyan Zhao and Guanzhong Xu and Luis Antonio Perez and Farhad Shahmohammadi and Tran Bui and Neil McKay and Vera Lychagina and Brett Elliott}, - year = {2019}, - URL = {https://dl.acm.org/citation.cfm?id=3360438}, - journal = {PVLDB}, - pages = {2022-2034}, - volume = {12(12)} -} - -@article{2023/AnyBlob, - author = {Durner, Dominik and Leis, Viktor and Neumann, Thomas}, - title = {Exploiting Cloud Object Storage for High-Performance Analytics}, - year = {2023}, - issue_date = {July 2023}, - publisher = {VLDB Endowment}, - volume = {16}, - number = {11}, - issn = {2150-8097}, - url = {https://doi.org/10.14778/3611479.3611486}, - doi = {10.14778/3611479.3611486}, - abstract = {Elasticity of compute and storage is crucial for analytical cloud database systems. All cloud vendors provide disaggregated object stores, which can be used as storage backend for analytical query engines. Until recently, local storage was unavoidable to process large tables efficiently due to the bandwidth limitations of the network infrastructure in public clouds. However, the gap between remote network and local NVMe bandwidth is closing, making cloud storage more attractive. This paper presents a blueprint for performing efficient analytics directly on cloud object stores. We derive cost- and performance-optimal retrieval configurations for cloud object stores with the first in-depth study of this foundational service in the context of analytical query processing. For achieving high retrieval performance, we present AnyBlob, a novel download manager for query engines that optimizes throughput while minimizing CPU usage. We discuss the integration of high-performance data retrieval in query engines and demonstrate it by incorporating AnyBlob in our database system Umbra. Our experiments show that even without caching, Umbra with integrated AnyBlob achieves similar performance to state-of-the-art cloud data warehouses that cache data on local SSDs while improving resource elasticity.}, - journal = {Proc. VLDB Endow.}, - month = jul, - pages = {2769–2782}, - numpages = {14} -} - -@article{2023/SelectionPushdown, - author = {Li, Yinan and Lu, Jianan and Chandramouli, Badrish}, - title = {Selection Pushdown in Column Stores using Bit Manipulation Instructions}, - year = {2023}, - issue_date = {June 2023}, - publisher = {Association for Computing Machinery}, - address = {New York, NY, USA}, - volume = {1}, - number = {2}, - url = {https://doi.org/10.1145/3589323}, - doi = {10.1145/3589323}, - abstract = {Modern analytical database systems predominantly rely on column-oriented storage, which offers superior compression efficiency due to the nature of the columnar layout. This compression, however, creates challenges in decoding speed during query processing. Previous research has explored predicate pushdown on encoded values to avoid decoding, but these techniques are restricted to specific encoding schemes and predicates, limiting their practical use. In this paper, we propose a generic predicate pushdown approach that supports arbitrary predicates by leveraging selection pushdown to reduce decoding costs. At the core of our approach is a fast select operator capable of directly extracting selected encoded values without decoding, by using Bit Manipulation Instructions, an instruction set extension to the X86 architecture. We empirically evaluate the proposed techniques in the context of Apache Parquet using both micro-benchmarks and the TPC-H benchmark, and show that our techniques improve the query performance of Parquet by up to one order of magnitude with representative scan queries. Further experimentation using Apache Spark demonstrates speed improvements of up to 5.5X even for end-to-end queries involving complex joins.}, - journal = {Proc. ACM Manag. Data}, - month = jun, - articleno = {178}, - numpages = {26}, - keywords = {BMI, column stores, compression, predicate pushdown} -} - -@inproceedings{2021/JSONTiles, - author = {Durner, Dominik and Leis, Viktor and Neumann, Thomas}, - title = {JSON Tiles: Fast Analytics on Semi-Structured Data}, - year = {2021}, - isbn = {9781450383431}, - publisher = {Association for Computing Machinery}, - address = {New York, NY, USA}, - url = {https://doi.org/10.1145/3448016.3452809}, - doi = {10.1145/3448016.3452809}, - abstract = {Developers often prefer flexibility over upfront schema design, making semi-structured data formats such as JSON increasingly popular. Large amounts of JSON data are therefore stored and analyzed by relational database systems. In existing systems, however, JSON's lack of a fixed schema results in slow analytics. In this paper, we present JSON tiles, which, without losing the flexibility of JSON, enables relational systems to perform analytics on JSON data at native speed. JSON tiles automatically detects the most important keys and extracts them transparently - often achieving scan performance similar to columnar storage. At the same time, JSON tiles is capable of handling heterogeneous and changing data. Furthermore, we automatically collect statistics that enable the query optimizer to find good execution plans. Our experimental evaluation compares against state-of-the-art systems and research proposals and shows that our approach is both robust and efficient.}, - booktitle = {Proceedings of the 2021 International Conference on Management of Data}, - pages = {445–458}, - numpages = {14}, - keywords = {storage, semi-structured data, scan, olap, jsonb, json, analytics}, - location = {Virtual Event, China}, - series = {SIGMOD '21} -} - -@article{2024/MementoFilter, - author = {Eslami, Navid and Dayan, Niv}, - title = {Memento Filter: A Fast, Dynamic, and Robust Range Filter}, - year = {2024}, - issue_date = {December 2024}, - publisher = {Association for Computing Machinery}, - address = {New York, NY, USA}, - volume = {2}, - number = {6}, - url = {https://doi.org/10.1145/3698820}, - doi = {10.1145/3698820}, - abstract = {Range filters are probabilistic data structures that answer approximate range emptiness queries. They aid in avoiding processing empty range queries and have use cases in many application domains such as key-value stores and social web analytics. However, current range filters do not support dynamically changing and growing datasets. Moreover, several of these designs also exhibit impractically high false positive rates under correlated workloads, which are common in practice. These impediments restrict the applicability of range filters across a wide range of use cases.We introduce Memento filter, the first range filter to simultaneously offer dynamicity, fast operations, and a robust false positive rate for any workload. Memento filter partitions the key universe and clusters its keys according to this partitioning. For each cluster, it stores a fingerprint and a list of key suffixes contiguously. The encoding of these lists makes them amenable to existing dynamic filter structures. Due to the one-to-one mapping from keys to suffixes, Memento filter supports inserts and deletes and can even expand to accommodate a growing dataset.We implement Memento filter on top of a Rank-and-Select Quotient filter and InfiniFilter and demonstrate that it achieves a competitive false positive rate and performance with the state of the art while also providing dynamicity. Due to its dynamicity, Memento filter is the first range filter applicable to B-Trees. We showcase this by integrating Memento filter into WiredTiger, a B-Tree-based key-value store, significantly boosting its performance for mixed workloads.}, - journal = {Proc. ACM Manag. Data}, - month = dec, - articleno = {244}, - numpages = {27}, - keywords = {data growth, dynamic data structure, range filter, scalability} -} \ No newline at end of file diff --git a/docs/references.md b/docs/references.md index 93f29f3b43c..662290005da 100644 --- a/docs/references.md +++ b/docs/references.md @@ -6,28 +6,32 @@ one another to send links to the same papers over and over again). ## Compression & Encodings -* The FastLanes Compression Layout: Decoding > 100 Billion Integers per Second with Scalar - Code {cite:p}`2023/FastLanes` -* Accelerating GPU Data Processing using FastLanes Compression {cite:p}`2024/FastLanesGPU` -* ALP: Adaptive Lossless floating-Point Compression {cite:p}`2023/ALP` -* FSST: fast random access string compression {cite:p}`2020/FSST` +* [The FastLanes Compression Layout: Decoding > 100 Billion Integers per Second with Scalar + Code](https://doi.org/10.14778/3598581.3598587) — Afroozeh, Boncz. PVLDB 16(9), 2023. +* [Accelerating GPU Data Processing using FastLanes + Compression](https://doi.org/10.1145/3662010.3663450) — Afroozeh, Felius, Boncz. DaMoN 2024. +* [ALP: Adaptive Lossless floating-Point Compression](https://doi.org/10.1145/3626717) — + Afroozeh, Kuffo, Boncz. Proc. ACM Manag. Data 1(4), 2023. +* [FSST: fast random access string compression](https://doi.org/10.14778/3407790.3407851) — + Boncz, Neumann, Leis. PVLDB 13(12), 2020. ## Scanning & Compute -* Selection Pushdown in Column Stores using Bit Manipulation Instructions {cite:p}`2023/SelectionPushdown` -* JSON Tiles: Fast Analytics on Semi-Structured Data {cite:p}`2021/JSONTiles` -* Memento Filter: A Fast, Dynamic, and Robust Range Filter {cite:p}`2024/MementoFilter` +* [Selection Pushdown in Column Stores using Bit Manipulation + Instructions](https://doi.org/10.1145/3589323) — Li, Lu, Chandramouli. Proc. ACM Manag. Data 1(2), 2023. +* [JSON Tiles: Fast Analytics on Semi-Structured Data](https://doi.org/10.1145/3448016.3452809) — + Durner, Leis, Neumann. SIGMOD 2021. +* [Memento Filter: A Fast, Dynamic, and Robust Range + Filter](https://doi.org/10.1145/3698820) — Eslami, Dayan. Proc. ACM Manag. Data 2(6), 2024. ## Columnar File Formats -* BtrBlocks: A High-Performance In-Memory Columnar Storage Format {cite:p}`2023/BtrBlocks` -* Procella: Unifying serving and analytical data at YouTube {cite:p}`2019/Procella` +* [BtrBlocks: Efficient Columnar Compression for Data Lakes](https://doi.org/10.1145/3589263) — + Kuschewski, Sauerwein, Alhomssi, Leis. Proc. ACM Manag. Data 1(2), 2023. +* [Procella: Unifying serving and analytical data at + YouTube](https://dl.acm.org/citation.cfm?id=3360438) — Chattopadhyay et al. PVLDB 12(12), 2019. ## I/O & Cloud Storage -* Exploiting Cloud Object Storage for High-Performance Analytics {cite:p}`2023/AnyBlob` - ---- - -```{bibliography} references.bib -``` +* [Exploiting Cloud Object Storage for High-Performance + Analytics](https://doi.org/10.14778/3611479.3611486) — Durner, Leis, Neumann. PVLDB 16(11), 2023. diff --git a/docs/specs/editions.md b/docs/specs/editions.md new file mode 100644 index 00000000000..ee5f494861f --- /dev/null +++ b/docs/specs/editions.md @@ -0,0 +1,386 @@ +# Editions + +Vortex files contain several kinds of serialized **component**: array encodings, layout encodings, extension dtypes, and +aggregate functions. An **edition** is a named set of their concrete wire IDs. It controls what a writer may put in a +file and, once frozen, identifies its origin library or project and minimum version: the earliest release of that origin +that recognizes every ID in the set. + +Editions belong to independently versioned families and are cumulative within a family. Each edition includes all +components from the preceding edition in that family, plus any additions. A writer selects at most one edition from +each family and may use the union of their component IDs. For example, selecting `core2026.08.3` and +`tensor2026.04.0` allows the core components together with tensor arrays and dtypes. Every family names the origin +library or project whose release versions its editions use. + +The first frozen edition, `core2025.05.0`, contains the components that Vortex `0.36.0` could write. This marks the +start of the Vortex file format's stability guarantee. Every Vortex release from `0.36.0` onward can read +`core2025.05.0`, and later frozen `core` editions extend that guarantee to newer components. + +When a writer selects only frozen editions from one origin, the highest of their minimum library versions is the +earliest release guaranteed to read the resulting file. With multiple origins, the file requires the recorded minimum +version of each. Editions without minimum library versions are drafts and carry no guarantee about their future +compatibility. + +## What an edition contains + +An edition records every component by kind and wire ID. IDs are unique within a kind, but not across kinds: a layout +named `vortex.flat` and an array encoding with the same ID are distinct components. The writer therefore builds and +enforces a separate allowlist for each kind: + +| Kind | What it identifies | Used at | +|-------------|--------------------------------------------|------------------------------| +| `array` | a serialized array representation | array serialization | +| `layout` | the footer's layout tree | layout serialization context | +| `dtype` | extension dtypes nested in the file schema | file writer | +| `aggregate` | zone maps in zoned layouts | layout writer context | + +Writing a component that is absent from the selected editions fails the write. This rule applies to every kind, +including aggregates. Although a zone map is only an optimization and could be dropped, doing so would silently change +the writer's configured pruning behavior. + +Only aggregates that would actually be written are checked. If a column's dtype cannot support an aggregate, the writer +omits it and there is no edition violation. + +An empty allowlist permits no components. Collectively, the selected editions must declare every serialized array ID, +layout encoding, extension dtype, and aggregate function that the writer writes. An array serializer may expose several +wire IDs for one in-memory encoding. The serializer chooses the representation, and the serialization context rejects +the write if the chosen wire ID is not declared by the selected editions. + +For example, `core2026.08.0` declares the aggregate functions that the default writer may store in zone maps: `min`, +`max`, `bounded_min`, `bounded_max`, `nan_count`, and `null_count`. It does not declare `sum`, because the writer does +not store sums in zone maps. File-level statistics use a fixed legacy field for sums rather than a serialized aggregate +function ID, so this allowlist does not apply to them. + +Optional Vortex modules enable their own edition families alongside `core`. Tensor support enables +`tensor2026.04.0`, for example, while Zstd buffer wrapping enables `zstd2026.02.0`. + +## Resolving an unknown-component error + +An unknown-ID error means that the reader does not recognize a serialized component in the file. Find the component's +kind and ID in the [registry](#edition-registry): + +1. **It belongs to a frozen edition.** Upgrade the edition's named origin to at least its minimum library version. +2. **It belongs to a draft edition.** No released reader is guaranteed to support it. Use a build that registers the + component or ask the file's producer which build to use. +3. **It is not in the registry.** The file contains a custom, third-party, or experimental component outside the + editions system. Ask the producer for its implementation and register it with the reader's session. + +Tools that only inspect or copy data can opt in to `allow_unknown`. Unknown array encodings, layout encodings, and +extension dtypes are then preserved as inert representations. An unknown aggregate function disables the affected +zone-map pruning rather than causing the file to be rejected. + +## Writing with an edition + +By default, the Vortex facade targets `core2026.08.3`. A new encoding or serialization feature that is +still evolving gets a new draft edition; later additions create later editions rather than changing an already +published feature set. Each feature advances through its own independently versioned family until it is ready to join +the shared `preview` family. Preview components remain opt-in until they are ready to join `core`. Components supplied +by an optional plugin belong to that plugin's family, such as `tensor`, `zstd`, `spatial`, or `json`. + +Edition configuration belongs to the writer's Vortex session. Registering an edition makes its declaration available to +the session; enabling it allows the writer to use its components. Enabling another edition in the same family replaces +the previous selection. + +You can change the default configuration to: + +- **Target an older `core` edition** when the file must remain readable by an older Vortex deployment. +- **Enable another family** to use components outside `core`. Vortex currently defines `preview`, `tensor`, `zstd`, + `spatial`, and `json` in addition to `core`. + +Sessions created without the Vortex facade must register and enable their editions before writing files. + +For experimental or custom components that do not belong to an edition, the Rust writer exposes +`disable_editions()`. This disables every edition check for that write: every array representation registered in the +session is eligible for compression and serialization, while layouts, extension dtypes, and aggregate functions are +unrestricted. It does not register missing readers, so files written this way have no edition compatibility guarantee. + +Compression and edition compatibility are separate. Compressors produce current in-memory arrays and do not select a +wire ID. The writer maps each allowed serialized ID to its current in-memory encoding and restricts the default +BtrBlocks compressor to schemes producing those encodings. Custom compressors remain unrestricted, with serialization +providing the final compatibility boundary when edition enforcement is enabled. At that boundary, the array plugin +produces an ID, metadata, buffers, and children. The serialization context interns the returned ID and fails the write +if the selected editions do not permit it. A serializer may emit a historical ID when the value satisfies that ID's +frozen contract, but it does not inspect the edition allowlist. Without disabling edition enforcement, a custom layout +or compressor therefore cannot bypass the final wire-ID check. + +## How editions change + +A frozen edition never changes: neither its membership list nor the meaning of its component IDs may be altered. +Introducing a new serialized object or a reader-visible revision requires a new edition; it is never added +retroactively to an existing edition. A component supplied by an optional plugin creates that edition in the plugin's +independently versioned family. + +Core-maintained objects do not enter `core` directly. When an object is ready for users to try and its wire format is +believed complete, it enters a draft edition in an independently versioned family. Publishing that edition is a +format-stability commitment, not the start of format design: the serialized contract should change only when absolutely +necessary to resolve a problem found during testing. After successful testing, the same object ID and wire contract move +into a new `preview` edition for broad opt-in use, and later into a new `core` edition for use by default. + +A new stable `core` or plugin edition may freeze in the release of its origin project in which it first ships. Until +that release is cut, its version is not known and the declaration keeps `min_library_version: None`. After the release +is cut, the declaration is updated with that newly released version, usually during development of the next release. +This backfills the documented minimum library version; it does not delay the freeze or its read-forever compatibility +guarantee. + +A component may later be deprecated, meaning that writers stop using it. Readers must continue to support it, so +deprecation does not invalidate existing files. + +Writer behavior evolves independently from the in-memory representation. A change that an old reader must distinguish +uses a new serialized ID, even when the new deserializer produces the same in-memory array. A serializer may continue +emitting the older ID for values that satisfy its frozen contract; the selected editions validate the ID it emits. + +## How serialized components evolve + +Editions govern serialized components, not in-memory representations. An in-memory representation may gain capabilities +or be replaced without changing an edition. Each in-memory array plugin owns the mapping between that representation and +its wire history: + +- the serialized IDs its deserializer recognizes; +- one serializer that returns the appropriate lossless variant as an ID, metadata, buffers, and children; and +- a deserializer that receives the exact ID found in the file and constructs the current in-memory representation. + +An in-memory representation often has one serialized ID equal to its in-memory encoding ID, but this is only the simple +case. Editions constrain the ID stored in the file, because that is what an old reader can recognize. + +### Reader-visible evolution requires a new ID + +Any new form that an old reader does not already understand uses a new serialized ID. This includes additive metadata or +children when an old reader would accept the ID but reject or misinterpret the new combination. The ID is the capability +tag: readers do not consult the edition or negotiate a separate version while decoding an array. + +Keeping an ID is safe only when the emitted representation remains within that ID's existing frozen contract. A writer +may choose a different but already-valid encoding of the same contract, and a reader may fix a bug or normalize the old +form into a newer in-memory structure. Neither action expands what the wire ID means. + +A new wire ID does not normally require a second in-memory array. The current plugin registers every historical ID, +serializes the current value under the oldest allowed lossless one, and deserializes all of them into the current type. +The old ID remains registered forever. If the compressor and serializer cannot preserve one common in-memory +representation and losslessly downgrade it, the change instead needs a new in-memory array, compressor, and +deserializer. + +Name successive incompatible revisions by appending a version to the same base name: `vortex.foo`, `vortex.foo_v2`, +`vortex.foo_v3`. Do not give successor versions descriptive names. A linear naming scheme keeps the component's +serialized history unambiguous. + +#### Example: multi-part decimals + +`vortex.decimal_byte_parts` entered `core2025.05.0` with each decimal value represented by one signed integer child. Its +metadata includes `lower_part_count`, but readers of this component require that field to be zero. Suppose the in-memory +representation gains support for wide decimals, represented by a signed most-significant part and one or more unsigned +64-bit lower parts: + +- The serializer first tries to construct the old single-signed-child form. If every value can be + represented that way, it emits `vortex.decimal_byte_parts` with `lower_part_count = 0`, even if + the current in-memory array has lower-part children. +- An array that cannot be collapsed into that old form losslessly uses the new + `vortex.decimal_byte_parts_v2` component, initially staged in a draft edition. +- A new reader deserializes both IDs into the same in-memory representation. An older reader reports + `vortex.decimal_byte_parts_v2` as unknown instead of trying to decode a wire format it does not support. +- When targeting an edition that permits only the old ID, serializing a value that can be collapsed succeeds; an + irreducibly multi-part value fails because no lossless downgrade exists. + +#### Example: Pco 8-bit integers + +The historical `vortex.pco` contract does not include `i8` or `u8`; readers implementing that contract must not be +sent an 8-bit Pco payload under the familiar ID. Adding 8-bit support keeps one current in-memory `Pco` array but adds +`vortex.pco.v2` as a serialized component: + +- The single Pco serializer emits `vortex.pco` for the primitive types covered by the old contract, even when both IDs + are permitted. +- For `i8` or `u8`, the earliest lossless form is `vortex.pco.v2`. A target edition without that ID rejects the write. +- The current deserializer registers both IDs. When given `vortex.pco`, it still rejects an 8-bit dtype; understanding + the v2 payload does not silently broaden the frozen v1 contract. +- The Pco compression scheme can sample and construct 8-bit Pco arrays without consulting editions. Wire selection + remains the serializer's responsibility. + +If writing an older edition must succeed for every input, its compression policy must choose an in-memory encoding +whose serializer has a permitted lossless form. It must not disguise the newer Pco form with the old ID. + +### Reading: deserialize into the current representation + +Every component in a frozen edition remains readable. Its deserializer may convert old data directly into the current +in-memory representation rather than preserving a parallel legacy representation. For example, a `vortex.alp` array with +interior patches is read as a `Patched` array around a patch-free ALP array. Similarly, old zone maps, including +`vortex.stats` layouts, are read by the machinery used for modern `vortex.zoned` layouts. + +Readers do not negotiate versions. They resolve the component ID, pass that exact ID to its deserializer, and either +construct the current in-memory array or report an +[unknown-component error](#resolving-an-unknown-component-error). + +A current deserializer must preserve each historical ID's contract. Recognizing a newer ID does not authorize it to +accept the newer metadata, child shape, dtype coverage, or buffer interpretation when the file carries an older ID. + +A file contains its array ID, dtype, metadata, children, and buffers. A newer plugin may be registered under both +`vortex.foo` and `vortex.foo_v2`, but an older build is registered only under `vortex.foo`. This is what guarantees that +the older build rejects a v2 file before interpreting its contents. + +### Writing: validate the selected component and writer behavior + +For each in-memory array, the writer calls its plugin's single serializer. The serializer owns the versioning logic and +returns the appropriate lossless variant. It may change metadata, buffers, and children without constructing a legacy +in-memory array. Returning `None` means the array cannot be serialized. The serialization context then interns the +returned ID, failing the write if that ID is not permitted by the selected editions. + +This selection happens recursively after compression. Compressor output therefore remains an in-memory concern: a +compressor does not label its array with an edition or choose a wire version. Layouts, extension dtypes, and aggregates +perform their analogous compatibility checks at their own serialization boundaries. + +### What this means for each kind + +- **Arrays.** The array serialization context permits only wire IDs from the selected editions. The in-memory array's + serializer chooses its lossless representation, and the context rejects it if its ID is not permitted. +- **Layouts.** The layout strategy builds the layout tree at write time. When targeting an older edition, it must use + structures available in that edition, such as plain chunked data in place of newer auxiliary layouts. +- **Extension dtypes.** Before writing any bytes, the file writer recursively validates every extension dtype in the + schema. Readers resolve serialized dtype IDs against the session's dtype registry. +- **Aggregate functions.** Zone maps serialize aggregate function IDs and their options. A zone map containing a + function outside the selected editions fails the write. With `allow_unknown`, readers disable a zone map whose + aggregate function they do not recognize; ignoring a zone map only reduces pruning and does not affect correctness. + +## The `preview` family + +The additive `preview` family is the shared opt-in set for core-maintained components whose serialized contracts have +survived independent testing but are not yet available to the default core writer. Preview currently contains no +components. Adding the first component will create a later preview edition; unrelated work remains in independent +families until it meets the preview compatibility bar. + +## Independently versioned component families + +Components that are ready for focused testing but are not yet ready for the shared preview set advance through their own +families. Optional modules use families such as `tensor`, `zstd`, `spatial`, and `json`. Each family can evolve without +coupling its chronology or selection to unrelated components. + +The wire format is expected to be complete when its first draft edition is published and should change only when +necessary to resolve an issue discovered during testing. If a correction changes what readers must understand, give the +corrected representation a new ID and add a later edition to the same family. Once testing establishes that an object is +ready for broad opt-in use, promote that same ID and serialized contract into a new `preview` edition. Later adoption by +the default writer promotes it into a new `core` edition. + +The default writer does not emit a component merely because its reader understands it. Users opt in by enabling the +edition containing that component. + +## Declaring, freezing, and the edition records + +The default declarations live in `vortex-edition/src/declarations/`, while optional-module declarations live in their +owning crates. Each declared edition is exported as a TOML record under `vortex/editions/`, grouped by family. A record +names the origin library or project whose releases `min_library_version` refers to. Draft records omit that field and +carry no read-forever guarantee. Regenerate the records by running: + +```sh +cargo run -p xtask -- generate-editions +``` + +Changing the declarations follows the edition's lifecycle: + +1. **Create a new family and edition for every new object.** Never add an unrelated serialized + object or reader-visible revision to an existing family. A revision advances the family that + owns its earlier ID. +2. **Publish test-ready work as a draft.** When an object is ready to be tried and its format is + believed complete, give it a wire ID and add it to a draft edition in its family. Change that + format only when necessary to resolve an issue found during testing; a reader-visible + correction gets another ID and a later edition in the same family. +3. **Promote the tested contract to preview.** Once it is ready for broad opt-in use, add the same + object ID and wire contract to a new additive `preview` edition. Promotion must not redesign + the format. +4. **Promote the adopted contract to core.** Once it is ready for use by the default writer, add + the same object ID and wire contract to a new `core` edition with `min_library_version: None` + and regenerate its draft record, then ship it in a release. The edition freezes as part of that + release. Its minimum library version cannot be populated yet because the release version is not + known until the release is cut. +5. **Backfill the released version.** After cutting the release, set `min_library_version` to that + newly released Vortex version — the version that first shipped readers for every member — and + regenerate the records, converting the draft record into a frozen record. This update usually + lands during development of the next release, but it documents the freeze that already + happened; it does not freeze the edition later. +6. **Never touch it again.** A frozen record is immutable: CI (`cargo run -p xtask -- check-editions`) rejects any + change that edits, renames, unfreezes, + or deletes a frozen record, and rejects new editions that do not extend their family's + chronology. To change what writers may emit, declare the next edition instead. + +## Edition registry + +Registry entries list the edition in which each component first appeared. Later editions in the same family inherit all +earlier components. + +### Frozen `core` editions + +#### `core2025.05.0` + +Minimum library version: `0.36.0`. + +- `array`: `fastlanes.bitpacked`, `fastlanes.for`, `vortex.alp`, `vortex.alprd`, `vortex.bool`, + `vortex.bytebool`, `vortex.chunked`, `vortex.constant`, `vortex.datetimeparts`, `vortex.decimal`, + `vortex.decimal_byte_parts`, `vortex.dict`, `vortex.ext`, `vortex.fsst`, `vortex.list`, + `vortex.null`, `vortex.primitive`, `vortex.runend`, `vortex.sparse`, `vortex.struct`, + `vortex.varbin`, `vortex.varbinview`, `vortex.zigzag` +- `layout`: `vortex.chunked`, `vortex.dict`, `vortex.flat`, `vortex.stats`, `vortex.struct` +- `dtype`: `vortex.date`, `vortex.time`, `vortex.timestamp` + +#### `core2025.06.0` + +Minimum library version: `0.40.0`. + +- `array`: `vortex.pco`, `vortex.sequence`, `vortex.zstd` + +#### `core2025.10.0` + +Minimum library version: `0.54.0`. + +- `array`: `fastlanes.rle`, `vortex.fixed_size_list`, `vortex.listview`, `vortex.masked` + +#### `core2026.08.0` + +Minimum library version: `0.84.0`. + +- `layout`: `vortex.zoned` +- `aggregate`: `vortex.bounded_max`, `vortex.bounded_min`, `vortex.max`, `vortex.min`, + `vortex.nan_count`, `vortex.null_count` + +#### `core2026.08.1` + +Minimum library version: `0.84.0`. + +- `array`: `vortex.onpair` + +#### `core2026.08.2` + +Minimum library version: `0.85.0`. + +- `array`: `vortex.map` + +#### `core2026.08.3` + +Minimum library version: `0.85.0`. + +- `array`: `vortex.parquet.variant`, `vortex.variant` +- `dtype`: `vortex.uuid` + +### Editions without a frozen guarantee + +These editions have no minimum library version. Evolving features advance through new draft editions in their own +families. Their formats are expected to remain compatible unless a defect is serious enough to block promotion into +core. Optional plugin families state their own policy. + +#### `preview2026.08.0` + +This edition currently adds no components. + +#### `tensor2026.04.0` + +- `array`: `vortex.tensor.cosine_similarity`, `vortex.tensor.inner_product`, `vortex.tensor.l2_norm`, + `vortex.tensor.l2_normalize` +- `dtype`: `vortex.tensor.fixed_shape_tensor`, `vortex.tensor.vector` + +#### `zstd2026.02.0` + +- `array`: `vortex.zstd_buffers` + +#### `spatial2026.08.0` + +- `dtype`: `vortex.st.box`, `vortex.st.linestring`, `vortex.st.multilinestring`, + `vortex.st.multipoint`, `vortex.st.multipolygon`, `vortex.st.point`, `vortex.st.polygon`, + `vortex.st.wkb` +- `aggregate`: `vortex.st.aabb` + +#### `json2026.08.0` + +- `dtype`: `vortex.json` diff --git a/docs/specs/file-format.md b/docs/specs/file-format.md index 425294f2d99..4a977365cb8 100644 --- a/docs/specs/file-format.md +++ b/docs/specs/file-format.md @@ -17,7 +17,9 @@ definition that allows efficiently querying the layout. Other considerations for the Vortex file format include: -* Backwards compatibility, and (coming soon) forwards compatibility. +* Backwards compatibility, and (coming soon) forwards compatibility. The set of encodings a + writer may put in a file — and the resulting read-compatibility promise — is governed by + [Editions](/specs/editions). * Fine-grained encryption. * Efficient access for both local disk and cloud storage. * Minimal overhead reading few columns or rows from wide or long arrays. @@ -52,6 +54,13 @@ The postscript contains the locations of: 2. a `layout` segment containing the root `Layout` 3. a `statistics` segment containing file-level per-field statistics (e.g., minima and maxima of each field/column, for whole-file pruning) 4. a `footer` segment containing a dictionary-encoded _segment map_, and other shared configuration such as compression and encryption schemes +5. up to 16 user-defined `metadata` segments, each identified by a unique, non-empty UTF-8 key of at most 64 bytes + +The postscript carries a locator (offset, length, and alignment) for each metadata segment that is +present; a file written without user metadata (and any file predating this feature) carries none. +Readers do not load the opaque metadata values by default. Opt-in metadata reads resolve each +locator separately, allowing values outside the initial file-tail read to be fetched without reading +the intervening file contents. :::{literalinclude} ../../vortex-flatbuffers/flatbuffers/vortex-file/footer.fbs :start-after: [postscript] diff --git a/docs/specs/index.md b/docs/specs/index.md index 045ccb0710d..12fab8da430 100644 --- a/docs/specs/index.md +++ b/docs/specs/index.md @@ -8,6 +8,7 @@ maxdepth: 2 --- file-format +editions ipc-format dtype-format scalar-format diff --git a/docs/user-guide/duckdb.md b/docs/user-guide/duckdb.md index 3fc0b03be9b..3a8f6ec6b5e 100644 --- a/docs/user-guide/duckdb.md +++ b/docs/user-guide/duckdb.md @@ -1,75 +1,73 @@ # DuckDB -Vortex is a [core extension](https://duckdb.org/docs/stable/core_extensions/vortex) shipped with -DuckDB, available from DuckDB 1.4.2+ on Linux and macOS (amd64, arm64). +Vortex [extension](https://duckdb.org/docs/stable/core_extensions/vortex) is +available from DuckDB 1.4.2+ on Linux and macOS (amd64, arm64). Windows support +[is planned](https://github.com/vortex-data/vortex/issues/9569). ## Setup ```sql -INSTALL vortex; -LOAD vortex; +INSTALL vortex; LOAD vortex; ``` -## Reading Vortex Files - -Use the `read_vortex` function to query a Vortex file: +## Reading files ```sql +SELECT * FROM 'data.vortex'; +# this syntax supports arguments like hive_partitioning=true SELECT * FROM read_vortex('data.vortex'); ``` -Filters and projections are pushed down into Vortex, so only the columns and rows needed by the -query are read and decompressed. +## Writing files ```sql -SELECT name, age -FROM read_vortex('data.vortex') -WHERE age > 30; +COPY (SELECT * FROM my_table) TO 'output.vortex'; ``` -:::{note} -Direct file path syntax (`SELECT * FROM 'data.vortex'`) is coming in an upcoming DuckDB release. -::: +Make sure to call `LOAD vortex` before writing any files. Duckdb writes CSV +files by default, so if Vortex extension wasn't loaded prior to `COPY`, +`output.vortex` will be a CSV file. If you want to prevent this from happening, +you can use `COPY ... TO 'output.vortex' (FORMAT vortex)` syntax instead which +will fail if Vortex is not loaded. -## Writing Vortex Files +## Python client -Export data to Vortex using the `COPY` statement. The `FORMAT vortex` clause is required — -without it, DuckDB defaults to CSV. +```python +import duckdb -```sql -COPY (SELECT * FROM my_table) TO 'output.vortex' (FORMAT vortex); -``` +duckdb.sql("INSTALL vortex") +duckdb.sql("LOAD vortex") -## Extension Options +result = duckdb.sql("SELECT * FROM 'data.vortex' WHERE age > 30") +result.show() +``` -### `vortex_filesystem` +## Object storage secrets -Controls which filesystem implementation is used for reading and writing Vortex files. +Vortex does not use DuckDB's secret storage. It uses environment variables via +[`object_store`](https://docs.rs/object_store/latest/object_store) instead. -| Value | Description | -| -------------------- | ---------------------------------------------------------------------------------------- | -| `'vortex'` (default) | Uses Vortex's built-in object store filesystem. Supports `file://` and `s3://` schemes. | -| `'duckdb'` | Uses DuckDB's built-in filesystem, including any filesystem extensions such as `httpfs`. | +If you want to read or write from an [S3 +bucket](https://docs.rs/object_store/latest/object_store/aws/struct.AmazonS3Builder.html), +set the following variables: -```sql -SET vortex_filesystem = 'duckdb'; +```sh +export AWS_ACCESS_KEY_ID="your_access_key" +export AWS_SECRET_ACCESS_KEY="your_secret_key" +export AWS_REGION="your_bucket_region" ``` -Use `'duckdb'` when you want to leverage DuckDB's filesystem extensions (e.g., `httpfs` for HTTP -or S3 access with DuckDB's credential management). Use `'vortex'` (the default) for direct -object store access via Vortex's own S3 integration, which reads credentials from environment -variables. - -## Python - -The DuckDB Python client works with `read_vortex` the same way: +For [Google Cloud Storage](https://docs.rs/object_store/latest/object_store/gcp/struct.GoogleCloudStorageBuilder.html), +use a key file variable: -```python -import duckdb +```sh +export GOOGLE_APPLICATION_CREDENTIALS="my service account JSON key file" +``` -duckdb.sql("INSTALL vortex") -duckdb.sql("LOAD vortex") +For [Azure Blob Storage](https://docs.rs/object_store/latest/object_store/azure/struct.MicrosoftAzureBuilder.html), +use -result = duckdb.sql("SELECT * FROM read_vortex('data.vortex') WHERE age > 30") -result.show() +```sh +export AZURE_STORAGE_ACCOUNT="my account" +export AZURE_STORAGE_KEY="my key" ``` diff --git a/docs/user-guide/pandas.md b/docs/user-guide/pandas.md index 4fdc13acbd6..ce080b4f759 100644 --- a/docs/user-guide/pandas.md +++ b/docs/user-guide/pandas.md @@ -38,11 +38,11 @@ convert: ... {'name': 'Mikhail', 'age': 57}, ... ]) >>> struct_arr.to_pandas() - age name - 0 25 Joseph - 1 31 Narendra - 2 33 Angela - 3 57 Mikhail + name age +0 Joseph 25 +1 Narendra 31 +2 Angela 33 +3 Mikhail 57 ``` {func}`~vortex.array` converts from a Pandas DataFrame into a Vortex array: diff --git a/docs/user-guide/pyarrow.md b/docs/user-guide/pyarrow.md index 770ac6385b1..300ed1344f5 100644 --- a/docs/user-guide/pyarrow.md +++ b/docs/user-guide/pyarrow.md @@ -97,9 +97,9 @@ Struct arrays convert to Arrow tables with {meth}`.Array.to_arrow_table`: ... ]) >>> struct_arr.to_arrow_table() pyarrow.Table -age: int64 name: string +age: int64 ---- -age: [[25,31,33,57]] name: [["Joseph","Narendra","Angela","Mikhail"]] +age: [[25,31,33,57]] ``` diff --git a/docs/user-guide/spark.md b/docs/user-guide/spark.md index 5b1f9b44b8d..0410b709aa6 100644 --- a/docs/user-guide/spark.md +++ b/docs/user-guide/spark.md @@ -6,11 +6,11 @@ connector is published to Maven Central in two flavors: - `dev.vortex:vortex-spark_2.13` for Spark 4.x (Scala 2.13) - `dev.vortex:vortex-spark_2.12` for Spark 3.5.x (Scala 2.12) -Use the `all` classifier JAR (e.g. `vortex-spark_2.13-0.78.0-all.jar`). It is self-contained: -it bundles the Vortex JNI bindings, native libraries for Linux (x86_64 and aarch64) and macOS -(aarch64), and relocates its Arrow, Guava, and Jackson dependencies to avoid classpath -conflicts with Spark. The thin (unclassified) JAR does not work on its own because it -references relocated classes that only ship in the `all` JAR. +Use the `all` classifier JAR (e.g. `vortex-spark_2.13-0.85.0-all.jar`). It is self-contained: +it bundles the Vortex JNI bindings, native libraries for Linux (x86_64 and aarch64), macOS +(aarch64), and Windows (x86_64), and relocates its Arrow, Guava, and Jackson dependencies to +avoid classpath conflicts with Spark. The thin (unclassified) JAR does not work on its own +because it references relocated classes that only ship in the `all` JAR. ## Getting Vortex into Spark @@ -18,7 +18,7 @@ For `spark-shell`, `spark-submit`, or `pyspark`, pass the `all` JAR with `--jars accepts either a local path or a URL, so you can point directly at Maven Central: ```shell -spark-shell --jars https://repo1.maven.org/maven2/dev/vortex/vortex-spark_2.13/0.78.0/vortex-spark_2.13-0.78.0-all.jar +spark-shell --jars https://repo1.maven.org/maven2/dev/vortex/vortex-spark_2.13/0.85.0/vortex-spark_2.13-0.85.0-all.jar ``` Or equivalently when building a session programmatically, e.g. in PySpark: @@ -26,13 +26,13 @@ Or equivalently when building a session programmatically, e.g. in PySpark: ```python spark = ( SparkSession.builder - .config("spark.jars", "/path/to/vortex-spark_2.13-0.78.0-all.jar") + .config("spark.jars", "/path/to/vortex-spark_2.13-0.85.0-all.jar") .getOrCreate() ) ``` ```{note} -`--packages dev.vortex:vortex-spark_2.13:0.78.0` does not work: `--packages` cannot select +`--packages dev.vortex:vortex-spark_2.13:0.85.0` does not work: `--packages` cannot select the `all` classifier and resolves the thin JAR, which fails at runtime with `NoClassDefFoundError: dev/vortex/relocated/...`. ``` @@ -44,22 +44,22 @@ format name `vortex` — no session configuration is required. To depend on the connector from a JVM project, add the `all` classifier to the dependency: -````{tab} Gradle (Kotlin) +Gradle (Kotlin): + ```kotlin -implementation("dev.vortex:vortex-spark_2.13:0.78.0:all") +implementation("dev.vortex:vortex-spark_2.13:0.85.0:all") ``` -```` -````{tab} Maven +Maven: + ```xml dev.vortex vortex-spark_2.13 - 0.78.0 + 0.85.0 all ``` -```` ## Reading Vortex Files @@ -79,6 +79,11 @@ partition per file. Column pruning is pushed down — only the columns referenced by the query are read from the file. +Filter pushdown is also supported. Comparisons (`=`, `<>`, `>`, `>=`, `<`, `<=`), `IS NULL`, +`IS NOT NULL`, `IN`, the string predicates `STARTS_WITH`, `ENDS_WITH` and `CONTAINS`, and a bare +boolean column are evaluated during the native scan. Anything else, including predicates on +partition columns, is returned to Spark for post-scan evaluation. + ## Writing Vortex Files ```java @@ -187,6 +192,7 @@ INSERT INTO vortex.`/path/to/data` VALUES (1, 'Alice', 20); | `TimestampNTZType` | Timestamp (microseconds, no timezone) | | `ArrayType` | List | | `StructType` | Struct | +| `MapType` | Map | ## S3 Support @@ -198,3 +204,17 @@ Dataset df = spark.read() .option("path", "s3://bucket/path/to/data") .load(); ``` + +## Azure Support + +Azure Data Lake Storage paths work the same way: + +```java +Dataset df = spark.read() + .format("vortex") + .option("path", "abfss://container@account.dfs.core.windows.net/path/to/data") + .load(); +``` + +Account keys are read from the Hadoop configuration: any `fs.azure.account.key*` entry is +forwarded to the native reader as the storage account key. diff --git a/docs/user-guide/vortex-python.md b/docs/user-guide/vortex-python.md index 5f2b1236da9..2675d3b9748 100644 --- a/docs/user-guide/vortex-python.md +++ b/docs/user-guide/vortex-python.md @@ -7,17 +7,17 @@ the Rust API are not yet exposed. See the {doc}`/api/python/index` for the full ## Installation -````{tab} pip +With pip: + ```bash pip install vortex-data ``` -```` -````{tab} uv +With uv: + ```bash uv add vortex-data ``` -```` ## Creating Arrays @@ -49,7 +49,7 @@ A list of {class}`dict` produces a struct array. Missing values may appear at an ... None, ... ]) >>> arr.dtype -struct({"age": int(64, nullable=True), "name": utf8(nullable=True)}, nullable=True) +struct({"name": utf8(nullable=True), "age": int(64, nullable=True)}, nullable=True) ``` {func}`~vortex.array` also accepts {class}`pyarrow.Array`, {class}`pyarrow.Table`, diff --git a/encodings/alp/Cargo.toml b/encodings/alp/Cargo.toml index 13ccc67d63b..d870a0fe0cf 100644 --- a/encodings/alp/Cargo.toml +++ b/encodings/alp/Cargo.toml @@ -17,17 +17,16 @@ version = { workspace = true } workspace = true [dependencies] +alp = { workspace = true } itertools = { workspace = true } num-traits = { workspace = true } prost = { workspace = true } -rustc-hash = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } vortex-fastlanes = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } -vortex-utils = { workspace = true } [dev-dependencies] divan = { workspace = true } diff --git a/encodings/alp/benches/alp_compress.rs b/encodings/alp/benches/alp_compress.rs index 6e769ffc22e..4af17372e22 100644 --- a/encodings/alp/benches/alp_compress.rs +++ b/encodings/alp/benches/alp_compress.rs @@ -12,6 +12,7 @@ use rand::rngs::StdRng; use vortex_alp::ALPFloat; use vortex_alp::ALPRDFloat; use vortex_alp::RDEncoder; +use vortex_alp::RDEncoderExt; use vortex_alp::alp_encode; use vortex_alp::decompress_into_array; use vortex_array::Canonical; @@ -116,12 +117,12 @@ fn decompress_alp(bencher: Bencher, args: (usize, f64 const RD_BENCH_ARGS: &[(usize, f64)] = &[ // length, fraction_patch + (2_000, 0.0), + (2_000, 0.01), + (2_000, 0.1), (10_000, 0.0), (10_000, 0.01), (10_000, 0.1), - (100_000, 0.0), - (100_000, 0.01), - (100_000, 0.1), ]; fn make_rd_array(n: usize, fraction_patch: f64) -> PrimitiveArray { diff --git a/encodings/alp/src/alp/array.rs b/encodings/alp/src/alp/array.rs index 5b0984dce7a..3e2010a1ad5 100644 --- a/encodings/alp/src/alp/array.rs +++ b/encodings/alp/src/alp/array.rs @@ -25,6 +25,7 @@ use vortex_array::array_slots; use vortex_array::arrays::Primitive; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; use vortex_array::dtype::PType; use vortex_array::patches::PatchSlotIndices; use vortex_array::patches::Patches; @@ -200,12 +201,16 @@ impl VTable for ALP { #[array_slots(ALP)] pub struct ALPSlots { /// The ALP-encoded values array. + #[slot(0)] pub encoded: ArrayRef, /// The indices of exception values that could not be ALP-encoded. + #[slot(1)] pub patch_indices: Option, /// The exception values that could not be ALP-encoded. + #[slot(2)] pub patch_values: Option, /// Chunk offsets for the patch indices/values. + #[slot(3)] pub patch_chunk_offsets: Option, } @@ -309,7 +314,10 @@ impl ALPData { } /// Validate that any patches provided are valid for the ALPArray. - fn validate_patches(patches: &Patches, encoded: &ArrayRef) -> VortexResult<()> { + fn validate_patches( + patches: &Patches, + encoded: &ArrayRef, + ) -> VortexResult<()> { vortex_ensure!( patches.array_len() == encoded.len(), "patches array_len != encoded len: {} != {}", diff --git a/encodings/alp/src/alp/compress.rs b/encodings/alp/src/alp/compress.rs index 819fc4b3dc4..fcb572c3159 100644 --- a/encodings/alp/src/alp/compress.rs +++ b/encodings/alp/src/alp/compress.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use ::alp::ENCODE_CHUNK_SIZE; use itertools::Itertools; use vortex_array::ArrayRef; use vortex_array::ArrayView; @@ -8,6 +9,7 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::NativePType; use vortex_array::dtype::PType; use vortex_array::patches::Patches; use vortex_array::validity::Validity; @@ -67,14 +69,33 @@ fn alp_encode_components_typed( ctx: &mut ExecutionCtx, ) -> VortexResult<(Exponents, ArrayRef, Option)> where - T: ALPFloat, + T: ALPFloat + NativePType, + T::ALPInt: NativePType, { let values_slice = values.as_slice::(); - let (exponents, encoded, exceptional_positions, exceptional_values, mut chunk_offsets) = - T::encode(values_slice, exponents); - - let encoded_array = PrimitiveArray::new(encoded, values.validity()?).into_array(); + // Encode straight into a Vortex buffer, so that the encoded array owns its values rather than + // adopting a `Vec` allocated by the encoder. + let mut encoded = BufferMut::::with_capacity(values_slice.len()); + + // Estimate capacity to be one patch per 32 values. + let mut exceptional_positions = Vec::with_capacity(values_slice.len() / 32); + let mut exceptional_values = Vec::with_capacity(values_slice.len() / 32); + + // There's exactly one offset per chunk. + let mut chunk_offsets = Vec::with_capacity(values_slice.len().div_ceil(ENCODE_CHUNK_SIZE)); + let exponents = ::alp::encode_into( + values_slice, + exponents, + &mut encoded.spare_capacity_mut()[..values_slice.len()], + &mut exceptional_positions, + &mut exceptional_values, + &mut chunk_offsets, + ); + // SAFETY: `encode_into` initializes exactly one element per value. + unsafe { encoded.set_len(values_slice.len()) }; + + let encoded_array = PrimitiveArray::new(encoded.freeze(), values.validity()?).into_array(); let validity = values .array() @@ -82,31 +103,28 @@ where .execute_mask(values.array().len(), ctx)?; // exceptional_positions may contain exceptions at invalid positions (which contain garbage // data). We remove null exceptions in order to keep the Patches small. - let (valid_exceptional_positions, valid_exceptional_values): (Buffer, Buffer) = - match validity { - Mask::AllTrue(_) => (exceptional_positions, exceptional_values), - Mask::AllFalse(_) => { - // no valid positions, ergo nothing worth patching - (Buffer::empty(), Buffer::empty()) - } - Mask::Values(is_valid) => { - let (pos, vals): (BufferMut, BufferMut) = exceptional_positions - .into_iter() - .zip_eq(exceptional_values) - .filter(|(index, _)| { - let is_valid = is_valid.value(*index as usize); - if !is_valid { - let patch_chunk = *index as usize / 1024; - for chunk_idx in (patch_chunk + 1)..chunk_offsets.len() { - chunk_offsets[chunk_idx] -= 1; - } - } - is_valid - }) - .unzip(); - (pos.freeze(), vals.freeze()) - } - }; + let (valid_exceptional_positions, valid_exceptional_values): (Vec, Vec) = match validity + { + Mask::AllTrue(_) => (exceptional_positions, exceptional_values), + Mask::AllFalse(_) => { + // no valid positions, ergo nothing worth patching + (Vec::new(), Vec::new()) + } + Mask::Values(is_valid) => exceptional_positions + .into_iter() + .zip_eq(exceptional_values) + .filter(|(index, _)| { + let is_valid = is_valid.value(*index as usize); + if !is_valid { + let patch_chunk = *index as usize / ENCODE_CHUNK_SIZE; + for chunk_idx in (patch_chunk + 1)..chunk_offsets.len() { + chunk_offsets[chunk_idx] -= 1; + } + } + is_valid + }) + .unzip(), + }; let patches = if valid_exceptional_positions.is_empty() { None } else { @@ -116,14 +134,15 @@ where Validity::NonNullable }; let valid_exceptional_values = - PrimitiveArray::new(valid_exceptional_values, patches_validity).into_array(); + PrimitiveArray::new(Buffer::from(valid_exceptional_values), patches_validity) + .into_array(); Some(Patches::new( values_slice.len(), 0, - valid_exceptional_positions.into_array(), + Buffer::from(valid_exceptional_positions).into_array(), valid_exceptional_values, - Some(chunk_offsets.into_array()), + Some(Buffer::from(chunk_offsets).into_array()), )?) }; Ok((exponents, encoded_array, patches)) @@ -309,7 +328,7 @@ mod tests { let decoded_val = decoded.as_slice::()[idx]; let original_val = original.as_slice::()[idx]; assert!( - decoded_val.is_eq(original_val), + NativePType::is_eq(decoded_val, original_val), "Expected {original_val} but got {decoded_val}" ); } @@ -645,7 +664,7 @@ mod tests { let decoded_val = decoded.as_slice::()[idx]; let original_val = values[idx]; assert!( - decoded_val.is_eq(original_val), + NativePType::is_eq(decoded_val, original_val), "At index {idx}: Expected {original_val} but got {decoded_val}" ); } diff --git a/encodings/alp/src/alp/compute/between.rs b/encodings/alp/src/alp/compute/between.rs index dbb6fec9a22..1a24a2cf328 100644 --- a/encodings/alp/src/alp/compute/between.rs +++ b/encodings/alp/src/alp/compute/between.rs @@ -12,6 +12,7 @@ use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::NativeDType; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; +use vortex_array::scalar::PValue; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::between::BetweenOptions; use vortex_array::scalar_fn::fns::between::BetweenReduce; @@ -64,7 +65,7 @@ fn between_impl( ) -> VortexResult where Scalar: From, - ::ALPInt: NativeDType + Debug, + ::ALPInt: NativeDType + NativePType + Into + Debug, { let exponents = array.exponents(); @@ -99,7 +100,7 @@ where ) } -fn encode_lower_bound( +fn encode_lower_bound( lower: T, exponents: Exponents, strict: StrictComparison, @@ -116,7 +117,7 @@ fn encode_lower_bound( ) } -fn encode_upper_bound( +fn encode_upper_bound( upper: T, exponents: Exponents, strict: StrictComparison, diff --git a/encodings/alp/src/alp/compute/compare.rs b/encodings/alp/src/alp/compute/compare.rs index 83e38e19bda..8e59fb23f7c 100644 --- a/encodings/alp/src/alp/compute/compare.rs +++ b/encodings/alp/src/alp/compute/compare.rs @@ -70,7 +70,7 @@ impl CompareKernel for ALP { /// We can compare a scalar to an ALPArray by encoding the scalar into the ALP domain and comparing /// the encoded value to the encoded values in the ALPArray. There are fixups when the value doesn't /// encode into the ALP domain. -fn alp_scalar_compare>( +fn alp_scalar_compare>( alp: ArrayView, value: F, operator: CompareOperator, @@ -178,7 +178,7 @@ mod tests { session }); - fn test_alp_compare>( + fn test_alp_compare>( alp: ArrayView, value: F, operator: CompareOperator, diff --git a/encodings/alp/src/alp/mod.rs b/encodings/alp/src/alp/mod.rs index 9992a9ddbe8..04f463c6493 100644 --- a/encodings/alp/src/alp/mod.rs +++ b/encodings/alp/src/alp/mod.rs @@ -1,17 +1,14 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::fmt::Display; -use std::fmt::Formatter; -use std::mem::size_of; -use std::mem::transmute; -use std::mem::transmute_copy; +//! Vortex array support for classic ALP, on top of the `alp` crate's encoder. -use itertools::Itertools; -use num_traits::CheckedSub; -use num_traits::Float; -use num_traits::PrimInt; -use num_traits::ToPrimitive; +pub use ::alp::ALPFloat; +pub use ::alp::Exponents; +pub use array::*; +pub use compress::alp_encode; +pub use decompress::decompress_into_array; +use vortex_session::VortexSession; mod array; mod compress; @@ -23,6 +20,10 @@ mod rules; pub(crate) use plugin::ALPPatchedPlugin; +pub(crate) fn initialize(session: &VortexSession) { + rules::initialize(session); +} + #[cfg(test)] mod tests { use prost::Message; @@ -30,49 +31,8 @@ mod tests { use vortex_array::patches::PatchesMetadata; use vortex_array::test_harness::check_metadata; - use crate::alp::ALPFloat; - use crate::alp::Exponents; use crate::alp::array::ALPMetadata; - // The allocation-free estimate must match a full encode + estimate for every candidate - // exponent pair, so `find_best_exponents` picks the same exponents and compression is unchanged. - fn check_estimate_matches(values: &[T]) { - for e in 0..T::MAX_EXPONENT { - for f in 0..e { - let exp = Exponents { e, f }; - let lightweight = T::estimate_encoded_size_for_exponents(values, exp); - let (_, encoded, _, patches, _) = T::encode(values, Some(exp)); - let full = T::estimate_encoded_size(&encoded, &patches); - assert_eq!( - lightweight, - full, - "mismatch at e={e}, f={f}, len={}", - values.len() - ); - } - } - } - - #[test] - fn estimate_for_exponents_matches_full_encode() { - // Clean 2-decimal values (mostly kept), repeating decimals (many patches), large - // magnitudes, constants, and a single element. - let mut f64s: Vec = (0..200).map(|i| i as f64 / 100.0).collect(); - f64s.extend((0..60).map(|i| i as f64 / 7.0)); - f64s.extend([1e17, -1e17, 0.0, 123.0]); - check_estimate_matches(&f64s); - check_estimate_matches::(&[123.456; 5]); - check_estimate_matches::(&[42.0]); - // Every value patches at every exponent -> exercises the all-patched branch. - check_estimate_matches::(&[1.0 / 3.0; 8]); - - let mut f32s: Vec = (0..200).map(|i| i as f32 / 100.0).collect(); - f32s.extend((0..60).map(|i| i as f32 / 7.0)); - f32s.extend([1e9, -1e9, 0.0, 123.0]); - check_estimate_matches(&f32s); - check_estimate_matches::(&[1.0 / 3.0; 8]); - } - #[cfg_attr(miri, ignore)] #[test] fn test_alp_metadata() { @@ -94,478 +54,3 @@ mod tests { ); } } - -pub use array::*; -pub use compress::alp_encode; -pub use decompress::decompress_into_array; -use vortex_array::dtype::NativePType; -use vortex_array::scalar::PValue; -use vortex_buffer::Buffer; -use vortex_buffer::BufferMut; -use vortex_session::VortexSession; - -const SAMPLE_SIZE: usize = 32; - -pub(crate) fn initialize(session: &VortexSession) { - rules::initialize(session); -} - -#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)] -pub struct Exponents { - pub e: u8, - pub f: u8, -} - -impl Display for Exponents { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "e: {}, f: {}", self.e, self.f) - } -} - -mod private { - pub trait Sealed {} - - impl Sealed for f32 {} - impl Sealed for f64 {} -} - -/// Widen a running `(min, max)` bound to include `value`, seeding it on the first value. -fn update_bounds(bounds: &mut Option<(I, I)>, value: I) { - *bounds = Some(bounds.map_or((value, value), |(min, max)| { - (min.min(value), max.max(value)) - })); -} - -pub trait ALPFloat: private::Sealed + Float + Display + NativePType { - type ALPInt: PrimInt + Display + ToPrimitive + Copy + NativePType + Into; - - const FRACTIONAL_BITS: u8; - const MAX_EXPONENT: u8; - const SWEET: Self; - const F10: &'static [Self]; - const IF10: &'static [Self]; - - /// Round to the nearest floating integer by shifting in and out of the low precision range. - #[inline] - fn fast_round(self) -> Self { - (self + Self::SWEET) - Self::SWEET - } - - /// Equivalent to calling `as` to cast the primitive float to the target integer type. - fn as_int(self) -> Self::ALPInt; - - /// Convert from the integer type back to the float type using `as`. - fn from_int(n: Self::ALPInt) -> Self; - - fn find_best_exponents(values: &[Self]) -> Exponents { - let mut best_exp = Exponents { e: 0, f: 0 }; - let mut best_nbytes: usize = usize::MAX; - - let sample = (values.len() > SAMPLE_SIZE).then(|| { - values - .iter() - .step_by(values.len() / SAMPLE_SIZE) - .cloned() - .collect_vec() - }); - let sample = sample.as_deref().unwrap_or(values); - - for e in (0..Self::MAX_EXPONENT).rev() { - for f in 0..e { - let exp = Exponents { e, f }; - let size = Self::estimate_encoded_size_for_exponents(sample, exp); - if size < best_nbytes { - best_nbytes = size; - best_exp = exp; - } else if size == best_nbytes && e - f < best_exp.e - best_exp.f { - best_exp = exp; - } - } - } - - best_exp - } - - /// Size estimate for `values` under `exponents` matching a full [`Self::encode`] plus - /// [`Self::estimate_encoded_size`], but without the per-candidate allocations. - fn estimate_encoded_size_for_exponents(values: &[Self], exponents: Exponents) -> usize { - // `kept` is the (min, max) over values that round-trip exactly (kept inline by `encode`); - // `all` is the (min, max) over every encoded value. `encode` fills patched slots in-range, - // so its emitted range is `kept`, except with all values patched (no fill) where `all` wins. - let mut kept: Option<(Self::ALPInt, Self::ALPInt)> = None; - let mut all: Option<(Self::ALPInt, Self::ALPInt)> = None; - let mut patch_count = 0usize; - - for &value in values { - let encoded = Self::encode_single_unchecked(value, exponents); - update_bounds(&mut all, encoded); - if Self::decode_single(encoded, exponents).is_eq(value) { - update_bounds(&mut kept, encoded); - } else { - patch_count += 1; - } - } - - let range = if patch_count == values.len() { - all - } else { - kept - }; - - let bits_per_encoded = range - .and_then(|(min, max)| max.checked_sub(&min)) - .and_then(|range_size| range_size.to_u64()) - .and_then(|range_size| { - range_size - .checked_ilog2() - .map(|bits| (bits + 1) as usize) - .or(Some(0)) - }) - .unwrap_or(size_of::() * 8); - - let encoded_bytes = (values.len() * bits_per_encoded).div_ceil(8); - // each patch is a value + a position - // in practice, patch positions are in [0, u16::MAX] because of how we chunk - let patch_bytes = patch_count * (size_of::() + size_of::()); - - encoded_bytes + patch_bytes - } - - #[inline] - fn estimate_encoded_size(encoded: &[Self::ALPInt], patches: &[Self]) -> usize { - let bits_per_encoded = encoded - .iter() - .minmax() - .into_option() - // estimating bits per encoded value assuming frame-of-reference + bitpacking-without-patches - .and_then(|(min, max)| max.checked_sub(min)) - .and_then(|range_size: ::ALPInt| range_size.to_u64()) - .and_then(|range_size| { - range_size - .checked_ilog2() - .map(|bits| (bits + 1) as usize) - .or(Some(0)) - }) - .unwrap_or(size_of::() * 8); - - let encoded_bytes = (encoded.len() * bits_per_encoded).div_ceil(8); - // each patch is a value + a position - // in practice, patch positions are in [0, u16::MAX] because of how we chunk - let patch_bytes = patches.len() * (size_of::() + size_of::()); - - encoded_bytes + patch_bytes - } - - #[expect( - clippy::type_complexity, - reason = "tuple return type is appropriate for multiple encoding outputs" - )] - fn encode( - values: &[Self], - exponents: Option, - ) -> ( - Exponents, - Buffer, - Buffer, - Buffer, - BufferMut, - ) { - let exp = exponents.unwrap_or_else(|| Self::find_best_exponents(values)); - - let mut encoded_output = BufferMut::::with_capacity(values.len()); - - // Estimate capacity to be one patch per 32 values. - let mut patch_indices = BufferMut::::with_capacity(values.len() / 32); - let mut patch_values = BufferMut::::with_capacity(values.len() / 32); - - // There's exactly one offset per 1024 chunk. - let mut chunk_offsets = BufferMut::::with_capacity(values.len().div_ceil(1024)); - let mut fill_value: Option = None; - - for chunk in values.chunks(1024) { - chunk_offsets.push(patch_indices.len() as u64); - encode_chunk_unchecked( - chunk, - exp, - &mut encoded_output, - &mut patch_indices, - &mut patch_values, - &mut fill_value, - ); - } - - ( - exp, - encoded_output.freeze(), - patch_indices.freeze(), - patch_values.freeze(), - chunk_offsets, - ) - } - - #[inline] - fn encode_single(value: Self, exponents: Exponents) -> Option { - let encoded = Self::encode_single_unchecked(value, exponents); - let decoded = Self::decode_single(encoded, exponents); - if decoded.is_eq(value) { - return Some(encoded); - } - None - } - - fn encode_above(value: Self, exponents: Exponents) -> Self::ALPInt { - (value * Self::F10[exponents.e as usize] * Self::IF10[exponents.f as usize]) - .ceil() - .as_int() - } - - fn encode_below(value: Self, exponents: Exponents) -> Self::ALPInt { - (value * Self::F10[exponents.e as usize] * Self::IF10[exponents.f as usize]) - .floor() - .as_int() - } - - fn decode(encoded: &[Self::ALPInt], exponents: Exponents) -> Vec { - let mut values = Vec::with_capacity(encoded.len()); - for encoded in encoded { - values.push(Self::decode_single(*encoded, exponents)); - } - values - } - - fn decode_buffer(encoded: BufferMut, exponents: Exponents) -> BufferMut { - encoded.map_each_in_place(move |encoded| Self::decode_single(encoded, exponents)) - } - - fn decode_into(encoded: &[Self::ALPInt], exponents: Exponents, output: &mut [Self]) { - assert_eq!(encoded.len(), output.len()); - - for i in 0..encoded.len() { - output[i] = Self::decode_single(encoded[i], exponents) - } - } - - fn decode_slice_inplace(encoded: &mut [Self::ALPInt], exponents: Exponents) { - let decoded: &mut [Self] = unsafe { transmute(encoded) }; - decoded.iter_mut().for_each(|v| { - *v = Self::decode_single( - unsafe { transmute_copy::(v) }, - exponents, - ) - }) - } - - #[inline(always)] - fn decode_single(encoded: Self::ALPInt, exponents: Exponents) -> Self { - Self::from_int(encoded) * Self::F10[exponents.f as usize] * Self::IF10[exponents.e as usize] - } - - /// Encode single float value. The returned value might decode to a different value than passed in. - /// Consider using [`Self::encode_single`] if you want the checked version of this function. - #[inline(always)] - fn encode_single_unchecked(value: Self, exponents: Exponents) -> Self::ALPInt { - (value * Self::F10[exponents.e as usize] * Self::IF10[exponents.f as usize]) - .fast_round() - .as_int() - } -} - -#[expect( - clippy::cast_possible_truncation, - reason = "intentional truncation for ALP encoding" -)] -fn encode_chunk_unchecked( - chunk: &[T], - exp: Exponents, - encoded_output: &mut BufferMut, - patch_indices: &mut BufferMut, - patch_values: &mut BufferMut, - fill_value: &mut Option, -) { - let num_prev_encoded = encoded_output.len(); - let num_prev_patches = patch_indices.len(); - assert_eq!(patch_indices.len(), patch_values.len()); - let has_filled = fill_value.is_some(); - - // encode the chunk, counting the number of patches - let mut chunk_patch_count = 0; - encoded_output.extend_trusted(chunk.iter().map(|&v| { - let encoded = T::encode_single_unchecked(v, exp); - let decoded = T::decode_single(encoded, exp); - let neq = !decoded.is_eq(v) as usize; - chunk_patch_count += neq; - encoded - })); - let chunk_patch_count = chunk_patch_count; // immutable hereafter - assert_eq!(encoded_output.len(), num_prev_encoded + chunk.len()); - - if chunk_patch_count > 0 { - // we need to gather the patches for this chunk - // preallocate space for the patches (plus one because our loop may attempt to write one past the end) - patch_indices.reserve(chunk_patch_count + 1); - patch_values.reserve(chunk_patch_count + 1); - - // record the patches in this chunk - let patch_indices_mut = patch_indices.spare_capacity_mut(); - let patch_values_mut = patch_values.spare_capacity_mut(); - let mut chunk_patch_index = 0; - for i in num_prev_encoded..encoded_output.len() { - let decoded = T::decode_single(encoded_output[i], exp); - // write() is only safe to call more than once because the values are primitive (i.e., Drop is a no-op) - patch_indices_mut[chunk_patch_index].write(i as u64); - patch_values_mut[chunk_patch_index].write(chunk[i - num_prev_encoded]); - chunk_patch_index += !decoded.is_eq(chunk[i - num_prev_encoded]) as usize; - } - assert_eq!(chunk_patch_index, chunk_patch_count); - unsafe { - patch_indices.set_len(num_prev_patches + chunk_patch_count); - patch_values.set_len(num_prev_patches + chunk_patch_count); - } - } - - // find the first successfully encoded value (i.e., not patched) - // this is our fill value for missing values - if fill_value.is_none() && (num_prev_encoded + chunk_patch_count < encoded_output.len()) { - assert_eq!(num_prev_encoded, num_prev_patches); - for i in num_prev_encoded..encoded_output.len() { - if i >= patch_indices.len() || patch_indices[i] != i as u64 { - *fill_value = Some(encoded_output[i]); - break; - } - } - } - - // replace the patched values in the encoded array with the fill value - // for better downstream compression - if let Some(fill_value) = fill_value { - // handle the edge case where the first N >= 1 chunks are all patches - let start_patch = if !has_filled { 0 } else { num_prev_patches }; - for patch_idx in &patch_indices[start_patch..] { - encoded_output[*patch_idx as usize] = *fill_value; - } - } -} - -impl ALPFloat for f32 { - type ALPInt = i32; - const FRACTIONAL_BITS: u8 = 23; - const MAX_EXPONENT: u8 = 10; - const SWEET: Self = - (1 << Self::FRACTIONAL_BITS) as Self + (1 << (Self::FRACTIONAL_BITS - 1)) as Self; - - const F10: &'static [Self] = &[ - 1.0, - 10.0, - 100.0, - 1000.0, - 10000.0, - 100000.0, - 1000000.0, - 10000000.0, - 100000000.0, - 1000000000.0, - 10000000000.0, // 10^10 - ]; - const IF10: &'static [Self] = &[ - 1.0, - 0.1, - 0.01, - 0.001, - 0.0001, - 0.00001, - 0.000001, - 0.0000001, - 0.00000001, - 0.000000001, - 0.0000000001, // 10^-10 - ]; - - #[inline(always)] - #[expect( - clippy::cast_possible_truncation, - reason = "intentional float to int truncation for ALP encoding" - )] - fn as_int(self) -> Self::ALPInt { - self as _ - } - - #[inline(always)] - fn from_int(n: Self::ALPInt) -> Self { - n as _ - } -} - -impl ALPFloat for f64 { - type ALPInt = i64; - const FRACTIONAL_BITS: u8 = 52; - const MAX_EXPONENT: u8 = 18; // 10^18 is the maximum i64 - const SWEET: Self = - (1u64 << Self::FRACTIONAL_BITS) as Self + (1u64 << (Self::FRACTIONAL_BITS - 1)) as Self; - const F10: &'static [Self] = &[ - 1.0, - 10.0, - 100.0, - 1000.0, - 10000.0, - 100000.0, - 1000000.0, - 10000000.0, - 100000000.0, - 1000000000.0, - 10000000000.0, - 100000000000.0, - 1000000000000.0, - 10000000000000.0, - 100000000000000.0, - 1000000000000000.0, - 10000000000000000.0, - 100000000000000000.0, - 1000000000000000000.0, - 10000000000000000000.0, - 100000000000000000000.0, - 1000000000000000000000.0, - 10000000000000000000000.0, - 100000000000000000000000.0, // 10^23 - ]; - - const IF10: &'static [Self] = &[ - 1.0, - 0.1, - 0.01, - 0.001, - 0.0001, - 0.00001, - 0.000001, - 0.0000001, - 0.00000001, - 0.000000001, - 0.0000000001, - 0.00000000001, - 0.000000000001, - 0.0000000000001, - 0.00000000000001, - 0.000000000000001, - 0.0000000000000001, - 0.00000000000000001, - 0.000000000000000001, - 0.0000000000000000001, - 0.00000000000000000001, - 0.000000000000000000001, - 0.0000000000000000000001, - 0.00000000000000000000001, // 10^-23 - ]; - - #[inline(always)] - #[expect( - clippy::cast_possible_truncation, - reason = "intentional float to int truncation for ALP encoding" - )] - fn as_int(self) -> Self::ALPInt { - self as _ - } - - #[inline(always)] - fn from_int(n: Self::ALPInt) -> Self { - n as _ - } -} diff --git a/encodings/alp/src/alp/plugin.rs b/encodings/alp/src/alp/plugin.rs index c14133109d1..5756dc5a991 100644 --- a/encodings/alp/src/alp/plugin.rs +++ b/encodings/alp/src/alp/plugin.rs @@ -7,17 +7,17 @@ //! This enables zero-cost backward compatibility with previously written datasets. use vortex_array::Array; +use vortex_array::ArrayDeserialization; use vortex_array::ArrayId; use vortex_array::ArrayPlugin; use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; use vortex_array::ArrayVTable; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::Patched; -use vortex_array::buffer::BufferHandle; -use vortex_array::dtype::DType; -use vortex_array::serde::ArrayChildren; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_session::VortexSession; @@ -42,22 +42,29 @@ impl ArrayPlugin for ALPPatchedPlugin { &self, array: &ArrayRef, session: &VortexSession, - ) -> VortexResult>> { + ) -> VortexResult> { // Delegate to ALP's metadata serde - ALP.serialize(array, session) + ArrayPlugin::serialize(&ALP, array, session) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, session: &VortexSession, ) -> VortexResult { + vortex_ensure!( + parts.serialized_id == self.id(), + "ALP plugin does not recognize serialized ID {}", + parts.serialized_id, + ); let alp_array = Array::::try_from_parts(ArrayVTable::deserialize( - &ALP, dtype, len, metadata, buffers, children, session, + &ALP, + parts.dtype, + parts.len, + parts.metadata, + parts.buffers, + parts.children, + session, )?) .map_err(|_| vortex_err!("ALP plugin should only deserialize vortex.alp"))?; @@ -91,6 +98,7 @@ mod tests { use std::f64::consts::PI; use std::sync::LazyLock; + use vortex_array::ArrayDeserialization; use vortex_array::ArrayPlugin; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; @@ -133,7 +141,7 @@ mod tests { let array = alp_encoded.as_array(); - let metadata = SESSION.array_serialize(array)?.unwrap(); + let serialization = SESSION.array_serialize(array)?.unwrap(); let children = array.children(); let buffers = array .buffers() @@ -142,11 +150,14 @@ mod tests { .collect::>(); let deserialized = ALPPatchedPlugin.deserialize( - array.dtype(), - array.len(), - &metadata, - &buffers, - &children, + ArrayDeserialization::new( + ALPPatchedPlugin.id(), + array.dtype(), + array.len(), + &serialization.metadata, + &buffers, + &children, + ), &SESSION, )?; @@ -182,7 +193,7 @@ mod tests { let array = alp_encoded.as_array(); - let metadata = SESSION.array_serialize(array)?.unwrap(); + let serialization = SESSION.array_serialize(array)?.unwrap(); let children = array.children(); let buffers = array .buffers() @@ -191,11 +202,14 @@ mod tests { .collect::>(); let deserialized = ALPPatchedPlugin.deserialize( - array.dtype(), - array.len(), - &metadata, - &buffers, - &children, + ArrayDeserialization::new( + ALPPatchedPlugin.id(), + array.dtype(), + array.len(), + &serialization.metadata, + &buffers, + &children, + ), &SESSION, )?; @@ -213,7 +227,7 @@ mod tests { fn primitive_array_returns_error() { let array = PrimitiveArray::from_iter([1.0f64, 2.0, 3.0]).into_array(); - let metadata = SESSION.array_serialize(&array).unwrap().unwrap(); + let serialization = SESSION.array_serialize(&array).unwrap().unwrap(); let children = array.children(); let buffers = array .buffers() @@ -223,11 +237,14 @@ mod tests { // This panics because PrimitiveArray has no children and ALP requires encoded child. let _result = ALPPatchedPlugin.deserialize( - array.dtype(), - array.len(), - &metadata, - &buffers, - &children, + ArrayDeserialization::new( + ALPPatchedPlugin.id(), + array.dtype(), + array.len(), + &serialization.metadata, + &buffers, + &children, + ), &SESSION, ); } diff --git a/encodings/alp/src/alp_rd/array.rs b/encodings/alp/src/alp_rd/array.rs index dcaeb0fde79..b6ee50d7b1f 100644 --- a/encodings/alp/src/alp_rd/array.rs +++ b/encodings/alp/src/alp_rd/array.rs @@ -22,6 +22,7 @@ use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::IntoArray; use vortex_array::TypedArrayRef; +use vortex_array::array_slots; use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::buffer::BufferHandle; @@ -105,11 +106,12 @@ impl VTable for ALPRD { len: usize, slots: &[Option], ) -> VortexResult<()> { + let alprd_slots = ALPRDSlotsView::from_slots(slots); validate_parts( dtype, len, - left_parts_from_slots(slots), - right_parts_from_slots(slots), + alprd_slots.left_parts, + alprd_slots.right_parts, patches_from_slots(slots, data.patches_data.as_ref(), len).as_ref(), ) } @@ -230,17 +232,18 @@ impl VTable for ALPRD { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + ALPRDSlots::NAMES[idx].to_string() } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - let array = require_child!(array, array.left_parts(), 0 => Primitive); - let array = require_child!(array, array.right_parts(), 1 => Primitive); + let array = require_child!(array, array.left_parts(), ALPRDSlots::LEFT_PARTS => Primitive); + let array = + require_child!(array, array.right_parts(), ALPRDSlots::RIGHT_PARTS => Primitive); require_patches!( array, - LP_PATCH_INDICES_SLOT, - LP_PATCH_VALUES_SLOT, - LP_PATCH_CHUNK_OFFSETS_SLOT + ALPRDSlots::PATCH_INDICES, + ALPRDSlots::PATCH_VALUES, + ALPRDSlots::PATCH_CHUNK_OFFSETS ); let dtype = array.dtype().clone(); @@ -307,29 +310,30 @@ impl VTable for ALPRD { } } -/// The left (most significant) parts of the real-double encoded values. -pub(super) const LEFT_PARTS_SLOT: usize = 0; -/// The right (least significant) parts of the real-double encoded values. -pub(super) const RIGHT_PARTS_SLOT: usize = 1; -/// The indices of left-parts exception values that could not be dictionary-encoded. -pub(super) const LP_PATCH_SLOTS: PatchSlotIndices = PatchSlotIndices { - indices: LP_PATCH_INDICES_SLOT, - values: LP_PATCH_VALUES_SLOT, - chunk_offsets: LP_PATCH_CHUNK_OFFSETS_SLOT, +#[array_slots(ALPRD)] +pub struct ALPRDSlots { + /// The left (most significant) parts of the real-double encoded values. + #[slot(0)] + pub left_parts: ArrayRef, + /// The right (least significant) parts of the real-double encoded values. + #[slot(1)] + pub right_parts: ArrayRef, + /// The indices of left-parts exception values that could not be dictionary-encoded. + #[slot(2)] + pub patch_indices: Option, + /// The exception values for left-parts that could not be dictionary-encoded. + #[slot(3)] + pub patch_values: Option, + /// Chunk offsets for the left-parts patch indices/values. + #[slot(4)] + pub patch_chunk_offsets: Option, +} + +const LP_PATCH_SLOTS: PatchSlotIndices = PatchSlotIndices { + indices: ALPRDSlots::PATCH_INDICES, + values: ALPRDSlots::PATCH_VALUES, + chunk_offsets: ALPRDSlots::PATCH_CHUNK_OFFSETS, }; -pub(super) const LP_PATCH_INDICES_SLOT: usize = 2; -/// The exception values for left-parts that could not be dictionary-encoded. -pub(super) const LP_PATCH_VALUES_SLOT: usize = 3; -/// Chunk offsets for the left-parts patch indices/values. -pub(super) const LP_PATCH_CHUNK_OFFSETS_SLOT: usize = 4; -pub(super) const NUM_SLOTS: usize = 5; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = [ - "left_parts", - "right_parts", - "patch_indices", - "patch_values", - "patch_chunk_offsets", -]; #[derive(Clone, Debug)] pub struct ALPRDData { @@ -451,18 +455,6 @@ impl ALPRDData { } } -fn left_parts_from_slots(slots: &[Option]) -> &ArrayRef { - slots[LEFT_PARTS_SLOT] - .as_ref() - .vortex_expect("ALPRDArray left_parts slot") -} - -fn right_parts_from_slots(slots: &[Option]) -> &ArrayRef { - slots[RIGHT_PARTS_SLOT] - .as_ref() - .vortex_expect("ALPRDArray right_parts slot") -} - fn patches_from_slots( slots: &[Option], patches_data: Option<&PatchesData>, @@ -539,15 +531,7 @@ fn validate_parts( Ok(()) } -pub trait ALPRDArrayExt: TypedArrayRef { - fn left_parts(&self) -> &ArrayRef { - left_parts_from_slots(self.as_ref().slots()) - } - - fn right_parts(&self) -> &ArrayRef { - right_parts_from_slots(self.as_ref().slots()) - } - +pub trait ALPRDArrayExt: ALPRDArraySlotsExt { fn right_bit_width(&self) -> u8 { ALPRDData::right_bit_width(self) } @@ -601,6 +585,7 @@ mod test { use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; + use vortex_array::dtype::NativePType; use vortex_array::dtype::PType; use vortex_array::patches::PatchesMetadata; use vortex_array::test_harness::check_metadata; @@ -608,6 +593,7 @@ mod test { use super::ALPRDMetadata; use crate::ALPRDFloat; + use crate::RDEncoderExt; use crate::alp_rd; static SESSION: LazyLock = LazyLock::new(|| { @@ -619,7 +605,7 @@ mod test { #[rstest] #[case(vec![0.1f32.next_up(); 1024], 1.123_848_f32)] #[case(vec![0.1f64.next_up(); 1024], 1.123_848_591_110_992_f64)] - fn test_array_encode_with_nulls_and_patches( + fn test_array_encode_with_nulls_and_patches( #[case] reals: Vec, #[case] seed: T, ) { diff --git a/encodings/alp/src/alp_rd/compute/cast.rs b/encodings/alp/src/alp_rd/compute/cast.rs index 9a08f00e298..7129c27194f 100644 --- a/encodings/alp/src/alp_rd/compute/cast.rs +++ b/encodings/alp/src/alp_rd/compute/cast.rs @@ -10,6 +10,7 @@ use vortex_array::scalar_fn::fns::cast::CastReduce; use vortex_error::VortexResult; use crate::ALPRDArrayExt; +use crate::ALPRDArraySlotsExt; use crate::alp_rd::ALPRD; impl CastReduce for ALPRD { @@ -58,6 +59,7 @@ mod tests { use vortex_array::dtype::PType; use crate::RDEncoder; + use crate::RDEncoderExt; #[test] fn test_cast_alprd_f32_to_f64() { diff --git a/encodings/alp/src/alp_rd/compute/filter.rs b/encodings/alp/src/alp_rd/compute/filter.rs index 69c26c53e3b..62f4985546d 100644 --- a/encodings/alp/src/alp_rd/compute/filter.rs +++ b/encodings/alp/src/alp_rd/compute/filter.rs @@ -11,6 +11,7 @@ use vortex_mask::Mask; use crate::ALPRD; use crate::ALPRDArrayExt; +use crate::ALPRDArraySlotsExt; impl FilterKernel for ALPRD { fn filter( @@ -49,6 +50,7 @@ mod test { use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::compute::conformance::filter::test_filter_conformance; + use vortex_array::dtype::NativePType; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_mask::Mask; @@ -57,6 +59,7 @@ mod test { use crate::ALPRDArrayExt; use crate::ALPRDFloat; use crate::RDEncoder; + use crate::RDEncoderExt; static SESSION: LazyLock = LazyLock::new(|| { let session = array_session(); @@ -67,7 +70,7 @@ mod test { #[rstest] #[case(0.1f32, 0.2f32, 3e25f32)] #[case(0.1f64, 0.2f64, 3e100f64)] - fn test_filter(#[case] a: T, #[case] b: T, #[case] outlier: T) { + fn test_filter(#[case] a: T, #[case] b: T, #[case] outlier: T) { let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::new(buffer![a, b, outlier], Validity::NonNullable); let encoded = RDEncoder::new(&[a, b]).encode(array.as_view()); @@ -85,7 +88,11 @@ mod test { #[rstest] #[case(0.1f32, 0.2f32, 3e25f32)] #[case(0.1f64, 0.2f64, 3e100f64)] - fn test_filter_simple(#[case] a: T, #[case] b: T, #[case] outlier: T) { + fn test_filter_simple( + #[case] a: T, + #[case] b: T, + #[case] outlier: T, + ) { let mut ctx = SESSION.create_execution_ctx(); test_filter_conformance( &RDEncoder::new(&[a, b]) @@ -98,7 +105,7 @@ mod test { #[rstest] #[case(0.1f32, 3e25f32)] #[case(0.5f64, 1e100f64)] - fn test_filter_with_nulls(#[case] a: T, #[case] outlier: T) { + fn test_filter_with_nulls(#[case] a: T, #[case] outlier: T) { let mut ctx = SESSION.create_execution_ctx(); test_filter_conformance( &RDEncoder::new(&[a]) diff --git a/encodings/alp/src/alp_rd/compute/mask.rs b/encodings/alp/src/alp_rd/compute/mask.rs index 25bdc8f9ab4..2fc55aff988 100644 --- a/encodings/alp/src/alp_rd/compute/mask.rs +++ b/encodings/alp/src/alp_rd/compute/mask.rs @@ -4,23 +4,19 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::IntoArray; -use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; -use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::fns::mask::Mask as MaskExpr; use vortex_array::scalar_fn::fns::mask::MaskReduce; use vortex_error::VortexResult; use crate::ALPRD; use crate::ALPRDArrayExt; +use crate::ALPRDArraySlotsExt; impl MaskReduce for ALPRD { #[allow(clippy::disallowed_methods)] fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult> { - let masked_left_parts = MaskExpr.try_new_array( - array.left_parts().len(), - EmptyOptions, - [array.left_parts().clone(), mask.clone()], - )?; + let masked_left_parts = + MaskExpr::try_new(array.left_parts().clone(), mask.clone())?.into_array(); Ok(Some( ALPRD::try_new( array.dtype().as_nullable(), @@ -45,10 +41,12 @@ mod tests { use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::mask::test_mask_conformance; + use vortex_array::dtype::NativePType; use vortex_session::VortexSession; use crate::ALPRDFloat; use crate::RDEncoder; + use crate::RDEncoderExt; static SESSION: LazyLock = LazyLock::new(|| { let session = array_session(); @@ -59,7 +57,11 @@ mod tests { #[rstest] #[case(0.1f32, 0.2f32, 3e25f32)] #[case(0.1f64, 0.2f64, 3e100f64)] - fn test_mask_simple(#[case] a: T, #[case] b: T, #[case] outlier: T) { + fn test_mask_simple( + #[case] a: T, + #[case] b: T, + #[case] outlier: T, + ) { let mut ctx = SESSION.create_execution_ctx(); test_mask_conformance( &RDEncoder::new(&[a, b]) @@ -72,7 +74,7 @@ mod tests { #[rstest] #[case(0.1f32, 3e25f32)] #[case(0.5f64, 1e100f64)] - fn test_mask_with_nulls(#[case] a: T, #[case] outlier: T) { + fn test_mask_with_nulls(#[case] a: T, #[case] outlier: T) { let mut ctx = SESSION.create_execution_ctx(); test_mask_conformance( &RDEncoder::new(&[a]) diff --git a/encodings/alp/src/alp_rd/compute/mod.rs b/encodings/alp/src/alp_rd/compute/mod.rs index 696e622b873..c89ae057087 100644 --- a/encodings/alp/src/alp_rd/compute/mod.rs +++ b/encodings/alp/src/alp_rd/compute/mod.rs @@ -21,6 +21,7 @@ mod tests { use crate::ALPRDArray; use crate::RDEncoder; + use crate::RDEncoderExt; static SESSION: LazyLock = LazyLock::new(|| { let session = array_session(); diff --git a/encodings/alp/src/alp_rd/compute/take.rs b/encodings/alp/src/alp_rd/compute/take.rs index 18d5e2b5ad4..b02e189a79e 100644 --- a/encodings/alp/src/alp_rd/compute/take.rs +++ b/encodings/alp/src/alp_rd/compute/take.rs @@ -12,6 +12,7 @@ use vortex_error::VortexResult; use crate::ALPRD; use crate::ALPRDArrayExt; +use crate::ALPRDArraySlotsExt; impl TakeExecute for ALPRD { fn take( @@ -67,11 +68,13 @@ mod test { use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::compute::conformance::take::test_take_conformance; + use vortex_array::dtype::NativePType; use vortex_session::VortexSession; use crate::ALPRDArrayExt; use crate::ALPRDFloat; use crate::RDEncoder; + use crate::RDEncoderExt; static SESSION: LazyLock = LazyLock::new(|| { let session = array_session(); @@ -82,7 +85,7 @@ mod test { #[rstest] #[case(0.1f32, 0.2f32, 3e25f32)] #[case(0.1f64, 0.2f64, 3e100f64)] - fn test_take(#[case] a: T, #[case] b: T, #[case] outlier: T) { + fn test_take(#[case] a: T, #[case] b: T, #[case] outlier: T) { use vortex_array::IntoArray as _; use vortex_buffer::buffer; @@ -111,7 +114,11 @@ mod test { #[rstest] #[case(0.1f32, 0.2f32, 3e25f32)] #[case(0.1f64, 0.2f64, 3e100f64)] - fn take_with_nulls(#[case] a: T, #[case] b: T, #[case] outlier: T) { + fn take_with_nulls( + #[case] a: T, + #[case] b: T, + #[case] outlier: T, + ) { let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::from_iter([a, b, outlier]); let encoded = RDEncoder::new(&[a, b]).encode(array.as_view()); @@ -141,7 +148,11 @@ mod test { #[rstest] #[case(0.1f32, 0.2f32, 3e25f32)] #[case(0.1f64, 0.2f64, 3e100f64)] - fn test_take_conformance_alprd(#[case] a: T, #[case] b: T, #[case] outlier: T) { + fn test_take_conformance_alprd( + #[case] a: T, + #[case] b: T, + #[case] outlier: T, + ) { let mut ctx = SESSION.create_execution_ctx(); test_take_conformance( &RDEncoder::new(&[a, b]) @@ -154,7 +165,10 @@ mod test { #[rstest] #[case(0.1f32, 3e25f32)] #[case(0.5f64, 1e100f64)] - fn test_take_with_nulls_conformance(#[case] a: T, #[case] outlier: T) { + fn test_take_with_nulls_conformance( + #[case] a: T, + #[case] outlier: T, + ) { let mut ctx = SESSION.create_execution_ctx(); test_take_conformance( &RDEncoder::new(&[a]) diff --git a/encodings/alp/src/alp_rd/mod.rs b/encodings/alp/src/alp_rd/mod.rs index 07afe6b8adf..2cdeacd207d 100644 --- a/encodings/alp/src/alp_rd/mod.rs +++ b/encodings/alp/src/alp_rd/mod.rs @@ -1,494 +1,212 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Vortex array support for ALP-RD ("real doubles"), on top of the `alp` crate's encoder. + #![expect(clippy::cast_possible_truncation)] +pub use ::alp::ALPRDFloat; +pub use ::alp::RDEncoder; +use ::alp::alp_rd_apply_patches; +use ::alp::alp_rd_combine_codes_inplace; +use ::alp::alp_rd_combine_inplace; +use ::alp::alp_rd_dict_decode_inplace; +use ::alp::bit_width; pub use array::*; +use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::patches::Patches; -use vortex_array::validity::Validity; -use vortex_fastlanes::bitpack_compress::bitpack_encode_unchecked; - -mod array; -mod compute; -mod kernel; -mod ops; -mod rules; -mod slice; - -use std::ops::Shl; -use std::ops::Shr; - -use itertools::Itertools; -use num_traits::Float; -use num_traits::One; -use num_traits::PrimInt; -use rustc_hash::FxBuildHasher; -use vortex_array::ArrayView; use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::match_each_integer_ptype; +use vortex_array::patches::Patches; +use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_panic; +use vortex_fastlanes::bitpack_compress::bitpack_encode_unchecked; use vortex_session::VortexSession; -use vortex_utils::aliases::hash_map::HashMap; use crate::match_each_alp_float_ptype; -macro_rules! bit_width { - ($value:expr) => { - if $value == 0 { - 1 - } else { - $value.ilog2().wrapping_add(1) as usize - } - }; -} - -/// Max number of bits to cut from the MSB section of each float. -const CUT_LIMIT: usize = 16; - -const MAX_DICT_SIZE: u8 = 8; +mod array; +mod compute; +mod kernel; +mod ops; +mod rules; +mod slice; pub(crate) fn initialize(session: &VortexSession) { kernel::initialize(session); } -mod private { - pub trait Sealed {} - - impl Sealed for f32 {} - impl Sealed for f64 {} -} - -/// Main trait for ALP-RD encodable floating point numbers. -/// -/// Like the paper, we limit this to the IEEE7 754 single-precision (`f32`) and double-precision -/// (`f64`) floating point types. -pub trait ALPRDFloat: private::Sealed + Float + Copy + NativePType { - /// The unsigned integer type with the same bit-width as the floating-point type. - type UINT: NativePType + PrimInt + One + Copy; - - /// Number of bits the value occupies in registers. - const BITS: usize = size_of::() * 8; - - /// Bit-wise transmute from the unsigned integer type to the floating-point type. - fn from_bits(bits: Self::UINT) -> Self; - - /// Bit-wise transmute into the unsigned integer type. - fn to_bits(value: Self) -> Self::UINT; - - /// Truncating conversion from the unsigned integer type to `u16`. - fn to_u16(bits: Self::UINT) -> u16; - - /// Type-widening conversion from `u16` to the unsigned integer type. - fn from_u16(value: u16) -> Self::UINT; -} - -impl ALPRDFloat for f64 { - type UINT = u64; - - fn from_bits(bits: Self::UINT) -> Self { - f64::from_bits(bits) - } - - fn to_bits(value: Self) -> Self::UINT { - value.to_bits() - } - - fn to_u16(bits: Self::UINT) -> u16 { - bits as u16 - } - - fn from_u16(value: u16) -> Self::UINT { - value as u64 - } -} - -impl ALPRDFloat for f32 { - type UINT = u32; - - fn from_bits(bits: Self::UINT) -> Self { - f32::from_bits(bits) - } - - fn to_bits(value: Self) -> Self::UINT { - value.to_bits() - } - - fn to_u16(bits: Self::UINT) -> u16 { - bits as u16 - } - - fn from_u16(value: u16) -> Self::UINT { - value as u32 - } -} - -/// Encoder for ALP-RD ("real doubles") values. -/// -/// The encoder calculates its parameters from a single sample of floating-point values, -/// and then can be applied to many vectors. -/// -/// ALP-RD uses the algorithm outlined in Section 3.4 of the paper. The crux of it is that the front -/// (most significant) bits of many double vectors tend to be the same, i.e. most doubles in a -/// vector often use the same exponent and front bits. Compression proceeds by finding the best -/// prefix of up to 16 bits that can be collapsed into a dictionary of -/// up to 8 elements. Each double can then be broken into the front/left `L` bits, which neatly -/// bit-packs down to 1-3 bits per element (depending on the actual dictionary size). -/// The remaining `R` bits naturally bit-pack. -/// -/// In the ideal case, this scheme allows us to store a sequence of doubles in 49 bits-per-value. -/// -/// Our implementation draws on the MIT-licensed [C++ implementation] provided by the original authors. -/// -/// [C++ implementation]: https://github.com/cwida/ALP/blob/main/include/alp/rd.hpp -pub struct RDEncoder { - right_bit_width: u8, - codes: Vec, -} - -impl RDEncoder { - /// Build a new encoder from a sample of doubles. - pub fn new(sample: &[T]) -> Self - where - T: ALPRDFloat + NativePType, - T::UINT: NativePType, - { - let dictionary = find_best_dictionary::(sample); - - let mut codes = vec![0; dictionary.dictionary.len()]; - dictionary.dictionary.into_iter().for_each(|(bits, code)| { - // write the reverse mapping into the codes vector. - codes[code as usize] = bits - }); - - Self { - right_bit_width: dictionary.right_bit_width, - codes, - } - } - - /// Build a new encoder from known parameters. - pub fn from_parts(right_bit_width: u8, codes: Vec) -> Self { - Self { - right_bit_width, - codes, - } - } - +/// Encoding of Vortex arrays with an ALP-RD [`RDEncoder`]. +pub trait RDEncoderExt { /// Encode a set of floating point values with ALP-RD. /// - /// Each value will be split into a left and right component, which are compressed individually. + /// Each value is split into a left and right component, which are compressed individually. // TODO(joe): make fallible - pub fn encode(&self, array: ArrayView<'_, Primitive>) -> ALPRDArray { - match_each_alp_float_ptype!(array.ptype(), |P| { self.encode_generic::

(array) }) - } - - fn encode_generic(&self, array: ArrayView<'_, Primitive>) -> ALPRDArray - where - T: ALPRDFloat + NativePType, - T::UINT: NativePType, - { - assert!( - !self.codes.is_empty(), - "codes lookup table must be populated before RD encoding" - ); - - let doubles = array.as_slice::(); - - let mut left_parts: BufferMut = BufferMut::with_capacity(doubles.len()); - let mut right_parts: BufferMut = BufferMut::with_capacity(doubles.len()); - let mut exceptions_pos: BufferMut = BufferMut::with_capacity(doubles.len() / 4); - let mut exceptions: BufferMut = BufferMut::with_capacity(doubles.len() / 4); - - // mask for right-parts - let right_mask = T::UINT::one().shl(self.right_bit_width as _) - T::UINT::one(); - let max_code = self.codes.len() - 1; - let left_bit_width = bit_width!(max_code); - - for v in doubles.iter().copied() { - right_parts.push(T::to_bits(v) & right_mask); - left_parts.push(::to_u16( - T::to_bits(v).shr(self.right_bit_width as _), - )); - } + fn encode(&self, array: ArrayView<'_, Primitive>) -> ALPRDArray; +} - // dict-encode the left-parts, keeping track of exceptions - for (idx, left) in left_parts.iter_mut().enumerate() { - // TODO: revisit if we need to change the branch order for perf. - if let Some(code) = self.codes.iter().position(|v| *v == *left) { - *left = code as u16; - } else { - exceptions.push(*left); - exceptions_pos.push(idx as _); +impl RDEncoderExt for RDEncoder { + fn encode(&self, array: ArrayView<'_, Primitive>) -> ALPRDArray { + match_each_alp_float_ptype!(array.ptype(), |P| { encode_generic::

(self, array) }) + } +} - *left = 0u16; - } - } +fn encode_generic(encoder: &RDEncoder, array: ArrayView<'_, Primitive>) -> ALPRDArray +where + T: ALPRDFloat + NativePType, + T::UINT: NativePType, +{ + let (left_parts, right_parts, exceptions_pos, exceptions) = + encoder.split_parts(array.as_slice::()); + let len = left_parts.len(); + + // Bit-pack down the encoded left-parts array that have been dictionary encoded. + let primitive_left = PrimitiveArray::new( + Buffer::from(left_parts), + array + .validity() + .vortex_expect("ALP RD validity should be derivable"), + ); + // SAFETY: by construction, all values in left_parts can be packed to left_bit_width. + let packed_left = unsafe { + bitpack_encode_unchecked(primitive_left, encoder.left_bit_width()) + .vortex_expect("bitpack_encode_unchecked should succeed for left parts") + .into_array() + }; - // Bit-pack down the encoded left-parts array that have been dictionary encoded. - let primitive_left = PrimitiveArray::new( - left_parts, - array - .validity() - .vortex_expect("ALP RD validity should be derivable"), - ); - // SAFETY: by construction, all values in left_parts can be packed to left_bit_width. - let packed_left = unsafe { - bitpack_encode_unchecked(primitive_left, left_bit_width as _) - .vortex_expect("bitpack_encode_unchecked should succeed for left parts") - .into_array() - }; + let primitive_right = PrimitiveArray::new(Buffer::from(right_parts), Validity::NonNullable); + // SAFETY: by construction, all values in right_parts are right_bit_width + leading zeros. + let packed_right = unsafe { + bitpack_encode_unchecked(primitive_right, encoder.right_bit_width()) + .vortex_expect("bitpack_encode_unchecked should succeed for right parts") + .into_array() + }; - let primitive_right = PrimitiveArray::new(right_parts, Validity::NonNullable); - // SAFETY: by construction, all values in right_parts are right_bit_width + leading zeros. - let packed_right = unsafe { - bitpack_encode_unchecked(primitive_right, self.right_bit_width as _) - .vortex_expect("bitpack_encode_unchecked should succeed for right parts") + // Patches for the left-parts that could not be dictionary encoded. + let exceptions = (!exceptions_pos.is_empty()).then(|| { + let max_exc_pos = exceptions_pos.last().copied().unwrap_or_default(); + let bw = bit_width(max_exc_pos); + + let exc_pos_array = + PrimitiveArray::new(Buffer::from(exceptions_pos), Validity::NonNullable); + // SAFETY: We calculate bw such that it is wide enough to hold the largest position index. + let packed_pos = unsafe { + bitpack_encode_unchecked(exc_pos_array, bw) + .vortex_expect("bitpack_encode_unchecked should succeed for exception positions") .into_array() }; - // Bit-pack the dict-encoded left-parts - // Bit-pack the right-parts - // Patches for exceptions. - let exceptions = (!exceptions_pos.is_empty()).then(|| { - let max_exc_pos = exceptions_pos.last().copied().unwrap_or_default(); - let bw = bit_width!(max_exc_pos) as u8; - - let exc_pos_array = PrimitiveArray::new(exceptions_pos, Validity::NonNullable); - // SAFETY: We calculate bw such that it is wide enough to hold the largest position index. - let packed_pos = unsafe { - bitpack_encode_unchecked(exc_pos_array, bw) - .vortex_expect( - "bitpack_encode_unchecked should succeed for exception positions", - ) - .into_array() - }; - - Patches::new( - doubles.len(), - 0, - packed_pos, - exceptions.into_array(), - // TODO(0ax1): handle chunk offsets - None, - ) - .vortex_expect("Patches construction in encode") - }); - - ALPRD::try_new( - DType::Primitive(T::PTYPE, packed_left.dtype().nullability()), - packed_left, - Buffer::::copy_from(&self.codes), - packed_right, - self.right_bit_width, - exceptions, + Patches::new( + len, + 0, + packed_pos, + Buffer::from(exceptions).into_array(), + // TODO(0ax1): handle chunk offsets + None, ) - .vortex_expect("ALPRDArray construction in encode") - } + .vortex_expect("Patches construction in encode") + }); + + ALPRD::try_new( + DType::Primitive(T::PTYPE, packed_left.dtype().nullability()), + packed_left, + Buffer::::copy_from(encoder.codes()), + packed_right, + encoder.right_bit_width(), + exceptions, + ) + .vortex_expect("ALPRDArray construction in encode") } /// Decode ALP-RD encoded values back into their original floating point format. /// +/// The decoded values are written over `right_parts`, which shares its layout with the floats they +/// decode to. +/// /// # Panics /// /// Panics if `left_parts` and `right_parts` differ in length. -pub fn alp_rd_decode( +pub fn alp_rd_decode( mut left_parts: BufferMut, left_parts_dict: &[u16], right_bit_width: u8, - right_parts: BufferMut, + mut right_parts: BufferMut, left_parts_patches: Option, ctx: &mut ExecutionCtx, -) -> VortexResult> { +) -> VortexResult> +where + T::UINT: NativePType, +{ if left_parts.len() != right_parts.len() { vortex_panic!("alp_rd_decode: left_parts.len != right_parts.len"); } - let shift = right_bit_width as usize; - if let Some(patches) = left_parts_patches { // Patched path: some left-part codes map to exception values that live outside // the dictionary. We must dictionary-decode first, then overwrite the exceptions, // before we can combine with right-parts. + alp_rd_dict_decode_inplace(left_parts.as_mut_slice(), left_parts_dict); - // Dictionary-decode every code in-place (code → actual left bit-pattern). - for code in left_parts.iter_mut() { - *code = left_parts_dict[*code as usize]; - } - - // Overwrite exception positions with their true left bit-patterns. let indices = patches.indices().clone().execute::(ctx)?; let patch_values = patches.values().clone().execute::(ctx)?; - alp_rd_apply_patches(&mut left_parts, &indices, &patch_values, patches.offset()); + match_each_integer_ptype!(indices.ptype(), |I| { + alp_rd_apply_patches( + left_parts.as_mut_slice(), + indices.as_slice::(), + patch_values.as_slice::(), + patches.offset(), + ) + }); - // Reconstruct floats by shifting each decoded left value into the MSBs - // and OR-ing with the corresponding right value. alp_rd_combine_inplace::( - right_parts, - |right, &left| { - *right = (::from_u16(left) << shift) | *right; - }, + right_parts.as_mut_slice(), left_parts.as_ref(), - ) + right_bit_width, + ); } else { - // Non-patched fast path: every code maps through the dictionary, so we can - // pre-shift the entire dictionary once and reduce the per-element hot loop to - // a single table lookup + OR. - let mut shifted_dict = [T::UINT::default(); MAX_DICT_SIZE as usize]; - for (i, &entry) in left_parts_dict.iter().enumerate() { - shifted_dict[i] = ::from_u16(entry) << shift; - } - - // Each element: look up the pre-shifted left value by code, OR with right-parts. - alp_rd_combine_inplace::( - right_parts, - |right, &code| { - // SAFETY: codes are bounded by dict size (< left_parts_dict.len() <= MAX_DICT_SIZE). - *right = unsafe { *shifted_dict.get_unchecked(code as usize) } | *right; - }, + // Non-patched fast path: every code maps through the dictionary, so the hot loop is a + // lookup into the pre-shifted dictionary and an OR. + alp_rd_combine_codes_inplace::( + right_parts.as_mut_slice(), left_parts.as_ref(), - ) + left_parts_dict, + right_bit_width, + ); } -} -/// Apply patches to the decoded left-parts values. -fn alp_rd_apply_patches( - values: &mut BufferMut, - indices: &PrimitiveArray, - patch_values: &PrimitiveArray, - offset: usize, -) { - match_each_integer_ptype!(indices.ptype(), |T| { - indices - .as_slice::() - .iter() - .copied() - .map(|idx| idx - offset as T) - .zip(patch_values.as_slice::().iter()) - .for_each(|(idx, v)| values[idx as usize] = *v); - }) -} - -/// Zip `right_parts` with `left_data`, apply `combine_fn` per element, then reinterpret the -/// buffer from `T::UINT` to `T` (same bit-width: u32↔f32, u64↔f64). -fn alp_rd_combine_inplace( - mut right_parts: BufferMut, - combine_fn: impl Fn(&mut T::UINT, &u16), - left_data: &[u16], -) -> VortexResult> { - for (right, left) in right_parts.as_mut_slice().iter_mut().zip(left_data.iter()) { - combine_fn(right, left); - } // SAFETY: all bit patterns of T::UINT are valid T (u32↔f32 or u64↔f64). Ok(unsafe { right_parts.transmute::() }.freeze()) } -/// Find the best "cut point" for a set of floating point values such that we can -/// cast them all to the relevant value instead. -fn find_best_dictionary(samples: &[T]) -> ALPRDDictionary { - let mut best_est_size = f64::MAX; - let mut best_dict = ALPRDDictionary::default(); - for p in 1..=16 { - let candidate_right_bw = (T::BITS - p) as u8; - let (dictionary, exception_count) = - build_left_parts_dictionary::(samples, candidate_right_bw, MAX_DICT_SIZE); - let estimated_size = estimate_compression_size( - dictionary.right_bit_width, - dictionary.left_bit_width, - exception_count, - samples.len(), - ); - if estimated_size < best_est_size { - best_est_size = estimated_size; - best_dict = dictionary; - } - } - - best_dict -} +#[cfg(test)] +mod tests { + use vortex_array::arrays::PrimitiveArray; -/// Build dictionary of the leftmost bits. -fn build_left_parts_dictionary( - samples: &[T], - right_bw: u8, - max_dict_size: u8, -) -> (ALPRDDictionary, usize) { - assert!( - right_bw >= (T::BITS - CUT_LIMIT) as _, - "left-parts must be <= 16 bits" - ); + use super::ALPRDArrayExt; + use super::RDEncoder; + use super::RDEncoderExt; - // Count the number of occurrences of each left bit pattern - let mut counts = HashMap::new(); - samples - .iter() - .copied() - .map(|v| ::to_u16(T::to_bits(v).shr(right_bw as _))) - .for_each(|item| *counts.entry(item).or_default() += 1); + #[test] + fn rd_encoder_uses_deterministic_codes_for_tied_prefixes() { + let values: Vec<_> = (1_000_i64..2_000) + .map(|row_id| ((row_id * 53) % 36_000) as f64 / 100.0 - 180.0) + .collect(); + let array = PrimitiveArray::from_iter(values.iter().copied()); - // Sorted counts: sort by negative count so that heavy hitters sort first. - let mut sorted_bit_counts: Vec<(u16, usize)> = counts.into_iter().collect_vec(); - sorted_bit_counts.sort_by_key(|(_, count)| count.wrapping_neg()); + let encoded = RDEncoder::new(&values).encode(array.as_view()); - // Assign the most-frequently occurring left-bits as dictionary codes, up to `dict_size`... - let mut dictionary = HashMap::with_capacity_and_hasher(max_dict_size as _, FxBuildHasher); - let mut code = 0u16; - while code < (max_dict_size as _) && (code as usize) < sorted_bit_counts.len() { - let (bits, _) = sorted_bit_counts[code as usize]; - dictionary.insert(bits, code); - code += 1; + assert_eq!(encoded.right_bit_width(), 53); + assert_eq!( + encoded.left_parts_dictionary().as_slice(), + &[514, 1538, 515, 1539, 513, 1537, 512, 1536] + ); } - - // ...and the rest are exceptions. - let exception_count: usize = sorted_bit_counts - .iter() - .skip(code as _) - .map(|(_, count)| *count) - .sum(); - - // Left bit-width is determined based on the actual dictionary size. - let max_code = dictionary.len() - 1; - let left_bw = bit_width!(max_code) as u8; - - ( - ALPRDDictionary { - dictionary, - right_bit_width: right_bw, - left_bit_width: left_bw, - }, - exception_count, - ) -} - -/// Estimate the bits-per-value when using these compression settings. -fn estimate_compression_size( - right_bw: u8, - left_bw: u8, - exception_count: usize, - sample_n: usize, -) -> f64 { - const EXC_POSITION_SIZE: usize = 16; // two bytes for exception position. - const EXC_SIZE: usize = 16; // two bytes for each exception (up to 16 front bits). - - let exceptions_size = exception_count * (EXC_POSITION_SIZE + EXC_SIZE); - (right_bw as f64) + (left_bw as f64) + ((exceptions_size as f64) / (sample_n as f64)) -} - -/// The ALP-RD dictionary, encoding the "left parts" and their dictionary encoding. -#[derive(Debug, Default)] -struct ALPRDDictionary { - /// Items in the dictionary are bit patterns, along with their 16-bit encoding. - dictionary: HashMap, - /// The (compressed) left bit width. This is after bit-packing the dictionary codes. - left_bit_width: u8, - /// The right bit width. This is the bit-packed width of each of the "real double" values. - right_bit_width: u8, } diff --git a/encodings/alp/src/alp_rd/ops.rs b/encodings/alp/src/alp_rd/ops.rs index 0fc0dc40fa1..edb2fb21186 100644 --- a/encodings/alp/src/alp_rd/ops.rs +++ b/encodings/alp/src/alp_rd/ops.rs @@ -11,6 +11,7 @@ use vortex_error::VortexResult; use crate::ALPRD; use crate::ALPRDArrayExt; +use crate::ALPRDArraySlotsExt; impl OperationsVTable for ALPRD { fn scalar_at( @@ -71,12 +72,14 @@ mod test { use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; + use vortex_array::dtype::NativePType; use vortex_array::scalar::Scalar; use vortex_session::VortexSession; use crate::ALPRDArrayExt; use crate::ALPRDFloat; use crate::RDEncoder; + use crate::RDEncoderExt; static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session(); @@ -87,7 +90,7 @@ mod test { #[rstest] #[case(0.1f32, 0.2f32, 3e25f32)] #[case(0.1f64, 0.2f64, 3e100f64)] - fn test_slice(#[case] a: T, #[case] b: T, #[case] outlier: T) { + fn test_slice(#[case] a: T, #[case] b: T, #[case] outlier: T) { let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::from_iter([a, b, outlier]); let encoded = RDEncoder::new(&[a, b]).encode(array.as_view()); @@ -99,7 +102,7 @@ mod test { #[rstest] #[case(0.1f32, 0.2f32, 3e25f32)] #[case(0.1f64, 0.2f64, 3e100f64)] - fn test_scalar_at>( + fn test_scalar_at>( #[case] a: T, #[case] b: T, #[case] outlier: T, diff --git a/encodings/alp/src/alp_rd/slice.rs b/encodings/alp/src/alp_rd/slice.rs index c36e138ff84..5f81bdc7b5d 100644 --- a/encodings/alp/src/alp_rd/slice.rs +++ b/encodings/alp/src/alp_rd/slice.rs @@ -11,6 +11,7 @@ use vortex_array::arrays::slice::SliceKernel; use vortex_error::VortexResult; use crate::ALPRDArrayExt; +use crate::ALPRDArraySlotsExt; use crate::alp_rd::ALPRD; impl SliceKernel for ALPRD { diff --git a/encodings/alp/src/lib.rs b/encodings/alp/src/lib.rs index 1286f3e6e2b..160ab568cd4 100644 --- a/encodings/alp/src/lib.rs +++ b/encodings/alp/src/lib.rs @@ -16,8 +16,6 @@ //! //! [paper]: https://ir.cwi.nl/pub/33334/33334.pdf -pub use alp::*; -pub use alp_rd::*; use vortex_array::ArrayVTable; use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::fns::nan_count::NanCount; @@ -26,6 +24,9 @@ use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::session::ArraySessionExt; use vortex_session::VortexSession; +pub use self::alp::*; +pub use self::alp_rd::*; + mod alp; mod alp_rd; diff --git a/encodings/bytebool/src/array.rs b/encodings/bytebool/src/array.rs index fe00bd5fe4d..faa1fea81fa 100644 --- a/encodings/bytebool/src/array.rs +++ b/encodings/bytebool/src/array.rs @@ -19,6 +19,7 @@ use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::IntoArray; use vortex_array::TypedArrayRef; +use vortex_array::array_slots; use vortex_array::arrays::BoolArray; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; @@ -72,7 +73,8 @@ impl VTable for ByteBool { len: usize, slots: &[Option], ) -> VortexResult<()> { - let validity = child_to_validity(slots[VALIDITY_SLOT].as_ref(), dtype.nullability()); + let validity = + child_to_validity(slots[ByteBoolSlots::VALIDITY].as_ref(), dtype.nullability()); ByteBoolData::validate(data.buffer(), &validity, dtype, len) } @@ -153,7 +155,7 @@ impl VTable for ByteBool { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + ByteBoolSlots::NAMES[idx].to_string() } fn reduce_parent( @@ -174,10 +176,12 @@ impl VTable for ByteBool { } } -/// The validity bitmap indicating which elements are non-null. -pub(super) const VALIDITY_SLOT: usize = 0; -pub(super) const NUM_SLOTS: usize = 1; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"]; +#[array_slots(ByteBool)] +pub struct ByteBoolSlots { + /// The validity bitmap indicating which elements are non-null. + #[slot(0)] + pub validity: Option, +} #[derive(Clone, Debug)] pub struct ByteBoolData { @@ -190,10 +194,11 @@ impl Display for ByteBoolData { } } -pub trait ByteBoolArrayExt: TypedArrayRef { - fn validity(&self) -> Validity { +pub trait ByteBoolArrayExt: TypedArrayRef + ByteBoolArraySlotsExt { + /// Returns the [`Validity`] derived from the validity slot. + fn bytebool_validity(&self) -> Validity { child_to_validity( - self.as_ref().slots()[VALIDITY_SLOT].as_ref(), + self.as_ref().slots()[ByteBoolSlots::VALIDITY].as_ref(), self.as_ref().dtype().nullability(), ) } @@ -302,7 +307,7 @@ impl ByteBoolData { impl ValidityVTable for ByteBool { fn validity(array: ArrayView<'_, ByteBool>) -> VortexResult { - Ok(ByteBoolArrayExt::validity(&array)) + Ok(array.bytebool_validity()) } } diff --git a/encodings/datetime-parts/Cargo.toml b/encodings/datetime-parts/Cargo.toml index 52894fc7c07..3bbdde465dd 100644 --- a/encodings/datetime-parts/Cargo.toml +++ b/encodings/datetime-parts/Cargo.toml @@ -26,6 +26,12 @@ vortex-mask = { workspace = true } vortex-session = { workspace = true } [dev-dependencies] +divan = { workspace = true } +rand = { workspace = true } rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } vortex-error = { workspace = true } + +[[bench]] +name = "split_temporal" +harness = false diff --git a/encodings/datetime-parts/benches/split_temporal.rs b/encodings/datetime-parts/benches/split_temporal.rs new file mode 100644 index 00000000000..213192899fd --- /dev/null +++ b/encodings/datetime-parts/benches/split_temporal.rs @@ -0,0 +1,58 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use rand::RngExt; +use rand::SeedableRng as _; +use rand::rngs::StdRng; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::TemporalArray; +use vortex_array::extension::datetime::TimeUnit; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_datetime_parts::split_temporal; +use vortex_session::VortexSession; + +fn main() { + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + +const BENCH_ARGS: &[(usize, TimeUnit)] = &[ + (65_536, TimeUnit::Seconds), + (65_536, TimeUnit::Milliseconds), + (65_536, TimeUnit::Microseconds), + (65_536, TimeUnit::Nanoseconds), +]; + +#[divan::bench(args = BENCH_ARGS)] +fn split(bencher: Bencher, args: (usize, TimeUnit)) { + let (n, unit) = args; + let divisor: i64 = match unit { + TimeUnit::Seconds => 1, + TimeUnit::Milliseconds => 1_000, + TimeUnit::Microseconds => 1_000_000, + TimeUnit::Nanoseconds => 1_000_000_000, + TimeUnit::Days => unreachable!(), + }; + let mut rng = StdRng::seed_from_u64(0); + let timestamps = Buffer::from_iter((0..n).map(|_| { + rng.random_range(1_500_000_000i64..1_800_000_000) * divisor + rng.random_range(0..divisor) + })); + let array = TemporalArray::new_timestamp( + PrimitiveArray::new(timestamps, Validity::NonNullable).into_array(), + unit, + Some("UTC".into()), + ); + + bencher + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) + .bench_values(|(array, mut ctx)| split_temporal(array, &mut ctx).unwrap()) +} diff --git a/encodings/datetime-parts/src/array.rs b/encodings/datetime-parts/src/array.rs index 42569da9728..30594c78367 100644 --- a/encodings/datetime-parts/src/array.rs +++ b/encodings/datetime-parts/src/array.rs @@ -212,10 +212,13 @@ impl VTable for DateTimeParts { #[array_slots(DateTimeParts)] pub struct DateTimePartsSlots { /// The days component of the datetime, stored as an integer array. + #[slot(0)] pub days: ArrayRef, /// The seconds component of the datetime (within the day). + #[slot(1)] pub seconds: ArrayRef, /// The sub-second component of the datetime. + #[slot(2)] pub subseconds: ArrayRef, } diff --git a/encodings/datetime-parts/src/compress.rs b/encodings/datetime-parts/src/compress.rs index a3ab3e6f23b..1198257ebae 100644 --- a/encodings/datetime-parts/src/compress.rs +++ b/encodings/datetime-parts/src/compress.rs @@ -1,18 +1,26 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::mem::MaybeUninit; + use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::TemporalArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::PType; +use vortex_array::extension::datetime::TimeUnit; use vortex_buffer::BufferMut; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use crate::timestamp; +/// All parts are stored as i32. seconds and subseconds always fit in i32. +/// For days we have a upper day limitation due to "jiff" validation so days +/// also fit in i32. pub struct TemporalParts { pub days: ArrayRef, pub seconds: ArrayRef, @@ -24,6 +32,11 @@ pub struct TemporalParts { /// Splitting the components by granularity creates more small values, which enables better /// cascading compression. pub fn split_temporal(array: TemporalArray, ctx: &mut ExecutionCtx) -> VortexResult { + let time_unit = array.temporal_metadata().time_unit(); + if matches!(time_unit, TimeUnit::Days) { + vortex_bail!("Cannot handle day-level data"); + } + let temporal_values = array .temporal_values() .clone() @@ -40,24 +53,121 @@ pub fn split_temporal(array: TemporalArray, ctx: &mut ExecutionCtx) -> VortexRes .execute::(ctx)?; let length = timestamps.len(); - let mut days = BufferMut::with_capacity(length); - let mut seconds = BufferMut::with_capacity(length); - let mut subseconds = BufferMut::with_capacity(length); - - for &ts in timestamps.as_slice::() { - let ts_parts = timestamp::split(ts, array.temporal_metadata().time_unit())?; - days.push(ts_parts.days); - seconds.push(ts_parts.seconds); - subseconds.push(ts_parts.subseconds); + + // If we don't [..length], compiler can't infer all 3 or 4 slices are the + // same length, and zip() in split_slice_* checks iterator boundaries. + let timestamps = ×tamps.as_slice::()[..length]; + + let mut days: BufferMut = BufferMut::with_capacity(timestamps.len()); + let mut seconds: BufferMut = BufferMut::with_capacity(timestamps.len()); + let days_ptr = &mut days.spare_capacity_mut()[..length]; + let seconds_ptr = &mut seconds.spare_capacity_mut()[..length]; + + if matches!(time_unit, TimeUnit::Seconds) { + split_slice_seconds(days_ptr, seconds_ptr, timestamps); + + // SAFETY: all items in [0; length) are filled in split_slice_seconds + unsafe { + days.set_len(length); + seconds.set_len(length); + } + + return Ok(TemporalParts { + days: PrimitiveArray::new(days.freeze(), temporal_values.validity()?).into_array(), + seconds: seconds.into_array(), + subseconds: ConstantArray::new(0, length).into_array(), + }); + } + + let mut subseconds: BufferMut = BufferMut::with_capacity(timestamps.len()); + let subseconds_ptr = &mut subseconds.spare_capacity_mut()[..length]; + + match time_unit { + TimeUnit::Nanoseconds => { + split_slice::<1_000_000_000>(days_ptr, seconds_ptr, subseconds_ptr, timestamps) + } + TimeUnit::Microseconds => { + split_slice::<1_000_000>(days_ptr, seconds_ptr, subseconds_ptr, timestamps) + } + TimeUnit::Milliseconds => { + split_slice::<1_000>(days_ptr, seconds_ptr, subseconds_ptr, timestamps) + } + _ => unreachable!("Handled before"), + }; + + // SAFETY: all items in [0; length) are filled in split_slice + unsafe { + days.set_len(length); + seconds.set_len(length); + subseconds.set_len(length); } Ok(TemporalParts { - days: PrimitiveArray::new(days, temporal_values.validity()?).into_array(), + days: PrimitiveArray::new(days.freeze(), temporal_values.validity()?).into_array(), seconds: seconds.into_array(), subseconds: subseconds.into_array(), }) } +#[inline] +fn split_slice( + days: &mut [MaybeUninit], + seconds: &mut [MaybeUninit], + subseconds: &mut [MaybeUninit], + timestamps: &[i64], +) { + // Computing chunks of 4 elements lets LLVM optimize stores into + // a 16-byte vector store per chunk. + let length = timestamps.len(); + let (timestamps, timestamps_rem) = timestamps.as_chunks::<4>(); + let (days, days_rem) = days[..length].as_chunks_mut::<4>(); + let (seconds, seconds_rem) = seconds[..length].as_chunks_mut::<4>(); + let (subseconds, subseconds_rem) = subseconds[..length].as_chunks_mut::<4>(); + + let chunks = timestamps.iter().zip(days).zip(seconds).zip(subseconds); + for (((timestamp, day), second), subsecond) in chunks { + let mut day_buf = [0i32; 4]; + let mut second_buf = [0i32; 4]; + let mut subsecond_buf = [0i32; 4]; + for k in 0..4 { + let parts = timestamp::split_with_divisor::(timestamp[k]); + day_buf[k] = parts.days; + second_buf[k] = parts.seconds; + subsecond_buf[k] = parts.subseconds; + } + for k in 0..4 { + day[k].write(day_buf[k]); + second[k].write(second_buf[k]); + subsecond[k].write(subsecond_buf[k]); + } + } + + let remainder = timestamps_rem + .iter() + .zip(days_rem) + .zip(seconds_rem) + .zip(subseconds_rem); + for (((&ts, day), second), subsecond) in remainder { + let parts = timestamp::split_with_divisor::(ts); + day.write(parts.days); + second.write(parts.seconds); + subsecond.write(parts.subseconds); + } +} + +#[inline] +fn split_slice_seconds( + days: &mut [MaybeUninit], + seconds: &mut [MaybeUninit], + timestamps: &[i64], +) { + for ((day, second), ts) in days.iter_mut().zip(seconds).zip(timestamps) { + let parts = timestamp::split_with_divisor::<1>(*ts); + day.write(parts.days); + second.write(parts.seconds); + } +} + #[cfg(test)] mod tests { use rstest::rstest; diff --git a/encodings/datetime-parts/src/compute/compare.rs b/encodings/datetime-parts/src/compute/compare.rs index 0791861df79..a4336e0d0fd 100644 --- a/encodings/datetime-parts/src/compute/compare.rs +++ b/encodings/datetime-parts/src/compute/compare.rs @@ -14,6 +14,7 @@ use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::binary::CompareKernel; use vortex_array::scalar_fn::fns::operators::CompareOperator; use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::validity::Validity; use vortex_error::VortexResult; use crate::array::DateTimeParts; @@ -175,26 +176,51 @@ fn compare_gt( fn compare_dtp( lhs: &ArrayRef, - rhs: i64, + rhs: i32, operator: CompareOperator, nullability: Nullability, ) -> VortexResult { // Since nullability is stripped from RHS and carried forward through nullability argument we want to incorporate it into lhs.dtype() that we cast rhs into - match ConstantArray::new(rhs, lhs.len()) - .into_array() - .cast(lhs.dtype().with_nullability(nullability)) - { - Ok(casted) => lhs.binary(casted, Operator::from(operator)), - // The narrowing cast failed. Therefore, we know lhs < rhs. + match Scalar::from(rhs).cast(&lhs.dtype().with_nullability(nullability)) { + Ok(casted) => lhs.binary( + ConstantArray::new(casted, lhs.len()).into_array(), + Operator::from(operator), + ), + // The narrowing cast failed, so rhs is either > or < every value in lhs. + // rhs positive => > all lhs + // rhs negative => < all lhs _ => { + let all_lhs_smaller = rhs > 0; let constant_value = match operator { - CompareOperator::Eq | CompareOperator::Gte | CompareOperator::Gt => false, - CompareOperator::NotEq | CompareOperator::Lte | CompareOperator::Lt => true, + CompareOperator::Eq => false, + CompareOperator::NotEq => true, + CompareOperator::Lt | CompareOperator::Lte => all_lhs_smaller, + CompareOperator::Gt | CompareOperator::Gte => !all_lhs_smaller, }; - Ok( - ConstantArray::new(Scalar::bool(constant_value, nullability), lhs.len()) - .into_array(), - ) + // If there are nulls in lhs, we need to propagate them + let validity = match nullability { + Nullability::NonNullable => Validity::NonNullable, + // lhs may be non-nullable while rhs contributes the nullability. + Nullability::Nullable => match lhs.validity()? { + Validity::NonNullable => Validity::AllValid, + validity => validity, + }, + }; + Ok(match validity { + Validity::NonNullable | Validity::AllValid => { + ConstantArray::new(Scalar::bool(constant_value, nullability), lhs.len()) + .into_array() + } + Validity::AllInvalid => { + ConstantArray::new(Scalar::null(DType::Bool(nullability)), lhs.len()) + .into_array() + } + Validity::Array(validity_array) => { + ConstantArray::new(Scalar::bool(constant_value, nullability), lhs.len()) + .into_array() + .mask(validity_array)? + } + }) } } } @@ -207,10 +233,16 @@ mod test { use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::sum::sum; use vortex_array::array_session; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::TemporalArray; + use vortex_array::assert_arrays_eq; use vortex_array::dtype::IntegerPType; use vortex_array::extension::datetime::TimeUnit; + use vortex_array::extension::datetime::Timestamp; + use vortex_array::extension::datetime::TimestampOptions; + use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::buffer; @@ -381,4 +413,101 @@ mod test { // `CompareOperator::Gt` and `CompareOperator::Gte` only cover the case of all lhs values // being larger. Therefore, these cases are not covered by unit tests. } + + #[test] + fn compare_date_time_parts_eq_out_of_range_seconds_constant() -> VortexResult<()> { + const QUERY_TIMESTAMP_MS: i64 = 1_786_576_785_621; + + let session = array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + let timestamp = Scalar::extension::( + TimestampOptions { + unit: TimeUnit::Milliseconds, + tz: None, + }, + QUERY_TIMESTAMP_MS.into(), + ); + let rhs = ConstantArray::new(timestamp, 1).into_array(); + let lhs = DateTimeParts::try_new( + rhs.dtype().clone(), + buffer![20_677i32].into_array(), + buffer![0u16].into_array(), + buffer![0u16].into_array(), + )?; + + let comparison = lhs.into_array().binary(rhs, Operator::Eq)?; + + assert_eq!(true_count(&comparison, &mut ctx), 0); + Ok(()) + } + + /// A days constant below the range of the narrowed storage type must compare as smaller than + /// every stored value, not larger. + #[test] + fn compare_date_time_parts_below_range_days_constant() -> VortexResult<()> { + // 1969-12-31, which splits to a day of `-1`: below the minimum of an unsigned days array. + const BEFORE_EPOCH_MS: i64 = -86_400_000; + + let session = array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + let timestamp = Scalar::extension::( + TimestampOptions { + unit: TimeUnit::Milliseconds, + tz: None, + }, + BEFORE_EPOCH_MS.into(), + ); + let rhs = ConstantArray::new(timestamp, 1).into_array(); + // 2026-08-08. Days narrow to `u16` because the column holds no pre-epoch value. + let lhs = DateTimeParts::try_new( + rhs.dtype().clone(), + buffer![20_677u16].into_array(), + buffer![0u16].into_array(), + buffer![0u16].into_array(), + )? + .into_array(); + + let lt = lhs.binary(rhs.clone(), Operator::Lt)?; + assert_eq!(true_count(<, &mut ctx), 0); + + let gt = lhs.binary(rhs, Operator::Gt)?; + assert_eq!(true_count(>, &mut ctx), 1); + Ok(()) + } + + /// A null lhs value stays null, even when the constant is outside the range of the narrowed + /// days storage type. + #[rstest] + #[case(Validity::AllInvalid, [None, None])] + #[case(Validity::from_iter([false, true]), [None, Some(true)])] + fn compare_date_time_parts_null_out_of_range_days_constant( + #[case] days_validity: Validity, + #[case] expected: [Option; 2], + ) -> VortexResult<()> { + let session = array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + // 1969-12-31, a day of `-1`: below the minimum of an unsigned days array. + let timestamp = Scalar::extension::( + TimestampOptions { + unit: TimeUnit::Milliseconds, + tz: None, + }, + (-86_400_000i64).into(), + ); + let rhs = ConstantArray::new(timestamp, 2).into_array(); + let lhs = DateTimeParts::try_new( + rhs.dtype().with_nullability(Nullability::Nullable), + PrimitiveArray::new(buffer![20_677u16, 20_678u16], days_validity).into_array(), + buffer![0u16, 0u16].into_array(), + buffer![0u16, 0u16].into_array(), + )? + .into_array(); + + let gt = lhs.binary(rhs, Operator::Gt)?; + assert_arrays_eq!(gt, BoolArray::from_iter(expected), &mut ctx); + Ok(()) + } } diff --git a/encodings/datetime-parts/src/compute/rules.rs b/encodings/datetime-parts/src/compute/rules.rs index 84f58310ac1..12546c23fd8 100644 --- a/encodings/datetime-parts/src/compute/rules.rs +++ b/encodings/datetime-parts/src/compute/rules.rs @@ -143,7 +143,7 @@ impl ArrayParentReduceRule for DTPComparisonPushDownRule { /// Try to extract the days value from a constant timestamp. /// Returns None if the constant is not a timestamp or has non-zero seconds/subseconds. -fn try_extract_days_constant(array: &ArrayRef) -> Option { +fn try_extract_days_constant(array: &ArrayRef) -> Option { let constant = array.as_constant()?; // Extract the timestamp value @@ -185,7 +185,6 @@ mod tests { use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::TemporalArray; - use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::extension::datetime::TimeUnit; use vortex_array::extension::datetime::TimestampOptions; use vortex_array::optimizer::ArrayOptimizer; @@ -282,9 +281,9 @@ mod tests { // Compare: dtp <= day 1 (midnight) let constant = midnight_constant(1, TimeUnit::Seconds, len); - let comparison = Binary - .try_new_array(len, Operator::Lte, [dtp.into_array(), constant]) - .unwrap(); + let comparison = Binary::try_new(dtp.into_array(), constant, Operator::Lte) + .unwrap() + .into_array(); // Optimize should push down to days let optimized = comparison.optimize().unwrap(); @@ -310,16 +309,17 @@ mod tests { let lower = midnight_constant(1, TimeUnit::Seconds, len); let upper = midnight_constant(3, TimeUnit::Seconds, len); - let between = Between - .try_new_array( - len, - BetweenOptions { - lower_strict: StrictComparison::NonStrict, - upper_strict: StrictComparison::NonStrict, - }, - [dtp.into_array(), lower, upper], - ) - .unwrap(); + let between = Between::try_new( + dtp.into_array(), + lower, + upper, + BetweenOptions { + lower_strict: StrictComparison::NonStrict, + upper_strict: StrictComparison::NonStrict, + }, + ) + .unwrap() + .into_array(); // Optimize should push down to days let optimized = between.optimize().unwrap(); @@ -337,9 +337,9 @@ mod tests { // Compare against non-midnight constant (day 1 at noon) let constant = non_midnight_constant(1, 43200, TimeUnit::Seconds, len); - let comparison = Binary - .try_new_array(len, Operator::Lte, [dtp.into_array(), constant]) - .unwrap(); + let comparison = Binary::try_new(dtp.into_array(), constant, Operator::Lte) + .unwrap() + .into_array(); // Optimize should NOT push down (constant has non-zero seconds) let optimized = comparison.optimize().unwrap(); @@ -372,9 +372,9 @@ mod tests { // Compare against midnight constant let constant = midnight_constant(1, TimeUnit::Seconds, len); - let comparison = Binary - .try_new_array(len, Operator::Lte, [dtp.into_array(), constant]) - .unwrap(); + let comparison = Binary::try_new(dtp.into_array(), constant, Operator::Lte) + .unwrap() + .into_array(); // Should still compute correctly (just not optimized via pushdown) let optimized = comparison.optimize().unwrap(); diff --git a/encodings/datetime-parts/src/ops.rs b/encodings/datetime-parts/src/ops.rs index ab8fa362cb7..86e22ea9fe0 100644 --- a/encodings/datetime-parts/src/ops.rs +++ b/encodings/datetime-parts/src/ops.rs @@ -37,24 +37,24 @@ impl OperationsVTable for DateTimeParts { return Ok(Scalar::null(DType::Extension(ext))); } - let days: i64 = array + let days: i32 = array .days() .execute_scalar(index, ctx)? .as_primitive() - .as_::() - .vortex_expect("days fits in i64"); - let seconds: i64 = array + .as_::() + .vortex_expect("days fits in i32"); + let seconds: i32 = array .seconds() .execute_scalar(index, ctx)? .as_primitive() - .as_::() - .vortex_expect("seconds fits in i64"); - let subseconds: i64 = array + .as_::() + .vortex_expect("seconds fits in i32"); + let subseconds: i32 = array .subseconds() .execute_scalar(index, ctx)? .as_primitive() - .as_::() - .vortex_expect("subseconds fits in i64"); + .as_::() + .vortex_expect("subseconds fits in i32"); let ts = timestamp::combine( TimestampParts { diff --git a/encodings/datetime-parts/src/timestamp.rs b/encodings/datetime-parts/src/timestamp.rs index 6cafeb530d5..2aa10d4fac6 100644 --- a/encodings/datetime-parts/src/timestamp.rs +++ b/encodings/datetime-parts/src/timestamp.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use num_traits::AsPrimitive; use vortex_array::extension::datetime::TimeUnit; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -11,13 +12,13 @@ pub const SECONDS_PER_DAY: i64 = 86_400; // 24 * 60 * 60 /// Parts of a Unix timestamp (time since 1970-01-01 00:00:00 UTC). /// /// Broken down into: -/// * Days since epoch +/// * Days since epoch (max range is "jiff"'s range) /// * Seconds within day (0-86399) /// * Subseconds (range depends on `TimeUnit`) pub struct TimestampParts { - pub days: i64, - pub seconds: i64, - pub subseconds: i64, + pub days: i32, + pub seconds: i32, + pub subseconds: i32, } /// Splits a Unix timestamp into its component parts. @@ -31,22 +32,29 @@ pub struct TimestampParts { /// /// Returns an error if `time_unit` is days, which cannot be split. pub fn split(timestamp: i64, time_unit: TimeUnit) -> VortexResult { - let divisor = match time_unit { - TimeUnit::Nanoseconds => 1_000_000_000, - TimeUnit::Microseconds => 1_000_000, - TimeUnit::Milliseconds => 1_000, - TimeUnit::Seconds => 1, + Ok(match time_unit { + TimeUnit::Nanoseconds => split_with_divisor::<1_000_000_000>(timestamp), + TimeUnit::Microseconds => split_with_divisor::<1_000_000>(timestamp), + TimeUnit::Milliseconds => split_with_divisor::<1_000>(timestamp), + TimeUnit::Seconds => split_with_divisor::<1>(timestamp), TimeUnit::Days => vortex_bail!("Cannot handle day-level data"), - }; - - let ticks_per_day = SECONDS_PER_DAY * divisor; - Ok(TimestampParts { - days: timestamp / ticks_per_day, - seconds: (timestamp % ticks_per_day) / divisor, - subseconds: (timestamp % ticks_per_day) % divisor, }) } +/// Split a Unix timestamp into parts using a constant DIVISOR. +#[inline] +pub(crate) fn split_with_divisor(timestamp: i64) -> TimestampParts { + let days = timestamp / (SECONDS_PER_DAY * DIVISOR); + let total_seconds = timestamp / DIVISOR; + let seconds = total_seconds - days * SECONDS_PER_DAY; + let subseconds = timestamp - total_seconds * DIVISOR; + TimestampParts { + days: days.as_(), + seconds: seconds.as_(), + subseconds: subseconds.as_(), + } +} + /// Combines timestamp parts back into a Unix timestamp. /// /// # Arguments @@ -66,13 +74,48 @@ pub fn combine(ts_parts: TimestampParts, time_unit: TimeUnit) -> i64 { TimeUnit::Days => vortex_panic!("Cannot handle day-level data"), }; - ts_parts.days * SECONDS_PER_DAY * divisor + ts_parts.seconds * divisor + ts_parts.subseconds + ts_parts.days as i64 * SECONDS_PER_DAY * divisor + + ts_parts.seconds as i64 * divisor + + ts_parts.subseconds as i64 } #[cfg(test)] mod tests { + use rstest::rstest; + use super::*; + #[rstest] + #[case(TimeUnit::Seconds, 1)] + #[case(TimeUnit::Milliseconds, 1_000)] + #[case(TimeUnit::Microseconds, 1_000_000)] + #[case(TimeUnit::Nanoseconds, 1_000_000_000)] + fn test_split_divmod(#[case] unit: TimeUnit, #[case] divisor: i64) -> VortexResult<()> { + let ticks_per_day = SECONDS_PER_DAY * divisor; + for ts in [ + 0, + 1, + -1, + divisor - 1, + 3723 * divisor + divisor / 2, + -(3723 * divisor + divisor / 2), + ticks_per_day - 1, + ticks_per_day, + -ticks_per_day, + ] { + let parts = split(ts, unit)?; + + let days: i32 = (ts / ticks_per_day).as_(); + let seconds: i32 = ((ts % ticks_per_day) / divisor).as_(); + let subseconds: i32 = ((ts % ticks_per_day) % divisor).as_(); + assert_eq!(parts.days, days); + assert_eq!(parts.seconds, seconds); + assert_eq!(parts.subseconds, subseconds); + assert_eq!(combine(parts, unit), ts); + } + Ok(()) + } + #[test] fn test_split_seconds() { // 1970-01-02 01:02:03 UTC diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs index 303e9d3247b..5ae1bf0101e 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs @@ -10,7 +10,7 @@ use vortex_array::scalar_fn::fns::cast::CastReduce; use vortex_error::VortexResult; use crate::DecimalByteParts; -use crate::decimal_byte_parts::DecimalBytePartsArrayExt; +use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; impl CastReduce for DecimalByteParts { fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs index af6e98bddc2..3044bd6e605 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs @@ -25,7 +25,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::DecimalByteParts; -use crate::decimal_byte_parts::DecimalBytePartsArrayExt; +use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; use crate::decimal_byte_parts::compute::compare::Sign::Positive; impl CompareKernel for DecimalByteParts { diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs index a1b50ad788a..a47a6ed846b 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs @@ -10,7 +10,7 @@ use vortex_error::VortexResult; use vortex_mask::Mask; use crate::DecimalByteParts; -use crate::decimal_byte_parts::DecimalBytePartsArrayExt; +use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; impl FilterReduce for DecimalByteParts { fn filter(array: ArrayView<'_, Self>, mask: &Mask) -> VortexResult> { DecimalByteParts::try_new( diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs index f4da528a9b3..065bc5e0051 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs @@ -11,7 +11,7 @@ use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use crate::DecimalByteParts; -use crate::decimal_byte_parts::DecimalBytePartsArrayExt; +use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; /// DecimalByteParts-specific is_constant kernel. /// diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs index 0f007519894..e7dc95af84f 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs @@ -4,23 +4,17 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::IntoArray; -use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; -use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::fns::mask::Mask as MaskExpr; use vortex_array::scalar_fn::fns::mask::MaskReduce; use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::DecimalByteParts; -use crate::decimal_byte_parts::DecimalBytePartsArrayExt; +use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; impl MaskReduce for DecimalByteParts { fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult> { - let masked_msp = MaskExpr.try_new_array( - array.msp().len(), - EmptyOptions, - [array.msp().clone(), mask.clone()], - )?; + let masked_msp = MaskExpr::try_new(array.msp().clone(), mask.clone())?.into_array(); Ok(Some( DecimalByteParts::try_new( masked_msp, diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs index 0a5477e93cf..6c2d0dabb31 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs @@ -16,6 +16,7 @@ mod tests { use vortex_array::VortexSessionExecute; use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; + use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array; use vortex_array::compute::conformance::consistency::test_array_consistency; use vortex_array::dtype::DecimalDType; use vortex_buffer::buffer; @@ -74,4 +75,22 @@ mod tests { let ctx = &mut array_session().create_execution_ctx(); test_array_consistency(&array.into_array(), ctx); } + + #[rstest] + #[case::decimal_i32(DecimalByteParts::try_new( + buffer![100i32, 200, 300, 400, 500].into_array(), + DecimalDType::new(10, 2) + ).unwrap())] + #[case::decimal_nullable_i64(DecimalByteParts::try_new( + PrimitiveArray::from_option_iter([Some(1000i64), None, Some(3000), Some(4000), None]).into_array(), + DecimalDType::new(19, 4) + ).unwrap())] + #[case::decimal_negative(DecimalByteParts::try_new( + buffer![-100i32, -200, 300, -400, 500].into_array(), + DecimalDType::new(10, 2) + ).unwrap())] + fn test_decimal_byte_parts_binary_numeric(#[case] array: DecimalBytePartsArray) { + let ctx = &mut array_session().create_execution_ctx(); + test_binary_numeric_array(&array.into_array(), ctx); + } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs index 807b660ec73..7a18f7bf91b 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs @@ -10,7 +10,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::DecimalByteParts; -use crate::decimal_byte_parts::DecimalBytePartsArrayExt; +use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; impl TakeExecute for DecimalByteParts { fn take( diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs index 864cd604885..d5b0024f5b7 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -21,7 +21,7 @@ use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::IntoArray; -use vortex_array::TypedArrayRef; +use vortex_array::array_slots; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::buffer::BufferHandle; @@ -90,9 +90,7 @@ impl VTable for DecimalByteParts { let Some(decimal_dtype) = dtype.as_decimal_opt() else { vortex_bail!("expected decimal dtype, got {}", dtype) }; - let msp = slots[MSP_SLOT] - .as_ref() - .vortex_expect("DecimalBytePartsArray msp slot"); + let msp = DecimalBytePartsSlotsView::from_slots(slots).msp; DecimalBytePartsData::validate(msp, *decimal_dtype, dtype, len) } @@ -158,7 +156,7 @@ impl VTable for DecimalByteParts { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + DecimalBytePartsSlots::NAMES[idx].to_string() } fn reduce_parent( @@ -174,10 +172,12 @@ impl VTable for DecimalByteParts { } } -/// The most significant parts of the decimal values. -pub(super) const MSP_SLOT: usize = 0; -pub(super) const NUM_SLOTS: usize = 1; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["msp"]; +#[array_slots(DecimalByteParts)] +pub struct DecimalBytePartsSlots { + /// The most significant parts of the decimal values. + #[slot(0)] + pub msp: ArrayRef, +} /// This array encodes decimals as between 1-4 columns of primitive typed children. /// The most significant part (msp) sorting the most significant decimal bits. @@ -202,16 +202,6 @@ pub struct DecimalBytePartsDataParts { pub msp: ArrayRef, } -pub trait DecimalBytePartsArrayExt: TypedArrayRef { - fn msp(&self) -> &ArrayRef { - self.as_ref().slots()[MSP_SLOT] - .as_ref() - .vortex_expect("DecimalBytePartsArray msp slot") - } -} - -impl> DecimalBytePartsArrayExt for T {} - impl DecimalBytePartsData { pub fn validate( msp: &ArrayRef, diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/rules.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/rules.rs index 7fafa0d5289..d4052a4bed8 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/rules.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/rules.rs @@ -15,7 +15,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::DecimalByteParts; -use crate::decimal_byte_parts::DecimalBytePartsArrayExt; +use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; pub(super) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&DecimalBytePartsFilterPushDownRule), diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/slice.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/slice.rs index a9c022d2051..14807421c73 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/slice.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/slice.rs @@ -11,7 +11,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::DecimalByteParts; -use crate::decimal_byte_parts::DecimalBytePartsArrayExt; +use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; impl SliceReduce for DecimalByteParts { fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { diff --git a/encodings/experimental/onpair/README.md b/encodings/experimental/onpair/README.md deleted file mode 100644 index 9628c006201..00000000000 --- a/encodings/experimental/onpair/README.md +++ /dev/null @@ -1,39 +0,0 @@ -# Vortex OnPair - -A Vortex Encoding for Binary and Utf8 data that uses the -[OnPair][onpair] short-string compression algorithm. OnPair is a -dictionary-based encoder with fast per-row random access. - -The trainer / encoder lives in the standalone [`onpair`][onpair-crate] -crate; this crate wraps the resulting column as a Vortex array with -cascading-compressor support on every integer child. - -## Compute - -Like the FSST encoding, this crate provides `cast` and `filter` -pushdown. Other operators fall back to ordinary decompression. - -## Default Configuration - -The default training preset is **dict-12**: 12 bits per token, -dictionary capped at 4 096 entries. Token codes are stored as a -`PrimitiveArray`; downstream `FastLanes::BitPacking` losslessly -narrows the child to exactly `bits`-bit codes on disk. - -## Layout - -- Buffer 0 — `dict_bytes`: dictionary blob built by the OnPair trainer, - padded with `MAX_TOKEN_SIZE` trailing zero bytes so the over-copy - decoder can read 16 bytes past the last token. -- Slot 0 — `dict_offsets`: `PrimitiveArray`, len `dict_size + 1`. -- Slot 1 — `codes`: `PrimitiveArray`, length `total_tokens`. -- Slot 2 — `codes_offsets`: `PrimitiveArray`, length `num_rows + 1`. -- Slot 3 — `uncompressed_lengths`: integer `PrimitiveArray`, length - `num_rows`. -- Slot 4 — optional validity child. - -All four integer slot children flow through the standard cascading -compressor pipeline (FoR / BitPacking / RunEnd / etc.). - -[onpair]: https://arxiv.org/abs/2508.02280 -[onpair-crate]: https://github.com/spiraldb/onpair diff --git a/encodings/experimental/onpair/goldenfiles/onpair.metadata b/encodings/experimental/onpair/goldenfiles/onpair.metadata deleted file mode 100644 index e96baf1a0ab..00000000000 --- a/encodings/experimental/onpair/goldenfiles/onpair.metadata +++ /dev/null @@ -1 +0,0 @@ -  (08 \ No newline at end of file diff --git a/encodings/experimental/onpair/src/canonical.rs b/encodings/experimental/onpair/src/canonical.rs deleted file mode 100644 index 20a058e7f90..00000000000 --- a/encodings/experimental/onpair/src/canonical.rs +++ /dev/null @@ -1,117 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors -// -//! Convert an [`OnPairArray`] to its canonical `VarBinViewArray` by handing -//! the materialised parts to `onpair::decompress_into`. -//! -//! [`OnPairArray`]: crate::OnPairArray - -use std::sync::Arc; - -use num_traits::AsPrimitive; -use onpair::Parts; -use vortex_array::ArrayRef; -use vortex_array::ArrayView; -use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::arrays::PrimitiveArray; -use vortex_array::arrays::VarBinViewArray; -use vortex_array::arrays::varbinview::build_views::BinaryView; -use vortex_array::arrays::varbinview::build_views::MAX_BUFFER_LEN; -use vortex_array::arrays::varbinview::build_views::build_views; -use vortex_array::match_each_integer_ptype; -use vortex_buffer::Buffer; -use vortex_buffer::ByteBuffer; -use vortex_buffer::ByteBufferMut; -use vortex_error::VortexResult; -use vortex_error::vortex_ensure; - -use crate::OnPair; -use crate::OnPairArraySlotsExt; -use crate::decode::code_boundary_at; -use crate::decode::collect_widened; - -pub(super) fn canonicalize_onpair( - array: ArrayView<'_, OnPair>, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let (buffers, views) = onpair_decode_views(array, 0, ctx)?; - let validity = array.array().validity()?; - Ok(unsafe { - VarBinViewArray::new_unchecked(views, Arc::from(buffers), array.dtype().clone(), validity) - .into_array() - }) -} - -pub(crate) fn onpair_decode_views( - array: ArrayView<'_, OnPair>, - start_buf_index: u32, - ctx: &mut ExecutionCtx, -) -> VortexResult<(Vec, Buffer)> { - let lengths = array - .uncompressed_lengths() - .clone() - .execute::(ctx)?; - - let total_size: usize = match_each_integer_ptype!(lengths.ptype(), |P| { - lengths - .as_slice::

() - .iter() - .map(|&l| AsPrimitive::::as_(l)) - .sum() - }); - - // `codes_offsets` holds the per-row code boundaries and may itself be a - // sliced or filtered view of the original. Its first and last entries - // bound the contiguous run of `codes` belonging to the rows present in - // this array: `slice` keeps the full `codes` child and only narrows - // `codes_offsets` (so `code_start > 0` and/or `code_end < codes.len()`), - // while `filter` rebuilds both children so the window is the whole stream. - // OnPair has no `TakeExecute`, so a reordering take is served from the - // canonical `VarBinView` and never reaches this path. We only need those - // two boundaries, so point-look them up rather than decoding every offset. - let codes_offsets = array.codes_offsets(); - let code_start = code_boundary_at(codes_offsets, 0, ctx)?; - let code_end = code_boundary_at(codes_offsets, array.len(), ctx)?; - vortex_ensure!( - code_start <= code_end, - "OnPair codes_offsets must be nondecreasing" - ); - vortex_ensure!( - code_end <= array.codes().len(), - "OnPair codes_offsets end {} exceeds codes len {}", - code_end, - array.codes().len() - ); - - // Slice the `codes` child to that window *before* unpacking it, so a sliced - // array materialises only its own codes rather than the whole column's. The - // contiguous decoder walks `codes` in order and never reads the per-row - // boundaries, so an empty boundary slice is sound. - let codes = collect_widened::(&array.codes().slice(code_start..code_end)?, ctx)?; - let dict_offsets = collect_widened::(array.dict_offsets(), ctx)?; - - let mut out_bytes = ByteBufferMut::with_capacity(total_size); - let written = onpair::decompress_into( - Parts { - dict_bytes: array.dict_bytes().as_slice(), - dict_offsets: dict_offsets.as_slice(), - bits: array.bits(), - codes: codes.as_slice(), - }, - out_bytes.spare_capacity_mut(), - ); - debug_assert_eq!(written, total_size); - // SAFETY: `decompress_into` initialised exactly `written` bytes of the - // spare capacity reserved above. - unsafe { out_bytes.set_len(written) }; - - match_each_integer_ptype!(lengths.ptype(), |P| { - Ok(build_views( - start_buf_index, - MAX_BUFFER_LEN, - out_bytes, - lengths.as_slice::

(), - )) - }) -} diff --git a/encodings/experimental/onpair/src/compress.rs b/encodings/experimental/onpair/src/compress.rs deleted file mode 100644 index 3616ec7e717..00000000000 --- a/encodings/experimental/onpair/src/compress.rs +++ /dev/null @@ -1,193 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Train + compress entry points for the OnPair encoding. - -use onpair::Config; -use onpair::Offset; -use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::arrays::ConstantArray; -use vortex_array::arrays::VarBinViewArray; -use vortex_array::arrays::varbinview::BinaryView; -use vortex_array::buffer::BufferHandle; -use vortex_array::validity::Validity; -use vortex_buffer::Alignment; -use vortex_buffer::Buffer; -use vortex_buffer::BufferMut; -use vortex_buffer::ByteBuffer; -use vortex_buffer::ByteBufferMut; -use vortex_error::VortexExpect; -use vortex_error::VortexResult; -use vortex_error::vortex_err; -use vortex_mask::AllOr; - -use crate::OnPair; -use crate::OnPairArray; - -/// Default OnPair training configuration: 12-bit codes ("dict-12"). -pub const DEFAULT_DICT12_CONFIG: Config = Config { - seed: Some(42), - ..onpair::DEFAULT_CONFIG -}; - -fn onpair_compress_varbinview( - array: VarBinViewArray, - config: Config, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - O: Offset, -{ - let len = array.len(); - let mask = array.validity()?.execute_mask(len, ctx)?; - if mask.all_false() { - return OnPair::try_new( - array.dtype().clone(), - BufferHandle::new_host(ByteBuffer::empty()), - ConstantArray::new(0, len).into_array(), - ConstantArray::new(0u16, len).into_array(), - ConstantArray::new(0u32, len + 1).into_array(), - ConstantArray::new(0i32, len).into_array(), - Validity::AllInvalid, - 9, - ); - } - - let mut flat: Vec = Vec::with_capacity(len * 16); - let mut offsets: Vec = Vec::with_capacity(len + 1); - let mut uncompressed_lengths: BufferMut = BufferMut::with_capacity(len); - offsets.push(O::from_usize(0)); - let views = array.views(); - let buffers = array - .data_buffers() - .as_ref() - .iter() - .map(|b| b.as_host()) - .collect::>(); - - match mask.bit_buffer() { - AllOr::All => { - for view in views { - let bytes = view_bytes(view, &buffers); - flat.extend_from_slice(bytes); - offsets.push(O::from_usize(flat.len())); - uncompressed_lengths - .push(i32::try_from(view.len()).vortex_expect("must fit in i32")); - } - } - AllOr::None => { - unreachable!("all_false() should have been caught earlier"); - } - AllOr::Some(validity) => { - for (view, valid) in views.iter().zip(validity.iter()) { - if valid { - let bytes = view_bytes(view, &buffers); - flat.extend_from_slice(bytes); - offsets.push(O::from_usize(flat.len())); - uncompressed_lengths - .push(i32::try_from(view.len()).vortex_expect("must fit in i32")); - } else { - offsets.push(O::from_usize(flat.len())); - uncompressed_lengths.push(0); - } - } - } - } - - let column = onpair::compress(&flat, &offsets, config) - .map_err(|e| vortex_err!("OnPair compress failed: {e}"))?; - let bits = column.bits; - let dict_bytes = dict_bytes_to_buffer(column.dict_bytes); - let codes_offsets = - build_codes_offsets(&column.codes, &column.dict_offsets, &offsets)?.into_array(); - let codes = Buffer::from(column.codes).into_array(); - let dict_offsets = Buffer::from(column.dict_offsets).into_array(); - - let uncompressed_lengths = uncompressed_lengths.into_array(); - - OnPair::try_new( - array.dtype().clone(), - dict_bytes, - dict_offsets, - codes, - codes_offsets, - uncompressed_lengths, - array.validity()?, - bits, - ) -} - -fn view_bytes<'a>(view: &'a BinaryView, buffers: &'a [&ByteBuffer]) -> &'a [u8] { - if view.is_inlined() { - view.as_inlined().value() - } else { - let view_ref = view.as_view(); - &buffers[view_ref.buffer_index as usize][view_ref.as_range()] - } -} - -/// Lift compressed dictionary bytes into the Vortex buffer slot. -fn dict_bytes_to_buffer(dict_bytes: Vec) -> BufferHandle { - // Pad the dictionary blob with MAX_TOKEN_SIZE zero bytes so the - // over-copy decoder can issue a fixed 16-byte load for every token - // without risking an OOB read on the last entry. - // - // Align dict_bytes to 8 bytes so the segment that ultimately holds the - // OnPair tree starts at an 8-aligned in-memory address. Without this - // anchor, the per-buffer padding the serializer inserts is only - // *relative* to the segment start; if the segment lands at a u8-aligned - // heap address, downstream `PrimitiveArray::deserialize` panics - // with `Misaligned buffer cannot be used to build PrimitiveArray of u32`. - let mut padded = ByteBufferMut::with_capacity_aligned( - dict_bytes.len() + onpair::MAX_TOKEN_SIZE, - Alignment::new(8), - ); - padded.extend_from_slice(&dict_bytes); - unsafe { padded.push_n_unchecked(0, dict_bytes.len() + onpair::MAX_TOKEN_SIZE - padded.len()) }; - BufferHandle::new_host(padded.freeze()) -} - -/// Reconstruct the per-row `codes_offsets` from the flat `codes`, the -/// dictionary `dict_offsets` (token byte lengths) and the per-row decoded byte -/// boundaries. Returns `nrows + 1` cumulative code counts (`u32`). -// TODO(joe): can we compute this while compressing the array, yes but a worse API. -fn build_codes_offsets( - codes: &[u16], - dict_offsets: &[u32], - row_byte_offsets: &[O], -) -> VortexResult> { - let nrows = row_byte_offsets.len() - 1; - let mut codes_offsets = BufferMut::with_capacity(nrows + 1); - codes_offsets.push(0u32); - let mut decoded_bytes: u64 = 0; - let mut code_idx: usize = 0; - for r in 0..nrows { - let target = row_byte_offsets[r + 1] - .to_usize() - .ok_or_else(|| vortex_err!("OnPair row byte offset does not fit usize"))? - as u64; - while decoded_bytes < target { - let code = codes[code_idx] as usize; - decoded_bytes += u64::from(dict_offsets[code + 1] - dict_offsets[code]); - code_idx += 1; - } - codes_offsets.push( - u32::try_from(code_idx) - .map_err(|_| vortex_err!("OnPair: code boundary {code_idx} does not fit u32"))?, - ); - } - Ok(codes_offsets.freeze()) -} - -/// Compress any [`ArrayRef`] whose canonical form is a string array, by first -/// canonicalising to `VarBinViewArray`. -pub fn onpair_compress( - array: &ArrayRef, - config: Config, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let view = array.clone().execute::(ctx)?; - onpair_compress_varbinview::(view, config, ctx) -} diff --git a/encodings/experimental/onpair/src/compute/cast.rs b/encodings/experimental/onpair/src/compute/cast.rs deleted file mode 100644 index 93e1fdd8f8a..00000000000 --- a/encodings/experimental/onpair/src/compute/cast.rs +++ /dev/null @@ -1,43 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use vortex_array::ArrayRef; -use vortex_array::ArrayView; -use vortex_array::IntoArray; -use vortex_array::dtype::DType; -use vortex_array::scalar_fn::fns::cast::CastReduce; -use vortex_error::VortexResult; - -use crate::OnPair; -use crate::OnPairArraySlotsExt; - -/// Cast between `Utf8` and `Binary` (or adjust nullability) without touching -/// any of the encoded payload — we only rewrap into a new outer DType. -impl CastReduce for OnPair { - fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { - if !array.dtype().eq_ignore_nullability(dtype) { - return Ok(None); - } - let validity = array.array().validity()?; - let Some(new_validity) = - validity.trivially_cast_nullability(dtype.nullability(), array.array().len())? - else { - return Ok(None); - }; - Ok(Some( - unsafe { - OnPair::new_unchecked( - dtype.clone(), - array.dict_bytes_handle().clone(), - array.dict_offsets().clone(), - array.codes().clone(), - array.codes_offsets().clone(), - array.uncompressed_lengths().clone(), - new_validity, - array.bits(), - ) - } - .into_array(), - )) - } -} diff --git a/encodings/experimental/onpair/src/compute/compare.rs b/encodings/experimental/onpair/src/compute/compare.rs deleted file mode 100644 index 939a1a2fb42..00000000000 --- a/encodings/experimental/onpair/src/compute/compare.rs +++ /dev/null @@ -1,145 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use vortex_array::ArrayRef; -use vortex_array::ArrayView; -use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::arrays::BoolArray; -use vortex_array::arrays::ConstantArray; -use vortex_array::builtins::ArrayBuiltins; -use vortex_array::dtype::DType; -use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::binary::CompareKernel; -use vortex_array::scalar_fn::fns::operators::CompareOperator; -use vortex_buffer::BitBuffer; -use vortex_error::VortexResult; - -use crate::OnPair; -use crate::OnPairArraySlotsExt; - -impl CompareKernel for OnPair { - fn compare( - lhs: ArrayView<'_, Self>, - rhs: &ArrayRef, - operator: CompareOperator, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - let Some(constant) = rhs.as_constant() else { - return Ok(None); - }; - let is_empty = match constant.dtype() { - DType::Utf8(_) => constant.as_utf8().is_empty(), - DType::Binary(_) => constant.as_binary().is_empty(), - _ => return Ok(None), - }; - if is_empty != Some(true) { - return Ok(None); - } - - let lengths = lhs.uncompressed_lengths(); - let buffer = match operator { - // every value is greater than an empty string - CompareOperator::Gte => BitBuffer::new_set(lhs.len()), - // no value is less than an empty string - CompareOperator::Lt => BitBuffer::new_unset(lhs.len()), - _ => lengths - .binary( - ConstantArray::new(Scalar::zero_value(lengths.dtype()), lengths.len()) - .into_array(), - operator.into(), - )? - .execute(ctx)?, - }; - Ok(Some( - BoolArray::new( - buffer, - lhs.validity()? - .union_nullability(constant.dtype().nullability()), - ) - .into_array(), - )) - } -} - -#[cfg(test)] -mod tests { - use std::sync::LazyLock; - - use rstest::rstest; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::array_session; - use vortex_array::arrays::BoolArray; - use vortex_array::arrays::ConstantArray; - use vortex_array::arrays::VarBinArray; - use vortex_array::assert_arrays_eq; - use vortex_array::builtins::ArrayBuiltins; - use vortex_array::dtype::DType; - use vortex_array::dtype::Nullability; - use vortex_array::scalar::Scalar; - use vortex_array::scalar_fn::fns::operators::Operator; - use vortex_error::VortexResult; - use vortex_session::VortexSession; - - use crate::compress::DEFAULT_DICT12_CONFIG; - use crate::compress::onpair_compress; - - static SESSION: LazyLock = LazyLock::new(array_session); - - #[cfg_attr(miri, ignore)] - #[rstest] - #[case(Operator::Eq, [true, false, true, false])] - #[case(Operator::NotEq, [false, true, false, true])] - #[case(Operator::Gt, [false, true, false, true])] - #[case(Operator::Gte, [true, true, true, true])] - #[case(Operator::Lt, [false, false, false, false])] - #[case(Operator::Lte, [true, false, true, false])] - fn compare_empty_string(#[case] op: Operator, #[case] expected: [bool; 4]) -> VortexResult<()> { - let input = VarBinArray::from_iter( - [Some(""), Some("a"), Some(""), Some("bbb")], - DType::Utf8(Nullability::NonNullable), - ); - let mut ctx = SESSION.create_execution_ctx(); - let arr = onpair_compress(input.as_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?.into_array(); - - let result = arr - .binary(ConstantArray::new("", input.len()).into_array(), op)? - .execute::(&mut ctx)?; - assert_arrays_eq!(&result, &BoolArray::from_iter(expected), &mut ctx); - Ok(()) - } - - #[cfg_attr(miri, ignore)] - #[test] - fn compare_empty_string_nullable() -> VortexResult<()> { - let input = VarBinArray::from_iter( - [Some(""), None, Some("x")], - DType::Utf8(Nullability::Nullable), - ); - let mut ctx = SESSION.create_execution_ctx(); - let arr = onpair_compress(input.as_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?.into_array(); - - let eq_empty = arr - .clone() - .binary(ConstantArray::new("", arr.len()).into_array(), Operator::Eq)? - .execute::(&mut ctx)?; - assert_arrays_eq!( - &eq_empty, - &BoolArray::from_iter([Some(true), None, Some(false)]), - &mut ctx - ); - - let null_rhs = - ConstantArray::new(Scalar::null(DType::Utf8(Nullability::Nullable)), arr.len()); - let eq_null = arr - .binary(null_rhs.into_array(), Operator::Eq)? - .execute::(&mut ctx)?; - assert_arrays_eq!( - &eq_null, - &BoolArray::from_iter([None::, None, None]), - &mut ctx - ); - Ok(()) - } -} diff --git a/encodings/experimental/onpair/src/decode.rs b/encodings/experimental/onpair/src/decode.rs deleted file mode 100644 index fd0b02b0284..00000000000 --- a/encodings/experimental/onpair/src/decode.rs +++ /dev/null @@ -1,43 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors -// -//! Helpers for turning [`OnPair`](super::OnPair) slot children into the inputs the upstream -//! `onpair` decoder consumes. - -use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; -use vortex_array::arrays::PrimitiveArray; -use vortex_array::builtins::ArrayBuiltins; -use vortex_array::dtype::DType; -use vortex_array::dtype::NativePType; -use vortex_buffer::Buffer; -use vortex_error::VortexResult; -use vortex_error::vortex_err; - -/// Canonicalise a slot child to the decoder's native primitive width. -pub(crate) fn collect_widened( - arr: &ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult> { - let dtype = DType::Primitive(T::PTYPE, arr.dtype().nullability()); - Ok(arr - .cast(dtype)? - .execute::(ctx)? - .into_buffer::()) -} - -/// Read one `codes_offsets` boundary by point lookup. This decodes at most a -/// single chunk of the child — never the whole per-row offsets array — so the -/// callers that only need a row window (`scalar_at`, the canonical decode's -/// start/end bounds) don't pay to materialise every boundary. -pub(crate) fn code_boundary_at( - codes_offsets: &ArrayRef, - index: usize, - ctx: &mut ExecutionCtx, -) -> VortexResult { - codes_offsets - .execute_scalar(index, ctx)? - .as_primitive() - .as_::() - .ok_or_else(|| vortex_err!("OnPair codes_offsets[{index}] is null")) -} diff --git a/encodings/fastlanes/Cargo.toml b/encodings/fastlanes/Cargo.toml index 4a187f8faae..9085390b67b 100644 --- a/encodings/fastlanes/Cargo.toml +++ b/encodings/fastlanes/Cargo.toml @@ -58,11 +58,6 @@ name = "compute_between" harness = false required-features = ["_test-harness"] -[[bench]] -name = "bit_transpose" -harness = false -required-features = ["_test-harness"] - [[bench]] name = "bitpack_compare" harness = false diff --git a/encodings/fastlanes/benches/bit_transpose.rs b/encodings/fastlanes/benches/bit_transpose.rs deleted file mode 100644 index 08c3ffb12e5..00000000000 --- a/encodings/fastlanes/benches/bit_transpose.rs +++ /dev/null @@ -1,310 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use divan::Bencher; -use vortex_fastlanes::bit_transpose::scalar::transpose_bits_scalar; -use vortex_fastlanes::bit_transpose::scalar::untranspose_bits_scalar; - -fn main() { - divan::main(); -} - -/// Generate deterministic test data. -#[expect(clippy::cast_possible_truncation)] -fn generate_test_data(seed: usize) -> [u8; 128] { - let mut data = [0u8; 128]; - for (i, byte) in data.iter_mut().enumerate() { - *byte = seed.wrapping_mul(17).wrapping_add(i).wrapping_mul(31) as u8; - } - data -} - -const BATCH_SIZE: usize = 1000; - -// ============================================================================ -// Transpose: single array -// ============================================================================ - -#[divan::bench] -fn transpose_scalar(bencher: Bencher) { - let input = generate_test_data(42); - - bencher - .with_inputs(|| (&input, [0u8; 128])) - .bench_refs(|(input, output)| { - transpose_bits_scalar(input, output); - }); -} - -// ============================================================================ -// Transpose: throughput (1000 arrays) -// ============================================================================ - -#[divan::bench] -fn transpose_scalar_throughput(bencher: Bencher) { - let inputs: Vec<[u8; 128]> = (0..BATCH_SIZE).map(generate_test_data).collect(); - - bencher - .with_inputs(|| (&inputs, vec![[0u8; 128]; BATCH_SIZE])) - .bench_refs(|(inputs, outputs)| { - for (input, output) in inputs.iter().zip(outputs.iter_mut()) { - transpose_bits_scalar(input, output); - } - }); -} - -// ============================================================================ -// Untranspose: single array -// ============================================================================ - -#[divan::bench] -fn untranspose_scalar(bencher: Bencher) { - let input = generate_test_data(42); - - bencher - .with_inputs(|| (&input, [0u8; 128])) - .bench_refs(|(input, output)| { - untranspose_bits_scalar(input, output); - }); -} - -// ============================================================================ -// Untranspose: throughput (1000 arrays) -// ============================================================================ - -#[divan::bench] -fn untranspose_scalar_throughput(bencher: Bencher) { - let inputs: Vec<[u8; 128]> = (0..BATCH_SIZE).map(generate_test_data).collect(); - - bencher - .with_inputs(|| (&inputs, vec![[0u8; 128]; BATCH_SIZE])) - .bench_refs(|(inputs, outputs)| { - for (input, output) in inputs.iter().zip(outputs.iter_mut()) { - untranspose_bits_scalar(input, output); - } - }); -} - -// ============================================================================ -// x86_64 benchmarks -// ============================================================================ - -#[cfg(target_arch = "x86_64")] -mod x86 { - use divan::Bencher; - use vortex_fastlanes::bit_transpose::x86::has_bmi2; - use vortex_fastlanes::bit_transpose::x86::has_vbmi; - use vortex_fastlanes::bit_transpose::x86::transpose_bits_bmi2; - use vortex_fastlanes::bit_transpose::x86::transpose_bits_vbmi; - use vortex_fastlanes::bit_transpose::x86::untranspose_bits_bmi2; - use vortex_fastlanes::bit_transpose::x86::untranspose_bits_vbmi; - - use super::BATCH_SIZE; - use super::generate_test_data; - - // --- Transpose: single array --- - - #[divan::bench] - fn transpose_bmi2(bencher: Bencher) { - if !has_bmi2() { - return; - } - - let input = generate_test_data(42); - - bencher - .with_inputs(|| (&input, [0u8; 128])) - .bench_refs(|(input, output)| { - unsafe { transpose_bits_bmi2(input, output) }; - }); - } - - #[divan::bench] - fn transpose_vbmi(bencher: Bencher) { - if !has_vbmi() { - return; - } - - let input = generate_test_data(42); - - bencher - .with_inputs(|| (&input, [0u8; 128])) - .bench_refs(|(input, output)| { - unsafe { transpose_bits_vbmi(input, output) }; - }); - } - - // --- Untranspose: single array --- - - #[divan::bench] - fn untranspose_bmi2(bencher: Bencher) { - if !has_bmi2() { - return; - } - - let input = generate_test_data(42); - - bencher - .with_inputs(|| (&input, [0u8; 128])) - .bench_refs(|(input, output)| { - unsafe { untranspose_bits_bmi2(input, output) }; - }); - } - - #[divan::bench] - fn untranspose_vbmi(bencher: Bencher) { - if !has_vbmi() { - return; - } - - let input = generate_test_data(42); - - bencher - .with_inputs(|| (&input, [0u8; 128])) - .bench_refs(|(input, output)| { - unsafe { untranspose_bits_vbmi(input, output) }; - }); - } - - // --- Transpose: throughput (1000 arrays) --- - - #[divan::bench] - fn transpose_bmi2_throughput(bencher: Bencher) { - if !has_bmi2() { - return; - } - - let inputs: Vec<[u8; 128]> = (0..BATCH_SIZE).map(generate_test_data).collect(); - - bencher - .with_inputs(|| (&inputs, vec![[0u8; 128]; BATCH_SIZE])) - .bench_refs(|(inputs, outputs)| { - for (input, output) in inputs.iter().zip(outputs.iter_mut()) { - unsafe { transpose_bits_bmi2(input, output) }; - } - }); - } - - #[divan::bench] - fn transpose_vbmi_throughput(bencher: Bencher) { - if !has_vbmi() { - return; - } - - let inputs: Vec<[u8; 128]> = (0..BATCH_SIZE).map(generate_test_data).collect(); - - bencher - .with_inputs(|| (&inputs, vec![[0u8; 128]; BATCH_SIZE])) - .bench_refs(|(inputs, outputs)| { - for (input, output) in inputs.iter().zip(outputs.iter_mut()) { - unsafe { transpose_bits_vbmi(input, output) }; - } - }); - } - - // --- Untranspose: throughput (1000 arrays) --- - - #[divan::bench] - fn untranspose_bmi2_throughput(bencher: Bencher) { - if !has_bmi2() { - return; - } - - let inputs: Vec<[u8; 128]> = (0..BATCH_SIZE).map(generate_test_data).collect(); - - bencher - .with_inputs(|| (&inputs, vec![[0u8; 128]; BATCH_SIZE])) - .bench_refs(|(inputs, outputs)| { - for (input, output) in inputs.iter().zip(outputs.iter_mut()) { - unsafe { untranspose_bits_bmi2(input, output) }; - } - }); - } - - #[divan::bench] - fn untranspose_vbmi_throughput(bencher: Bencher) { - if !has_vbmi() { - return; - } - - let inputs: Vec<[u8; 128]> = (0..BATCH_SIZE).map(generate_test_data).collect(); - - bencher - .with_inputs(|| (&inputs, vec![[0u8; 128]; BATCH_SIZE])) - .bench_refs(|(inputs, outputs)| { - for (input, output) in inputs.iter().zip(outputs.iter_mut()) { - unsafe { untranspose_bits_vbmi(input, output) }; - } - }); - } -} - -// ============================================================================ -// aarch64 benchmarks -// ============================================================================ - -#[cfg(target_arch = "aarch64")] -mod aarch64 { - use vortex_fastlanes::bit_transpose::aarch64::transpose_bits_neon; - use vortex_fastlanes::bit_transpose::aarch64::untranspose_bits_neon; - - use super::BATCH_SIZE; - use super::Bencher; - use super::generate_test_data; - - // --- Transpose: single array --- - - #[divan::bench] - fn transpose_neon(bencher: Bencher) { - let input = generate_test_data(42); - - bencher - .with_inputs(|| (&input, [0u8; 128])) - .bench_refs(|(input, output)| { - unsafe { transpose_bits_neon(input, output) }; - }); - } - - // --- Untranspose: single array --- - - #[divan::bench] - fn untranspose_neon(bencher: Bencher) { - let input = generate_test_data(42); - - bencher - .with_inputs(|| (&input, [0u8; 128])) - .bench_refs(|(input, output)| { - unsafe { untranspose_bits_neon(input, output) }; - }); - } - - // --- Transpose: throughput (1000 arrays) --- - - #[divan::bench] - fn transpose_neon_throughput(bencher: Bencher) { - let inputs: Vec<[u8; 128]> = (0..BATCH_SIZE).map(generate_test_data).collect(); - - bencher - .with_inputs(|| (&inputs, vec![[0u8; 128]; BATCH_SIZE])) - .bench_refs(|(inputs, outputs)| { - for (input, output) in inputs.iter().zip(outputs.iter_mut()) { - unsafe { transpose_bits_neon(input, output) }; - } - }); - } - - // --- Untranspose: throughput (1000 arrays) --- - - #[divan::bench] - fn untranspose_neon_throughput(bencher: Bencher) { - let inputs: Vec<[u8; 128]> = (0..BATCH_SIZE).map(generate_test_data).collect(); - - bencher - .with_inputs(|| (&inputs, vec![[0u8; 128]; BATCH_SIZE])) - .bench_refs(|(inputs, outputs)| { - for (input, output) in inputs.iter().zip(outputs.iter_mut()) { - unsafe { untranspose_bits_neon(input, output) }; - } - }); - } -} diff --git a/encodings/fastlanes/benches/canonicalize_bench.rs b/encodings/fastlanes/benches/canonicalize_bench.rs index 30ee0756960..7d8d7abb659 100644 --- a/encodings/fastlanes/benches/canonicalize_bench.rs +++ b/encodings/fastlanes/benches/canonicalize_bench.rs @@ -95,11 +95,13 @@ fn canonical_into_non_nullable( bencher .with_inputs(|| { let chunked = ChunkedArray::from_iter(chunks.clone()).into_array(); - let primitive_builder = PrimitiveBuilder::::with_capacity( + let ctx = SESSION.create_execution_ctx(); + let primitive_builder = PrimitiveBuilder::::with_capacity_in( chunked.dtype().nullability(), chunk_len * chunk_count, + ctx.allocator(), ); - (chunked, primitive_builder, SESSION.create_execution_ctx()) + (chunked, primitive_builder, ctx) }) .bench_refs(|(chunked, primitive_builder, ctx)| { chunked @@ -174,11 +176,13 @@ fn canonical_into_nullable( bencher .with_inputs(|| { let chunked = ChunkedArray::from_iter(chunks.clone()).into_array(); - let primitive_builder = PrimitiveBuilder::::with_capacity( + let ctx = SESSION.create_execution_ctx(); + let primitive_builder = PrimitiveBuilder::::with_capacity_in( chunked.dtype().nullability(), chunk_len * chunk_count, + ctx.allocator(), ); - (chunked, primitive_builder, SESSION.create_execution_ctx()) + (chunked, primitive_builder, ctx) }) .bench_refs(|(chunked, primitive_builder, ctx)| { chunked diff --git a/encodings/fastlanes/benches/cast_bitpacked.rs b/encodings/fastlanes/benches/cast_bitpacked.rs index baeff20f5e1..62c6c5533c3 100644 --- a/encodings/fastlanes/benches/cast_bitpacked.rs +++ b/encodings/fastlanes/benches/cast_bitpacked.rs @@ -50,10 +50,10 @@ const U32: DType = DType::Primitive(PType::U32, Nullability::NonNullable); const ARGS: &[(usize, usize, f64)] = &[ (65_536, 1, 0.00), (65_536, 1, 0.01), - (65_536, 4, 0.00), - (65_536, 4, 0.01), - (262_144, 1, 0.00), - (262_144, 1, 0.01), + (4_096, 4, 0.00), + (4_096, 4, 0.01), + (81_920, 1, 0.00), + (81_920, 1, 0.01), ]; /// Build a single bit-packed `u16` chunk. Most values fit in `bit_width` bits; `fraction_patched` diff --git a/encodings/fastlanes/src/bit_transpose.rs b/encodings/fastlanes/src/bit_transpose.rs new file mode 100644 index 00000000000..b9761ee20cf --- /dev/null +++ b/encodings/fastlanes/src/bit_transpose.rs @@ -0,0 +1,223 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::mem; +use std::mem::MaybeUninit; + +use vortex_buffer::Alignment; +use vortex_buffer::BitBuffer; +use vortex_buffer::BufferMut; +use vortex_buffer::ByteBuffer; +use vortex_error::VortexExpect; + +use crate::FL_CHUNK_SIZE; + +/// Transposes `bits` into FastLanes order, padding the result to whole 1,024-bit chunks. +/// +/// Bit `i` of the result holds logical bit [`fastlanes::transpose(i)`](fastlanes::transpose), the +/// same element order [`fastlanes::Transpose::transpose`] gives a value vector. To address the +/// result by logical index instead, invert that with `crate::untranspose_idx`. +pub fn transpose_bitbuffer(bits: BitBuffer) -> BitBuffer { + bits_op(bits, fastlanes::transpose_bits::) +} + +/// Untransposes whole 1,024-bit FastLanes chunks back into sequential bit order. +pub fn untranspose_bitbuffer(bits: BitBuffer) -> BitBuffer { + assert!( + bits.len().is_multiple_of(FL_CHUNK_SIZE), + "Transposed BitBuffer length must be a multiple of {FL_CHUNK_SIZE}" + ); + bits_op(bits, fastlanes::untranspose_bits) +} + +fn bits_op(bits: BitBuffer, op: F) -> BitBuffer { + // Normalize to offset 0 with no bytes outside the logical range, so that byte-buffer chunk + // boundaries line up with logical 1,024-bit chunks regardless of how the buffer was sliced. + let bits = bits.sliced(); + let (_offset, len, bytes) = bits.into_inner(); + + if len.is_multiple_of(FL_CHUNK_SIZE) && bytes.is_aligned(Alignment::of::()) { + match bytes.try_into_mut() { + Ok(mut bytes_mut) => { + let (chunks, _) = bytes_mut.as_chunks_mut::<128>(); + let mut tmp = [0u64; 16]; + for chunk in chunks { + // SAFETY: the buffer is u64-aligned and chunks start at 128-byte multiples. + let chunk_u64 = + unsafe { mem::transmute::<&mut [u8; 128], &mut [u64; 16]>(chunk) }; + op(chunk_u64, &mut tmp); + chunk_u64.copy_from_slice(&tmp); + } + BitBuffer::new(bytes_mut.freeze().into_byte_buffer(), len) + } + Err(bytes) => bits_op_with_copy(bytes, len, op), + } + } else { + bits_op_with_copy(bytes, len, op) + } +} + +fn bits_op_with_copy( + bytes: ByteBuffer, + len: usize, + op: F, +) -> BitBuffer { + let output_len = bytes.len().div_ceil(8).next_multiple_of(16); + let mut output = BufferMut::::with_capacity(output_len); + let (input_chunks, input_trailer) = bytes.as_chunks::<128>(); + // Bound to the requested `output_len`: `spare_capacity_mut` may expose extra over-aligned + // capacity, which would otherwise split into spurious trailing chunks and make `last_mut` + // below target a chunk past the data we actually initialize. + let (output_chunks, _) = unsafe { + mem::transmute::<&mut [MaybeUninit], &mut [u64]>( + &mut output.spare_capacity_mut()[..output_len], + ) + } + .as_chunks_mut::<16>(); + + for (input, output) in input_chunks.iter().zip(output_chunks.iter_mut()) { + op(&load_chunk_unaligned(input), output); + } + + if !input_trailer.is_empty() { + let mut padded_input = [0u8; 128]; + padded_input[0..input_trailer.len()].clone_from_slice(input_trailer); + op( + &load_chunk_unaligned(&padded_input), + output_chunks + .last_mut() + .vortex_expect("Output wasn't a multiple of 128 bytes"), + ); + } + + unsafe { output.set_len(output_len) }; + BitBuffer::new( + output.freeze().into_byte_buffer(), + len.next_multiple_of(FL_CHUNK_SIZE), + ) +} + +/// Loads a 128-byte chunk into a `[u64; 16]` without requiring the input to be 8-byte aligned. +/// +/// Native endianness is required: this must produce the same words as the aligned in-place path, +/// which reinterprets the buffer as host-endian `u64`s. +#[allow(clippy::host_endian_bytes)] +fn load_chunk_unaligned(chunk: &[u8; 128]) -> [u64; 16] { + let mut words = [0u64; 16]; + let (bytes, _) = chunk.as_chunks::<8>(); + for (word, bytes) in words.iter_mut().zip(bytes) { + *word = u64::from_ne_bytes(*bytes); + } + words +} + +#[cfg(test)] +mod tests { + use vortex_buffer::BitBuffer; + use vortex_buffer::BitBufferMut; + use vortex_buffer::ByteBuffer; + use vortex_buffer::ByteBufferMut; + + use super::*; + + fn make_validity_bits(num_bits: usize) -> BitBuffer { + let mut builder = BitBufferMut::with_capacity(num_bits); + for i in 0..num_bits { + builder.append(i % 3 != 0); + } + builder.freeze() + } + + fn force_copy_path(bits: BitBuffer) -> (BitBuffer, ByteBuffer) { + let (offset, len, bytes) = bits.into_inner(); + let extra_ref = bytes.clone(); + (BitBuffer::new_with_offset(bytes, len, offset), extra_ref) + } + + #[test] + fn transpose_padding_copy_produces_same_bits() { + let bits = make_validity_bits(500); + let transposed = transpose_bitbuffer(bits.clone()); + assert_eq!(transposed.len(), 1024); + let untransposed = untranspose_bitbuffer(transposed); + assert_eq!(untransposed.slice(0..500), bits) + } + + #[test] + fn transpose_inplace_and_copy_produce_same_bits() { + let bits = make_validity_bits(2048); + + let inplace_result = transpose_bitbuffer(bits.clone()); + + let (bits_shared, _hold) = force_copy_path(bits); + let copy_result = transpose_bitbuffer(bits_shared); + + assert_eq!(inplace_result.len(), copy_result.len()); + assert_eq!(inplace_result, copy_result); + } + + #[test] + fn transpose_bitbuffer_roundtrip_non_aligned() { + let original_len = 1500; + let bits = make_validity_bits(original_len); + + let transposed = transpose_bitbuffer(bits.clone()); + let roundtripped = untranspose_bitbuffer(transposed); + assert_eq!(bits, roundtripped.slice(0..original_len)); + } + + /// Regression: the copy path split the over-aligned spare capacity into extra 128-byte + /// chunks and wrote the padded remainder via `last_mut()`, which landed past the requested + /// length and left the real trailing chunk uninitialized. Whether the surplus capacity + /// produced an extra chunk depended on the allocation address, so we repeat across sizes to + /// defeat that luck; the fix bounds the spare slice so the result no longer depends on it. + #[test] + fn transpose_copy_path_survives_overallocation() { + for original_len in [129, 500, 1500, 9999] { + let bits = make_validity_bits(original_len); + for _ in 0..64 { + let transposed = transpose_bitbuffer(bits.clone()); + let roundtripped = untranspose_bitbuffer(transposed); + assert_eq!( + roundtripped.slice(0..original_len), + bits, + "len={original_len}" + ); + } + } + } + + /// Regression: buffers that are not 8-byte aligned (e.g. views into file segments) must not + /// panic and must produce the same result as the aligned path. + #[test] + fn untranspose_unaligned_buffer() { + let bits = make_validity_bits(2048); + let transposed = transpose_bitbuffer(bits.clone()); + let expected = untranspose_bitbuffer(transposed.clone()); + + // Rebuild the transposed bytes at an odd offset within a larger allocation so the + // resulting buffer is misaligned for u64. + let (_, _, transposed_bytes) = transposed.sliced().into_inner(); + let mut shifted = ByteBufferMut::with_capacity(transposed_bytes.len() + 1); + shifted.push(0xFF); + shifted.extend_from_slice(&transposed_bytes); + let misaligned = shifted.freeze().slice(1..transposed_bytes.len() + 1); + assert!(!misaligned.is_aligned(Alignment::of::())); + + let untransposed = untranspose_bitbuffer(BitBuffer::new(misaligned, 2048)); + assert_eq!(untransposed, expected); + assert_eq!(untransposed.slice(0..2048), bits); + } + + /// Regression: a bit-offset view of transposed chunks (a sliced validity) must untranspose + /// only the viewed chunks rather than assuming the buffer starts at a chunk boundary. + #[test] + fn untranspose_chunk_aligned_bit_offset() { + let bits = make_validity_bits(3 * FL_CHUNK_SIZE); + let transposed = transpose_bitbuffer(bits.clone()); + + let view = transposed.slice(FL_CHUNK_SIZE..3 * FL_CHUNK_SIZE); + let untransposed = untranspose_bitbuffer(view); + assert_eq!(untransposed, bits.slice(FL_CHUNK_SIZE..3 * FL_CHUNK_SIZE)); + } +} diff --git a/encodings/fastlanes/src/bit_transpose/aarch64.rs b/encodings/fastlanes/src/bit_transpose/aarch64.rs deleted file mode 100644 index 85dbb1e0690..00000000000 --- a/encodings/fastlanes/src/bit_transpose/aarch64.rs +++ /dev/null @@ -1,299 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#![cfg(target_arch = "aarch64")] - -use core::arch::aarch64::uint64x2_t; -use core::arch::aarch64::vandq_u64; -use core::arch::aarch64::vdupq_n_u64; -use core::arch::aarch64::veorq_u64; -use core::arch::aarch64::vgetq_lane_u64; -use core::arch::aarch64::vld1q_u8; -use core::arch::aarch64::vld1q_u8_x4; -use core::arch::aarch64::vorrq_u8; -use core::arch::aarch64::vqtbl4q_u8; -use core::arch::aarch64::vreinterpretq_u8_u64; -use core::arch::aarch64::vreinterpretq_u64_u8; -use core::arch::aarch64::vshlq_n_u64; -use core::arch::aarch64::vshrq_n_u64; -use core::arch::aarch64::vst1q_u8; - -use crate::bit_transpose::BASE_PATTERN_FIRST; -use crate::bit_transpose::BASE_PATTERN_SECOND; -use crate::bit_transpose::TRANSPOSE_2X2; -use crate::bit_transpose::TRANSPOSE_4X4; -use crate::bit_transpose::TRANSPOSE_8X8; - -/// Gather indices for the first half from input[0..64]. -/// Each group needs 4 bytes at stride 16 (the low half of the stride pattern). -/// Layout: [`g0_from_lo(4` bytes), pad(4 bytes), `g1_from_lo(4` bytes), pad(4 bytes), ...] -/// Two groups per 16-byte NEON register. -static GATHER_FIRST_LO: [[u8; 16]; 4] = [ - // Groups 0,1 from BASE_PATTERN_FIRST: bases 0, 8 - [ - 0, 16, 32, 48, 0xFF, 0xFF, 0xFF, 0xFF, 8, 24, 40, 56, 0xFF, 0xFF, 0xFF, 0xFF, - ], - // Groups 2,3: bases 4, 12 - [ - 4, 20, 36, 52, 0xFF, 0xFF, 0xFF, 0xFF, 12, 28, 44, 60, 0xFF, 0xFF, 0xFF, 0xFF, - ], - // Groups 4,5: bases 2, 10 - [ - 2, 18, 34, 50, 0xFF, 0xFF, 0xFF, 0xFF, 10, 26, 42, 58, 0xFF, 0xFF, 0xFF, 0xFF, - ], - // Groups 6,7: bases 6, 14 - [ - 6, 22, 38, 54, 0xFF, 0xFF, 0xFF, 0xFF, 14, 30, 46, 62, 0xFF, 0xFF, 0xFF, 0xFF, - ], -]; - -/// Gather indices for the first half from input[64..128]. -/// These fill in bytes 4-7 of each u64 (the high half of the stride pattern). -static GATHER_FIRST_HI: [[u8; 16]; 4] = [ - // Groups 0,1: bases 0, 8 (offset by -64 since table starts at input[64]) - [ - 0xFF, 0xFF, 0xFF, 0xFF, 0, 16, 32, 48, 0xFF, 0xFF, 0xFF, 0xFF, 8, 24, 40, 56, - ], - // Groups 2,3: bases 4, 12 - [ - 0xFF, 0xFF, 0xFF, 0xFF, 4, 20, 36, 52, 0xFF, 0xFF, 0xFF, 0xFF, 12, 28, 44, 60, - ], - // Groups 4,5: bases 2, 10 - [ - 0xFF, 0xFF, 0xFF, 0xFF, 2, 18, 34, 50, 0xFF, 0xFF, 0xFF, 0xFF, 10, 26, 42, 58, - ], - // Groups 6,7: bases 6, 14 - [ - 0xFF, 0xFF, 0xFF, 0xFF, 6, 22, 38, 54, 0xFF, 0xFF, 0xFF, 0xFF, 14, 30, 46, 62, - ], -]; - -/// Gather indices for the second half from input[0..64]. -/// Uses `BASE_PATTERN_SECOND`: bases [1, 9, 5, 13, 3, 11, 7, 15] -static GATHER_SECOND_LO: [[u8; 16]; 4] = [ - [ - 1, 17, 33, 49, 0xFF, 0xFF, 0xFF, 0xFF, 9, 25, 41, 57, 0xFF, 0xFF, 0xFF, 0xFF, - ], - [ - 5, 21, 37, 53, 0xFF, 0xFF, 0xFF, 0xFF, 13, 29, 45, 61, 0xFF, 0xFF, 0xFF, 0xFF, - ], - [ - 3, 19, 35, 51, 0xFF, 0xFF, 0xFF, 0xFF, 11, 27, 43, 59, 0xFF, 0xFF, 0xFF, 0xFF, - ], - [ - 7, 23, 39, 55, 0xFF, 0xFF, 0xFF, 0xFF, 15, 31, 47, 63, 0xFF, 0xFF, 0xFF, 0xFF, - ], -]; - -/// Gather indices for the second half from input[64..128]. -static GATHER_SECOND_HI: [[u8; 16]; 4] = [ - [ - 0xFF, 0xFF, 0xFF, 0xFF, 1, 17, 33, 49, 0xFF, 0xFF, 0xFF, 0xFF, 9, 25, 41, 57, - ], - [ - 0xFF, 0xFF, 0xFF, 0xFF, 5, 21, 37, 53, 0xFF, 0xFF, 0xFF, 0xFF, 13, 29, 45, 61, - ], - [ - 0xFF, 0xFF, 0xFF, 0xFF, 3, 19, 35, 51, 0xFF, 0xFF, 0xFF, 0xFF, 11, 27, 43, 59, - ], - [ - 0xFF, 0xFF, 0xFF, 0xFF, 7, 23, 39, 55, 0xFF, 0xFF, 0xFF, 0xFF, 15, 31, 47, 63, - ], -]; - -/// 8x8 byte transpose (scatter) permutation split into 4 × 16-byte chunks for NEON TBL. -/// Input layout: [g0b0..g0b7, g1b0..g1b7, ..., g7b0..g7b7] (64 bytes, group-major) -/// Output layout: [g0b0,g1b0,..,g7b0, g0b1,g1b1,..,g7b1, ...] (64 bytes, row-major) -/// Same permutation as x86 `SCATTER_8X8`, split for 16-byte NEON registers. -static SCATTER_8X8_NEON: [[u8; 16]; 4] = [ - [0, 8, 16, 24, 32, 40, 48, 56, 1, 9, 17, 25, 33, 41, 49, 57], - [2, 10, 18, 26, 34, 42, 50, 58, 3, 11, 19, 27, 35, 43, 51, 59], - [4, 12, 20, 28, 36, 44, 52, 60, 5, 13, 21, 29, 37, 45, 53, 61], - [6, 14, 22, 30, 38, 46, 54, 62, 7, 15, 23, 31, 39, 47, 55, 63], -]; - -/// Perform 8x8 bit transpose on two u64s packed in a `uint64x2_t`. -#[expect(unsafe_op_in_unsafe_fn)] -unsafe fn bit_transpose_8x8_neon(mut v: uint64x2_t) -> uint64x2_t { - let mask1 = vdupq_n_u64(TRANSPOSE_2X2); - let t = vandq_u64(veorq_u64(v, vshrq_n_u64::<7>(v)), mask1); - v = veorq_u64(veorq_u64(v, t), vshlq_n_u64::<7>(t)); - - let mask2 = vdupq_n_u64(TRANSPOSE_4X4); - let t = vandq_u64(veorq_u64(v, vshrq_n_u64::<14>(v)), mask2); - v = veorq_u64(veorq_u64(v, t), vshlq_n_u64::<14>(t)); - - let mask3 = vdupq_n_u64(TRANSPOSE_8X8); - let t = vandq_u64(veorq_u64(v, vshrq_n_u64::<28>(v)), mask3); - veorq_u64(veorq_u64(v, t), vshlq_n_u64::<28>(t)) -} - -/// Transpose 1024 bits using ARM NEON with TBL-based vectorized gather and scatter. -/// -/// Uses `vqtbl4q_u8` to gather bytes from the 128-byte input in parallel, -/// avoiding scalar byte-by-byte loads. Then uses `vqtbl4q_u8` again to perform -/// the 8x8 byte transpose for scatter. This is the NEON analog of x86 VBMI's -/// `vpermb`/`vpermi2b` byte permutation instructions. -/// -/// # Safety -/// Requires `AArch64` with NEON (always available on `AArch64`). -#[expect(unsafe_op_in_unsafe_fn)] -#[inline(never)] -pub unsafe fn transpose_bits_neon(input: &[u8; 128], output: &mut [u8; 128]) { - // Load all 128 input bytes into two uint8x16x4_t tables (64 bytes each) - let tbl_lo = vld1q_u8_x4(input.as_ptr()); - let tbl_hi = vld1q_u8_x4(input.as_ptr().add(64)); - - // Load scatter permutation indices (4 × 16 bytes) - let scatter0 = vld1q_u8(SCATTER_8X8_NEON[0].as_ptr()); - let scatter1 = vld1q_u8(SCATTER_8X8_NEON[1].as_ptr()); - let scatter2 = vld1q_u8(SCATTER_8X8_NEON[2].as_ptr()); - let scatter3 = vld1q_u8(SCATTER_8X8_NEON[3].as_ptr()); - - // Process first 64 output bytes (8 groups from BASE_PATTERN_FIRST) - // Gather and bit-transpose all 4 pairs, then scatter the full 64 bytes - let mut buf = [0u8; 64]; - for (i, (gather_lo, gather_high)) in [ - (GATHER_FIRST_LO, GATHER_FIRST_HI), - (GATHER_SECOND_LO, GATHER_SECOND_HI), - ] - .iter() - .enumerate() - { - for pair in 0..4 { - let idx_lo = vld1q_u8(gather_lo[pair].as_ptr()); - let idx_hi = vld1q_u8(gather_high[pair].as_ptr()); - - let from_lo = vqtbl4q_u8(tbl_lo, idx_lo); - let from_hi = vqtbl4q_u8(tbl_hi, idx_hi); - let gathered = vorrq_u8(from_lo, from_hi); - - let v = bit_transpose_8x8_neon(vreinterpretq_u64_u8(gathered)); - vst1q_u8(buf.as_mut_ptr().add(pair * 16), vreinterpretq_u8_u64(v)); - } - - // Load the 64-byte result as a TBL table and apply 8x8 byte transpose - let result_tbl = vld1q_u8_x4(buf.as_ptr()); - vst1q_u8( - output.as_mut_ptr().add(i * 64), - vqtbl4q_u8(result_tbl, scatter0), - ); - vst1q_u8( - output.as_mut_ptr().add(i * 64 + 16), - vqtbl4q_u8(result_tbl, scatter1), - ); - vst1q_u8( - output.as_mut_ptr().add(i * 64 + 32), - vqtbl4q_u8(result_tbl, scatter2), - ); - vst1q_u8( - output.as_mut_ptr().add(i * 64 + 48), - vqtbl4q_u8(result_tbl, scatter3), - ); - } -} - -/// Untranspose 1024 bits using ARM NEON with TBL-based vectorized operations. -/// -/// # Safety -/// Requires `AArch64` with NEON (always available on `AArch64`). -#[expect(unsafe_op_in_unsafe_fn)] -#[inline(never)] -pub unsafe fn untranspose_bits_neon(input: &[u8; 128], output: &mut [u8; 128]) { - // Load scatter indices (SCATTER_8X8 is self-inverse, so same table un-scatters) - let scatter0 = vld1q_u8(SCATTER_8X8_NEON[0].as_ptr()); - let scatter1 = vld1q_u8(SCATTER_8X8_NEON[1].as_ptr()); - let scatter2 = vld1q_u8(SCATTER_8X8_NEON[2].as_ptr()); - let scatter3 = vld1q_u8(SCATTER_8X8_NEON[3].as_ptr()); - - // Each iteration un-scatters the 64-byte input block to group-major order - let mut buf = [0u8; 64]; - for (i, base_pattern) in [BASE_PATTERN_FIRST, BASE_PATTERN_SECOND].iter().enumerate() { - let in_tbl = vld1q_u8_x4(input.as_ptr().add(i * 64)); - vst1q_u8(buf.as_mut_ptr(), vqtbl4q_u8(in_tbl, scatter0)); - vst1q_u8(buf.as_mut_ptr().add(16), vqtbl4q_u8(in_tbl, scatter1)); - vst1q_u8(buf.as_mut_ptr().add(32), vqtbl4q_u8(in_tbl, scatter2)); - vst1q_u8(buf.as_mut_ptr().add(48), vqtbl4q_u8(in_tbl, scatter3)); - - // Bit-transpose each pair and scatter to stride-16 output - for pair in 0..4 { - let base_group_0 = pair * 2; - let base_group_1 = pair * 2 + 1; - - let gathered = vld1q_u8(buf.as_ptr().add(pair * 16)); - let v = bit_transpose_8x8_neon(vreinterpretq_u64_u8(gathered)); - - let result_0 = vgetq_lane_u64::<0>(v); - let result_1 = vgetq_lane_u64::<1>(v); - - let out_base_0 = base_pattern[base_group_0]; - let out_base_1 = base_pattern[base_group_1]; - for i in 0..8 { - output[out_base_0 + i * 16] = (result_0 >> (i * 8)) as u8; - output[out_base_1 + i * 16] = (result_1 >> (i * 8)) as u8; - } - } - } -} - -#[cfg(test)] -mod tests { - use crate::bit_transpose::aarch64::transpose_bits_neon; - use crate::bit_transpose::aarch64::untranspose_bits_neon; - use crate::bit_transpose::generate_test_data; - use crate::bit_transpose::transpose_bits_baseline; - use crate::bit_transpose::untranspose_bits_baseline; - - #[test] - fn test_neon_matches_baseline() { - for seed in [0, 42, 123, 255] { - let input = generate_test_data(seed); - let mut baseline_out = [0u8; 128]; - let mut tbl_out = [0u8; 128]; - - transpose_bits_baseline(&input, &mut baseline_out); - unsafe { transpose_bits_neon(&input, &mut tbl_out) }; - - assert_eq!( - baseline_out, tbl_out, - "NEON TBL transpose doesn't match baseline for seed {seed}" - ); - } - } - - #[test] - fn test_neon_roundtrip() { - for seed in [0, 42, 123, 255] { - let input = generate_test_data(seed); - let mut transposed = [0u8; 128]; - let mut roundtrip = [0u8; 128]; - - unsafe { - transpose_bits_neon(&input, &mut transposed); - untranspose_bits_neon(&transposed, &mut roundtrip); - } - - assert_eq!( - input, roundtrip, - "NEON TBL roundtrip failed for seed {seed}" - ); - } - } - - #[test] - fn test_untranspose_neon_matches_baseline() { - for seed in [0, 42, 123, 255] { - let input = generate_test_data(seed); - let mut baseline_out = [0u8; 128]; - let mut tbl_out = [0u8; 128]; - - untranspose_bits_baseline(&input, &mut baseline_out); - unsafe { untranspose_bits_neon(&input, &mut tbl_out) }; - - assert_eq!( - baseline_out, tbl_out, - "NEON TBL untranspose doesn't match baseline for seed {seed}" - ); - } - } -} diff --git a/encodings/fastlanes/src/bit_transpose/mod.rs b/encodings/fastlanes/src/bit_transpose/mod.rs deleted file mode 100644 index 26e8b595195..00000000000 --- a/encodings/fastlanes/src/bit_transpose/mod.rs +++ /dev/null @@ -1,197 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Fast implementations of the `FastLanes` 1024-bit transpose. -//! -//! The `FastLanes` transpose is a fixed permutation of 1024 bits (128 bytes) that -//! enables SIMD parallelism for encodings like delta and RLE. This module provides -//! optimized implementations for different x86 SIMD instruction sets. -//! -//! The key insight is that each output byte is formed by extracting the SAME bit -//! position from 8 different input bytes at stride 16. The input byte groups follow -//! the `FL_ORDER` permutation pattern. - -#[cfg(feature = "_test-harness")] -pub mod aarch64; -#[cfg(feature = "_test-harness")] -pub mod scalar; -#[cfg(feature = "_test-harness")] -pub mod x86; - -#[cfg(not(feature = "_test-harness"))] -mod aarch64; -#[cfg(not(feature = "_test-harness"))] -mod scalar; -#[cfg(not(feature = "_test-harness"))] -mod x86; - -mod validity; - -pub use validity::*; -use vortex_buffer::CpuKernel; - -/// Signature shared by all 1024-bit transpose kernels. -type TransposeKernel = unsafe fn(&[u8; 128], &mut [u8; 128]); - -/// Base indices for the first 64 output bytes (lanes 0-7). -/// Each entry indicates the starting input byte index for that output byte group. -/// Pattern: [0*2, 4*2, 2*2, 6*2, 1*2, 5*2, 3*2, 7*2] = [0, 8, 4, 12, 2, 10, 6, 14] -const BASE_PATTERN_FIRST: [usize; 8] = [0, 8, 4, 12, 2, 10, 6, 14]; - -/// Base indices for the second 64 output bytes (lanes 8-15). -/// Pattern: first pattern + 1 = [1, 9, 5, 13, 3, 11, 7, 15] -const BASE_PATTERN_SECOND: [usize; 8] = [1, 9, 5, 13, 3, 11, 7, 15]; - -/// Masks for transposing 8x8 bit blocks. -const TRANSPOSE_2X2: u64 = 0x00AA_00AA_00AA_00AA; -const TRANSPOSE_4X4: u64 = 0x0000_CCCC_0000_CCCC; -const TRANSPOSE_8X8: u64 = 0x0000_0000_F0F0_F0F0; - -/// Transpose 1024-bits into FastLanes layout. -/// -/// Dispatch to the best available implementation, selected once on first call. -#[inline] -pub fn transpose_bits(input: &[u8; 128], output: &mut [u8; 128]) { - static KERNEL: CpuKernel = CpuKernel::new(|| { - #[cfg(target_arch = "x86_64")] - { - // VBMI is fastest - if x86::has_vbmi() { - return x86::transpose_bits_vbmi; - } - if x86::has_bmi2() { - return x86::transpose_bits_bmi2; - } - } - // NEON is architecturally guaranteed on aarch64, so it needs no probe. - #[cfg(target_arch = "aarch64")] - return aarch64::transpose_bits_neon; - // The aarch64 arm above returns unconditionally, making this portable default - // unreachable there. - #[allow(unreachable_code)] - { - scalar::transpose_bits_scalar - } - }); - // SAFETY: the selector only returns kernels that are safe or whose required CPU - // features were probed before selection. - unsafe { KERNEL.get()(input, output) } -} - -/// Untranspose 1024-bits from FastLanes layout. -/// -/// Dispatch untranspose to the best available implementation, selected once on first call. -#[inline] -pub fn untranspose_bits(input: &[u8; 128], output: &mut [u8; 128]) { - static KERNEL: CpuKernel = CpuKernel::new(|| { - #[cfg(target_arch = "x86_64")] - { - // VBMI is fastest - if x86::has_vbmi() { - return x86::untranspose_bits_vbmi; - } - if x86::has_bmi2() { - return x86::untranspose_bits_bmi2; - } - } - // NEON is architecturally guaranteed on aarch64, so it needs no probe. - #[cfg(target_arch = "aarch64")] - return aarch64::untranspose_bits_neon; - // The aarch64 arm above returns unconditionally, making this portable default - // unreachable there. - #[allow(unreachable_code)] - { - scalar::untranspose_bits_scalar - } - }); - // SAFETY: the selector only returns kernels that are safe or whose required CPU - // features were probed before selection. - unsafe { KERNEL.get()(input, output) } -} - -#[cfg(test)] -#[expect(clippy::cast_possible_truncation)] -fn generate_test_data(seed: u8) -> [u8; 128] { - let mut data = [0u8; 128]; - for (i, byte) in data.iter_mut().enumerate() { - *byte = seed.wrapping_mul(17).wrapping_add(i as u8).wrapping_mul(31); - } - data -} - -#[cfg(test)] -pub fn transpose_bits_baseline(input: &[u8; 128], output: &mut [u8; 128]) { - for in_bit in 0..1024 { - let out_bit = fastlanes::transpose(in_bit); - let in_byte = in_bit / 8; - let in_bit_pos = in_bit % 8; - let out_byte = out_bit / 8; - let out_bit_pos = out_bit % 8; - let bit_val = (input[in_byte] >> in_bit_pos) & 1; - output[out_byte] |= bit_val << out_bit_pos; - } -} - -#[cfg(test)] -pub fn untranspose_bits_baseline(input: &[u8; 128], output: &mut [u8; 128]) { - for out_bit in 0..1024 { - let in_bit = fastlanes::transpose(out_bit); - let in_byte = in_bit / 8; - let in_bit_pos = in_bit % 8; - let out_byte = out_bit / 8; - let out_bit_pos = out_bit % 8; - let bit_val = (input[in_byte] >> in_bit_pos) & 1; - output[out_byte] |= bit_val << out_bit_pos; - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_transpose_baseline_roundtrip() { - let input = generate_test_data(42); - let mut transposed = [0u8; 128]; - let mut roundtrip = [0u8; 128]; - - transpose_bits_baseline(&input, &mut transposed); - untranspose_bits_baseline(&transposed, &mut roundtrip); - - assert_eq!(input, roundtrip); - } - - #[test] - fn test_dispatch_matches_baseline() { - for seed in [0, 42, 123, 255] { - let input = generate_test_data(seed); - let mut baseline_out = [0u8; 128]; - let mut out = [0u8; 128]; - - transpose_bits_baseline(&input, &mut baseline_out); - transpose_bits(&input, &mut out); - - assert_eq!( - baseline_out, out, - "best dispatch doesn't match baseline for seed {seed}" - ); - } - } - - #[test] - fn test_untranspose_dispatch_matches_baseline() { - for seed in [0, 42, 123, 255] { - let input = generate_test_data(seed); - let mut baseline_out = [0u8; 128]; - let mut out = [0u8; 128]; - - untranspose_bits_baseline(&input, &mut baseline_out); - untranspose_bits(&input, &mut out); - - assert_eq!( - baseline_out, out, - "best untranspose dispatch doesn't match baseline for seed {seed}" - ); - } - } -} diff --git a/encodings/fastlanes/src/bit_transpose/scalar.rs b/encodings/fastlanes/src/bit_transpose/scalar.rs deleted file mode 100644 index 425e74087ab..00000000000 --- a/encodings/fastlanes/src/bit_transpose/scalar.rs +++ /dev/null @@ -1,219 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use crate::bit_transpose::BASE_PATTERN_FIRST; -use crate::bit_transpose::BASE_PATTERN_SECOND; -use crate::bit_transpose::TRANSPOSE_2X2; -use crate::bit_transpose::TRANSPOSE_4X4; -use crate::bit_transpose::TRANSPOSE_8X8; - -/// Fast scalar transpose using the 8x8 bit matrix transpose algorithm. -/// -/// This version uses 64-bit gather + parallel bit operations instead of -/// extracting bits one by one. Typically 5-10x faster than the basic scalar version. -#[inline(never)] -#[allow(dead_code)] -pub fn transpose_bits_scalar(input: &[u8; 128], output: &mut [u8; 128]) { - // Helper to perform 8x8 bit transpose on a u64 (each byte becomes a row) - fn transpose_8x8(mut x: u64) -> u64 { - // Step 1: Transpose 2x2 bit blocks - let t = (x ^ (x >> 7)) & TRANSPOSE_2X2; - x = x ^ t ^ (t << 7); - // Step 2: Transpose 4x4 bit blocks - let t = (x ^ (x >> 14)) & TRANSPOSE_4X4; - x = x ^ t ^ (t << 14); - // Step 3: Transpose 8x8 bit blocks - let t = (x ^ (x >> 28)) & TRANSPOSE_8X8; - x ^ t ^ (t << 28) - } - - // Helper to gather 8 bytes at stride 16 into a u64 - fn gather(input: &[u8; 128], base: usize) -> u64 { - u64::from(input[base]) - | (u64::from(input[base + 16]) << 8) - | (u64::from(input[base + 32]) << 16) - | (u64::from(input[base + 48]) << 24) - | (u64::from(input[base + 64]) << 32) - | (u64::from(input[base + 80]) << 40) - | (u64::from(input[base + 96]) << 48) - | (u64::from(input[base + 112]) << 56) - } - - // Process first half (8 base groups, fully unrolled) - let r0 = transpose_8x8(gather(input, BASE_PATTERN_FIRST[0])); - let r1 = transpose_8x8(gather(input, BASE_PATTERN_FIRST[1])); - let r2 = transpose_8x8(gather(input, BASE_PATTERN_FIRST[2])); - let r3 = transpose_8x8(gather(input, BASE_PATTERN_FIRST[3])); - let r4 = transpose_8x8(gather(input, BASE_PATTERN_FIRST[4])); - let r5 = transpose_8x8(gather(input, BASE_PATTERN_FIRST[5])); - let r6 = transpose_8x8(gather(input, BASE_PATTERN_FIRST[6])); - let r7 = transpose_8x8(gather(input, BASE_PATTERN_FIRST[7])); - - // Write first 64 output bytes (unrolled) - for bit_pos in 0..8 { - output[bit_pos * 8] = (r0 >> (bit_pos * 8)) as u8; - output[bit_pos * 8 + 1] = (r1 >> (bit_pos * 8)) as u8; - output[bit_pos * 8 + 2] = (r2 >> (bit_pos * 8)) as u8; - output[bit_pos * 8 + 3] = (r3 >> (bit_pos * 8)) as u8; - output[bit_pos * 8 + 4] = (r4 >> (bit_pos * 8)) as u8; - output[bit_pos * 8 + 5] = (r5 >> (bit_pos * 8)) as u8; - output[bit_pos * 8 + 6] = (r6 >> (bit_pos * 8)) as u8; - output[bit_pos * 8 + 7] = (r7 >> (bit_pos * 8)) as u8; - } - - // Process second half - let r0 = transpose_8x8(gather(input, BASE_PATTERN_SECOND[0])); - let r1 = transpose_8x8(gather(input, BASE_PATTERN_SECOND[1])); - let r2 = transpose_8x8(gather(input, BASE_PATTERN_SECOND[2])); - let r3 = transpose_8x8(gather(input, BASE_PATTERN_SECOND[3])); - let r4 = transpose_8x8(gather(input, BASE_PATTERN_SECOND[4])); - let r5 = transpose_8x8(gather(input, BASE_PATTERN_SECOND[5])); - let r6 = transpose_8x8(gather(input, BASE_PATTERN_SECOND[6])); - let r7 = transpose_8x8(gather(input, BASE_PATTERN_SECOND[7])); - - for bit_pos in 0..8 { - output[64 + bit_pos * 8] = (r0 >> (bit_pos * 8)) as u8; - output[64 + bit_pos * 8 + 1] = (r1 >> (bit_pos * 8)) as u8; - output[64 + bit_pos * 8 + 2] = (r2 >> (bit_pos * 8)) as u8; - output[64 + bit_pos * 8 + 3] = (r3 >> (bit_pos * 8)) as u8; - output[64 + bit_pos * 8 + 4] = (r4 >> (bit_pos * 8)) as u8; - output[64 + bit_pos * 8 + 5] = (r5 >> (bit_pos * 8)) as u8; - output[64 + bit_pos * 8 + 6] = (r6 >> (bit_pos * 8)) as u8; - output[64 + bit_pos * 8 + 7] = (r7 >> (bit_pos * 8)) as u8; - } -} - -/// Fast scalar untranspose using the 8x8 bit matrix transpose algorithm. -#[inline(never)] -#[allow(dead_code)] -pub fn untranspose_bits_scalar(input: &[u8; 128], output: &mut [u8; 128]) { - fn transpose_8x8(mut x: u64) -> u64 { - let t = (x ^ (x >> 7)) & TRANSPOSE_2X2; - x = x ^ t ^ (t << 7); - let t = (x ^ (x >> 14)) & TRANSPOSE_4X4; - x = x ^ t ^ (t << 14); - let t = (x ^ (x >> 28)) & TRANSPOSE_8X8; - x ^ t ^ (t << 28) - } - - fn gather_transposed(input: &[u8; 128], base_group: usize, offset: usize) -> u64 { - let mut result: u64 = 0; - for bit_pos in 0..8 { - result |= u64::from(input[offset + bit_pos * 8 + base_group]) << (bit_pos * 8); - } - result - } - - fn scatter(output: &mut [u8; 128], base: usize, val: u64) { - output[base] = val as u8; - output[base + 16] = (val >> 8) as u8; - output[base + 32] = (val >> 16) as u8; - output[base + 48] = (val >> 24) as u8; - output[base + 64] = (val >> 32) as u8; - output[base + 80] = (val >> 40) as u8; - output[base + 96] = (val >> 48) as u8; - output[base + 112] = (val >> 56) as u8; - } - - // First half (unrolled) - let r0 = transpose_8x8(gather_transposed(input, 0, 0)); - let r1 = transpose_8x8(gather_transposed(input, 1, 0)); - let r2 = transpose_8x8(gather_transposed(input, 2, 0)); - let r3 = transpose_8x8(gather_transposed(input, 3, 0)); - let r4 = transpose_8x8(gather_transposed(input, 4, 0)); - let r5 = transpose_8x8(gather_transposed(input, 5, 0)); - let r6 = transpose_8x8(gather_transposed(input, 6, 0)); - let r7 = transpose_8x8(gather_transposed(input, 7, 0)); - - scatter(output, BASE_PATTERN_FIRST[0], r0); - scatter(output, BASE_PATTERN_FIRST[1], r1); - scatter(output, BASE_PATTERN_FIRST[2], r2); - scatter(output, BASE_PATTERN_FIRST[3], r3); - scatter(output, BASE_PATTERN_FIRST[4], r4); - scatter(output, BASE_PATTERN_FIRST[5], r5); - scatter(output, BASE_PATTERN_FIRST[6], r6); - scatter(output, BASE_PATTERN_FIRST[7], r7); - - // Second half - let r0 = transpose_8x8(gather_transposed(input, 0, 64)); - let r1 = transpose_8x8(gather_transposed(input, 1, 64)); - let r2 = transpose_8x8(gather_transposed(input, 2, 64)); - let r3 = transpose_8x8(gather_transposed(input, 3, 64)); - let r4 = transpose_8x8(gather_transposed(input, 4, 64)); - let r5 = transpose_8x8(gather_transposed(input, 5, 64)); - let r6 = transpose_8x8(gather_transposed(input, 6, 64)); - let r7 = transpose_8x8(gather_transposed(input, 7, 64)); - - scatter(output, BASE_PATTERN_SECOND[0], r0); - scatter(output, BASE_PATTERN_SECOND[1], r1); - scatter(output, BASE_PATTERN_SECOND[2], r2); - scatter(output, BASE_PATTERN_SECOND[3], r3); - scatter(output, BASE_PATTERN_SECOND[4], r4); - scatter(output, BASE_PATTERN_SECOND[5], r5); - scatter(output, BASE_PATTERN_SECOND[6], r6); - scatter(output, BASE_PATTERN_SECOND[7], r7); -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::bit_transpose::generate_test_data; - - #[test] - fn test_scalar_matches_baseline() { - for seed in [0, 42, 123, 255] { - let input = generate_test_data(seed); - let mut baseline_out = [0u8; 128]; - let mut fast_out = [0u8; 128]; - - transpose_bits_scalar(&input, &mut baseline_out); - transpose_bits_scalar(&input, &mut fast_out); - - assert_eq!( - baseline_out, fast_out, - "scalar_fast transpose doesn't match baseline for seed {seed}" - ); - } - } - - #[test] - fn test_scalar_roundtrip() { - for seed in [0, 42, 123, 255] { - let input = generate_test_data(seed); - let mut transposed = [0u8; 128]; - let mut roundtrip = [0u8; 128]; - - transpose_bits_scalar(&input, &mut transposed); - untranspose_bits_scalar(&transposed, &mut roundtrip); - - assert_eq!( - input, roundtrip, - "scalar_fast roundtrip failed for seed {seed}" - ); - } - } - - #[test] - fn test_all_zeros() { - let input = [0u8; 128]; - let mut output = [0xFFu8; 128]; - - transpose_bits_scalar(&input, &mut output); - assert_eq!(output, [0u8; 128]); - - untranspose_bits_scalar(&input, &mut output); - assert_eq!(output, [0u8; 128]); - } - - #[test] - fn test_all_ones() { - let input = [0xFFu8; 128]; - let mut output = [0u8; 128]; - - transpose_bits_scalar(&input, &mut output); - assert_eq!(output, [0xFFu8; 128]); - - untranspose_bits_scalar(&input, &mut output); - assert_eq!(output, [0xFFu8; 128]); - } -} diff --git a/encodings/fastlanes/src/bit_transpose/validity.rs b/encodings/fastlanes/src/bit_transpose/validity.rs deleted file mode 100644 index a6b2ec43cbc..00000000000 --- a/encodings/fastlanes/src/bit_transpose/validity.rs +++ /dev/null @@ -1,215 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::mem; -use std::mem::MaybeUninit; - -use vortex_array::Canonical; -use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::arrays::BoolArray; -use vortex_array::validity::Validity; -use vortex_buffer::BitBuffer; -use vortex_buffer::ByteBuffer; -use vortex_buffer::ByteBufferMut; -use vortex_error::VortexExpect; -use vortex_error::VortexResult; - -use crate::bit_transpose::transpose_bits; -use crate::bit_transpose::untranspose_bits; - -pub fn transpose_validity(validity: &Validity, ctx: &mut ExecutionCtx) -> VortexResult { - match validity { - Validity::Array(mask) => { - let bools = mask - .clone() - .execute::(ctx)? - .into_bool() - .into_bit_buffer(); - - Ok(Validity::Array( - BoolArray::new(transpose_bitbuffer(bools), Validity::NonNullable).into_array(), - )) - } - v @ Validity::AllValid | v @ Validity::AllInvalid | v @ Validity::NonNullable => { - Ok(v.clone()) - } - } -} - -pub fn transpose_bitbuffer(bits: BitBuffer) -> BitBuffer { - let (offset, len, bytes) = bits.into_inner(); - - if bytes.len().is_multiple_of(128) { - match bytes.try_into_mut() { - Ok(mut bytes_mut) => { - // We can ignore the spare trailer capacity that can be an artifact of allocator as we requested 128 multiple chunks - let (chunks, _) = bytes_mut.as_chunks_mut::<128>(); - let mut tmp = [0u8; 128]; - for chunk in chunks { - transpose_bits(chunk, &mut tmp); - chunk.copy_from_slice(&tmp); - } - BitBuffer::new_with_offset(bytes_mut.freeze().into_byte_buffer(), len, offset) - } - Err(bytes) => bits_op_with_copy(bytes, len, offset, transpose_bits), - } - } else { - bits_op_with_copy(bytes, len, offset, transpose_bits) - } -} - -pub fn untranspose_validity(validity: &Validity, ctx: &mut ExecutionCtx) -> VortexResult { - match validity { - Validity::Array(mask) => { - let bools = mask - .clone() - .execute::(ctx)? - .into_bool() - .into_bit_buffer(); - - Ok(Validity::Array( - BoolArray::new(untranspose_bitbuffer(bools), Validity::NonNullable).into_array(), - )) - } - v @ Validity::AllValid | v @ Validity::AllInvalid | v @ Validity::NonNullable => { - Ok(v.clone()) - } - } -} - -pub fn untranspose_bitbuffer(bits: BitBuffer) -> BitBuffer { - assert!( - bits.inner().len().is_multiple_of(128), - "Transpose BitBuffer byte length must be a multiple of 128" - ); - let (offset, len, bytes) = bits.into_inner(); - match bytes.try_into_mut() { - Ok(mut bytes_mut) => { - let (chunks, _) = bytes_mut.as_chunks_mut::<128>(); - let mut tmp = [0u8; 128]; - for chunk in chunks { - untranspose_bits(chunk, &mut tmp); - chunk.copy_from_slice(&tmp); - } - BitBuffer::new_with_offset(bytes_mut.freeze().into_byte_buffer(), len, offset) - } - Err(bytes) => bits_op_with_copy(bytes, len, offset, untranspose_bits), - } -} - -fn bits_op_with_copy( - bytes: ByteBuffer, - len: usize, - offset: usize, - op: F, -) -> BitBuffer { - let output_len = bytes.len().next_multiple_of(128); - let mut output = ByteBufferMut::with_capacity(output_len); - let (input_chunks, input_trailer) = bytes.as_chunks::<128>(); - // Bound to the requested `output_len`: `spare_capacity_mut` may expose extra over-aligned - // capacity, which would otherwise split into spurious trailing chunks and make `last_mut` - // below target a chunk past the data we actually initialize. - let (output_chunks, _) = output.spare_capacity_mut()[..output_len].as_chunks_mut::<128>(); - - for (input, output) in input_chunks.iter().zip(output_chunks.iter_mut()) { - op(input, unsafe { - mem::transmute::<&mut [MaybeUninit; 128], &mut [u8; 128]>(output) - }); - } - - if !input_trailer.is_empty() { - let mut padded_input = [0u8; 128]; - padded_input[0..input_trailer.len()].clone_from_slice(input_trailer); - op(&padded_input, unsafe { - mem::transmute::<&mut [MaybeUninit; 128], &mut [u8; 128]>( - output_chunks - .last_mut() - .vortex_expect("Output wasn't a multiple of 128 bytes"), - ) - }); - } - - unsafe { output.set_len(output_len) }; - BitBuffer::new_with_offset( - output.freeze().into_byte_buffer(), - len.next_multiple_of(1024), - offset, - ) -} - -#[cfg(test)] -mod tests { - use vortex_buffer::BitBuffer; - use vortex_buffer::BitBufferMut; - use vortex_buffer::ByteBuffer; - - use super::*; - - fn make_validity_bits(num_bits: usize) -> BitBuffer { - let mut builder = BitBufferMut::with_capacity(num_bits); - for i in 0..num_bits { - builder.append(i % 3 != 0); - } - builder.freeze() - } - - fn force_copy_path(bits: BitBuffer) -> (BitBuffer, ByteBuffer) { - let (offset, len, bytes) = bits.into_inner(); - let extra_ref = bytes.clone(); - (BitBuffer::new_with_offset(bytes, len, offset), extra_ref) - } - - #[test] - fn transpose_padding_copy_produces_same_bits() { - let bits = make_validity_bits(500); - let transposed = transpose_bitbuffer(bits.clone()); - assert_eq!(transposed.len(), 1024); - let untransposed = untranspose_bitbuffer(transposed); - assert_eq!(untransposed.slice(0..500), bits) - } - - #[test] - fn transpose_inplace_and_copy_produce_same_bits() { - let bits = make_validity_bits(2048); - - let inplace_result = transpose_bitbuffer(bits.clone()); - - let (bits_shared, _hold) = force_copy_path(bits); - let copy_result = transpose_bitbuffer(bits_shared); - - assert_eq!(inplace_result.len(), copy_result.len()); - assert_eq!(inplace_result, copy_result); - } - - #[test] - fn transpose_validity_roundtrip_non_aligned() { - let original_len = 1500; - let bits = make_validity_bits(original_len); - - let transposed = transpose_bitbuffer(bits.clone()); - let roundtripped = untranspose_bitbuffer(transposed); - assert_eq!(bits, roundtripped.slice(0..original_len)); - } - - /// Regression: the copy path split the over-aligned spare capacity into extra 128-byte - /// chunks and wrote the padded remainder via `last_mut()`, which landed past the requested - /// length and left the real trailing chunk uninitialized. Whether the surplus capacity - /// produced an extra chunk depended on the allocation address, so we repeat across sizes to - /// defeat that luck; the fix bounds the spare slice so the result no longer depends on it. - #[test] - fn transpose_copy_path_survives_overallocation() { - for original_len in [129, 500, 1500, 9999] { - let bits = make_validity_bits(original_len); - for _ in 0..64 { - let transposed = transpose_bitbuffer(bits.clone()); - let roundtripped = untranspose_bitbuffer(transposed); - assert_eq!( - roundtripped.slice(0..original_len), - bits, - "len={original_len}" - ); - } - } - } -} diff --git a/encodings/fastlanes/src/bit_transpose/x86.rs b/encodings/fastlanes/src/bit_transpose/x86.rs deleted file mode 100644 index 9eaef248d28..00000000000 --- a/encodings/fastlanes/src/bit_transpose/x86.rs +++ /dev/null @@ -1,715 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors -#![cfg(target_arch = "x86_64")] - -use core::arch::x86_64::__m512i; -use core::arch::x86_64::_mm512_and_si512; -use core::arch::x86_64::_mm512_loadu_si512; -use core::arch::x86_64::_mm512_permutex2var_epi8; -use core::arch::x86_64::_mm512_permutexvar_epi8; -use core::arch::x86_64::_mm512_set1_epi64; -use core::arch::x86_64::_mm512_slli_epi64; -use core::arch::x86_64::_mm512_srli_epi64; -use core::arch::x86_64::_mm512_storeu_si512; -use core::arch::x86_64::_mm512_xor_si512; -use core::arch::x86_64::_pdep_u64; -use core::arch::x86_64::_pext_u64; -use std::is_x86_feature_detected; - -use crate::bit_transpose::BASE_PATTERN_FIRST; -use crate::bit_transpose::BASE_PATTERN_SECOND; -use crate::bit_transpose::TRANSPOSE_2X2; -use crate::bit_transpose::TRANSPOSE_4X4; -use crate::bit_transpose::TRANSPOSE_8X8; - -/// Check if BMI2 is available. -#[inline] -#[must_use] -pub fn has_bmi2() -> bool { - is_x86_feature_detected!("bmi2") -} - -/// Check if AVX-512 VBMI is available (for byte permutation). -#[inline] -#[must_use] -pub fn has_vbmi() -> bool { - is_x86_feature_detected!("avx512vbmi") -} - -/// Transpose 1024 bits using BMI2 PEXT instruction. -/// -/// PEXT extracts bits at positions specified by a mask into contiguous low bits. -/// Fully unrolled for ~12% better performance vs looped version. -/// -/// # Safety -/// Requires BMI2 support. Check with `has_bmi2()` before calling. -#[target_feature(enable = "bmi2")] -#[inline(never)] -#[expect(clippy::too_many_lines)] -pub unsafe fn transpose_bits_bmi2(input: &[u8; 128], output: &mut [u8; 128]) { - // Helper to gather 8 bytes at stride 16 into a u64 - fn gather(input: &[u8; 128], base: usize) -> u64 { - (input[base] as u64) - | ((input[base + 16] as u64) << 8) - | ((input[base + 32] as u64) << 16) - | ((input[base + 48] as u64) << 24) - | ((input[base + 64] as u64) << 32) - | ((input[base + 80] as u64) << 40) - | ((input[base + 96] as u64) << 48) - | ((input[base + 112] as u64) << 56) - } - - // Gather all 16 groups (fully unrolled) - // First half: BASE_PATTERN_FIRST = [0, 8, 4, 12, 2, 10, 6, 14] - let g0 = gather(input, 0); - let g1 = gather(input, 8); - let g2 = gather(input, 4); - let g3 = gather(input, 12); - let g4 = gather(input, 2); - let g5 = gather(input, 10); - let g6 = gather(input, 6); - let g7 = gather(input, 14); - // Second half: BASE_PATTERN_SECOND = [1, 9, 5, 13, 3, 11, 7, 15] - let g8 = gather(input, 1); - let g9 = gather(input, 9); - let g10 = gather(input, 5); - let g11 = gather(input, 13); - let g12 = gather(input, 3); - let g13 = gather(input, 11); - let g14 = gather(input, 7); - let g15 = gather(input, 15); - - // Masks for each bit position - let m0: u64 = 0x0101_0101_0101_0101; - let m1: u64 = 0x0202_0202_0202_0202; - let m2: u64 = 0x0404_0404_0404_0404; - let m3: u64 = 0x0808_0808_0808_0808; - let m4: u64 = 0x1010_1010_1010_1010; - let m5: u64 = 0x2020_2020_2020_2020; - let m6: u64 = 0x4040_4040_4040_4040; - let m7: u64 = 0x8080_8080_8080_8080; - - // First half - 64 PEXT operations (fully unrolled) - output[0] = _pext_u64(g0, m0) as u8; - output[1] = _pext_u64(g1, m0) as u8; - output[2] = _pext_u64(g2, m0) as u8; - output[3] = _pext_u64(g3, m0) as u8; - output[4] = _pext_u64(g4, m0) as u8; - output[5] = _pext_u64(g5, m0) as u8; - output[6] = _pext_u64(g6, m0) as u8; - output[7] = _pext_u64(g7, m0) as u8; - output[8] = _pext_u64(g0, m1) as u8; - output[9] = _pext_u64(g1, m1) as u8; - output[10] = _pext_u64(g2, m1) as u8; - output[11] = _pext_u64(g3, m1) as u8; - output[12] = _pext_u64(g4, m1) as u8; - output[13] = _pext_u64(g5, m1) as u8; - output[14] = _pext_u64(g6, m1) as u8; - output[15] = _pext_u64(g7, m1) as u8; - output[16] = _pext_u64(g0, m2) as u8; - output[17] = _pext_u64(g1, m2) as u8; - output[18] = _pext_u64(g2, m2) as u8; - output[19] = _pext_u64(g3, m2) as u8; - output[20] = _pext_u64(g4, m2) as u8; - output[21] = _pext_u64(g5, m2) as u8; - output[22] = _pext_u64(g6, m2) as u8; - output[23] = _pext_u64(g7, m2) as u8; - output[24] = _pext_u64(g0, m3) as u8; - output[25] = _pext_u64(g1, m3) as u8; - output[26] = _pext_u64(g2, m3) as u8; - output[27] = _pext_u64(g3, m3) as u8; - output[28] = _pext_u64(g4, m3) as u8; - output[29] = _pext_u64(g5, m3) as u8; - output[30] = _pext_u64(g6, m3) as u8; - output[31] = _pext_u64(g7, m3) as u8; - output[32] = _pext_u64(g0, m4) as u8; - output[33] = _pext_u64(g1, m4) as u8; - output[34] = _pext_u64(g2, m4) as u8; - output[35] = _pext_u64(g3, m4) as u8; - output[36] = _pext_u64(g4, m4) as u8; - output[37] = _pext_u64(g5, m4) as u8; - output[38] = _pext_u64(g6, m4) as u8; - output[39] = _pext_u64(g7, m4) as u8; - output[40] = _pext_u64(g0, m5) as u8; - output[41] = _pext_u64(g1, m5) as u8; - output[42] = _pext_u64(g2, m5) as u8; - output[43] = _pext_u64(g3, m5) as u8; - output[44] = _pext_u64(g4, m5) as u8; - output[45] = _pext_u64(g5, m5) as u8; - output[46] = _pext_u64(g6, m5) as u8; - output[47] = _pext_u64(g7, m5) as u8; - output[48] = _pext_u64(g0, m6) as u8; - output[49] = _pext_u64(g1, m6) as u8; - output[50] = _pext_u64(g2, m6) as u8; - output[51] = _pext_u64(g3, m6) as u8; - output[52] = _pext_u64(g4, m6) as u8; - output[53] = _pext_u64(g5, m6) as u8; - output[54] = _pext_u64(g6, m6) as u8; - output[55] = _pext_u64(g7, m6) as u8; - output[56] = _pext_u64(g0, m7) as u8; - output[57] = _pext_u64(g1, m7) as u8; - output[58] = _pext_u64(g2, m7) as u8; - output[59] = _pext_u64(g3, m7) as u8; - output[60] = _pext_u64(g4, m7) as u8; - output[61] = _pext_u64(g5, m7) as u8; - output[62] = _pext_u64(g6, m7) as u8; - output[63] = _pext_u64(g7, m7) as u8; - - // Second half - 64 PEXT operations (fully unrolled) - output[64] = _pext_u64(g8, m0) as u8; - output[65] = _pext_u64(g9, m0) as u8; - output[66] = _pext_u64(g10, m0) as u8; - output[67] = _pext_u64(g11, m0) as u8; - output[68] = _pext_u64(g12, m0) as u8; - output[69] = _pext_u64(g13, m0) as u8; - output[70] = _pext_u64(g14, m0) as u8; - output[71] = _pext_u64(g15, m0) as u8; - output[72] = _pext_u64(g8, m1) as u8; - output[73] = _pext_u64(g9, m1) as u8; - output[74] = _pext_u64(g10, m1) as u8; - output[75] = _pext_u64(g11, m1) as u8; - output[76] = _pext_u64(g12, m1) as u8; - output[77] = _pext_u64(g13, m1) as u8; - output[78] = _pext_u64(g14, m1) as u8; - output[79] = _pext_u64(g15, m1) as u8; - output[80] = _pext_u64(g8, m2) as u8; - output[81] = _pext_u64(g9, m2) as u8; - output[82] = _pext_u64(g10, m2) as u8; - output[83] = _pext_u64(g11, m2) as u8; - output[84] = _pext_u64(g12, m2) as u8; - output[85] = _pext_u64(g13, m2) as u8; - output[86] = _pext_u64(g14, m2) as u8; - output[87] = _pext_u64(g15, m2) as u8; - output[88] = _pext_u64(g8, m3) as u8; - output[89] = _pext_u64(g9, m3) as u8; - output[90] = _pext_u64(g10, m3) as u8; - output[91] = _pext_u64(g11, m3) as u8; - output[92] = _pext_u64(g12, m3) as u8; - output[93] = _pext_u64(g13, m3) as u8; - output[94] = _pext_u64(g14, m3) as u8; - output[95] = _pext_u64(g15, m3) as u8; - output[96] = _pext_u64(g8, m4) as u8; - output[97] = _pext_u64(g9, m4) as u8; - output[98] = _pext_u64(g10, m4) as u8; - output[99] = _pext_u64(g11, m4) as u8; - output[100] = _pext_u64(g12, m4) as u8; - output[101] = _pext_u64(g13, m4) as u8; - output[102] = _pext_u64(g14, m4) as u8; - output[103] = _pext_u64(g15, m4) as u8; - output[104] = _pext_u64(g8, m5) as u8; - output[105] = _pext_u64(g9, m5) as u8; - output[106] = _pext_u64(g10, m5) as u8; - output[107] = _pext_u64(g11, m5) as u8; - output[108] = _pext_u64(g12, m5) as u8; - output[109] = _pext_u64(g13, m5) as u8; - output[110] = _pext_u64(g14, m5) as u8; - output[111] = _pext_u64(g15, m5) as u8; - output[112] = _pext_u64(g8, m6) as u8; - output[113] = _pext_u64(g9, m6) as u8; - output[114] = _pext_u64(g10, m6) as u8; - output[115] = _pext_u64(g11, m6) as u8; - output[116] = _pext_u64(g12, m6) as u8; - output[117] = _pext_u64(g13, m6) as u8; - output[118] = _pext_u64(g14, m6) as u8; - output[119] = _pext_u64(g15, m6) as u8; - output[120] = _pext_u64(g8, m7) as u8; - output[121] = _pext_u64(g9, m7) as u8; - output[122] = _pext_u64(g10, m7) as u8; - output[123] = _pext_u64(g11, m7) as u8; - output[124] = _pext_u64(g12, m7) as u8; - output[125] = _pext_u64(g13, m7) as u8; - output[126] = _pext_u64(g14, m7) as u8; - output[127] = _pext_u64(g15, m7) as u8; -} - -/// Untranspose 1024 bits using BMI2 PDEP instruction. -/// -/// Structured per-output-group: for each group of 8 output bytes at stride 16, -/// PDEP 8 input bytes into different bit positions, OR in registers, then -/// scatter-store once. Each output byte is written exactly once (no read-modify-write). -/// -/// # Safety -/// Requires BMI2 support. Check with `has_bmi2()` before calling. -#[target_feature(enable = "bmi2")] -#[inline(never)] -#[expect(clippy::too_many_lines)] -pub unsafe fn untranspose_bits_bmi2(input: &[u8; 128], output: &mut [u8; 128]) { - // Helper: scatter a u64 to 8 output bytes at stride 16 - fn scatter(output: &mut [u8; 128], base: usize, val: u64) { - output[base] = val as u8; - output[base + 16] = (val >> 8) as u8; - output[base + 32] = (val >> 16) as u8; - output[base + 48] = (val >> 24) as u8; - output[base + 64] = (val >> 32) as u8; - output[base + 80] = (val >> 40) as u8; - output[base + 96] = (val >> 48) as u8; - output[base + 112] = (val >> 56) as u8; - } - - // Masks for each bit position - let m0: u64 = 0x0101_0101_0101_0101; - let m1: u64 = 0x0202_0202_0202_0202; - let m2: u64 = 0x0404_0404_0404_0404; - let m3: u64 = 0x0808_0808_0808_0808; - let m4: u64 = 0x1010_1010_1010_1010; - let m5: u64 = 0x2020_2020_2020_2020; - let m6: u64 = 0x4040_4040_4040_4040; - let m7: u64 = 0x8080_8080_8080_8080; - - // For each output group, the input bytes that contribute are at - // input[bit_pos * 8 + group_idx] for bit_pos 0..8. - // PDEP deposits the 8 bits of the input byte into the bit_pos position - // of each byte in the u64. - - // First half: 8 groups using BASE_PATTERN_FIRST - // Group 0 (base=0): input bytes [0, 8, 16, 24, 32, 40, 48, 56] - let v = _pdep_u64(input[0] as u64, m0) - | _pdep_u64(input[8] as u64, m1) - | _pdep_u64(input[16] as u64, m2) - | _pdep_u64(input[24] as u64, m3) - | _pdep_u64(input[32] as u64, m4) - | _pdep_u64(input[40] as u64, m5) - | _pdep_u64(input[48] as u64, m6) - | _pdep_u64(input[56] as u64, m7); - scatter(output, 0, v); - - // Group 1 (base=8) - let v = _pdep_u64(input[1] as u64, m0) - | _pdep_u64(input[9] as u64, m1) - | _pdep_u64(input[17] as u64, m2) - | _pdep_u64(input[25] as u64, m3) - | _pdep_u64(input[33] as u64, m4) - | _pdep_u64(input[41] as u64, m5) - | _pdep_u64(input[49] as u64, m6) - | _pdep_u64(input[57] as u64, m7); - scatter(output, 8, v); - - // Group 2 (base=4) - let v = _pdep_u64(input[2] as u64, m0) - | _pdep_u64(input[10] as u64, m1) - | _pdep_u64(input[18] as u64, m2) - | _pdep_u64(input[26] as u64, m3) - | _pdep_u64(input[34] as u64, m4) - | _pdep_u64(input[42] as u64, m5) - | _pdep_u64(input[50] as u64, m6) - | _pdep_u64(input[58] as u64, m7); - scatter(output, 4, v); - - // Group 3 (base=12) - let v = _pdep_u64(input[3] as u64, m0) - | _pdep_u64(input[11] as u64, m1) - | _pdep_u64(input[19] as u64, m2) - | _pdep_u64(input[27] as u64, m3) - | _pdep_u64(input[35] as u64, m4) - | _pdep_u64(input[43] as u64, m5) - | _pdep_u64(input[51] as u64, m6) - | _pdep_u64(input[59] as u64, m7); - scatter(output, 12, v); - - // Group 4 (base=2) - let v = _pdep_u64(input[4] as u64, m0) - | _pdep_u64(input[12] as u64, m1) - | _pdep_u64(input[20] as u64, m2) - | _pdep_u64(input[28] as u64, m3) - | _pdep_u64(input[36] as u64, m4) - | _pdep_u64(input[44] as u64, m5) - | _pdep_u64(input[52] as u64, m6) - | _pdep_u64(input[60] as u64, m7); - scatter(output, 2, v); - - // Group 5 (base=10) - let v = _pdep_u64(input[5] as u64, m0) - | _pdep_u64(input[13] as u64, m1) - | _pdep_u64(input[21] as u64, m2) - | _pdep_u64(input[29] as u64, m3) - | _pdep_u64(input[37] as u64, m4) - | _pdep_u64(input[45] as u64, m5) - | _pdep_u64(input[53] as u64, m6) - | _pdep_u64(input[61] as u64, m7); - scatter(output, 10, v); - - // Group 6 (base=6) - let v = _pdep_u64(input[6] as u64, m0) - | _pdep_u64(input[14] as u64, m1) - | _pdep_u64(input[22] as u64, m2) - | _pdep_u64(input[30] as u64, m3) - | _pdep_u64(input[38] as u64, m4) - | _pdep_u64(input[46] as u64, m5) - | _pdep_u64(input[54] as u64, m6) - | _pdep_u64(input[62] as u64, m7); - scatter(output, 6, v); - - // Group 7 (base=14) - let v = _pdep_u64(input[7] as u64, m0) - | _pdep_u64(input[15] as u64, m1) - | _pdep_u64(input[23] as u64, m2) - | _pdep_u64(input[31] as u64, m3) - | _pdep_u64(input[39] as u64, m4) - | _pdep_u64(input[47] as u64, m5) - | _pdep_u64(input[55] as u64, m6) - | _pdep_u64(input[63] as u64, m7); - scatter(output, 14, v); - - // Second half: 8 groups using BASE_PATTERN_SECOND - // Group 0 (base=1) - let v = _pdep_u64(input[64] as u64, m0) - | _pdep_u64(input[72] as u64, m1) - | _pdep_u64(input[80] as u64, m2) - | _pdep_u64(input[88] as u64, m3) - | _pdep_u64(input[96] as u64, m4) - | _pdep_u64(input[104] as u64, m5) - | _pdep_u64(input[112] as u64, m6) - | _pdep_u64(input[120] as u64, m7); - scatter(output, 1, v); - - // Group 1 (base=9) - let v = _pdep_u64(input[65] as u64, m0) - | _pdep_u64(input[73] as u64, m1) - | _pdep_u64(input[81] as u64, m2) - | _pdep_u64(input[89] as u64, m3) - | _pdep_u64(input[97] as u64, m4) - | _pdep_u64(input[105] as u64, m5) - | _pdep_u64(input[113] as u64, m6) - | _pdep_u64(input[121] as u64, m7); - scatter(output, 9, v); - - // Group 2 (base=5) - let v = _pdep_u64(input[66] as u64, m0) - | _pdep_u64(input[74] as u64, m1) - | _pdep_u64(input[82] as u64, m2) - | _pdep_u64(input[90] as u64, m3) - | _pdep_u64(input[98] as u64, m4) - | _pdep_u64(input[106] as u64, m5) - | _pdep_u64(input[114] as u64, m6) - | _pdep_u64(input[122] as u64, m7); - scatter(output, 5, v); - - // Group 3 (base=13) - let v = _pdep_u64(input[67] as u64, m0) - | _pdep_u64(input[75] as u64, m1) - | _pdep_u64(input[83] as u64, m2) - | _pdep_u64(input[91] as u64, m3) - | _pdep_u64(input[99] as u64, m4) - | _pdep_u64(input[107] as u64, m5) - | _pdep_u64(input[115] as u64, m6) - | _pdep_u64(input[123] as u64, m7); - scatter(output, 13, v); - - // Group 4 (base=3) - let v = _pdep_u64(input[68] as u64, m0) - | _pdep_u64(input[76] as u64, m1) - | _pdep_u64(input[84] as u64, m2) - | _pdep_u64(input[92] as u64, m3) - | _pdep_u64(input[100] as u64, m4) - | _pdep_u64(input[108] as u64, m5) - | _pdep_u64(input[116] as u64, m6) - | _pdep_u64(input[124] as u64, m7); - scatter(output, 3, v); - - // Group 5 (base=11) - let v = _pdep_u64(input[69] as u64, m0) - | _pdep_u64(input[77] as u64, m1) - | _pdep_u64(input[85] as u64, m2) - | _pdep_u64(input[93] as u64, m3) - | _pdep_u64(input[101] as u64, m4) - | _pdep_u64(input[109] as u64, m5) - | _pdep_u64(input[117] as u64, m6) - | _pdep_u64(input[125] as u64, m7); - scatter(output, 11, v); - - // Group 6 (base=7) - let v = _pdep_u64(input[70] as u64, m0) - | _pdep_u64(input[78] as u64, m1) - | _pdep_u64(input[86] as u64, m2) - | _pdep_u64(input[94] as u64, m3) - | _pdep_u64(input[102] as u64, m4) - | _pdep_u64(input[110] as u64, m5) - | _pdep_u64(input[118] as u64, m6) - | _pdep_u64(input[126] as u64, m7); - scatter(output, 7, v); - - // Group 7 (base=15) - let v = _pdep_u64(input[71] as u64, m0) - | _pdep_u64(input[79] as u64, m1) - | _pdep_u64(input[87] as u64, m2) - | _pdep_u64(input[95] as u64, m3) - | _pdep_u64(input[103] as u64, m4) - | _pdep_u64(input[111] as u64, m5) - | _pdep_u64(input[119] as u64, m6) - | _pdep_u64(input[127] as u64, m7); - scatter(output, 15, v); -} - -// Static permutation tables for VBMI gather operations -static GATHER_FIRST: [u8; 64] = [ - // Gather bytes at stride 16 for first 8 groups (bases from BASE_PATTERN_FIRST) - // Group 0: base=0 - 0, 16, 32, 48, 64, 80, 96, 112, // Group 1: base=8 - 8, 24, 40, 56, 72, 88, 104, 120, // Group 2: base=4 - 4, 20, 36, 52, 68, 84, 100, 116, // Group 3: base=12 - 12, 28, 44, 60, 76, 92, 108, 124, // Group 4: base=2 - 2, 18, 34, 50, 66, 82, 98, 114, // Group 5: base=10 - 10, 26, 42, 58, 74, 90, 106, 122, // Group 6: base=6 - 6, 22, 38, 54, 70, 86, 102, 118, // Group 7: base=14 - 14, 30, 46, 62, 78, 94, 110, 126, -]; - -static GATHER_SECOND: [u8; 64] = [ - // Gather bytes at stride 16 for second 8 groups (bases from BASE_PATTERN_SECOND) - // Group 0: base=1 - 1, 17, 33, 49, 65, 81, 97, 113, // Group 1: base=9 - 9, 25, 41, 57, 73, 89, 105, 121, // Group 2: base=5 - 5, 21, 37, 53, 69, 85, 101, 117, // Group 3: base=13 - 13, 29, 45, 61, 77, 93, 109, 125, // Group 4: base=3 - 3, 19, 35, 51, 67, 83, 99, 115, // Group 5: base=11 - 11, 27, 43, 59, 75, 91, 107, 123, // Group 6: base=7 - 7, 23, 39, 55, 71, 87, 103, 119, // Group 7: base=15 - 15, 31, 47, 63, 79, 95, 111, 127, -]; - -// 8x8 byte transpose permutation for scatter phase -// Input: [g0b0..g0b7, g1b0..g1b7, ..., g7b0..g7b7] (8 groups of 8 bytes) -// Output: [g0b0,g1b0,..,g7b0, g0b1,g1b1,..,g7b1, ...] (8 rows of 8 bytes) -static SCATTER_8X8: [u8; 64] = [ - 0, 8, 16, 24, 32, 40, 48, 56, // byte 0 from each group - 1, 9, 17, 25, 33, 41, 49, 57, // byte 1 from each group - 2, 10, 18, 26, 34, 42, 50, 58, // byte 2 from each group - 3, 11, 19, 27, 35, 43, 51, 59, // byte 3 from each group - 4, 12, 20, 28, 36, 44, 52, 60, // byte 4 from each group - 5, 13, 21, 29, 37, 45, 53, 61, // byte 5 from each group - 6, 14, 22, 30, 38, 46, 54, 62, // byte 6 from each group - 7, 15, 23, 31, 39, 47, 55, 63, // byte 7 from each group -]; - -/// Transpose 1024 bits using AVX-512 VBMI for vectorized gather and scatter. -/// -/// Uses vpermi2b to gather bytes from stride-16 positions in parallel, -/// and vpermb for the final 8x8 byte transpose to output format. -/// -/// # Safety -/// Requires AVX-512F, AVX-512BW, and AVX-512VBMI support. -#[target_feature(enable = "avx512f", enable = "avx512bw", enable = "avx512vbmi")] -#[inline(never)] -#[expect(clippy::cast_possible_wrap)] -#[expect(clippy::cast_ptr_alignment)] -#[expect(unsafe_op_in_unsafe_fn)] -pub unsafe fn transpose_bits_vbmi(input: &[u8; 128], output: &mut [u8; 128]) { - // Load all 128 input bytes into two ZMM registers - let in_lo = _mm512_loadu_si512(input.as_ptr().cast::<__m512i>()); - let in_hi = _mm512_loadu_si512(input.as_ptr().add(64).cast::<__m512i>()); - - // Load permutation indices (static tables) - let idx_first = _mm512_loadu_si512(GATHER_FIRST.as_ptr().cast::<__m512i>()); - let idx_second = _mm512_loadu_si512(GATHER_SECOND.as_ptr().cast::<__m512i>()); - let idx_scatter = _mm512_loadu_si512(SCATTER_8X8.as_ptr().cast::<__m512i>()); - - // Masks for 8x8 bit transpose - let mask1 = _mm512_set1_epi64(TRANSPOSE_2X2 as i64); - let mask2 = _mm512_set1_epi64(TRANSPOSE_4X4 as i64); - let mask3 = _mm512_set1_epi64(TRANSPOSE_8X8 as i64); - - // Process first half - let gathered = _mm512_permutex2var_epi8(in_lo, idx_first, in_hi); - - // 8x8 bit transpose on all 8 groups in parallel - let mut v = gathered; - let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<7>(v)), mask1); - v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<7>(t)); - let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<14>(v)), mask2); - v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<14>(t)); - let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<28>(v)), mask3); - v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<28>(t)); - - // 8x8 byte transpose for scatter using vpermb - let scattered = _mm512_permutexvar_epi8(idx_scatter, v); - _mm512_storeu_si512(output.as_mut_ptr().cast::<__m512i>(), scattered); - - // Process second half - let gathered = _mm512_permutex2var_epi8(in_lo, idx_second, in_hi); - - let mut v = gathered; - let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<7>(v)), mask1); - v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<7>(t)); - let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<14>(v)), mask2); - v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<14>(t)); - let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<28>(v)), mask3); - v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<28>(t)); - - let scattered = _mm512_permutexvar_epi8(idx_scatter, v); - _mm512_storeu_si512(output.as_mut_ptr().add(64).cast::<__m512i>(), scattered); -} - -/// Untranspose 1024 bits using AVX-512 VBMI for vectorized scatter. -/// -/// # Safety -/// Requires AVX-512F, AVX-512BW, and AVX-512VBMI support. -#[target_feature(enable = "avx512f", enable = "avx512bw", enable = "avx512vbmi")] -#[inline(never)] -#[expect(clippy::cast_possible_wrap)] -#[expect(clippy::cast_ptr_alignment)] -#[expect(unsafe_op_in_unsafe_fn)] -pub unsafe fn untranspose_bits_vbmi(input: &[u8; 128], output: &mut [u8; 128]) { - // For untranspose, we gather consecutive bytes from transposed layout, - // then scatter back to stride-16 positions - - // Gather indices for first half - collect 8 bytes per group from transposed layout - // In transposed layout, bytes for group 0 are at: [0, 8, 16, 24, 32, 40, 48, 56] - let gather_indices: [u8; 64] = [ - 0, 8, 16, 24, 32, 40, 48, 56, // Group 0 - 1, 9, 17, 25, 33, 41, 49, 57, // Group 1 - 2, 10, 18, 26, 34, 42, 50, 58, // Group 2 - 3, 11, 19, 27, 35, 43, 51, 59, // Group 3 - 4, 12, 20, 28, 36, 44, 52, 60, // Group 4 - 5, 13, 21, 29, 37, 45, 53, 61, // Group 5 - 6, 14, 22, 30, 38, 46, 54, 62, // Group 6 - 7, 15, 23, 31, 39, 47, 55, 63, // Group 7 - ]; - - let in_first = _mm512_loadu_si512(input.as_ptr().cast::<__m512i>()); - let idx = _mm512_loadu_si512(gather_indices.as_ptr().cast::<__m512i>()); - let gathered = _mm512_permutexvar_epi8(idx, in_first); - - // 8x8 bit transpose - let mask1 = _mm512_set1_epi64(TRANSPOSE_2X2 as i64); - let mask2 = _mm512_set1_epi64(TRANSPOSE_4X4 as i64); - let mask3 = _mm512_set1_epi64(TRANSPOSE_8X8 as i64); - - let mut v = gathered; - let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<7>(v)), mask1); - v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<7>(t)); - - let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<14>(v)), mask2); - v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<14>(t)); - - let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<28>(v)), mask3); - v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<28>(t)); - - // Scatter to output at stride 16 - need to use scalar stores for now - // (AVX-512 scatter is available but complex for this pattern) - let mut result = [0u64; 8]; - _mm512_storeu_si512(result.as_mut_ptr().cast::<__m512i>(), v); - - for base_group in 0..8 { - let out_base = BASE_PATTERN_FIRST[base_group]; - for i in 0..8 { - output[out_base + i * 16] = (result[base_group] >> (i * 8)) as u8; - } - } - - // Second half - let in_second = _mm512_loadu_si512(input.as_ptr().add(64).cast::<__m512i>()); - let gathered = _mm512_permutexvar_epi8(idx, in_second); - - let mut v = gathered; - let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<7>(v)), mask1); - v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<7>(t)); - - let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<14>(v)), mask2); - v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<14>(t)); - - let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<28>(v)), mask3); - v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<28>(t)); - - _mm512_storeu_si512(result.as_mut_ptr().cast::<__m512i>(), v); - - for base_group in 0..8 { - let out_base = BASE_PATTERN_SECOND[base_group]; - for i in 0..8 { - output[out_base + i * 16] = (result[base_group] >> (i * 8)) as u8; - } - } -} - -#[cfg(test)] -mod tests { - use crate::bit_transpose::generate_test_data; - use crate::bit_transpose::transpose_bits_baseline; - use crate::bit_transpose::x86::has_bmi2; - use crate::bit_transpose::x86::has_vbmi; - use crate::bit_transpose::x86::transpose_bits_bmi2; - use crate::bit_transpose::x86::transpose_bits_vbmi; - use crate::bit_transpose::x86::untranspose_bits_bmi2; - use crate::bit_transpose::x86::untranspose_bits_vbmi; - - #[test] - fn test_bmi2_matches_baseline() { - if !has_bmi2() { - return; - } - - for seed in [0, 42, 123, 255] { - let input = generate_test_data(seed); - let mut baseline_out = [0u8; 128]; - let mut bmi2_out = [0u8; 128]; - - transpose_bits_baseline(&input, &mut baseline_out); - unsafe { transpose_bits_bmi2(&input, &mut bmi2_out) }; - - assert_eq!( - baseline_out, bmi2_out, - "BMI2 transpose doesn't match baseline for seed {seed}" - ); - } - } - - #[test] - fn test_bmi2_roundtrip() { - if !has_bmi2() { - return; - } - - for seed in [0, 42, 123, 255] { - let input = generate_test_data(seed); - let mut transposed = [0u8; 128]; - let mut roundtrip = [0u8; 128]; - - unsafe { - transpose_bits_bmi2(&input, &mut transposed); - untranspose_bits_bmi2(&transposed, &mut roundtrip); - } - - assert_eq!(input, roundtrip, "BMI2 roundtrip failed for seed {seed}"); - } - } - - #[test] - fn test_vbmi_matches_baseline() { - if !has_vbmi() { - return; - } - - for seed in [0, 42, 123, 255] { - let input = generate_test_data(seed); - let mut baseline_out = [0u8; 128]; - let mut vbmi_out = [0u8; 128]; - - transpose_bits_baseline(&input, &mut baseline_out); - unsafe { transpose_bits_vbmi(&input, &mut vbmi_out) }; - - assert_eq!( - baseline_out, vbmi_out, - "VBMI transpose doesn't match baseline for seed {seed}" - ); - } - } - - #[test] - fn test_vbmi_roundtrip() { - if !has_vbmi() { - return; - } - - for seed in [0, 42, 123, 255] { - let input = generate_test_data(seed); - let mut transposed = [0u8; 128]; - let mut roundtrip = [0u8; 128]; - - unsafe { - transpose_bits_vbmi(&input, &mut transposed); - untranspose_bits_vbmi(&transposed, &mut roundtrip); - } - - assert_eq!(input, roundtrip, "VBMI roundtrip failed for seed {seed}"); - } - } -} diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs index a393db6ecc8..4fe91aad4ff 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs @@ -512,15 +512,21 @@ mod test { let chunked = ChunkedArray::from_iter(chunks).into_array(); let into_ca = chunked.clone().execute::(&mut ctx)?; - let mut primitive_builder = - PrimitiveBuilder::::with_capacity(chunked.dtype().nullability(), 10 * 100); + let mut primitive_builder = PrimitiveBuilder::::with_capacity_in( + chunked.dtype().nullability(), + 10 * 100, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); chunked.append_to_builder(&mut primitive_builder, &mut ctx)?; let ca_into = primitive_builder.finish(); assert_arrays_eq!(into_ca, ca_into, &mut ctx); - let mut primitive_builder = - PrimitiveBuilder::::with_capacity(chunked.dtype().nullability(), 10 * 100); + let mut primitive_builder = PrimitiveBuilder::::with_capacity_in( + chunked.dtype().nullability(), + 10 * 100, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); chunked.append_to_builder(&mut primitive_builder, &mut ctx)?; let ca_into = primitive_builder.finish(); diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs index 692e7dcdd7f..0684aea5e6a 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs @@ -22,6 +22,7 @@ use vortex_error::VortexResult; use crate::BitPacked; use crate::BitPackedArrayExt; +use crate::FL_CHUNK_SIZE; use crate::unpack_iter::BitPacked as BitPackedUnpack; use crate::unpack_iter::BitUnpackedChunks; @@ -39,7 +40,11 @@ pub fn unpack_primitive_array( array: ArrayView<'_, BitPacked>, ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut builder = PrimitiveBuilder::with_capacity(array.dtype().nullability(), array.len()); + let mut builder = PrimitiveBuilder::with_capacity_in( + array.dtype().nullability(), + array.len(), + ctx.allocator(), + ); unpack_into_primitive_builder::(array, &mut builder, ctx)?; assert_eq!(builder.len(), array.len()); Ok(builder.finish_into_primitive()) @@ -102,7 +107,7 @@ where F: BitPackedUnpack, T: NativePType, M: Fn(F) -> T, - D: FnOnce(&mut BitUnpackedChunks, &mut [MaybeUninit], &M), + D: FnOnce(&mut BitUnpackedChunks<'_, F>, &mut [MaybeUninit], &M), { if array.is_empty() { return Ok(()); @@ -119,7 +124,8 @@ where // SAFETY: `decode` writes a value to every slot in this range. let uninit_slice = unsafe { uninit_range.slice_uninit_mut(0, len) }; - let mut chunks = array.unpacked_chunks::()?; + let mut scratch = [const { MaybeUninit::::uninit() }; FL_CHUNK_SIZE]; + let mut chunks = array.unpacked_chunks::(&mut scratch)?; decode(&mut chunks, uninit_slice, &map); if let Some(patches) = array.patches() { @@ -292,22 +298,24 @@ mod tests { Validity::NonNullable, )?; - let mut listview_builder = ListViewBuilder::::with_capacity( + let mut listview_builder = ListViewBuilder::::with_capacity_in( Arc::clone(&element_dtype), Nullability::NonNullable, 0, 0, + ctx.allocator(), ); list.clone() .into_array() .append_to_builder(&mut listview_builder, &mut ctx)?; assert_arrays_eq!(listview_builder.finish(), list, &mut ctx); - let mut list_builder = ListBuilder::::with_capacity( + let mut list_builder = ListBuilder::::with_capacity_in( Arc::clone(&element_dtype), Nullability::NonNullable, 0, 0, + ctx.allocator(), ); list.clone() .into_array() @@ -321,8 +329,13 @@ mod tests { Buffer::from_iter(std::iter::repeat_n(64u32, 48)).into_array(), Validity::NonNullable, )?; - let mut list_builder = - ListBuilder::::with_capacity(element_dtype, Nullability::NonNullable, 0, 0); + let mut list_builder = ListBuilder::::with_capacity_in( + element_dtype, + Nullability::NonNullable, + 0, + 0, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); listview .into_array() .append_to_builder(&mut list_builder, &mut ctx)?; @@ -450,7 +463,10 @@ mod tests { let empty: PrimitiveArray = PrimitiveArray::from_iter(Vec::::new()); let bitpacked = encode(&empty, 0); - let mut builder = PrimitiveBuilder::::new(Nullability::NonNullable); + let mut builder = PrimitiveBuilder::::new_in( + Nullability::NonNullable, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); unpack_map_into_builder( bitpacked.as_view(), &mut builder, @@ -479,7 +495,11 @@ mod tests { let bitpacked = encode(&array, 3); // Unpack into a new builder. - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::Nullable, 5); + let mut builder = PrimitiveBuilder::::with_capacity_in( + Nullability::Nullable, + 5, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); unpack_map_into_builder( bitpacked.as_view(), &mut builder, @@ -517,7 +537,11 @@ mod tests { ); // Unpack into a new builder. - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 100); + let mut builder = PrimitiveBuilder::::with_capacity_in( + Nullability::NonNullable, + 100, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); unpack_map_into_builder( bitpacked.as_view(), &mut builder, diff --git a/encodings/fastlanes/src/bitpacking/array/mod.rs b/encodings/fastlanes/src/bitpacking/array/mod.rs index 772de2ac641..05cbee8b3ef 100644 --- a/encodings/fastlanes/src/bitpacking/array/mod.rs +++ b/encodings/fastlanes/src/bitpacking/array/mod.rs @@ -3,6 +3,7 @@ use std::fmt::Display; use std::fmt::Formatter; +use std::mem::MaybeUninit; use fastlanes::BitPacking; use vortex_array::ArrayRef; @@ -29,6 +30,7 @@ pub mod bitpack_decompress; pub mod unpack_iter; use crate::BitPackedArray; +use crate::FL_CHUNK_SIZE; use crate::bitpack_compress::bitpack_encode; use crate::unpack_iter::BitPacked as BitPackedIter; use crate::unpack_iter::BitUnpackedChunks; @@ -36,12 +38,16 @@ use crate::unpack_iter::BitUnpackedChunks; #[array_slots(crate::BitPacked)] pub struct BitPackedSlots { /// The indices of exception values that don't fit in the bit-packed representation. + #[slot(0)] pub patch_indices: Option, /// The exception values that don't fit in the bit-packed representation. + #[slot(1)] pub patch_values: Option, /// Chunk offsets for the patch indices/values. + #[slot(2)] pub patch_chunk_offsets: Option, /// The validity bitmap indicating which elements are non-null. + #[slot(3)] pub validity_child: Option, } @@ -219,17 +225,18 @@ impl BitPackedData { } /// Accessor for bit unpacked chunks - pub fn unpacked_chunks( - &self, + pub fn unpacked_chunks<'a, T: BitPackedIter>( + &'a self, dtype: &DType, len: usize, - ) -> VortexResult> { + scratch: &'a mut [MaybeUninit; FL_CHUNK_SIZE], + ) -> VortexResult> { assert_eq!( T::PTYPE, self.ptype(dtype), "Requested type doesn't match the array ptype" ); - BitUnpackedChunks::try_new(self, len) + BitUnpackedChunks::try_new(self, len, scratch) } /// Bit-width of the packed values @@ -312,8 +319,16 @@ pub trait BitPackedArrayExt: BitPackedArraySlotsExt { } #[inline] - fn unpacked_chunks(&self) -> VortexResult> { - BitPackedData::unpacked_chunks::(self, self.as_ref().dtype(), self.as_ref().len()) + fn unpacked_chunks<'a, T: BitPackedIter>( + &'a self, + scratch: &'a mut [MaybeUninit; FL_CHUNK_SIZE], + ) -> VortexResult> { + BitPackedData::unpacked_chunks::( + self, + self.as_ref().dtype(), + self.as_ref().len(), + scratch, + ) } } diff --git a/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs b/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs index 3c77e146ad0..518495896e2 100644 --- a/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs +++ b/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs @@ -11,13 +11,13 @@ use lending_iterator::prelude::Item; #[gat(Item)] use lending_iterator::prelude::LendingIterator; use vortex_array::dtype::PhysicalPType; -use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use crate::BitPackedData; +use crate::FL_CHUNK_SIZE; -const CHUNK_SIZE: usize = 1024; +const CHUNK_SIZE: usize = FL_CHUNK_SIZE; /// Strategy trait for fastlanes unpacking operations pub trait UnpackStrategy { @@ -33,6 +33,7 @@ pub trait UnpackStrategy { pub struct BitPackingStrategy; impl> UnpackStrategy for BitPackingStrategy { + #[allow(clippy::inline_always)] #[inline(always)] unsafe fn unpack_chunk( &self, @@ -51,6 +52,8 @@ impl> UnpackStrategy for BitPackingStr /// /// The usual pattern of usage should follow /// ``` +/// use std::mem::MaybeUninit; +/// /// use lending_iterator::gat; /// use lending_iterator::prelude::Item; /// #[gat(Item)] @@ -64,22 +67,26 @@ impl> UnpackStrategy for BitPackingStr /// /// let mut ctx = vortex_array::array_session().create_execution_ctx(); /// let array = BitPackedData::encode(&buffer![2, 3, 4, 5].into_array(), 2, &mut ctx).unwrap(); -/// let mut unpacked_chunks: BitUnpackedChunks = array.unpacked_chunks().unwrap(); +/// let mut scratch = [const { MaybeUninit::::uninit() }; 1024]; +/// let mut unpacked_chunks: BitUnpackedChunks = +/// array.unpacked_chunks(&mut scratch).unwrap(); /// /// if let Some(header) = unpacked_chunks.initial() { /// // handle partial initial chunk /// } /// -/// let mut chunks_iter = unpacked_chunks.full_chunks(); -/// while let Some(chunk) = chunks_iter.next() { -/// // handle full bitpacked chunks of 1024 elements +/// { +/// let mut chunks_iter = unpacked_chunks.full_chunks(); +/// while let Some(chunk) = chunks_iter.next() { +/// // handle full bitpacked chunks of 1024 elements +/// } /// } /// /// if let Some(trailer) = unpacked_chunks.trailer() { /// // handle partial trailing chunk /// } /// ``` -pub struct UnpackedChunks> { +pub struct UnpackedChunks<'a, T: PhysicalPType, S: UnpackStrategy> { strategy: S, bit_width: usize, offset: usize, @@ -87,20 +94,25 @@ pub struct UnpackedChunks> { num_chunks: usize, // 0 indicates full chunk of CHUNK_SIZE last_chunk_length: usize, - packed: ByteBuffer, - buffer: [MaybeUninit; CHUNK_SIZE], + packed: &'a [T::Physical], + scratch: &'a mut [MaybeUninit; CHUNK_SIZE], } -pub type BitUnpackedChunks = UnpackedChunks; +pub type BitUnpackedChunks<'a, T> = UnpackedChunks<'a, T, BitPackingStrategy>; -impl BitUnpackedChunks { - pub fn try_new(array: &BitPackedData, len: usize) -> VortexResult { +impl<'a, T: BitPacked> BitUnpackedChunks<'a, T> { + pub fn try_new( + array: &'a BitPackedData, + len: usize, + scratch: &'a mut [MaybeUninit; CHUNK_SIZE], + ) -> VortexResult { Self::try_new_with_strategy( BitPackingStrategy, - array.packed().clone().unwrap_host(), + array.packed_slice::(), array.bit_width() as usize, array.offset() as usize, len, + scratch, ) } @@ -109,8 +121,8 @@ impl BitUnpackedChunks { let last_chunk_is_sliced = self.last_chunk_is_sliced() as usize; let first_chunk_is_sliced = self.first_chunk_is_sliced(); BitUnpackIterator::new( - buffer_as_slice(&self.packed), - &mut self.buffer, + self.packed, + self.scratch, self.bit_width, elems_per_chunk, self.num_chunks - last_chunk_is_sliced, @@ -119,41 +131,30 @@ impl BitUnpackedChunks { } } -impl> UnpackedChunks { +impl<'a, T: PhysicalPType, S: UnpackStrategy> UnpackedChunks<'a, T, S> { pub fn try_new_with_strategy( strategy: S, - packed: ByteBuffer, + packed: &'a [T::Physical], bit_width: usize, offset: usize, len: usize, + scratch: &'a mut [MaybeUninit; CHUNK_SIZE], ) -> VortexResult { - vortex_ensure!( - offset < CHUNK_SIZE, - "Invalid bit-packed offset {offset}, expected < {CHUNK_SIZE}" - ); - let elems_per_chunk = 128 * bit_width / size_of::(); - let num_chunks = (offset + len).div_ceil(CHUNK_SIZE); - - vortex_ensure!( - packed.len() / size_of::() == num_chunks * elems_per_chunk, - "Invalid packed length: got {}, expected {}", - packed.len() / size_of::(), - num_chunks * elems_per_chunk - ); - - let last_chunk_length = (offset + len) % CHUNK_SIZE; + let (num_chunks, last_chunk_length) = + validate_packed::(packed.len(), bit_width, offset, len)?; Ok(Self { strategy, bit_width, offset, len, - packed, - buffer: [const { MaybeUninit::::uninit() }; CHUNK_SIZE], num_chunks, last_chunk_length, + packed, + scratch, }) } + #[allow(clippy::inline_always)] #[inline(always)] const fn elems_per_chunk(&self) -> usize { 128 * self.bit_width / size_of::() @@ -162,8 +163,8 @@ impl> UnpackedChunks { /// Access first chunk of the array if the last chunk has fewer than 1024 due to slicing pub fn initial(&mut self) -> Option<&mut [T]> { (self.first_chunk_is_sliced() || self.num_chunks == 1).then(|| { - let chunk: &[T::Physical] = &buffer_as_slice(&self.packed)[..self.elems_per_chunk()]; - let dst: &mut [MaybeUninit] = &mut self.buffer; + let chunk: &[T::Physical] = &self.packed[..self.elems_per_chunk()]; + let dst: &mut [MaybeUninit] = self.scratch; let dst: &mut [T::Physical] = unsafe { mem::transmute(dst) }; let header_end_slice = if self.num_chunks == 1 { @@ -176,7 +177,7 @@ impl> UnpackedChunks { // 2. buffer is exactly CHUNK_SIZE. unsafe { self.strategy.unpack_chunk(self.bit_width, chunk, dst); - mem::transmute(&mut self.buffer[self.offset..][..header_end_slice]) + mem::transmute(&mut self.scratch[self.offset..][..header_end_slice]) } }) } @@ -235,14 +236,14 @@ impl> UnpackedChunks { } if self.num_chunks > 1 { - let packed_slice: &[T::Physical] = buffer_as_slice(&self.packed); + let packed_slice = self.packed; let elems_per_chunk = self.elems_per_chunk(); for i in self.full_chunks_range() { let chunk = &packed_slice[i * elems_per_chunk..][..elems_per_chunk]; unsafe { - let dst: &mut [T::Physical] = mem::transmute(&mut self.buffer[..]); + let dst: &mut [T::Physical] = mem::transmute(&mut self.scratch[..]); self.strategy.unpack_chunk(self.bit_width, chunk, dst); - let unpacked: &mut [T] = mem::transmute(&mut self.buffer[..]); + let unpacked: &mut [T] = mem::transmute(&mut self.scratch[..]); f(unpacked, local_idx..local_idx + CHUNK_SIZE); } local_idx += CHUNK_SIZE; @@ -258,34 +259,6 @@ impl> UnpackedChunks { debug_assert_eq!(local_idx, self.len); } - /// Walk every *packed* chunk in array order, yielding the raw packed FastLanes block and the - /// padded bit range it covers, without unpacking it. - /// - /// Unlike [`Self::for_each_unpacked_chunk`], this does not fill the scratch buffer: it hands - /// the still-packed block to the callback so fused kernels (e.g. compare) can unpack and - /// consume it in a single pass. Each yielded block holds exactly `elems_per_chunk` packed - /// values (the buffer is zero-padded out to a whole final chunk). - /// - /// The yielded range is in *padded* coordinates: block `c` covers - /// `[c * 1024, min((c + 1) * 1024, offset + len))`, so it includes the leading `offset` rows - /// that slicing skips. Block starts are therefore always 1024-aligned regardless of `offset`. - /// Callers must account for the array's `offset` when mapping a block's rows back to logical - /// output positions (e.g. by viewing the output buffer at a bit offset of `offset`). - pub(crate) fn for_each_packed_chunk(&self, mut f: F) - where - F: FnMut(&[T::Physical], Range), - { - let packed_slice: &[T::Physical] = buffer_as_slice(&self.packed); - let elems_per_chunk = self.elems_per_chunk(); - let padded_len = self.offset + self.len; - for chunk in 0..self.num_chunks { - let packed_chunk = &packed_slice[chunk * elems_per_chunk..][..elems_per_chunk]; - let start = chunk * CHUNK_SIZE; - let end = (start + CHUNK_SIZE).min(padded_len); - f(packed_chunk, start..end); - } - } - /// Unpack full chunks into output range starting at the given index. fn decode_full_chunks_into_at( &mut self, @@ -298,7 +271,7 @@ impl> UnpackedChunks { let mut local_idx = start_idx; - let packed_slice: &[T::Physical] = buffer_as_slice(&self.packed); + let packed_slice = self.packed; let elems_per_chunk = self.elems_per_chunk(); for i in self.full_chunks_range() { let chunk = &packed_slice[i * elems_per_chunk..][..elems_per_chunk]; @@ -322,16 +295,16 @@ impl> UnpackedChunks { /// Access last chunk of the array if the last chunk has fewer than 1024 due to slicing pub fn trailer(&mut self) -> Option<&mut [T]> { (self.last_chunk_is_sliced() && self.num_chunks > 1).then(|| { - let chunk: &[T::Physical] = &buffer_as_slice(&self.packed) + let chunk: &[T::Physical] = &self.packed [(self.num_chunks - 1) * self.elems_per_chunk()..][..self.elems_per_chunk()]; - let dst: &mut [MaybeUninit] = &mut self.buffer; + let dst: &mut [MaybeUninit] = self.scratch; let dst: &mut [T::Physical] = unsafe { mem::transmute(dst) }; // SAFETY: // 1. chunk is elems_per_chunk. // 2. buffer is exactly CHUNK_SIZE. unsafe { self.strategy.unpack_chunk(self.bit_width, chunk, dst); - mem::transmute(&mut self.buffer[..self.last_chunk_length]) + mem::transmute(&mut self.scratch[..self.last_chunk_length]) } }) } @@ -345,6 +318,50 @@ impl> UnpackedChunks { } } +/// Walk every packed chunk in array order without allocating an unpack scratch buffer. +pub(crate) fn for_each_packed_chunk( + packed: &[T::Physical], + bit_width: usize, + offset: usize, + len: usize, + mut f: F, +) -> VortexResult<()> +where + T: PhysicalPType, + F: FnMut(&[T::Physical], Range), +{ + let (num_chunks, _) = validate_packed::(packed.len(), bit_width, offset, len)?; + let elems_per_chunk = 128 * bit_width / size_of::(); + let padded_len = offset + len; + for chunk in 0..num_chunks { + let packed_chunk = &packed[chunk * elems_per_chunk..][..elems_per_chunk]; + let start = chunk * CHUNK_SIZE; + let end = (start + CHUNK_SIZE).min(padded_len); + f(packed_chunk, start..end); + } + Ok(()) +} + +fn validate_packed( + packed_len: usize, + bit_width: usize, + offset: usize, + len: usize, +) -> VortexResult<(usize, usize)> { + vortex_ensure!( + offset < CHUNK_SIZE, + "Invalid bit-packed offset {offset}, expected < {CHUNK_SIZE}" + ); + let elems_per_chunk = 128 * bit_width / size_of::(); + let num_chunks = (offset + len).div_ceil(CHUNK_SIZE); + vortex_ensure!( + packed_len == num_chunks * elems_per_chunk, + "Invalid packed length: got {packed_len}, expected {}", + num_chunks * elems_per_chunk + ); + Ok((num_chunks, (offset + len) % CHUNK_SIZE)) +} + /// Iterator over full chunks of bitpacked array that yields unpacked chunks one at a time pub struct BitUnpackIterator<'a, T: BitPacked + 'a> { packed: &'a [T::Physical], @@ -401,17 +418,6 @@ impl<'a, T: BitPacked + 'a> LendingIterator for BitUnpackIterator<'a, T> { } } -fn buffer_as_slice(buffer: &ByteBuffer) -> &[T] { - let packed_ptr: *const T = buffer.as_ptr().cast(); - // Return number of elements of type `T` packed in the buffer - let packed_len = buffer.len() / size_of::(); - - // SAFETY: as_slice points to buffer memory that outlives the lifetime of `self`. - // Unfortunately Rust cannot understand this, so we reconstruct the slice from raw parts - // to get it to reinterpret the lifetime. - unsafe { std::slice::from_raw_parts(packed_ptr, packed_len) } -} - fn write_map(src: &[T], dst: &mut [MaybeUninit], f: &mut impl FnMut(T) -> U) { for (dst, &src) in dst.iter_mut().zip(src.iter()) { dst.write(f(src)); diff --git a/encodings/fastlanes/src/bitpacking/compute/between.rs b/encodings/fastlanes/src/bitpacking/compute/between.rs index 1dd3b61dbf4..8010fa208d2 100644 --- a/encodings/fastlanes/src/bitpacking/compute/between.rs +++ b/encodings/fastlanes/src/bitpacking/compute/between.rs @@ -55,10 +55,10 @@ impl BetweenKernel for BitPacked { let result = match_each_integer_ptype!(arr_ptype, |T| { let lo: T = lower_prim .typed_value::() - .vortex_expect("between precondition strips null lower"); + .vortex_expect("the between short circuit strips a null lower bound"); let up: T = upper_prim .typed_value::() - .vortex_expect("between precondition strips null upper"); + .vortex_expect("the between short circuit strips a null upper bound"); between_constant_typed::(array, lo, up, options, nullability, ctx)? }); Ok(Some(result)) diff --git a/encodings/fastlanes/src/bitpacking/compute/cast.rs b/encodings/fastlanes/src/bitpacking/compute/cast.rs index f7cfc56acf0..cdbb8141cea 100644 --- a/encodings/fastlanes/src/bitpacking/compute/cast.rs +++ b/encodings/fastlanes/src/bitpacking/compute/cast.rs @@ -100,7 +100,11 @@ impl CastKernel for BitPacked { .cast_nullability(tgt_nullability, array.len(), ctx)?; let result = match_each_integer_ptype!(tgt, |T| { - let mut builder = PrimitiveBuilder::::with_capacity(tgt_nullability, array.len()); + let mut builder = PrimitiveBuilder::::with_capacity_in( + tgt_nullability, + array.len(), + ctx.allocator(), + ); match_each_integer_ptype!(src, |F| { unpack_map_into_builder::(array, &mut builder, ctx, |v: F| v.as_())?; }); diff --git a/encodings/fastlanes/src/bitpacking/compute/compare.rs b/encodings/fastlanes/src/bitpacking/compute/compare.rs index 51aca62a9aa..c9d6b815b0d 100644 --- a/encodings/fastlanes/src/bitpacking/compute/compare.rs +++ b/encodings/fastlanes/src/bitpacking/compute/compare.rs @@ -45,8 +45,8 @@ impl CompareKernel for BitPacked { return Ok(None); }; - // Adaptor strips null-constant RHS, and the binary scalar-fn coerce_args step has - // already promoted both sides to a common ptype. + // The adaptor strips a null-constant RHS, and binary expressions require both argument + // types to match before dispatch. let nullability = lhs.dtype().nullability() | rhs.dtype().nullability(); let lhs_ptype = lhs.dtype().as_ptype(); if constant_prim.ptype() != lhs_ptype { diff --git a/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs b/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs index baf3b0ede96..1259ed815fe 100644 --- a/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs +++ b/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs @@ -8,16 +8,15 @@ //! FastLanes [`BitPackingCompare::unchecked_unpack_cmp`] kernel, which compares each value against //! the constant *as it is unpacked*, accumulating the boolean results straight into a 1024-bit //! mask (`[u64; 16]`) in transposed FastLanes lane order - one register-resident word per lane, no -//! `[bool; 1024]` or `[T; 1024]` scratch. A single SIMD [`untranspose_bits`] per block then rotates +//! `[bool; 1024]` or `[T; 1024]` scratch. A single SIMD [`transpose_bits`] per block then rotates //! that mask into logical row order. //! -//! The packed blocks are walked through the regular [`crate::unpack_iter::BitUnpackedChunks`] -//! iterator (via [`crate::unpack_iter::BitUnpackedChunks::for_each_packed_chunk`]) rather than a -//! bespoke chunk loop, so chunk sizing and bounds live in one place. +//! The packed blocks are walked through [`crate::unpack_iter::for_each_packed_chunk`], so chunk +//! sizing and bounds live in one place without allocating an unpack scratch buffer. //! //! Slicing is handled by working in *padded* coordinates: bit `offset + i` holds element `i`. The //! output buffer is over-allocated to whole 1024-bit blocks, so every block - the sliced first -//! block, the body, and the trailing partial - untransposes straight into a 64-bit-word-aligned +//! block, the body, and the trailing partial - transposes straight into a 64-bit-word-aligned //! slot with no per-block temporary and only one shared scratch `[u64; 16]`. The leading `offset` //! garbage rows are represented as the final [`BitBuffer`] bit offset, which naturally handles //! sub-byte slices without copy-aligning. Inline patches are spliced in afterwards by overwriting @@ -29,7 +28,7 @@ use fastlanes::BitPacking; use fastlanes::BitPackingCompare; use fastlanes::FastLanesComparable; -use fastlanes::untranspose_bits; +use fastlanes::transpose_bits; use num_traits::AsPrimitive; use vortex_array::ArrayRef; use vortex_array::ArrayView; @@ -50,6 +49,7 @@ use super::stream_predicate::stream_predicate; use crate::BitPacked; use crate::BitPackedArrayExt; use crate::unpack_iter::BitPacked as BitPackedIter; +use crate::unpack_iter::for_each_packed_chunk; const CHUNK_SIZE: usize = 1024; const U64_BITS: usize = u64::BITS as usize; @@ -92,29 +92,31 @@ where let num_chunks = (offset + len).div_ceil(CHUNK_SIZE); let mut words: BufferMut = BufferMut::zeroed(num_chunks * WORDS_PER_CHUNK); - let chunks = array.unpacked_chunks::()?; { let words = words.as_mut_slice(); - let mut transposed = [0u64; WORDS_PER_CHUNK]; - chunks.for_each_packed_chunk(|packed_chunk, range| { - // Block starts are always 1024-aligned (padded coords), so the slot is a full block. - let out = words[range.start / U64_BITS..] - .first_chunk_mut::() - .vortex_expect("over-allocated buffer holds a full block per chunk"); - // SAFETY: `packed_chunk` holds exactly `128 * bit_width / size_of::()` packed - // elements and `bit_width <= U::T`, satisfying `unchecked_unpack_cmp`'s contract. The - // kernel assigns every word in `transposed`, so its previous contents are irrelevant. - unsafe { - <::Physical as BitPackingCompare>::unchecked_unpack_cmp::( - bit_width, - packed_chunk, - &mut transposed, - cmp, - rhs, - ); - } - untranspose_bits::<::Physical>(&transposed, out); - }); + let mut lane_major = [0u64; WORDS_PER_CHUNK]; + for_each_packed_chunk::( + array.packed_slice::<::Physical>(), + bit_width, + offset, + len, + |packed_chunk, range| { + // Block starts are always 1024-aligned (padded coords), so the slot is a full block. + let out = words[range.start / U64_BITS..] + .first_chunk_mut::() + .vortex_expect("over-allocated buffer holds a full block per chunk"); + // SAFETY: `packed_chunk` holds exactly `128 * bit_width / size_of::()` packed + // elements and `bit_width <= U::T`, satisfying `unchecked_unpack_cmp`'s contract. The + // kernel assigns every word in `transposed`, so its previous contents are irrelevant. + unsafe { + <::Physical as BitPackingCompare>::unchecked_unpack_cmp::< + T, + _, + >(bit_width, packed_chunk, &mut lane_major, cmp, rhs); + } + transpose_bits::<::Physical>(&lane_major, out); + }, + )?; } let mut bits = BitBufferMut::from_buffer(words.into_byte_buffer(), offset, len); diff --git a/encodings/fastlanes/src/bitpacking/compute/filter.rs b/encodings/fastlanes/src/bitpacking/compute/filter.rs index 1530c33874d..0b1b9422f86 100644 --- a/encodings/fastlanes/src/bitpacking/compute/filter.rs +++ b/encodings/fastlanes/src/bitpacking/compute/filter.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::mem::MaybeUninit; -use std::sync::Arc; use fastlanes::BitPacking; use vortex_array::ArrayRef; @@ -20,7 +19,7 @@ use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_mask::Mask; -use vortex_mask::MaskValues; +use vortex_mask::MaskValuesRef; use super::chunked_indices; use super::take::UNPACK_CHUNK_THRESHOLD; @@ -82,7 +81,7 @@ impl FilterKernel for BitPacked { let patches = array .patches() - .map(|patches| patches.filter(&Mask::Values(Arc::clone(values)), ctx)) + .map(|patches| patches.filter(&Mask::Values(MaskValuesRef::clone(values)), ctx)) .transpose()? .flatten(); @@ -109,12 +108,12 @@ impl FilterKernel for BitPacked { /// Returns a tuple of (values buffer, validity mask). fn filter_primitive_without_patches( array: ArrayView<'_, BitPacked>, - selection: &Arc, + selection: &MaskValuesRef, ) -> VortexResult<(Buffer, Validity)> { let values = filter_with_indices(array.data(), selection.indices()); let validity = array .validity()? - .filter(&Mask::Values(Arc::clone(selection)))?; + .filter(&Mask::Values(MaskValuesRef::clone(selection)))?; Ok((values.freeze(), validity)) } @@ -176,7 +175,7 @@ fn filter_with_indices( } #[cfg(test)] -mod test { +mod tests { use std::sync::LazyLock; use vortex_array::IntoArray as _; diff --git a/encodings/fastlanes/src/bitpacking/compute/is_constant.rs b/encodings/fastlanes/src/bitpacking/compute/is_constant.rs index 3b4e2b9770a..0ab01a635ba 100644 --- a/encodings/fastlanes/src/bitpacking/compute/is_constant.rs +++ b/encodings/fastlanes/src/bitpacking/compute/is_constant.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::mem::MaybeUninit; use std::ops::Range; use itertools::Itertools; @@ -55,7 +56,8 @@ fn bitpacked_is_constant( array: ArrayView<'_, BitPacked>, ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut bit_unpack_iterator = array.unpacked_chunks::()?; + let mut scratch = [const { MaybeUninit::::uninit() }; 1024]; + let mut bit_unpack_iterator = array.unpacked_chunks::(&mut scratch)?; let patches = array .patches() .map(|p| -> VortexResult<_> { @@ -87,37 +89,39 @@ fn bitpacked_is_constant( } let mut first_chunk_value = None; - let mut chunks_iter = bit_unpack_iterator.full_chunks(); - while let Some(chunk) = chunks_iter.next() { - if let Some((indices, patches, offset)) = &patches { - let chunk_len = chunk.len(); - apply_patches( - chunk, - current_idx..current_idx + chunk_len, - indices, - patches.as_slice::(), - *offset, - ) - } - - if !compute_is_constant::<_, WIDTH>(chunk) { - return Ok(false); - } + { + let mut chunks_iter = bit_unpack_iterator.full_chunks(); + while let Some(chunk) = chunks_iter.next() { + if let Some((indices, patches, offset)) = &patches { + let chunk_len = chunk.len(); + apply_patches( + chunk, + current_idx..current_idx + chunk_len, + indices, + patches.as_slice::(), + *offset, + ) + } - if let Some(chunk_value) = first_chunk_value { - if chunk_value != chunk[0] { + if !compute_is_constant::<_, WIDTH>(chunk) { return Ok(false); } - } else { - if let Some(header_value) = header_constant_value - && header_value != chunk[0] - { - return Ok(false); + + if let Some(chunk_value) = first_chunk_value { + if chunk_value != chunk[0] { + return Ok(false); + } + } else { + if let Some(header_value) = header_constant_value + && header_value != chunk[0] + { + return Ok(false); + } + first_chunk_value = Some(chunk[0]); } - first_chunk_value = Some(chunk[0]); - } - current_idx += chunk.len(); + current_idx += chunk.len(); + } } if let Some(trailer) = bit_unpack_iterator.trailer() { diff --git a/encodings/fastlanes/src/bitpacking/compute/stream_predicate.rs b/encodings/fastlanes/src/bitpacking/compute/stream_predicate.rs index a85cc358c60..9154ca736c1 100644 --- a/encodings/fastlanes/src/bitpacking/compute/stream_predicate.rs +++ b/encodings/fastlanes/src/bitpacking/compute/stream_predicate.rs @@ -15,6 +15,8 @@ //! [`BitBuffer::collect_bool`]: vortex_buffer::BitBuffer::collect_bool //! [`Patches`]: vortex_array::patches::Patches +use std::mem::MaybeUninit; + use num_traits::AsPrimitive; use vortex_array::ArrayRef; use vortex_array::ArrayView; @@ -32,6 +34,7 @@ use vortex_error::VortexResult; use crate::BitPacked; use crate::BitPackedArrayExt; +use crate::FL_CHUNK_SIZE; use crate::unpack_iter::BitPacked as BitPackedIter; /// Stream `predicate` over the unpacked values of a [`BitPackedArray`](crate::BitPackedArray), one FastLanes @@ -50,7 +53,8 @@ where let mut words: BufferMut = BufferMut::zeroed(len.div_ceil(u64::BITS as usize)); if len > 0 { - let mut chunks = array.unpacked_chunks::()?; + let mut scratch = [const { MaybeUninit::::uninit() }; FL_CHUNK_SIZE]; + let mut chunks = array.unpacked_chunks::(&mut scratch)?; let words = words.as_mut_slice(); if let Some(p) = array.patches() { diff --git a/encodings/fastlanes/src/bitpacking/compute/take.rs b/encodings/fastlanes/src/bitpacking/compute/take.rs index 77f6b9d3127..86e97623cf6 100644 --- a/encodings/fastlanes/src/bitpacking/compute/take.rs +++ b/encodings/fastlanes/src/bitpacking/compute/take.rs @@ -99,11 +99,10 @@ fn take_primitive( let packed = &packed[chunk_idx * chunk_len..][..chunk_len]; let mut have_unpacked = false; - let mut offset_chunk_iter = indices_within_chunk.chunks_exact(UNPACK_CHUNK_THRESHOLD); + let (offset_chunks, remainder) = indices_within_chunk.as_chunks::(); // this loop only runs if we have at least UNPACK_CHUNK_THRESHOLD offsets - for offset_chunk in &mut offset_chunk_iter { - assert_eq!(offset_chunk.len(), UNPACK_CHUNK_THRESHOLD); // let compiler know slice length + for offset_chunk in offset_chunks { if !have_unpacked { unsafe { let dst: &mut [MaybeUninit] = &mut unpacked; @@ -119,16 +118,16 @@ fn take_primitive( } // if we have a remainder (i.e., < UNPACK_CHUNK_THRESHOLD leftover offsets), we need to handle it - if !offset_chunk_iter.remainder().is_empty() { + if !remainder.is_empty() { if have_unpacked { // we already bulk unpacked this chunk, so we can just push the remaining elements - for &index in offset_chunk_iter.remainder() { + for &index in remainder { output.push(unsafe { unpacked[index].assume_init() }); } } else { // we had fewer than UNPACK_CHUNK_THRESHOLD offsets in the first place, // so we need to unpack each one individually - for &index in offset_chunk_iter.remainder() { + for &index in remainder { output.push(unsafe { bitpack_decompress::unpack_single_primitive::(packed, bit_width, index) }); diff --git a/encodings/fastlanes/src/bitpacking/plugin.rs b/encodings/fastlanes/src/bitpacking/plugin.rs index a621d085514..3ff07db7e5c 100644 --- a/encodings/fastlanes/src/bitpacking/plugin.rs +++ b/encodings/fastlanes/src/bitpacking/plugin.rs @@ -7,17 +7,17 @@ //! This enables zero-cost backward compatibility with previously written datasets. use vortex_array::Array; +use vortex_array::ArrayDeserialization; use vortex_array::ArrayId; use vortex_array::ArrayPlugin; use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; use vortex_array::ArrayVTable; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::Patched; -use vortex_array::buffer::BufferHandle; -use vortex_array::dtype::DType; -use vortex_array::serde::ArrayChildren; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_session::VortexSession; @@ -41,22 +41,29 @@ impl ArrayPlugin for BitPackedPatchedPlugin { &self, array: &ArrayRef, session: &VortexSession, - ) -> VortexResult>> { + ) -> VortexResult> { // delegate to BitPacked VTable for serialization - BitPacked.serialize(array, session) + ArrayPlugin::serialize(&BitPacked, array, session) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, session: &VortexSession, ) -> VortexResult { + vortex_ensure!( + parts.serialized_id == self.id(), + "BitPacked plugin does not recognize serialized ID {}", + parts.serialized_id, + ); let bitpacked = Array::::try_from_parts(ArrayVTable::deserialize( - &BitPacked, dtype, len, metadata, buffers, children, session, + &BitPacked, + parts.dtype, + parts.len, + parts.metadata, + parts.buffers, + parts.children, + session, )?) .map_err(|_| vortex_err!("BitPacked plugin should only deserialize fastlanes.bitpacked"))?; @@ -93,6 +100,7 @@ impl ArrayPlugin for BitPackedPatchedPlugin { mod tests { use std::sync::LazyLock; + use vortex_array::ArrayDeserialization; use vortex_array::ArrayPlugin; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; @@ -134,7 +142,7 @@ mod tests { let array = bitpacked.as_array(); - let metadata = SESSION.array_serialize(array)?.unwrap(); + let serialization = SESSION.array_serialize(array)?.unwrap(); let children = array.children(); let buffers = array .buffers() @@ -143,11 +151,14 @@ mod tests { .collect::>(); let deserialized = BitPackedPatchedPlugin.deserialize( - array.dtype(), - array.len(), - &metadata, - &buffers, - &children, + ArrayDeserialization::new( + BitPackedPatchedPlugin.id(), + array.dtype(), + array.len(), + &serialization.metadata, + &buffers, + &children, + ), &SESSION, )?; @@ -184,7 +195,7 @@ mod tests { let array = bitpacked.as_array(); - let metadata = SESSION.array_serialize(array)?.unwrap(); + let serialization = SESSION.array_serialize(array)?.unwrap(); let children = array.children(); let buffers = array .buffers() @@ -193,11 +204,14 @@ mod tests { .collect::>(); let deserialized = BitPackedPatchedPlugin.deserialize( - array.dtype(), - array.len(), - &metadata, - &buffers, - &children, + ArrayDeserialization::new( + BitPackedPatchedPlugin.id(), + array.dtype(), + array.len(), + &serialization.metadata, + &buffers, + &children, + ), &SESSION, )?; @@ -214,7 +228,7 @@ mod tests { fn primitive_array_returns_error() -> VortexResult<()> { let array = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); - let metadata = SESSION.array_serialize(&array)?.unwrap(); + let serialization = SESSION.array_serialize(&array)?.unwrap(); let children = array.children(); let buffers = array .buffers() @@ -223,11 +237,14 @@ mod tests { .collect::>(); let result = BitPackedPatchedPlugin.deserialize( - array.dtype(), - array.len(), - &metadata, - &buffers, - &children, + ArrayDeserialization::new( + BitPackedPatchedPlugin.id(), + array.dtype(), + array.len(), + &serialization.metadata, + &buffers, + &children, + ), &SESSION, ); diff --git a/encodings/fastlanes/src/delta/array/delta_compress.rs b/encodings/fastlanes/src/delta/array/delta_compress.rs index 7f00eb172dc..c2ef38ceb79 100644 --- a/encodings/fastlanes/src/delta/array/delta_compress.rs +++ b/encodings/fastlanes/src/delta/array/delta_compress.rs @@ -8,17 +8,22 @@ use fastlanes::Delta; use fastlanes::FastLanes; use fastlanes::Transpose; use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::dtype::NativePType; use vortex_array::match_each_unsigned_integer_ptype; +use vortex_array::validity::Validity; +use vortex_buffer::BitBufferMut; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use crate::FL_CHUNK_SIZE; -use crate::bit_transpose::transpose_validity; +use crate::bit_transpose::transpose_bitbuffer; use crate::fill_forward_nulls; + pub fn delta_compress( array: &PrimitiveArray, ctx: &mut ExecutionCtx, @@ -33,8 +38,29 @@ pub fn delta_compress( // corrupted delta values propagate through the cumulative sum during decompression. let filled = fill_forward_nulls(array.to_buffer::(), &validity, ctx)?; let (bases, deltas) = compress_primitive::(&filled); - // TODO(robert): This can be avoided if we add TransposedBoolArray that performs index translation when necessary. - let validity = transpose_validity(&validity, ctx)?; + let validity = match validity { + Validity::Array(mask) => { + let bits = mask.execute::(ctx)?.into_bit_buffer(); + let pad = bits.len().next_multiple_of(FL_CHUNK_SIZE) - bits.len(); + // Pad remainder bits as valid to match last-value remainder padding. + // `transpose_bitbuffer` uses the same element order as the value transpose, so + // the pad bits land on exactly the padded value slots; zero-filling them would + // mark those slots null and bitpacking would then skip their patches. + let bits = if pad == 0 { + bits + } else { + // `sliced` first so the copy covers only the logical range, not whatever + // wider buffer the mask was sliced out of. + let mut padded = BitBufferMut::copy_from(&bits.sliced()); + padded.append_n(true, pad); + padded.freeze() + }; + Validity::Array( + BoolArray::new(transpose_bitbuffer(bits), Validity::NonNullable).into_array(), + ) + } + validity => validity, + }; ( PrimitiveArray::new(bases, array.dtype().nullability().into()), PrimitiveArray::new(deltas, validity), @@ -83,8 +109,10 @@ where } // Pad the remainder to 1024 elements and process as a full chunk. - if !remainder.is_empty() { - let mut padded_chunk = [T::default(); FL_CHUNK_SIZE]; + if let Some(&last) = remainder.last() { + // Repeat the last value for padding to prevent a value-to-zero step from producing + // huge wrapping deltas in the padded tail (same rationale as RLE compression). + let mut padded_chunk = [last; FL_CHUNK_SIZE]; padded_chunk[..remainder.len()].copy_from_slice(remainder); process_chunk(&padded_chunk, &mut output_deltas[full_chunks.len()]); } @@ -99,18 +127,25 @@ where #[cfg(test)] mod tests { + use std::iter; use std::sync::LazyLock; use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; + use vortex_array::arrays::Bool; + use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; + use vortex_array::validity::Validity; use vortex_error::VortexExpect; use vortex_error::VortexResult; + use vortex_error::vortex_bail; use vortex_session::VortexSession; use crate::Delta; + use crate::FL_CHUNK_SIZE; + use crate::bit_transpose::untranspose_bitbuffer; use crate::bitpack_compress::bitpack_encode; use crate::delta::array::delta_decompress::delta_decompress; use crate::delta_compress; @@ -155,6 +190,114 @@ mod tests { Ok(()) } + /// Zero-padding the trailing chunk inflated delta span on unaligned monotone columns, + /// causing DeltaScheme to reject encoding. Pad positions must repeat the last value. + #[test] + fn remainder_pad_preserves_small_delta_span() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + for n in [1025usize, 2049] { + let array = PrimitiveArray::from_iter((0..n as u32).map(|i| 1000 + i)); + let (_bases, deltas) = delta_compress(&array, &mut ctx)?; + let d = deltas.as_slice::(); + let min = *d.iter().min().unwrap(); + let max = *d.iter().max().unwrap(); + assert!( + max - min <= 1, + "n={n}: delta span should stay O(1), got min={min} max={max}", + ); + assert!( + !d.iter().any(|&v| v > u32::MAX / 2), + "n={n}: padding must not produce wrapping deltas", + ); + } + Ok(()) + } + + /// Padding remainder validity with `true` must not change logical nulls, including leading + /// and trailing nulls in the unaligned tail. After untranspose, pad bits are valid and are + /// sliced off by `logical_len`. + /// + /// Which transposed slots the pad bits land on varies with the remainder length, so cover + /// several, plus an aligned length that pads nothing at all. + #[rstest] + #[case::one_row_remainder(1025)] + #[case::mid_chunk_remainder(1500)] + #[case::two_chunks_plus_one(2049)] + #[case::one_row_short_of_aligned(3071)] + #[case::already_aligned(2048)] + fn remainder_validity_pad_does_not_clobber_logical_nulls( + #[case] len: usize, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + // Nulls at both ends, so a leading null and a null inside the padded tail are covered. + let array = PrimitiveArray::from_option_iter( + iter::once(None) + .chain((1..len as i32 - 1).map(Some)) + .chain(iter::once(None)), + ); + assert_eq!(array.len(), len); + + let (bases, deltas) = delta_compress(&array, &mut ctx)?; + let padded_len = len.next_multiple_of(FL_CHUNK_SIZE); + assert_eq!(deltas.len(), padded_len); + + let Validity::Array(storage) = deltas.validity()? else { + vortex_bail!("expected array-backed storage validity") + }; + let sequential = + untranspose_bitbuffer(storage.execute::(&mut ctx)?.into_bit_buffer()); + assert_eq!(sequential.len(), padded_len); + for i in 0..len { + assert_eq!( + sequential.value(i), + array.is_valid(i, &mut ctx)?, + "logical validity changed at {i}" + ); + } + for i in len..padded_len { + assert!(sequential.value(i), "pad bit {i} should be valid"); + } + + let delta = Delta::try_new(bases.into_array(), deltas.into_array(), 0, len)?; + assert_eq!(delta.len(), len); + assert!(!delta.is_valid(0, &mut ctx)?); + assert!(!delta.is_valid(len - 1, &mut ctx)?); + assert_arrays_eq!(delta, array, &mut ctx); + Ok(()) + } + + /// The transposed validity must line up slot for slot with the transposed delta values: + /// storage slot `j` describes logical row `transpose(j)`, the position `Transpose::transpose` + /// put that row's delta in. `fastlanes` unified the element order of the bit and value + /// transposes; before that these two permutations disagreed. + #[test] + fn storage_validity_aligns_with_transposed_value_slots() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + // Chunk-aligned, so no remainder padding takes part. + const LEN: usize = 2 * FL_CHUNK_SIZE; + let valid = |i: usize| !i.is_multiple_of(3); + let array = + PrimitiveArray::from_option_iter((0..LEN).map(|i| valid(i).then_some(i as u32))); + + let (_bases, deltas) = delta_compress(&array, &mut ctx)?; + let Validity::Array(storage) = deltas.validity()? else { + vortex_bail!("expected array-backed storage validity") + }; + let bits = storage.execute::(&mut ctx)?.into_bit_buffer(); + + for chunk in 0..LEN / FL_CHUNK_SIZE { + let base = chunk * FL_CHUNK_SIZE; + for slot in 0..FL_CHUNK_SIZE { + assert_eq!( + bits.value(base + slot), + valid(base + fastlanes::transpose(slot)), + "chunk={chunk} slot={slot}" + ); + } + } + Ok(()) + } + /// Regression test: delta + bitpacked encoding must correctly round-trip nullable arrays /// where null positions contain arbitrary values. Without fill-forward, the delta cumulative /// sum propagates corrupted values from null positions. @@ -165,6 +308,11 @@ mod tests { (0u8..200).map(|i| (!(50..100).contains(&i)).then_some(i)), ); let (bases, deltas) = delta_compress(&array, &mut ctx)?; + let Validity::Array(storage_validity) = deltas.validity()? else { + vortex_bail!("test input should have array-backed validity") + }; + assert!(storage_validity.is::()); + let bitpacked_deltas = bitpack_encode(&deltas, 1, None, &mut ctx)?; let packed_delta = Delta::try_new( bases.into_array(), diff --git a/encodings/fastlanes/src/delta/array/delta_decompress.rs b/encodings/fastlanes/src/delta/array/delta_decompress.rs index 7dcbeb44950..3ccfede27c1 100644 --- a/encodings/fastlanes/src/delta/array/delta_decompress.rs +++ b/encodings/fastlanes/src/delta/array/delta_decompress.rs @@ -18,8 +18,8 @@ use vortex_buffer::BufferMut; use vortex_error::VortexResult; use crate::DeltaArray; -use crate::bit_transpose::untranspose_validity; use crate::delta::array::DeltaArrayExt; +use crate::delta::array::DeltaArraySlotsExt; pub fn delta_decompress( array: &DeltaArray, @@ -31,8 +31,7 @@ pub fn delta_decompress( let start = array.offset(); let end = start + array.len(); - let validity = untranspose_validity(&deltas.validity()?, ctx)?; - let validity = validity.slice(start..end)?; + let validity = array.validity()?; let original_ptype = deltas.ptype(); // Signed inputs are processed through their unsigned counterpart; `wrapping_add` on the diff --git a/encodings/fastlanes/src/delta/array/mod.rs b/encodings/fastlanes/src/delta/array/mod.rs index 7754bb37a59..444c66e11c5 100644 --- a/encodings/fastlanes/src/delta/array/mod.rs +++ b/encodings/fastlanes/src/delta/array/mod.rs @@ -7,27 +7,32 @@ use std::fmt::Formatter; use fastlanes::FastLanes; use vortex_array::ArrayRef; use vortex_array::TypedArrayRef; +use vortex_array::array_slots; use vortex_array::dtype::PType; use vortex_array::match_each_unsigned_integer_ptype; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; +use crate::Delta; + pub mod delta_compress; pub mod delta_decompress; -/// The base values for each block of deltas. -pub(super) const BASES_SLOT: usize = 0; -/// The delta-encoded values relative to the base values. -pub(super) const DELTAS_SLOT: usize = 1; -pub(super) const NUM_SLOTS: usize = 2; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["bases", "deltas"]; +#[array_slots(Delta)] +pub struct DeltaSlots { + /// The base values for each block of deltas. + #[slot(0)] + pub bases: ArrayRef, + /// The delta-encoded values relative to the base values. + #[slot(1)] + pub deltas: ArrayRef, +} /// A FastLanes-style delta-encoded array of primitive values. /// /// A DeltaArray comprises a sequence of _chunks_ each representing exactly 1,024 /// delta-encoded values. If the input array length is not a multiple of 1,024, the last chunk -/// is padded with zeros to fill a complete 1,024-element chunk. +/// is padded with the last value to fill a complete 1,024-element chunk. /// /// # Examples /// @@ -88,25 +93,13 @@ impl Display for DeltaData { } } -pub trait DeltaArrayExt: TypedArrayRef { - fn bases(&self) -> &ArrayRef { - self.as_ref().slots()[BASES_SLOT] - .as_ref() - .vortex_expect("DeltaArray bases slot") - } - - fn deltas(&self) -> &ArrayRef { - self.as_ref().slots()[DELTAS_SLOT] - .as_ref() - .vortex_expect("DeltaArray deltas slot") - } - +pub trait DeltaArrayExt: DeltaArraySlotsExt { fn offset(&self) -> usize { self.offset } } -impl> DeltaArrayExt for T {} +impl> DeltaArrayExt for T {} impl DeltaData { pub fn try_new(offset: usize) -> VortexResult { diff --git a/encodings/fastlanes/src/delta/compute/cast.rs b/encodings/fastlanes/src/delta/compute/cast.rs index c6de602bb0b..5fe84f62b62 100644 --- a/encodings/fastlanes/src/delta/compute/cast.rs +++ b/encodings/fastlanes/src/delta/compute/cast.rs @@ -12,6 +12,7 @@ use vortex_error::VortexResult; use crate::delta::Delta; use crate::delta::array::DeltaArrayExt; +use crate::delta::array::DeltaArraySlotsExt; impl CastReduce for Delta { fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { diff --git a/encodings/fastlanes/src/delta/mod.rs b/encodings/fastlanes/src/delta/mod.rs index 52ea9b33574..3cf2d4a6815 100644 --- a/encodings/fastlanes/src/delta/mod.rs +++ b/encodings/fastlanes/src/delta/mod.rs @@ -2,7 +2,10 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod array; +pub use array::DeltaArrayExt; +pub use array::DeltaArraySlotsExt; pub use array::DeltaData; +pub use array::DeltaSlots; pub use array::delta_compress::delta_compress; mod compute; diff --git a/encodings/fastlanes/src/delta/vtable/mod.rs b/encodings/fastlanes/src/delta/vtable/mod.rs index 6dcfedc0bc9..1f2e3823689 100644 --- a/encodings/fastlanes/src/delta/vtable/mod.rs +++ b/encodings/fastlanes/src/delta/vtable/mod.rs @@ -21,9 +21,7 @@ use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::PType; use vortex_array::serde::ArrayChildren; -use vortex_array::smallvec::smallvec; use vortex_array::vtable::VTable; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; @@ -32,10 +30,10 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::DeltaData; -use crate::delta::array::BASES_SLOT; -use crate::delta::array::DELTAS_SLOT; use crate::delta::array::DeltaArrayExt; -use crate::delta::array::SLOT_NAMES; +use crate::delta::array::DeltaArraySlotsExt; +use crate::delta::array::DeltaSlots; +use crate::delta::array::DeltaSlotsView; use crate::delta::array::delta_decompress::delta_decompress; use crate::delta::array::lane_count; use crate::delta_compress; @@ -87,13 +85,14 @@ impl VTable for Delta { len: usize, slots: &[Option], ) -> VortexResult<()> { - let bases = slots[BASES_SLOT] - .as_ref() - .vortex_expect("DeltaArray bases slot"); - let deltas = slots[DELTAS_SLOT] - .as_ref() - .vortex_expect("DeltaArray deltas slot"); - validate_parts(bases, deltas, data.offset, dtype, len) + let delta_slots = DeltaSlotsView::from_slots(slots); + validate_parts( + delta_slots.bases, + delta_slots.deltas, + data.offset, + dtype, + len, + ) } fn nbuffers(_array: ArrayView<'_, Self>) -> usize { @@ -125,7 +124,7 @@ impl VTable for Delta { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + DeltaSlots::NAMES[idx].to_string() } fn serialize( @@ -175,7 +174,7 @@ impl VTable for Delta { let deltas = children.get(1, dtype, deltas_len)?; let data = DeltaData::try_new(metadata.offset as usize)?; - let slots = smallvec![Some(bases), Some(deltas)]; + let slots = DeltaSlots { bases, deltas }.into_slots(); Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) } @@ -198,7 +197,7 @@ impl Delta { ) -> VortexResult { let dtype = bases.dtype().with_nullability(deltas.dtype().nullability()); let data = DeltaData::try_new(offset)?; - let slots = smallvec![Some(bases), Some(deltas)]; + let slots = DeltaSlots { bases, deltas }.into_slots(); Array::try_from_parts(ArrayParts::new(Delta, dtype, len, data).with_slots(slots)) } diff --git a/encodings/fastlanes/src/delta/vtable/slice.rs b/encodings/fastlanes/src/delta/vtable/slice.rs index 17b088dac92..be028b2f2ca 100644 --- a/encodings/fastlanes/src/delta/vtable/slice.rs +++ b/encodings/fastlanes/src/delta/vtable/slice.rs @@ -11,6 +11,7 @@ use vortex_array::arrays::slice::SliceReduce; use vortex_error::VortexResult; use crate::delta::array::DeltaArrayExt; +use crate::delta::array::DeltaArraySlotsExt; use crate::delta::array::lane_count; use crate::delta::vtable::Delta; diff --git a/encodings/fastlanes/src/delta/vtable/validity.rs b/encodings/fastlanes/src/delta/vtable/validity.rs index c8d8b7b4dc0..3e56e696d45 100644 --- a/encodings/fastlanes/src/delta/vtable/validity.rs +++ b/encodings/fastlanes/src/delta/vtable/validity.rs @@ -2,26 +2,223 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::ArrayView; -use vortex_array::VortexSessionExecute; -use vortex_array::legacy_session; +use vortex_array::IntoArray; use vortex_array::validity::Validity; use vortex_array::vtable::ValidityVTable; use vortex_error::VortexResult; use crate::Delta; -use crate::bit_transpose::untranspose_validity; +use crate::TransposedBool; use crate::delta::array::DeltaArrayExt; +use crate::delta::array::DeltaArraySlotsExt; impl ValidityVTable for Delta { - #[allow(clippy::disallowed_methods)] fn validity(array: ArrayView<'_, Delta>) -> VortexResult { let start = array.offset(); - let end = start + array.len(); + let stop = start + array.len(); + let validity = match array.deltas().validity()? { + Validity::Array(mask) => Validity::Array(TransposedBool::try_new(mask)?.into_array()), + validity => validity, + }; + validity.slice(start..stop) + } +} + +#[cfg(test)] +mod tests { + use std::ops::Range; + use std::sync::LazyLock; + + use rstest::rstest; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::SliceArray; + use vortex_array::assert_arrays_eq; + use vortex_array::validity::Validity; + use vortex_buffer::Alignment; + use vortex_buffer::BitBuffer; + use vortex_buffer::Buffer; + use vortex_buffer::ByteBufferMut; + use vortex_error::VortexResult; + use vortex_error::vortex_bail; + use vortex_session::VortexSession; + + use super::*; + use crate::TransposedBool; + use crate::delta::array::delta_compress::delta_compress; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + + /// Length not a multiple of 1,024, so the last chunk is padded. + const LEN: usize = 2500; + + fn nullable_primitive() -> PrimitiveArray { + PrimitiveArray::from_option_iter( + (0u32..LEN as u32).map(|value| (value % 3 != 0).then_some(value)), + ) + } + + fn expected_validity(range: Range) -> BoolArray { + BoolArray::from_iter(range.map(|value| value % 3 != 0)) + } + + /// Rebuilds `bits` one byte into a fresh allocation so the backing buffer is not u64-aligned, + /// like a bitmap viewed out of a file segment. + fn misalign_bits(bits: BitBuffer) -> BitBuffer { + let (_offset, len, bytes) = bits.sliced().into_inner(); + let mut shifted = ByteBufferMut::with_capacity(bytes.len() + 1); + shifted.push(0xFF); + shifted.extend_from_slice(&bytes); + let misaligned = shifted.freeze().slice(1..bytes.len() + 1); + assert!(!misaligned.is_aligned(Alignment::of::())); + BitBuffer::new(misaligned, len) + } + + #[test] + fn validity_is_lazy_for_cross_chunk_slice() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let primitive = nullable_primitive(); + let delta = Delta::try_from_primitive_array(&primitive, &mut ctx)?; + let sliced = delta.slice(1000..1050)?; + + let Validity::Array(validity) = sliced.validity()? else { + vortex_bail!("expected array-backed validity") + }; + assert!(validity.is::()); + assert_arrays_eq!(validity, expected_validity(1000..1050), &mut ctx); + Ok(()) + } + + /// Slicing a DeltaArray must slice its lazily-untransposed validity to the same logical + /// range, wherever the slice falls relative to 1,024-element chunk boundaries. + #[rstest] + #[case::within_first_chunk(10..1000)] + #[case::cross_chunk_boundary(1000..1050)] + #[case::chunk_aligned(1024..2048)] + #[case::tail_into_padded_chunk(2400..LEN)] + #[case::full_range(0..LEN)] + fn sliced_validity_matches_slice_range(#[case] range: Range) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let primitive = nullable_primitive(); + let delta = Delta::try_from_primitive_array(&primitive, &mut ctx)?; + let sliced = delta.slice(range.clone())?; + + let Validity::Array(validity) = sliced.validity()? else { + vortex_bail!("expected array-backed validity") + }; + assert!(validity.is::()); + assert_arrays_eq!(validity, expected_validity(range.clone()), &mut ctx); + assert_arrays_eq!(sliced, primitive.slice(range)?, &mut ctx); + Ok(()) + } + + /// A slice of a slice must compose the physical offsets before untransposing the validity. + #[test] + fn validity_of_nested_slice_composes_offsets() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let primitive = nullable_primitive(); + let delta = Delta::try_from_primitive_array(&primitive, &mut ctx)?; + let sliced = delta.slice(100..2400)?.slice(900..1500)?; + + let Validity::Array(validity) = sliced.validity()? else { + vortex_bail!("expected array-backed validity") + }; + assert!(validity.is::()); + assert_arrays_eq!(validity, expected_validity(1000..1600), &mut ctx); + assert_arrays_eq!(sliced, primitive.slice(1000..1600)?, &mut ctx); + Ok(()) + } + + /// Regression: the deltas' storage validity is not always a raw `Bool` array — slicing or a + /// file round-trip can leave it wrapped in a lazy encoding such as `vortex.slice`. The Delta + /// validity must accept it rather than bail. + #[test] + fn validity_handles_slice_encoded_storage_validity() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let primitive = PrimitiveArray::from_option_iter( + (0u32..2048).map(|value| (value % 3 != 0).then_some(value)), + ); + let (bases, deltas) = delta_compress(&primitive, &mut ctx)?; + + // Rebuild the deltas with a lazily slice-encoded validity, as produced when the deltas + // child is sliced and the validity encoding has no static slice reduction. + let Validity::Array(storage_validity) = deltas.validity()? else { + vortex_bail!("expected array-backed storage validity") + }; + let lazy_validity = SliceArray::try_new(storage_validity, 0..deltas.len())?.into_array(); + let deltas = PrimitiveArray::new(deltas.to_buffer::(), Validity::Array(lazy_validity)); + let delta = Delta::try_new(bases.into_array(), deltas.into_array(), 0, primitive.len())?; + + let Validity::Array(validity) = delta.validity()? else { + vortex_bail!("expected array-backed validity") + }; + assert_arrays_eq!( + validity, + BoolArray::from_iter((0u32..2048).map(|value| value % 3 != 0)), + &mut ctx + ); + assert_arrays_eq!(delta, primitive, &mut ctx); + Ok(()) + } + + /// Regression: the transposed validity bits may sit in a buffer that is not u64-aligned + /// (e.g. a view into a file segment). Reading the delta validity — whole or sliced — must + /// take the copying untranspose path instead of panicking on alignment. + #[test] + fn validity_from_unaligned_storage_buffer() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let primitive = nullable_primitive(); + let (bases, deltas) = delta_compress(&primitive, &mut ctx)?; + + // Rebuild the deltas with the same transposed validity bits in a misaligned buffer. + let Validity::Array(storage_validity) = deltas.validity()? else { + vortex_bail!("expected array-backed storage validity") + }; + let bits = storage_validity + .execute::(&mut ctx)? + .into_bit_buffer(); + let unaligned = BoolArray::new(misalign_bits(bits), Validity::NonNullable).into_array(); + let deltas = PrimitiveArray::new(deltas.to_buffer::(), Validity::Array(unaligned)); + let delta = Delta::try_new(bases.into_array(), deltas.into_array(), 0, primitive.len())?; + + let Validity::Array(validity) = delta.validity()? else { + vortex_bail!("expected array-backed validity") + }; + assert_arrays_eq!(validity, expected_validity(0..LEN), &mut ctx); + assert_arrays_eq!(delta, primitive, &mut ctx); + + let sliced = delta.slice(1000..1050)?; + let Validity::Array(sliced_validity) = sliced.validity()? else { + vortex_bail!("expected array-backed validity") + }; + assert_arrays_eq!(sliced_validity, expected_validity(1000..1050), &mut ctx); + Ok(()) + } + + /// Creating a DeltaArray from a primitive whose validity mask is backed by an unaligned bit + /// buffer must take the copying transpose path and round-trip losslessly. The length is a + /// whole number of chunks so that only the misalignment forces the copy. + #[test] + fn compress_primitive_with_unaligned_validity_buffer() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let len = 2048u32; + let bits = misalign_bits(BitBuffer::from_iter((0..len).map(|value| value % 3 != 0))); + let mask = BoolArray::new(bits, Validity::NonNullable).into_array(); + let primitive = + PrimitiveArray::new((0..len).collect::>(), Validity::Array(mask)); + let delta = Delta::try_from_primitive_array(&primitive, &mut ctx)?; - let validity = untranspose_validity( - &array.deltas().validity()?, - &mut legacy_session().create_execution_ctx(), - )?; - validity.slice(start..end) + let Validity::Array(validity) = delta.validity()? else { + vortex_bail!("expected array-backed validity") + }; + assert!(validity.is::()); + assert_arrays_eq!(validity, expected_validity(0..len as usize), &mut ctx); + assert_arrays_eq!(delta, primitive, &mut ctx); + Ok(()) } } diff --git a/encodings/fastlanes/src/for/array/for_compress.rs b/encodings/fastlanes/src/for/array/for_compress.rs index 9c2d5bbe423..62c23b7ba21 100644 --- a/encodings/fastlanes/src/for/array/for_compress.rs +++ b/encodings/fastlanes/src/for/array/for_compress.rs @@ -68,6 +68,7 @@ mod test { use super::*; use crate::BitPackedData; use crate::r#for::array::FoRArrayExt; + use crate::r#for::array::FoRArraySlotsExt; use crate::r#for::array::for_decompress::decompress; use crate::r#for::array::for_decompress::fused_decompress; diff --git a/encodings/fastlanes/src/for/array/for_decompress.rs b/encodings/fastlanes/src/for/array/for_decompress.rs index ffe9449f962..e41eb27a5be 100644 --- a/encodings/fastlanes/src/for/array/for_decompress.rs +++ b/encodings/fastlanes/src/for/array/for_decompress.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::mem::MaybeUninit; + use fastlanes::FoR; use num_traits::PrimInt; use num_traits::WrappingAdd; @@ -19,9 +21,11 @@ use vortex_error::VortexResult; use crate::BitPacked; use crate::BitPackedArrayExt; +use crate::FL_CHUNK_SIZE; use crate::FoRArray; use crate::bitpack_decompress; use crate::r#for::array::FoRArrayExt; +use crate::r#for::array::FoRArraySlotsExt; use crate::unpack_iter::UnpackStrategy; use crate::unpack_iter::UnpackedChunks; @@ -31,6 +35,7 @@ struct FoRStrategy { } impl + FoR> UnpackStrategy for FoRStrategy { + #[allow(clippy::inline_always)] #[inline(always)] unsafe fn unpack_chunk( &self, @@ -92,19 +97,22 @@ pub(crate) fn fused_decompress< .vortex_expect("cannot be null"); let strategy = FoRStrategy { reference: ref_ }; + let mut scratch = [const { MaybeUninit::::uninit() }; FL_CHUNK_SIZE]; // Create [`UnpackedChunks`] with FoR strategy. let mut unpacked = UnpackedChunks::try_new_with_strategy( strategy, - bp.packed().as_host().clone(), + bp.packed_slice::(), bp.bit_width() as usize, bp.offset() as usize, bp.len(), + &mut scratch, )?; - let mut builder = PrimitiveBuilder::::with_capacity( + let mut builder = PrimitiveBuilder::::with_capacity_in( for_.reference_scalar().dtype().nullability(), bp.len(), + ctx.allocator(), ); let mut uninit_range = builder.uninit_range(bp.len()); unsafe { diff --git a/encodings/fastlanes/src/for/array/mod.rs b/encodings/fastlanes/src/for/array/mod.rs index 691888ef5da..444ab152c23 100644 --- a/encodings/fastlanes/src/for/array/mod.rs +++ b/encodings/fastlanes/src/for/array/mod.rs @@ -6,19 +6,21 @@ use std::fmt::Formatter; use vortex_array::ArrayRef; use vortex_array::TypedArrayRef; +use vortex_array::array_slots; use vortex_array::dtype::PType; use vortex_array::scalar::Scalar; -use vortex_error::VortexExpect as _; use vortex_error::VortexResult; use vortex_error::vortex_ensure; pub mod for_compress; pub mod for_decompress; -/// The encoded array with the frame-of-reference (minimum value) subtracted. -pub(super) const ENCODED_SLOT: usize = 0; -pub(super) const NUM_SLOTS: usize = 1; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["encoded"]; +#[array_slots(crate::FoR)] +pub struct FoRSlots { + /// The encoded array with the frame-of-reference (minimum value) subtracted. + #[slot(0)] + pub encoded: ArrayRef, +} /// Frame of Reference (FoR) encoded array. /// @@ -29,13 +31,7 @@ pub struct FoRData { pub(super) reference: Scalar, } -pub trait FoRArrayExt: TypedArrayRef { - fn encoded(&self) -> &ArrayRef { - self.as_ref().slots()[ENCODED_SLOT] - .as_ref() - .vortex_expect("FoRArray encoded slot") - } - +pub trait FoRArrayExt: FoRArraySlotsExt { fn reference_scalar(&self) -> &Scalar { &self.reference } diff --git a/encodings/fastlanes/src/for/compute/cast.rs b/encodings/fastlanes/src/for/compute/cast.rs index e56188d9e00..d321ee4a3e3 100644 --- a/encodings/fastlanes/src/for/compute/cast.rs +++ b/encodings/fastlanes/src/for/compute/cast.rs @@ -11,6 +11,7 @@ use vortex_error::VortexResult; use crate::r#for::FoR; use crate::r#for::array::FoRArrayExt; +use crate::r#for::array::FoRArraySlotsExt; impl CastReduce for FoR { fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { // FoR only supports integer types diff --git a/encodings/fastlanes/src/for/compute/compare.rs b/encodings/fastlanes/src/for/compute/compare.rs index 384bdee8037..a5ccce3efc8 100644 --- a/encodings/fastlanes/src/for/compute/compare.rs +++ b/encodings/fastlanes/src/for/compute/compare.rs @@ -24,6 +24,7 @@ use vortex_error::VortexResult; use crate::FoR; use crate::r#for::array::FoRArrayExt; +use crate::r#for::array::FoRArraySlotsExt; impl CompareKernel for FoR { fn compare( diff --git a/encodings/fastlanes/src/for/compute/is_constant.rs b/encodings/fastlanes/src/for/compute/is_constant.rs index 29393884e95..d146e3e1d7a 100644 --- a/encodings/fastlanes/src/for/compute/is_constant.rs +++ b/encodings/fastlanes/src/for/compute/is_constant.rs @@ -11,7 +11,7 @@ use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use crate::FoR; -use crate::r#for::array::FoRArrayExt; +use crate::r#for::array::FoRArraySlotsExt; /// FoR-specific is_constant kernel. /// diff --git a/encodings/fastlanes/src/for/compute/is_sorted.rs b/encodings/fastlanes/src/for/compute/is_sorted.rs index 7cf43a93da4..92102190863 100644 --- a/encodings/fastlanes/src/for/compute/is_sorted.rs +++ b/encodings/fastlanes/src/for/compute/is_sorted.rs @@ -14,7 +14,7 @@ use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use crate::FoR; -use crate::r#for::array::FoRArrayExt; +use crate::r#for::array::FoRArraySlotsExt; #[derive(Debug)] pub(crate) struct FoRIsSortedKernel; @@ -68,7 +68,7 @@ mod test { use vortex_buffer::buffer; use crate::FoRData; - use crate::r#for::array::FoRArrayExt; + use crate::r#for::array::FoRArraySlotsExt; #[test] fn test_sorted() { diff --git a/encodings/fastlanes/src/for/compute/mod.rs b/encodings/fastlanes/src/for/compute/mod.rs index 168fcbc4192..fdb2a162e5c 100644 --- a/encodings/fastlanes/src/for/compute/mod.rs +++ b/encodings/fastlanes/src/for/compute/mod.rs @@ -17,6 +17,7 @@ use vortex_mask::Mask; use crate::FoR; use crate::r#for::array::FoRArrayExt; +use crate::r#for::array::FoRArraySlotsExt; impl TakeExecute for FoR { fn take( diff --git a/encodings/fastlanes/src/for/mod.rs b/encodings/fastlanes/src/for/mod.rs index a5d495e3578..6c88c9de825 100644 --- a/encodings/fastlanes/src/for/mod.rs +++ b/encodings/fastlanes/src/for/mod.rs @@ -3,7 +3,9 @@ mod array; pub use array::FoRArrayExt; +pub use array::FoRArraySlotsExt; pub use array::FoRData; +pub use array::FoRSlots; pub(crate) mod compute; diff --git a/encodings/fastlanes/src/for/vtable/mod.rs b/encodings/fastlanes/src/for/vtable/mod.rs index 28811bd9582..452abcd4562 100644 --- a/encodings/fastlanes/src/for/vtable/mod.rs +++ b/encodings/fastlanes/src/for/vtable/mod.rs @@ -25,7 +25,6 @@ use vortex_array::serde::ArrayChildren; use vortex_array::smallvec::smallvec; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTableFromChild; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -35,7 +34,8 @@ use vortex_session::registry::CachedId; use crate::FoRData; use crate::r#for::array::FoRArrayExt; -use crate::r#for::array::SLOT_NAMES; +use crate::r#for::array::FoRSlots; +use crate::r#for::array::FoRSlotsView; use crate::r#for::array::for_decompress::decompress; use crate::r#for::vtable::rules::PARENT_RULES; @@ -82,7 +82,7 @@ impl VTable for FoR { len: usize, slots: &[Option], ) -> VortexResult<()> { - let encoded = slots[0].as_ref().vortex_expect("FoRArray encoded slot"); + let encoded = FoRSlotsView::from_slots(slots).encoded; validate_parts(encoded.dtype(), encoded.len(), &data.reference, dtype, len) } @@ -107,7 +107,7 @@ impl VTable for FoR { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + FoRSlots::NAMES[idx].to_string() } fn serialize( diff --git a/encodings/fastlanes/src/for/vtable/operations.rs b/encodings/fastlanes/src/for/vtable/operations.rs index e549431f973..36dac998cbe 100644 --- a/encodings/fastlanes/src/for/vtable/operations.rs +++ b/encodings/fastlanes/src/for/vtable/operations.rs @@ -11,6 +11,7 @@ use vortex_error::VortexResult; use super::FoR; use crate::r#for::array::FoRArrayExt; +use crate::r#for::array::FoRArraySlotsExt; impl OperationsVTable for FoR { fn scalar_at( array: ArrayView<'_, FoR>, diff --git a/encodings/fastlanes/src/for/vtable/rules.rs b/encodings/fastlanes/src/for/vtable/rules.rs index 8d220a86d7f..2ac1a0ae3b5 100644 --- a/encodings/fastlanes/src/for/vtable/rules.rs +++ b/encodings/fastlanes/src/for/vtable/rules.rs @@ -14,6 +14,7 @@ use vortex_error::VortexResult; use crate::FoR; use crate::r#for::array::FoRArrayExt; +use crate::r#for::array::FoRArraySlotsExt; pub(super) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ // TODO: add BetweenReduceAdaptor(FoR) diff --git a/encodings/fastlanes/src/for/vtable/slice.rs b/encodings/fastlanes/src/for/vtable/slice.rs index 206c47ef509..da2e15ed55a 100644 --- a/encodings/fastlanes/src/for/vtable/slice.rs +++ b/encodings/fastlanes/src/for/vtable/slice.rs @@ -11,6 +11,7 @@ use vortex_error::VortexResult; use crate::FoR; use crate::r#for::array::FoRArrayExt; +use crate::r#for::array::FoRArraySlotsExt; impl SliceReduce for FoR { fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { diff --git a/encodings/fastlanes/src/for/vtable/validity.rs b/encodings/fastlanes/src/for/vtable/validity.rs index dd2d28802ae..d6555f2e62b 100644 --- a/encodings/fastlanes/src/for/vtable/validity.rs +++ b/encodings/fastlanes/src/for/vtable/validity.rs @@ -6,7 +6,7 @@ use vortex_array::ArrayView; use vortex_array::vtable::ValidityChild; use super::FoR; -use crate::r#for::array::FoRArrayExt; +use crate::r#for::array::FoRArraySlotsExt; impl ValidityChild for FoR { fn validity_child(array: ArrayView<'_, FoR>) -> ArrayRef { diff --git a/encodings/fastlanes/src/lib.rs b/encodings/fastlanes/src/lib.rs index e6192a9cdfb..43d83c6fc7f 100644 --- a/encodings/fastlanes/src/lib.rs +++ b/encodings/fastlanes/src/lib.rs @@ -30,6 +30,7 @@ pub use bitpacking::*; pub use delta::*; pub use r#for::*; pub use rle::*; +pub use transposed_bool::*; use vortex_array::ExecutionCtx; use vortex_array::arrays::BoolArray; use vortex_array::arrays::bool::BoolArrayExt; @@ -43,9 +44,27 @@ mod bitpacking; mod delta; mod r#for; mod rle; +mod transposed_bool; pub const FL_CHUNK_SIZE: usize = 1024; +/// Returns the position in a `FastLanes`-transposed chunk that holds logical element `idx`. +/// +/// The inverse of [`fastlanes::transpose`], which reads the other way round: it returns the +/// logical element held by a given transposed position. `fastlanes` exports only that direction, +/// so accessors that address a transposed chunk by logical index need this one. +/// +/// `fastlanes::transpose` composes `lane`, `order` and `row` as +/// `lane * 64 + FL_ORDER[order] * 8 + row`, so recovering them from the transposed position only +/// needs `FL_ORDER` inverted, and `FL_ORDER` is its own inverse. +pub(crate) const fn untranspose_idx(idx: usize) -> usize { + let lane = idx / 64; + let order = fastlanes::FL_ORDER[(idx % 64) / 8]; + let row = idx % 8; + + row * 128 + order * 16 + lane +} + use bitpacking::compute::is_constant::BitPackedIsConstantKernel; use r#for::compute::is_constant::FoRIsConstantKernel; use r#for::compute::is_sorted::FoRIsSortedKernel; @@ -70,6 +89,7 @@ pub fn initialize(session: &VortexSession) { session.arrays().register(Delta); session.arrays().register(FoR); session.arrays().register(RLE); + session.arrays().register(TransposedBool); bitpacking::initialize(session); r#for::initialize(session); rle::initialize(session); @@ -174,6 +194,16 @@ mod test { session }); + /// `untranspose_idx` derives the inverse of `fastlanes::transpose` from `FL_ORDER` being + /// its own inverse, so prove the round trip over a whole chunk in both directions. + #[test] + fn untranspose_idx_inverts_transpose() { + for idx in 0..FL_CHUNK_SIZE { + assert_eq!(untranspose_idx(fastlanes::transpose(idx)), idx, "idx={idx}"); + assert_eq!(fastlanes::transpose(untranspose_idx(idx)), idx, "idx={idx}"); + } + } + #[test] fn fill_forward_nulls_resets_at_chunk_boundary() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); diff --git a/encodings/fastlanes/src/rle/array/mod.rs b/encodings/fastlanes/src/rle/array/mod.rs index 7ddff474836..2805b999a4f 100644 --- a/encodings/fastlanes/src/rle/array/mod.rs +++ b/encodings/fastlanes/src/rle/array/mod.rs @@ -7,29 +7,33 @@ use std::fmt::Formatter; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::TypedArrayRef; -use vortex_error::VortexExpect as _; +use vortex_array::array_slots; use vortex_error::VortexResult; use vortex_error::vortex_ensure; pub mod rle_compress; pub mod rle_decompress; -/// Run values in the dictionary. -pub(super) const VALUES_SLOT: usize = 0; -/// Chunk-local indices from all chunks. The start of each chunk is looked up in `values_idx_offsets`. -pub(super) const INDICES_SLOT: usize = 1; -/// Index start positions of each value chunk. -/// -/// # Example -/// ```text -/// // Chunk 0: [10, 20] (starts at index 0) -/// // Chunk 1: [30, 40] (starts at index 2) -/// let values = [10, 20, 30, 40]; // Global values array -/// let values_idx_offsets = [0, 2]; // Chunk 0 starts at index 0, Chunk 1 starts at index 2 -/// ``` -pub(super) const VALUES_IDX_OFFSETS_SLOT: usize = 2; -pub(super) const NUM_SLOTS: usize = 3; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["values", "indices", "values_idx_offsets"]; +#[array_slots(crate::RLE)] +pub struct RLESlots { + /// Run values in the dictionary. + #[slot(0)] + pub values: ArrayRef, + /// Chunk-local indices from all chunks. The start of each chunk is looked up in `values_idx_offsets`. + #[slot(1)] + pub indices: ArrayRef, + /// Index start positions of each value chunk. + /// + /// # Example + /// ```text + /// // Chunk 0: [10, 20] (starts at index 0) + /// // Chunk 1: [30, 40] (starts at index 2) + /// let values = [10, 20, 30, 40]; // Global values array + /// let values_idx_offsets = [0, 2]; // Chunk 0 starts at index 0, Chunk 1 starts at index 2 + /// ``` + #[slot(2)] + pub values_idx_offsets: ArrayRef, +} #[derive(Clone, Debug)] pub struct RLEData { @@ -79,28 +83,7 @@ impl RLEData { } } -pub trait RLEArrayExt: TypedArrayRef { - #[inline] - fn values(&self) -> &ArrayRef { - self.as_ref().slots()[VALUES_SLOT] - .as_ref() - .vortex_expect("RLEArray values slot must be populated") - } - - #[inline] - fn indices(&self) -> &ArrayRef { - self.as_ref().slots()[INDICES_SLOT] - .as_ref() - .vortex_expect("RLEArray indices slot must be populated") - } - - #[inline] - fn values_idx_offsets(&self) -> &ArrayRef { - self.as_ref().slots()[VALUES_IDX_OFFSETS_SLOT] - .as_ref() - .vortex_expect("RLEArray values_idx_offsets slot must be populated") - } - +pub trait RLEArrayExt: RLEArraySlotsExt { /// Values index offset relative to the first chunk. /// /// Offsets in `values_idx_offsets` are absolute and need to be shifted @@ -160,6 +143,7 @@ mod tests { use crate::RLE; use crate::RLEData; use crate::rle::array::RLEArrayExt; + use crate::rle::array::RLEArraySlotsExt; use crate::test::SESSION; #[test] diff --git a/encodings/fastlanes/src/rle/array/rle_compress.rs b/encodings/fastlanes/src/rle/array/rle_compress.rs index b7e04a13699..f623b84d778 100644 --- a/encodings/fastlanes/src/rle/array/rle_compress.rs +++ b/encodings/fastlanes/src/rle/array/rle_compress.rs @@ -81,11 +81,15 @@ where // returned from `T::encode` are relative to the chunk. values_idx_offsets.push(value_count_acc as u64); - let value_count = NativeValue::::encode( - input, - unsafe { &mut *(rle_vals.as_mut_ptr() as *mut [_; FL_CHUNK_SIZE]) }, - rle_idxs, - ); + // SAFETY: all three arguments are `FL_CHUNK_SIZE`-element arrays, which is + // all `encode_unchecked` requires. + let value_count = unsafe { + NativeValue::::encode_unchecked( + input, + &mut *(rle_vals.as_mut_ptr() as *mut [_; FL_CHUNK_SIZE]), + rle_idxs, + ) + }; value_count_acc += value_count; }; @@ -171,6 +175,7 @@ mod tests { use super::*; use crate::rle::array::RLEArrayExt; + use crate::rle::array::RLEArraySlotsExt; static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session(); diff --git a/encodings/fastlanes/src/rle/array/rle_decompress.rs b/encodings/fastlanes/src/rle/array/rle_decompress.rs index 51766398070..0c841b7cba6 100644 --- a/encodings/fastlanes/src/rle/array/rle_decompress.rs +++ b/encodings/fastlanes/src/rle/array/rle_decompress.rs @@ -12,11 +12,13 @@ use vortex_array::match_each_unsigned_integer_ptype; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_error::vortex_panic; use crate::FL_CHUNK_SIZE; use crate::RLEArray; use crate::rle::RLEArrayExt; +use crate::rle::RLEArraySlotsExt; /// Decompresses an RLE array back into a primitive array. pub fn rle_decompress(array: &RLEArray, ctx: &mut ExecutionCtx) -> VortexResult { @@ -61,9 +63,27 @@ where let values = array.values().clone().execute::(ctx)?; let values = values.as_slice::(); + // The offsets come from (possibly untrusted) storage. Validate them once here so + // that the per-chunk slicing and the unchecked decodes below stay within `values`: + // offsets must be non-decreasing and span at most `values.len()` values. + vortex_ensure!( + values_idx_offsets.is_sorted(), + "RLE values_idx_offsets must be non-decreasing" + ); + if let (Some(&first), Some(&last)) = (values_idx_offsets.first(), values_idx_offsets.last()) { + vortex_ensure!( + last - first <= values.len() as u64, + "RLE values_idx_offsets span {} values but only {} are present", + last - first, + values.len() + ); + } + let indices = array.indices().clone().execute::(ctx)?; assert!(indices.len().is_multiple_of(FL_CHUNK_SIZE)); - let has_invalid = !indices.all_valid(ctx)?; + let indices_validity = indices.validity()?.execute_mask(indices.len(), ctx)?; + // `None` means every position is valid. + let validity_bits = (!indices_validity.all_true()).then(|| indices_validity.to_bit_buffer()); let (indices_sl, _) = indices.as_slice::().as_chunks::(); let chunk_start_idx = array.offset() / FL_CHUNK_SIZE; @@ -88,6 +108,10 @@ where }; let num_chunk_values = u16::try_from(next_value_idx_offset - value_idx_offset) .vortex_expect("There can be at most 1024 values in RLE chunk"); + vortex_ensure!( + num_chunk_values > 0, + "RLE chunk {chunk_idx} references no values" + ); // SAFETY: `MaybeUninit` and `T` have the same layout. let buffer_values: &mut [V; FL_CHUNK_SIZE] = unsafe { std::mem::transmute(chunk_out) }; @@ -97,20 +121,47 @@ where // access. The indices may contain values other than 0 when they // have been further compressed (e.g., as a masked constant). buffer_values.fill(chunk_values[0]); - } else if has_invalid { - // When the indices array has invalid (null) positions, those - // positions may contain arbitrary garbage values after further - // compression. Clamp all indices into [0, num_chunk_values) to - // prevent out-of-bounds access in the fastlanes decoder. - let mut sanitized: [u16; FL_CHUNK_SIZE] = [0; FL_CHUNK_SIZE]; - for (idx_out, idx) in sanitized.iter_mut().zip(chunk_indices) { - let idx: u16 = - NumCast::from(*idx).vortex_expect("RLE indices are always less than u16"); - *idx_out = idx.min(num_chunk_values - 1); - } - V::decode(chunk_values, &sanitized, buffer_values); } else { - V::decode(chunk_values, chunk_indices, buffer_values); + let chunk_start = chunk_idx * FL_CHUNK_SIZE; + let chunk_valid_count = validity_bits.as_ref().map_or(FL_CHUNK_SIZE, |bits| { + bits.count_range(chunk_start, chunk_start + FL_CHUNK_SIZE) + }); + if chunk_valid_count == FL_CHUNK_SIZE { + decode_chunk_checked( + chunk_values, + chunk_indices, + num_chunk_values, + chunk_idx, + buffer_values, + )?; + } else if chunk_valid_count == 0 { + // Entirely-null chunk: every position is masked out, so the indices are + // meaningless (possibly garbage after further compression) — skip the + // gather and emit an arbitrary in-bounds value. + buffer_values.fill(chunk_values[0]); + } else { + // Null positions may contain arbitrary garbage indices after further + // compression, so zero them out (their decoded values are masked by the + // validity). Valid positions must be genuinely in bounds and are still + // bound-checked before the unchecked gather: a corrupt index at a valid + // position is an error, never silently remapped. + let bits = validity_bits + .as_ref() + .vortex_expect("mixed-validity chunk implies a materialized mask"); + let mut sanitized: [u16; FL_CHUNK_SIZE] = [0; FL_CHUNK_SIZE]; + bits.slice(chunk_start..chunk_start + FL_CHUNK_SIZE) + .for_each_set_index(|i| { + sanitized[i] = NumCast::from(chunk_indices[i]) + .vortex_expect("RLE indices are always less than u16"); + }); + decode_chunk_checked( + chunk_values, + &sanitized, + num_chunk_values, + chunk_idx, + buffer_values, + )?; + } } } @@ -127,3 +178,201 @@ where array.validity()?, )) } + +/// Bound-checks every index in the chunk, then runs the unchecked fastlanes gather. +/// +/// The indices come from (possibly untrusted) storage: a single max-reduction over the +/// chunk vectorizes and keeps the bounds check out of the per-element decode loop. +fn decode_chunk_checked( + chunk_values: &[V], + chunk_indices: &[I; FL_CHUNK_SIZE], + num_chunk_values: u16, + chunk_idx: usize, + out: &mut [V; FL_CHUNK_SIZE], +) -> VortexResult<()> +where + V: RLE, + I: Copy + Into, +{ + let max_index: usize = chunk_indices + .iter() + .map(|idx| (*idx).into()) + .max() + .unwrap_or_default(); + vortex_ensure!( + max_index < num_chunk_values as usize, + "RLE index {max_index} out of bounds for chunk {chunk_idx} with {num_chunk_values} values" + ); + // SAFETY: just checked that every index in the chunk is below `num_chunk_values`, + // which the caller's offset validation bounds by `chunk_values.len()`. + unsafe { V::decode_unchecked(chunk_values, chunk_indices, out) }; + Ok(()) +} + +#[cfg(test)] +mod tests { + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + use vortex_error::VortexResult; + + use crate::FL_CHUNK_SIZE; + use crate::RLE; + use crate::rle::array::rle_decompress::rle_decompress; + use crate::test::SESSION; + + fn indices_with_oob(oob: u16) -> vortex_array::ArrayRef { + let mut indices = [0u16, 1] + .iter() + .cycle() + .take(FL_CHUNK_SIZE) + .copied() + .collect::>(); + indices[100] = oob; + PrimitiveArray::from_iter(indices).into_array() + } + + #[test] + fn test_decode_rejects_out_of_bounds_index() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::from_iter([10u32, 20]).into_array(); + let values_idx_offsets = PrimitiveArray::from_iter([0u64]).into_array(); + + // Index 999 points far beyond the 2 values of the only chunk. + let rle = RLE::try_new( + values, + indices_with_oob(999), + values_idx_offsets, + 0, + FL_CHUNK_SIZE, + )?; + assert!(rle_decompress(&rle, &mut ctx).is_err()); + Ok(()) + } + + #[test] + fn test_decode_rejects_index_into_next_chunk() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + // Two chunks with 2 values each: an index of 2 in chunk 0 is within `values` + // but out of bounds for the chunk, and must be rejected rather than silently + // reading chunk 1's values. + let values = PrimitiveArray::from_iter([10u32, 20, 30, 40]).into_array(); + let values_idx_offsets = PrimitiveArray::from_iter([0u64, 2]).into_array(); + let mut indices = [0u16, 1].repeat(FL_CHUNK_SIZE); + indices[100] = 2; + let indices = PrimitiveArray::from_iter(indices).into_array(); + + let rle = RLE::try_new(values, indices, values_idx_offsets, 0, 2 * FL_CHUNK_SIZE)?; + assert!(rle_decompress(&rle, &mut ctx).is_err()); + Ok(()) + } + + #[test] + fn test_decode_rejects_non_monotonic_offsets() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::from_iter([10u32, 20, 30, 40]).into_array(); + let values_idx_offsets = PrimitiveArray::from_iter([2u64, 0]).into_array(); + let indices = PrimitiveArray::from_iter([0u16, 1].repeat(FL_CHUNK_SIZE)).into_array(); + + let rle = RLE::try_new(values, indices, values_idx_offsets, 0, 2 * FL_CHUNK_SIZE)?; + assert!(rle_decompress(&rle, &mut ctx).is_err()); + Ok(()) + } + + #[test] + fn test_decode_rejects_offsets_beyond_values() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::from_iter([10u32, 20, 30, 40]).into_array(); + // The offset span (100) exceeds the number of values present (4). + let values_idx_offsets = PrimitiveArray::from_iter([0u64, 100]).into_array(); + let indices = PrimitiveArray::from_iter([0u16, 1].repeat(FL_CHUNK_SIZE)).into_array(); + + let rle = RLE::try_new(values, indices, values_idx_offsets, 0, 2 * FL_CHUNK_SIZE)?; + assert!(rle_decompress(&rle, &mut ctx).is_err()); + Ok(()) + } + + #[test] + fn test_decode_rejects_empty_chunk() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::from_iter([10u32, 20]).into_array(); + // Chunk 0 references no values: offsets [0, 0]. + let values_idx_offsets = PrimitiveArray::from_iter([0u64, 0]).into_array(); + let indices = PrimitiveArray::from_iter([0u16, 1].repeat(FL_CHUNK_SIZE)).into_array(); + + let rle = RLE::try_new(values, indices, values_idx_offsets, 0, 2 * FL_CHUNK_SIZE)?; + assert!(rle_decompress(&rle, &mut ctx).is_err()); + Ok(()) + } + + #[test] + fn test_decode_rejects_out_of_bounds_index_at_valid_position() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::from_iter([10u32, 20]).into_array(); + let values_idx_offsets = PrimitiveArray::from_iter([0u64]).into_array(); + + // Position 100 is VALID but holds an out-of-bounds index: this is corruption + // and must fail, even though other positions are null. + let mut indices = [0u16, 1].repeat(FL_CHUNK_SIZE / 2); + indices[100] = 999; + let mut validity = vec![true; FL_CHUNK_SIZE]; + validity[200] = false; + let indices = PrimitiveArray::new( + indices.into_iter().collect::>(), + Validity::from_iter(validity), + ) + .into_array(); + + let rle = RLE::try_new(values, indices, values_idx_offsets, 0, FL_CHUNK_SIZE)?; + assert!(rle_decompress(&rle, &mut ctx).is_err()); + Ok(()) + } + + #[test] + fn test_decode_tolerates_garbage_index_at_null_position() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::from_iter([10u32, 20]).into_array(); + let values_idx_offsets = PrimitiveArray::from_iter([0u64]).into_array(); + + // Position 100 is NULL and holds a garbage out-of-bounds index (e.g. after + // further compression of the indices): its value is masked, so decoding + // must succeed and every valid position must decode correctly. + let mut indices = [0u16, 1].repeat(FL_CHUNK_SIZE / 2); + indices[100] = 999; + let mut validity = vec![true; FL_CHUNK_SIZE]; + validity[100] = false; + let indices = PrimitiveArray::new( + indices.into_iter().collect::>(), + Validity::from_iter(validity), + ) + .into_array(); + + let rle = RLE::try_new(values, indices, values_idx_offsets, 0, FL_CHUNK_SIZE)?; + let decoded = rle_decompress(&rle, &mut ctx)?; + let expected = PrimitiveArray::from_option_iter( + (0..FL_CHUNK_SIZE).map(|i| (i != 100).then_some(if i % 2 == 0 { 10u32 } else { 20 })), + ); + assert_arrays_eq!(decoded, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_decode_in_bounds_indices_roundtrip() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::from_iter([10u32, 20]).into_array(); + let values_idx_offsets = PrimitiveArray::from_iter([0u64]).into_array(); + let indices = + PrimitiveArray::from_iter([0u16, 1].iter().cycle().take(FL_CHUNK_SIZE).copied()) + .into_array(); + + let rle = RLE::try_new(values, indices, values_idx_offsets, 0, FL_CHUNK_SIZE)?; + let decoded = rle_decompress(&rle, &mut ctx)?; + let expected = + PrimitiveArray::from_iter([10u32, 20].iter().cycle().take(FL_CHUNK_SIZE).copied()); + assert_arrays_eq!(decoded, expected, &mut ctx); + Ok(()) + } +} diff --git a/encodings/fastlanes/src/rle/compute/cast.rs b/encodings/fastlanes/src/rle/compute/cast.rs index 38e7927f8e6..a8aa7b199f0 100644 --- a/encodings/fastlanes/src/rle/compute/cast.rs +++ b/encodings/fastlanes/src/rle/compute/cast.rs @@ -12,6 +12,7 @@ use vortex_error::VortexResult; use crate::rle::RLE; use crate::rle::RLEArrayExt; +use crate::rle::RLEArraySlotsExt; impl CastReduce for RLE { fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { diff --git a/encodings/fastlanes/src/rle/kernel.rs b/encodings/fastlanes/src/rle/kernel.rs index 9dd05ffbae5..933e6e3d42a 100644 --- a/encodings/fastlanes/src/rle/kernel.rs +++ b/encodings/fastlanes/src/rle/kernel.rs @@ -18,6 +18,7 @@ use vortex_session::VortexSession; use crate::FL_CHUNK_SIZE; use crate::RLE; use crate::rle::RLEArrayExt; +use crate::rle::RLEArraySlotsExt; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); diff --git a/encodings/fastlanes/src/rle/mod.rs b/encodings/fastlanes/src/rle/mod.rs index 649d11f307c..25610742c19 100644 --- a/encodings/fastlanes/src/rle/mod.rs +++ b/encodings/fastlanes/src/rle/mod.rs @@ -3,7 +3,9 @@ mod array; pub use array::RLEArrayExt; +pub use array::RLEArraySlotsExt; pub use array::RLEData; +pub use array::RLESlots; mod compute; mod kernel; diff --git a/encodings/fastlanes/src/rle/vtable/mod.rs b/encodings/fastlanes/src/rle/vtable/mod.rs index cea655e53b4..a3f6f01fa5a 100644 --- a/encodings/fastlanes/src/rle/vtable/mod.rs +++ b/encodings/fastlanes/src/rle/vtable/mod.rs @@ -22,9 +22,7 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::serde::ArrayChildren; -use vortex_array::smallvec::smallvec; use vortex_array::vtable::VTable; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; @@ -32,11 +30,10 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::RLEData; -use crate::rle::array::INDICES_SLOT; use crate::rle::array::RLEArrayExt; -use crate::rle::array::SLOT_NAMES; -use crate::rle::array::VALUES_IDX_OFFSETS_SLOT; -use crate::rle::array::VALUES_SLOT; +use crate::rle::array::RLEArraySlotsExt; +use crate::rle::array::RLESlots; +use crate::rle::array::RLESlotsView; use crate::rle::array::rle_decompress::rle_decompress; use crate::rle::vtable::rules::RULES; @@ -93,16 +90,11 @@ impl VTable for RLE { len: usize, slots: &[Option], ) -> VortexResult<()> { + let rle_slots = RLESlotsView::from_slots(slots); validate_parts( - slots[VALUES_SLOT] - .as_ref() - .vortex_expect("RLEArray values slot must be populated"), - slots[INDICES_SLOT] - .as_ref() - .vortex_expect("RLEArray indices slot must be populated"), - slots[VALUES_IDX_OFFSETS_SLOT] - .as_ref() - .vortex_expect("RLEArray values_idx_offsets slot must be populated"), + rle_slots.values, + rle_slots.indices, + rle_slots.values_idx_offsets, data.offset, dtype, len, @@ -138,7 +130,7 @@ impl VTable for RLE { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + RLESlots::NAMES[idx].to_string() } fn serialize( @@ -195,7 +187,12 @@ impl VTable for RLE { usize::try_from(metadata.values_idx_offsets_len)?, )?; - let slots = smallvec![Some(values), Some(indices), Some(values_idx_offsets)]; + let slots = RLESlots { + values, + indices, + values_idx_offsets, + } + .into_slots(); let data = RLEData::try_new(metadata.offset as usize)?; Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) } @@ -219,7 +216,12 @@ impl RLE { length: usize, ) -> VortexResult { let dtype = DType::Primitive(values.dtype().as_ptype(), indices.dtype().nullability()); - let slots = smallvec![Some(values), Some(indices), Some(values_idx_offsets)]; + let slots = RLESlots { + values, + indices, + values_idx_offsets, + } + .into_slots(); let data = RLEData::try_new(offset)?; Array::try_from_parts(ArrayParts::new(RLE, dtype, length, data).with_slots(slots)) } @@ -236,7 +238,12 @@ impl RLE { length: usize, ) -> RLEArray { let dtype = DType::Primitive(values.dtype().as_ptype(), indices.dtype().nullability()); - let slots = smallvec![Some(values), Some(indices), Some(values_idx_offsets)]; + let slots = RLESlots { + values, + indices, + values_idx_offsets, + } + .into_slots(); let data = unsafe { RLEData::new_unchecked(offset) }; unsafe { Array::from_parts_unchecked(ArrayParts::new(RLE, dtype, length, data).with_slots(slots)) diff --git a/encodings/fastlanes/src/rle/vtable/operations.rs b/encodings/fastlanes/src/rle/vtable/operations.rs index b8f79db95a6..ca4d2d39545 100644 --- a/encodings/fastlanes/src/rle/vtable/operations.rs +++ b/encodings/fastlanes/src/rle/vtable/operations.rs @@ -11,6 +11,7 @@ use vortex_error::VortexResult; use super::RLE; use crate::FL_CHUNK_SIZE; use crate::rle::RLEArrayExt; +use crate::rle::RLEArraySlotsExt; impl OperationsVTable for RLE { fn scalar_at( diff --git a/encodings/fastlanes/src/rle/vtable/validity.rs b/encodings/fastlanes/src/rle/vtable/validity.rs index 917b43da996..603cc9db542 100644 --- a/encodings/fastlanes/src/rle/vtable/validity.rs +++ b/encodings/fastlanes/src/rle/vtable/validity.rs @@ -7,6 +7,7 @@ use vortex_error::VortexResult; use crate::rle::RLE; use crate::rle::RLEArrayExt; +use crate::rle::RLEArraySlotsExt; impl ValidityVTable for RLE { fn validity(array: ArrayView<'_, RLE>) -> VortexResult { diff --git a/encodings/fastlanes/src/transposed_bool.rs b/encodings/fastlanes/src/transposed_bool.rs new file mode 100644 index 00000000000..efb3443cdfc --- /dev/null +++ b/encodings/fastlanes/src/transposed_bool.rs @@ -0,0 +1,366 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hash; +use std::hash::Hasher; +use std::ops::Range; + +use vortex_array::Array; +use vortex_array::ArrayEq; +use vortex_array::ArrayHash; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::EqMode; +use vortex_array::ExecutionCtx; +use vortex_array::ExecutionResult; +use vortex_array::IntoArray; +use vortex_array::TypedArrayRef; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::slice::SliceReduce; +use vortex_array::arrays::slice::SliceReduceAdaptor; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::optimizer::rules::ParentRuleSet; +use vortex_array::scalar::Scalar; +use vortex_array::serde::ArrayChildren; +use vortex_array::smallvec::smallvec; +use vortex_array::validity::Validity; +use vortex_array::vtable::OperationsVTable; +use vortex_array::vtable::VTable; +use vortex_array::vtable::ValidityVTable; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::FL_CHUNK_SIZE; +use crate::bit_transpose::untranspose_bitbuffer; +use crate::untranspose_idx; + +/// A non-nullable boolean array stored in FastLanes-transposed order. +pub type TransposedBoolArray = Array; + +/// The array encoding for a boolean bitmap stored in FastLanes-transposed order. +#[derive(Clone, Debug)] +pub struct TransposedBool; + +/// The transposed bitmap, as a non-nullable boolean array covering whole 1,024-bit chunks. +const TRANSPOSED_SLOT: usize = 0; + +/// Per-array data for a [`TransposedBoolArray`]. +#[derive(Clone, Debug)] +pub struct TransposedBoolData { + offset: usize, +} + +impl Display for TransposedBoolData { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "offset: {}", self.offset) + } +} + +impl ArrayHash for TransposedBoolData { + fn array_hash(&self, state: &mut H, _accuracy: EqMode) { + self.offset.hash(state); + } +} + +impl ArrayEq for TransposedBoolData { + fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool { + self.offset == other.offset + } +} + +/// Accessors for a [`TransposedBoolArray`]. +pub trait TransposedBoolArrayExt: TypedArrayRef { + /// Returns the logical offset into the first transposed chunk. + fn offset(&self) -> usize { + self.deref().offset + } + + /// Returns the backing bitmap in FastLanes-transposed order. + fn transposed(&self) -> &ArrayRef { + self.as_ref().slots()[TRANSPOSED_SLOT] + .as_ref() + .vortex_expect("TransposedBoolArray transposed slot") + } +} + +impl> TransposedBoolArrayExt for T {} + +impl TransposedBool { + /// Creates an array from a boolean array already stored in FastLanes-transposed order. + /// + /// The `transposed` array may use any encoding (e.g. a lazy slice); it is canonicalized when + /// this array is executed. + /// + /// # Errors + /// + /// Returns an error if `transposed` is not a non-nullable boolean array containing complete + /// 1,024-bit chunks. + pub fn try_new(transposed: ArrayRef) -> VortexResult { + let len = transposed.len(); + Self::try_new_view(transposed, 0, len) + } + + fn try_new_view( + transposed: ArrayRef, + offset: usize, + len: usize, + ) -> VortexResult { + Array::try_from_parts( + ArrayParts::new( + TransposedBool, + DType::Bool(Nullability::NonNullable), + len, + TransposedBoolData { offset }, + ) + .with_slots(smallvec![Some(transposed)]), + ) + } +} + +impl VTable for TransposedBool { + type TypedArrayData = TransposedBoolData; + type OperationsVTable = Self; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("fastlanes.transposed_bool"); + *ID + } + + fn validate( + &self, + data: &Self::TypedArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure!( + dtype == &DType::Bool(Nullability::NonNullable), + "TransposedBoolArray must have non-nullable boolean dtype, got {dtype}" + ); + vortex_ensure!( + slots.len() == 1, + "TransposedBoolArray expects one slot, got {}", + slots.len() + ); + let transposed = slots[TRANSPOSED_SLOT] + .as_ref() + .vortex_expect("TransposedBoolArray transposed slot"); + vortex_ensure!( + transposed.dtype() == &DType::Bool(Nullability::NonNullable), + "TransposedBoolArray transposed child must be a non-nullable boolean array, got {}", + transposed.dtype() + ); + vortex_ensure!( + transposed.len().is_multiple_of(FL_CHUNK_SIZE), + "TransposedBoolArray transposed child length {} must be a multiple of {FL_CHUNK_SIZE}", + transposed.len() + ); + vortex_ensure!( + data.offset < FL_CHUNK_SIZE, + "TransposedBoolArray offset {} must be less than {FL_CHUNK_SIZE}", + data.offset + ); + let end = data + .offset + .checked_add(len) + .ok_or_else(|| vortex_error::vortex_err!("TransposedBoolArray range end overflow"))?; + vortex_ensure!( + end <= transposed.len(), + "TransposedBoolArray range {}..{} exceeds transposed child length {}", + data.offset, + end, + transposed.len() + ); + Ok(()) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + vortex_panic!("TransposedBoolArray buffer index {idx} out of bounds") + } + + fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option { + None + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_array::vtable::with_empty_buffers(self, array, buffers) + } + + fn serialize( + _array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + vortex_bail!("Cannot serialise TransposedBoolArray"); + } + + fn deserialize( + &self, + _dtype: &DType, + _len: usize, + _metadata: &[u8], + _buffers: &[BufferHandle], + _children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_bail!("Cannot deserialise TransposedBoolArray"); + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + match idx { + TRANSPOSED_SLOT => "transposed".to_string(), + _ => vortex_panic!("TransposedBoolArray slot index {idx} out of bounds"), + } + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let len = array.len(); + let offset = array.offset(); + let bits = array + .transposed() + .clone() + .execute::(ctx)? + .into_bit_buffer(); + let untransposed = BoolArray::new(untranspose_bitbuffer(bits), Validity::NonNullable); + Ok(ExecutionResult::done( + untransposed.slice(offset..offset + len)?, + )) + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + RULES.evaluate(array, parent, child_idx) + } +} + +impl OperationsVTable for TransposedBool { + fn scalar_at( + array: ArrayView<'_, TransposedBool>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let logical_index = array.offset() + index; + let chunk_start = logical_index / FL_CHUNK_SIZE * FL_CHUNK_SIZE; + let transposed_index = chunk_start + untranspose_idx(logical_index % FL_CHUNK_SIZE); + array.transposed().execute_scalar(transposed_index, ctx) + } +} + +impl ValidityVTable for TransposedBool { + fn validity(_array: ArrayView<'_, TransposedBool>) -> VortexResult { + Ok(Validity::NonNullable) + } +} + +impl SliceReduce for TransposedBool { + fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { + let physical_start = array.offset() + range.start; + let physical_stop = array.offset() + range.end; + let start_chunk = physical_start / FL_CHUNK_SIZE; + let stop_chunk = physical_stop.div_ceil(FL_CHUNK_SIZE); + let transposed = array + .transposed() + .slice(start_chunk * FL_CHUNK_SIZE..stop_chunk * FL_CHUNK_SIZE)?; + + Ok(Some( + TransposedBool::try_new_view(transposed, physical_start % FL_CHUNK_SIZE, range.len())? + .into_array(), + )) + } +} + +static RULES: ParentRuleSet = + ParentRuleSet::new(&[ParentRuleSet::lift(&SliceReduceAdaptor(TransposedBool))]); + +#[cfg(test)] +mod tests { + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::SliceArray; + use vortex_array::assert_arrays_eq; + use vortex_buffer::BitBuffer; + use vortex_error::VortexResult; + + use super::*; + use crate::bit_transpose::transpose_bitbuffer; + + fn test_bits() -> BitBuffer { + BitBuffer::from_iter((0..2 * FL_CHUNK_SIZE).map(|i| i % 3 != 0 && i % 11 != 0)) + } + + fn transposed_bool_array(bits: BitBuffer) -> ArrayRef { + BoolArray::new(transpose_bitbuffer(bits), Validity::NonNullable).into_array() + } + + #[test] + fn execute_full_array() -> VortexResult<()> { + let expected = test_bits(); + let array = TransposedBool::try_new(transposed_bool_array(expected.clone()))?; + let mut ctx = array_session().create_execution_ctx(); + + assert_arrays_eq!(array, BoolArray::from(expected), &mut ctx); + Ok(()) + } + + #[test] + fn slice_stays_lazy_and_translates_scalars() -> VortexResult<()> { + let expected = test_bits(); + let array = TransposedBool::try_new(transposed_bool_array(expected.clone()))?; + let sliced = array.slice(1000..1050)?; + assert!(sliced.is::()); + + let mut ctx = array_session().create_execution_ctx(); + for index in [0, 23, 49] { + assert_eq!( + sliced.execute_scalar(index, &mut ctx)?.as_bool().value(), + Some(expected.value(1000 + index)) + ); + } + assert_arrays_eq!( + sliced, + BoolArray::from(expected.slice(1000..1050)), + &mut ctx + ); + Ok(()) + } + + /// Regression: the transposed child may be lazily encoded (e.g. a `vortex.slice` wrapper), + /// which must be canonicalized at execution rather than rejected. + #[test] + fn execute_slice_encoded_child() -> VortexResult<()> { + let expected = test_bits(); + let child = transposed_bool_array(expected.clone()); + let lazy_slice = SliceArray::try_new(child, FL_CHUNK_SIZE..2 * FL_CHUNK_SIZE)?.into_array(); + let array = TransposedBool::try_new(lazy_slice)?; + + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!( + array, + BoolArray::from(expected.slice(FL_CHUNK_SIZE..2 * FL_CHUNK_SIZE)), + &mut ctx + ); + Ok(()) + } +} diff --git a/encodings/fsst/benches/chunked_dict_fsst_builder.rs b/encodings/fsst/benches/chunked_dict_fsst_builder.rs index 8c2c6c2e388..de25bf274b6 100644 --- a/encodings/fsst/benches/chunked_dict_fsst_builder.rs +++ b/encodings/fsst/benches/chunked_dict_fsst_builder.rs @@ -9,7 +9,7 @@ use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ChunkedArray; -use vortex_array::builders::builder_with_capacity; +use vortex_array::builders::builder_with_capacity_in; use vortex_array::dtype::NativePType; use vortex_error::VortexExpect; use vortex_fsst::test_utils::gen_dict_fsst_test_data; @@ -19,13 +19,14 @@ fn main() { divan::main(); } +// `unique_values` must stay <= `len`, since codes index the values array. const BENCH_ARGS: &[(usize, usize, usize)] = &[ - (1000, 10, 10), - (1000, 100, 10), - (1000, 1000, 10), - (1000, 10, 100), - (1000, 100, 100), - (1000, 1000, 100), + (250, 10, 2), + (250, 100, 2), + (250, 250, 2), + (250, 10, 5), + (250, 100, 5), + (250, 250, 5), ]; static SESSION: LazyLock = LazyLock::new(|| { @@ -56,7 +57,8 @@ fn chunked_dict_fsst_canonical_into( bencher .with_inputs(|| (&chunk, SESSION.create_execution_ctx())) .bench_refs(|(chunk, ctx)| { - let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count); + let mut builder = + builder_with_capacity_in(chunk.dtype(), len * chunk_count, ctx.allocator()); chunk .append_to_builder(builder.as_mut(), ctx) .vortex_expect("append failed"); diff --git a/encodings/fsst/benches/fsst_compress.rs b/encodings/fsst/benches/fsst_compress.rs index 5e03d1d003d..ea6d9f7c544 100644 --- a/encodings/fsst/benches/fsst_compress.rs +++ b/encodings/fsst/benches/fsst_compress.rs @@ -40,18 +40,18 @@ static SESSION: LazyLock = LazyLock::new(|| { // [(string_count, avg_len, unique_chars)] const BENCH_ARGS: &[(usize, usize, u8)] = &[ - (1_000, 4, 4), - (1_000, 4, 8), - (1_000, 16, 4), - (1_000, 16, 8), - (1_000, 64, 4), - (1_000, 64, 8), - (10_000, 4, 4), - (10_000, 4, 8), - (10_000, 16, 4), - (10_000, 16, 8), - (10_000, 64, 4), - (10_000, 64, 8), + (200, 4, 4), + (200, 4, 8), + (200, 16, 4), + (200, 16, 8), + (200, 64, 4), + (200, 64, 8), + (500, 4, 4), + (500, 4, 8), + (500, 16, 4), + (500, 16, 8), + (500, 64, 4), + (500, 64, 8), ]; #[divan::bench(args = BENCH_ARGS)] @@ -77,14 +77,6 @@ fn decompress_fsst(bencher: Bencher, (string_count, avg_len, unique_chars): (usi .bench_refs(|(encoded, ctx)| (*encoded).clone().execute::(ctx)) } -#[divan::bench(args = BENCH_ARGS)] -fn train_compressor(bencher: Bencher, (string_count, avg_len, unique_chars): (usize, usize, u8)) { - let array = generate_test_data(string_count, avg_len, unique_chars); - bencher - .with_inputs(|| (&array, SESSION.create_execution_ctx())) - .bench_refs(|(array, ctx)| fsst_train_compressor(array, ctx).unwrap()) -} - #[divan::bench(args = BENCH_ARGS)] fn pushdown_compare(bencher: Bencher, (string_count, avg_len, unique_chars): (usize, usize, u8)) { let array = generate_test_data(string_count, avg_len, unique_chars); @@ -139,15 +131,15 @@ fn canonicalize_compare( // [(chunk_size, string_count, avg_len, unique_chars)] const CHUNKED_BENCH_ARGS: &[(usize, usize, usize, u8)] = &[ - (1000, 100, 16, 4), - (1000, 100, 16, 16), - (1000, 100, 16, 64), - (1000, 50, 8, 4), - (1000, 50, 8, 16), - (1000, 50, 8, 64), - (10, 10_000, 4, 4), - (10, 10_000, 16, 4), - (10, 10_000, 64, 4), + (20, 100, 16, 4), + (20, 100, 16, 16), + (20, 100, 16, 64), + (20, 50, 8, 4), + (20, 50, 8, 16), + (20, 50, 8, 64), + (10, 150, 4, 4), + (10, 150, 16, 4), + (10, 150, 64, 4), ]; #[divan::bench(args = CHUNKED_BENCH_ARGS)] @@ -160,9 +152,10 @@ fn chunked_canonicalize_into( bencher .with_inputs(|| (&array, SESSION.create_execution_ctx())) .bench_refs(|(array, ctx)| { - let mut builder = VarBinViewBuilder::with_capacity( + let mut builder = VarBinViewBuilder::with_capacity_in( DType::Binary(Nullability::NonNullable), array.len(), + ctx.allocator().clone(), ); array.append_to_builder(&mut builder, ctx).unwrap(); builder.finish() diff --git a/encodings/fsst/benches/fsst_like.rs b/encodings/fsst/benches/fsst_like.rs index 414c1b7afaf..0737100602b 100644 --- a/encodings/fsst/benches/fsst_like.rs +++ b/encodings/fsst/benches/fsst_like.rs @@ -11,7 +11,6 @@ use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; -use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::scalar_fn::fns::like::Like; use vortex_array::scalar_fn::fns::like::LikeOptions; use vortex_fsst::FSSTArray; @@ -121,7 +120,7 @@ fn bench_like(bencher: Bencher, fsst: &FSSTArray, pattern: &str) { bencher .with_inputs(|| SESSION.create_execution_ctx()) .bench_refs(|ctx| { - Like.try_new_array(len, LikeOptions::default(), [arr.clone(), pattern.clone()]) + Like::try_new(arr.clone(), pattern.clone(), LikeOptions::default()) .unwrap() .into_array() .execute::(ctx) diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index de0cf3157ba..86b8fd0c34d 100644 --- a/encodings/fsst/src/array.rs +++ b/encodings/fsst/src/array.rs @@ -5,12 +5,14 @@ use std::fmt::Debug; use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hasher; +use std::mem::MaybeUninit; use std::sync::Arc; use std::sync::OnceLock; use fsst::Compressor; use fsst::Decompressor; use fsst::Symbol; +use num_traits::AsPrimitive; use prost::Message as _; use vortex_array::Array; use vortex_array::ArrayEq; @@ -20,31 +22,34 @@ use vortex_array::ArrayParts; use vortex_array::ArrayRef; use vortex_array::ArraySlots; use vortex_array::ArrayView; -use vortex_array::Canonical; use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; -use vortex_array::IntoArray; use vortex_array::TypedArrayRef; use vortex_array::VortexSessionExecute; +use vortex_array::array_slots; use vortex_array::arrays::VarBin; use vortex_array::arrays::VarBinArray; -use vortex_array::arrays::varbin::VarBinArrayExt; +use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_array::buffer::BufferHandle; use vortex_array::builders::ArrayBuilder; +use vortex_array::builders::VarBinBuilder; use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; +use vortex_array::dtype::OffsetBuilderPType; use vortex_array::dtype::PType; use vortex_array::legacy_session; +use vortex_array::match_each_integer_ptype; +use vortex_array::match_each_varbin_builder; use vortex_array::serde::ArrayChildren; -use vortex_array::smallvec::smallvec; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTable; use vortex_array::vtable::child_to_validity; use vortex_array::vtable::validity_to_child; use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -55,8 +60,10 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::canonical::FSST_DECODE_SLACK; +use crate::canonical::FsstDecodePlan; use crate::canonical::canonicalize_fsst; -use crate::canonical::fsst_decode_views; +use crate::canonical::fsst_decode_bytes; use crate::rules::RULES; /// A [`FSST`]-encoded Vortex array. @@ -78,25 +85,28 @@ impl FSSTMetadata { } } +/// The number of entries in a fully-populated FSST symbol table. +/// +/// Code 255 is reserved as the escape code, leaving 255 usable codes. [`Decompressor`] borrows +/// the symbol table as fixed-size arrays of exactly this length, so Vortex pads the symbols and +/// symbol lengths buffers out to it on construction. +pub const FSST_SYMBOL_TABLE_LEN: usize = 255; + impl ArrayHash for FSSTData { fn array_hash(&self, state: &mut H, precision: EqMode) { - self.symbol_table.symbols.array_hash(state, precision); - self.symbol_table - .symbol_lengths - .array_hash(state, precision); + self.padded_symbols().array_hash(state, precision); + self.padded_symbol_lengths().array_hash(state, precision); self.codes_bytes.as_host().array_hash(state, precision); } } impl ArrayEq for FSSTData { fn array_eq(&self, other: &Self, precision: EqMode) -> bool { - self.symbol_table - .symbols - .array_eq(&other.symbol_table.symbols, precision) + self.padded_symbols() + .array_eq(other.padded_symbols(), precision) && self - .symbol_table - .symbol_lengths - .array_eq(&other.symbol_table.symbol_lengths, precision) + .padded_symbol_lengths() + .array_eq(other.padded_symbol_lengths(), precision) && self .codes_bytes .as_host() @@ -133,8 +143,18 @@ impl VTable for FSST { fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { match idx { - 0 => BufferHandle::new_host(array.symbols().clone().into_byte_buffer()), - 1 => BufferHandle::new_host(array.symbol_lengths().clone().into_byte_buffer()), + 0 => BufferHandle::new_host( + array + .padded_symbols() + .slice(0..array.n_symbols()) + .into_byte_buffer(), + ), + 1 => BufferHandle::new_host( + array + .padded_symbol_lengths() + .slice(0..array.n_symbols()) + .into_byte_buffer(), + ), 2 => array.codes_bytes_handle().clone(), _ => vortex_panic!("FSSTArray buffer index {idx} out of bounds"), } @@ -172,9 +192,7 @@ impl VTable for FSST { array: ArrayView<'_, Self>, _session: &VortexSession, ) -> VortexResult>> { - let codes_offsets = array.as_ref().slots()[CODES_OFFSETS_SLOT] - .as_ref() - .vortex_expect("FSSTArray codes_offsets slot"); + let codes_offsets = array.codes_offsets(); Ok(Some( FSSTMetadata { uncompressed_lengths_ptype: array.uncompressed_lengths().dtype().as_ptype().into(), @@ -266,8 +284,8 @@ impl VTable for FSST { }; FSSTData::validate_parts( - &symbols, - &symbol_lengths, + symbols.as_slice(), + symbol_lengths.as_slice(), &codes_bytes, &codes_offsets, dtype.nullability(), @@ -276,11 +294,12 @@ impl VTable for FSST { len, &mut ctx, )?; - let slots = smallvec![ - Some(uncompressed_lengths), - Some(codes_offsets), - validity_to_child(&codes_validity, len), - ]; + let slots = FSSTSlots { + uncompressed_lengths, + codes_offsets, + codes_validity: validity_to_child(&codes_validity, len), + } + .into_slots(); let data = FSSTData::try_new(symbols, symbol_lengths, codes_bytes, len)?; return Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)); } @@ -292,7 +311,7 @@ impl VTable for FSST { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + FSSTSlots::NAMES[idx].to_string() } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { @@ -304,29 +323,34 @@ impl VTable for FSST { builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { + if let Some(result) = + match_each_varbin_builder!(builder, |builder| append_to_varbin(array, builder, ctx)) + { + return result; + } + + // The two arms here are every builder a `Utf8`/`Binary` dtype has: all four + // `VarBinBuilder` widths above, and `VarBinViewBuilder` below. There is deliberately no + // canonicalize-then-append fallback — it would decode to a `VarBinView` only for + // `VarBinView::append_to_builder` to reject the same remainder. let Some(builder) = builder.as_any_mut().downcast_mut::() else { - return array - .array() - .clone() - .execute::(ctx)? - .into_array() - .append_to_builder(builder, ctx); + vortex_bail!("append_to_builder for FSST requires a variable-binary builder") }; - // Decompress the whole block of data into a new buffer, and create some views - // from it instead. The new buffer lands after any pending in-progress - // buffer that push_buffer_and_adjusted_views will flush first. - let next_buffer_index = builder.completed_block_count() + u32::from(builder.in_progress()); - let (buffers, views) = fsst_decode_views(array, next_buffer_index, ctx)?; - - builder.push_buffer_and_adjusted_views( - &buffers, - &views, - array - .array() - .validity()? - .execute_mask(array.array().len(), ctx)?, - ); + // Decompress the whole block of data into a new buffer, which the builder adopts as a + // data buffer with views built over it in place. + let validity = array + .array() + .validity()? + .execute_mask(array.array().len(), ctx)?; + let (uncompressed_bytes, uncompressed_lens) = fsst_decode_bytes(array, ctx)?; + match_each_integer_ptype!(uncompressed_lens.ptype(), |P| { + builder.append_buffer_with_lengths( + uncompressed_bytes.freeze(), + uncompressed_lens.as_slice::

(), + &validity, + ) + }); Ok(()) } @@ -339,15 +363,53 @@ impl VTable for FSST { } } -/// Lengths of the original values before compression, can be compressed. -pub(crate) const UNCOMPRESSED_LENGTHS_SLOT: usize = 0; -/// The offsets array for the FSST-compressed codes. -pub(crate) const CODES_OFFSETS_SLOT: usize = 1; -/// The validity bitmap for the compressed codes. -pub(crate) const CODES_VALIDITY_SLOT: usize = 2; -pub(crate) const NUM_SLOTS: usize = 3; -pub(crate) const SLOT_NAMES: [&str; NUM_SLOTS] = - ["uncompressed_lengths", "codes_offsets", "codes_validity"]; +/// Decompresses the code stream straight into `builder`'s byte storage. +/// +/// The offsets are the running sum of the uncompressed lengths the array already stores, so the +/// only work beyond the bulk `decompress_into` is one prefix sum over them. +fn append_to_varbin( + array: ArrayView<'_, FSST>, + builder: &mut VarBinBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> +where + usize: AsPrimitive, +{ + let plan = FsstDecodePlan::new(array, ctx)?; + let validity = array + .array() + .validity()? + .execute_mask(array.array().len(), ctx)?; + let decompressor = array.decompressor(); + // Built once, outside the ptype match: the decoder is `#[inline(always)]`, so creating the + // closure inside each arm would stamp out a copy of the whole decode loop per length type. + let mut decode = |out: &mut [MaybeUninit]| plan.decode_into(&decompressor, out); + match_each_integer_ptype!(plan.lengths.ptype(), |P| { + // SAFETY: `decode_into` initializes exactly the prefix whose length it returns. + unsafe { + builder.append_decoded( + plan.total_size, + FSST_DECODE_SLACK, + plan.lengths.as_slice::

(), + &validity, + &mut decode, + ) + } + }) +} + +#[array_slots(FSST)] +pub struct FSSTSlots { + /// Lengths of the original values before compression, can be compressed. + #[slot(0)] + pub uncompressed_lengths: ArrayRef, + /// The offsets array for the FSST-compressed codes. + #[slot(1)] + pub codes_offsets: ArrayRef, + /// The validity bitmap for the compressed codes. + #[slot(2)] + pub codes_validity: Option, +} /// The inner data for an FSST-compressed array. /// @@ -371,8 +433,7 @@ impl Display for FSSTData { write!( f, "len: {}, nsymbols: {}", - self.len, - self.symbol_table.symbols.len() + self.len, self.symbol_table.n_symbols ) } } @@ -380,8 +441,8 @@ impl Display for FSSTData { impl Debug for FSSTData { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("FSSTArray") - .field("symbols", &self.symbol_table.symbols) - .field("symbol_lengths", &self.symbol_table.symbol_lengths) + .field("symbols", &self.symbols()) + .field("symbol_lengths", &self.symbol_lengths()) .field("codes_bytes_len", &self.codes_bytes.len()) .field("len", &self.len) .field("uncompressed_lengths", &"") @@ -391,27 +452,131 @@ impl Debug for FSSTData { } } -pub(crate) struct FSSTSymbolTable { - symbols: Buffer, - symbol_lengths: Buffer, +pub struct FSSTSymbolTable { + /// Symbols padded out to [`FSST_SYMBOL_TABLE_LEN`] entries, zero-filled past `n_symbols`, + /// so that a [`Decompressor`] can borrow them without copying. + padded_symbols: Buffer, + /// Symbol lengths padded out to [`FSST_SYMBOL_TABLE_LEN`] entries, zero-filled past + /// `n_symbols`. + padded_symbol_lengths: Buffer, + /// The number of populated symbols. Entries at or past this index are padding. + n_symbols: usize, /// Memoized compressor used for push-down of compute by compressing the RHS. compressor: OnceLock, } impl FSSTSymbolTable { - fn new(symbols: Buffer, symbol_lengths: Buffer) -> Self { + /// Builds a symbol table, padding `symbols` and `symbol_lengths` out to + /// [`FSST_SYMBOL_TABLE_LEN`] entries if they are shorter. + /// + /// `n_symbols` is the number of populated entries; everything past it is padding. Buffers + /// longer than [`FSST_SYMBOL_TABLE_LEN`] are truncated; callers are expected to have rejected + /// them already (see [`FSSTData::try_new`]). + pub fn new(symbols: Buffer, symbol_lengths: Buffer, n_symbols: usize) -> Self { Self { - symbols, - symbol_lengths, + padded_symbols: pad_symbol_table(symbols, Symbol::ZERO), + padded_symbol_lengths: pad_symbol_table(symbol_lengths, 0), + n_symbols: n_symbols.min(FSST_SYMBOL_TABLE_LEN), compressor: OnceLock::new(), } } - fn compressor(&self) -> &Compressor { - self.compressor.get_or_init(|| { - Compressor::rebuild_from(self.symbols.as_slice(), self.symbol_lengths.as_slice()) + /// Builds a symbol table, padding `symbols` and `symbol_lengths` out to + /// [`FSST_SYMBOL_TABLE_LEN`] entries if they are shorter. + /// + /// `n_symbols` is the number of populated entries; everything past it is padding. Buffers + /// longer than [`FSST_SYMBOL_TABLE_LEN`] are truncated; callers are expected to have rejected + /// them already (see [`FSSTData::try_new`]). + pub fn new_padded( + padded_symbols: Buffer, + padded_symbol_lengths: Buffer, + n_symbols: usize, + ) -> VortexResult { + vortex_ensure!( + padded_symbols.len() == FSST_SYMBOL_TABLE_LEN + && padded_symbol_lengths.len() == FSST_SYMBOL_TABLE_LEN, + InvalidArgument: "padded symbol table must have exactly {FSST_SYMBOL_TABLE_LEN} entries, found {} symbols and {} symbol lengths", + padded_symbols.len(), + padded_symbol_lengths.len() + ); + vortex_ensure!( + n_symbols <= FSST_SYMBOL_TABLE_LEN, + InvalidArgument: "n_symbols must be <= {FSST_SYMBOL_TABLE_LEN}, found {n_symbols}" + ); + Ok(Self { + padded_symbols, + padded_symbol_lengths, + n_symbols, + compressor: OnceLock::new(), }) } + + /// The populated symbols, excluding the padding added by [`Self::new`]. + fn symbols(&self) -> &[Symbol] { + &self.padded_symbols.as_slice()[..self.n_symbols] + } + + /// The populated symbol lengths, excluding the padding added by [`Self::new`]. + fn symbol_lengths(&self) -> &[u8] { + &self.padded_symbol_lengths.as_slice()[..self.n_symbols] + } + + /// The symbols buffer, padded to exactly [`FSST_SYMBOL_TABLE_LEN`] entries with + /// [`Symbol::ZERO`]. + fn padded_symbols(&self) -> &Buffer { + &self.padded_symbols + } + + /// The symbol lengths buffer, padded to exactly [`FSST_SYMBOL_TABLE_LEN`] entries with zeros. + fn padded_symbol_lengths(&self) -> &Buffer { + &self.padded_symbol_lengths + } + + /// Borrow the padded symbol table as the fixed-size arrays expected by [`Decompressor`]. + /// + /// Both buffers are padded to [`FSST_SYMBOL_TABLE_LEN`] on construction, so this is a + /// length check and a pointer cast rather than a copy. + fn decompressor(&self) -> Decompressor<'_> { + const PADDED: &str = "FSST symbol table is padded to FSST_SYMBOL_TABLE_LEN entries"; + let symbols = self + .padded_symbols + .as_slice() + .first_chunk::() + .vortex_expect(PADDED); + let symbol_lengths = self + .padded_symbol_lengths + .as_slice() + .first_chunk::() + .vortex_expect(PADDED); + Decompressor::new(symbols, symbol_lengths) + } + + fn compressor(&self) -> &Compressor { + self.compressor + .get_or_init(|| Compressor::rebuild_from(self.symbols(), self.symbol_lengths())) + } +} + +/// Returns `buffer` resized to exactly [`FSST_SYMBOL_TABLE_LEN`] entries, filling any tail with +/// `pad`. Buffers that are already the right length are returned untouched. +fn pad_symbol_table(buffer: Buffer, pad: T) -> Buffer { + if buffer.len() == FSST_SYMBOL_TABLE_LEN { + return buffer; + } + padded_symbol_table(buffer.as_slice(), pad) +} + +/// Copies `values` into a buffer of exactly [`FSST_SYMBOL_TABLE_LEN`] entries, filling the tail +/// with `pad`. +/// +/// FSST symbol tables are stored padded so that [`FSSTData::decompressor`] can borrow them as the +/// fixed-size arrays [`Decompressor`] requires, without copying. +pub(crate) fn padded_symbol_table(values: &[T], pad: T) -> Buffer { + let populated = values.len().min(FSST_SYMBOL_TABLE_LEN); + let mut padded = BufferMut::with_capacity(FSST_SYMBOL_TABLE_LEN); + padded.extend_from_slice(&values[..populated]); + padded.push_n(pad, FSST_SYMBOL_TABLE_LEN - populated); + padded.freeze() } #[derive(Clone, Debug)] @@ -433,8 +598,8 @@ impl FSST { ) -> VortexResult { let len = codes.len(); FSSTData::validate_parts_from_codes( - &symbols, - &symbol_lengths, + symbols.as_slice(), + symbol_lengths.as_slice(), &codes, &uncompressed_lengths, &dtype, @@ -449,7 +614,7 @@ impl FSST { }) } - pub(crate) fn try_new_with_symbol_table( + pub fn try_new_with_symbol_table( dtype: DType, symbol_table: Arc, codes: VarBinArray, @@ -458,8 +623,8 @@ impl FSST { ) -> VortexResult { let len = codes.len(); FSSTData::validate_parts_from_codes( - &symbol_table.symbols, - &symbol_table.symbol_lengths, + symbol_table.symbols(), + symbol_table.symbol_lengths(), &codes, &uncompressed_lengths, &dtype, @@ -512,8 +677,8 @@ impl FSST { )?; FSSTData::validate_parts_from_codes( - symbols, - symbol_lengths, + symbols.as_slice(), + symbol_lengths.as_slice(), &codes, &uncompressed_lengths, dtype, @@ -545,16 +710,17 @@ impl FSST { impl FSSTData { fn make_slots(codes: &VarBinArray, uncompressed_lengths: &ArrayRef) -> ArraySlots { - smallvec![ - Some(uncompressed_lengths.clone()), - Some(codes.offsets().clone()), - validity_to_child( + FSSTSlots { + uncompressed_lengths: uncompressed_lengths.clone(), + codes_offsets: codes.offsets().clone(), + codes_validity: validity_to_child( &codes .validity() .vortex_expect("FSST codes validity should be derivable"), codes.len(), ), - ] + } + .into_slots() } /// Build FSST data from a set of `symbols`, `symbol_lengths`, and compressed codes bytes. @@ -567,6 +733,10 @@ impl FSSTData { /// corresponds either to a symbol or to the "escape code" (which tells the decoder to /// emit the following byte without doing a table lookup). /// + /// `symbols` and `symbol_lengths` hold only the populated entries; they are padded out to + /// [`FSST_SYMBOL_TABLE_LEN`] entries so that [`Self::decompressor`] can borrow them without + /// copying. + /// /// The offsets and validity for the codes are stored in the array's slots, not here. /// Use [`FSSTArrayExt::codes()`] to reconstruct a full `VarBinArray`. pub fn try_new( @@ -575,11 +745,23 @@ impl FSSTData { codes_bytes: BufferHandle, len: usize, ) -> VortexResult { - // SAFETY: all components validated above + vortex_ensure!( + symbols.len() == symbol_lengths.len(), + InvalidArgument: "symbols and symbol_lengths arrays must have same length, found {} and {}", + symbols.len(), + symbol_lengths.len() + ); + vortex_ensure!( + symbols.len() <= FSST_SYMBOL_TABLE_LEN, + InvalidArgument: "symbols array must have length <= {FSST_SYMBOL_TABLE_LEN}, found {}", + symbols.len() + ); + let n_symbols = symbols.len(); + // SAFETY: the symbol table shape is validated above. + let symbol_table = Arc::new(FSSTSymbolTable::new(symbols, symbol_lengths, n_symbols)); unsafe { - Ok(Self::new_unchecked( - symbols, - symbol_lengths, + Ok(Self::new_unchecked_with_symbol_table( + symbol_table, codes_bytes, len, )) @@ -593,16 +775,14 @@ impl FSSTData { slots: &[Option], ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - let codes_offsets = slots[CODES_OFFSETS_SLOT] - .as_ref() - .vortex_expect("FSSTArray codes_offsets slot"); + let fsst_slots = FSSTSlotsView::from_slots(slots); Self::validate_parts( - &self.symbol_table.symbols, - &self.symbol_table.symbol_lengths, + self.symbol_table.symbols(), + self.symbol_table.symbol_lengths(), &self.codes_bytes, - codes_offsets, + fsst_slots.codes_offsets, dtype.nullability(), - uncompressed_lengths_from_slots(slots), + fsst_slots.uncompressed_lengths, dtype, len, ctx, @@ -612,8 +792,8 @@ impl FSSTData { /// Validate using the decomposed components (codes bytes + offsets + nullability). #[expect(clippy::too_many_arguments)] fn validate_parts( - symbols: &Buffer, - symbol_lengths: &Buffer, + symbols: &[Symbol], + symbol_lengths: &[u8], codes_bytes: &BufferHandle, codes_offsets: &ArrayRef, codes_nullability: Nullability, @@ -627,15 +807,15 @@ impl FSSTData { "FSST arrays must be Binary or Utf8, found {dtype}" ); - if symbols.len() > 255 { - vortex_bail!(InvalidArgument: "symbols array must have length <= 255"); + if symbols.len() > FSST_SYMBOL_TABLE_LEN { + vortex_bail!(InvalidArgument: "symbols array must have length <= {FSST_SYMBOL_TABLE_LEN}"); } if symbols.len() != symbol_lengths.len() { vortex_bail!(InvalidArgument: "symbols and symbol_lengths arrays must have same length"); } - Self::validate_symbol_lengths(symbol_lengths.as_slice())?; + Self::validate_symbol_lengths(symbol_lengths)?; // codes_offsets.len() - 1 == number of elements let codes_len = codes_offsets.len().saturating_sub(1); @@ -706,8 +886,8 @@ impl FSSTData { /// Validate using a VarBinArray for the codes (convenience for construction paths). fn validate_parts_from_codes( - symbols: &Buffer, - symbol_lengths: &Buffer, + symbols: &[Symbol], + symbol_lengths: &[u8], codes: &VarBinArray, uncompressed_lengths: &ArrayRef, dtype: &DType, @@ -727,16 +907,6 @@ impl FSSTData { ) } - pub(crate) unsafe fn new_unchecked( - symbols: Buffer, - symbol_lengths: Buffer, - codes_bytes: BufferHandle, - len: usize, - ) -> Self { - let symbol_table = Arc::new(FSSTSymbolTable::new(symbols, symbol_lengths)); - unsafe { Self::new_unchecked_with_symbol_table(symbol_table, codes_bytes, len) } - } - pub(crate) unsafe fn new_unchecked_with_symbol_table( symbol_table: Arc, codes_bytes: BufferHandle, @@ -760,16 +930,39 @@ impl FSSTData { } /// Access the symbol table array. - pub fn symbols(&self) -> &Buffer { - &self.symbol_table.symbols + /// + /// This is the populated prefix of the symbol table; the padding added on construction to + /// reach [`FSST_SYMBOL_TABLE_LEN`] is not included. Use [`Self::padded_symbols`] to get the + /// whole buffer. + pub fn symbols(&self) -> &[Symbol] { + &self.symbol_table.padded_symbols().as_slice()[0..self.symbol_table.n_symbols] } /// Access the symbol lengths array. - pub fn symbol_lengths(&self) -> &Buffer { - &self.symbol_table.symbol_lengths + /// + /// As with [`Self::symbols`], this excludes the padding added on construction. + pub fn symbol_lengths(&self) -> &[u8] { + &self.symbol_table.padded_symbol_lengths().as_slice()[0..self.symbol_table.n_symbols] + } + + /// The whole symbols buffer, padded to exactly [`FSST_SYMBOL_TABLE_LEN`] entries with + /// [`Symbol::ZERO`]. Entries at or past [`Self::n_symbols`] are padding. + pub fn padded_symbols(&self) -> &Buffer { + self.symbol_table.padded_symbols() + } + + /// The whole symbol lengths buffer, padded to exactly [`FSST_SYMBOL_TABLE_LEN`] entries with + /// zeros. Entries at or past [`Self::n_symbols`] are padding. + pub fn padded_symbol_lengths(&self) -> &Buffer { + self.symbol_table.padded_symbol_lengths() } - pub(crate) fn symbol_table(&self) -> Arc { + /// The number of populated entries in the symbol table. + pub fn n_symbols(&self) -> usize { + self.symbol_table.n_symbols + } + + pub fn symbol_table(&self) -> Arc { Arc::clone(&self.symbol_table) } @@ -786,7 +979,7 @@ impl FSSTData { /// Build a [`Decompressor`] that can be used to decompress values from /// this array. pub fn decompressor(&self) -> Decompressor<'_> { - Decompressor::new(self.symbols().as_slice(), self.symbol_lengths().as_slice()) + self.symbol_table.decompressor() } /// Retrieves the FSST compressor. @@ -795,17 +988,7 @@ impl FSSTData { } } -fn uncompressed_lengths_from_slots(slots: &[Option]) -> &ArrayRef { - slots[UNCOMPRESSED_LENGTHS_SLOT] - .as_ref() - .vortex_expect("FSSTArray uncompressed_lengths slot") -} - -pub trait FSSTArrayExt: TypedArrayRef { - fn uncompressed_lengths(&self) -> &ArrayRef { - uncompressed_lengths_from_slots(self.as_ref().slots()) - } - +pub trait FSSTArrayExt: FSSTArraySlotsExt { fn uncompressed_lengths_dtype(&self) -> &DType { self.uncompressed_lengths().dtype() } @@ -813,14 +996,9 @@ pub trait FSSTArrayExt: TypedArrayRef { /// Reconstruct a [`VarBinArray`] for the compressed codes by combining the bytes /// from [`FSSTData`] with the offsets and validity stored in the array's slots. fn codes(&self) -> VarBinArray { - let offsets = self.as_ref().slots()[CODES_OFFSETS_SLOT] - .as_ref() - .vortex_expect("FSSTArray codes_offsets slot") - .clone(); - let validity = child_to_validity( - self.as_ref().slots()[CODES_VALIDITY_SLOT].as_ref(), - self.as_ref().dtype().nullability(), - ); + let offsets = self.codes_offsets().clone(); + let validity = + child_to_validity(self.codes_validity(), self.as_ref().dtype().nullability()); let codes_bytes = self.codes_bytes_handle().clone(); // SAFETY: components were validated at construction time. unsafe { @@ -844,7 +1022,7 @@ impl> FSSTArrayExt for T {} impl ValidityVTable for FSST { fn validity(array: ArrayView<'_, FSST>) -> VortexResult { Ok(child_to_validity( - array.slots()[CODES_VALIDITY_SLOT].as_ref(), + array.codes_validity(), array.dtype().nullability(), )) } @@ -855,6 +1033,7 @@ mod test { use fsst::Compressor; use fsst::Symbol; use prost::Message; + use vortex_array::ArrayDeserialization; use vortex_array::ArrayPlugin; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; @@ -865,14 +1044,21 @@ mod test { use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::test_harness::check_metadata; + use vortex_array::vtable::VTable as _; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_error::vortex_err; use crate::FSST; + use crate::array::FSST_SYMBOL_TABLE_LEN; use crate::array::FSSTArrayExt; + use crate::array::FSSTArraySlotsExt; + use crate::array::FSSTData; use crate::array::FSSTMetadata; + use crate::array::FSSTSymbolTable; + use crate::array::padded_symbol_table; use crate::fsst_compress; + use crate::fsst_train_compressor; #[test] fn slice_reuses_initialized_compressor() -> VortexResult<()> { @@ -898,6 +1084,177 @@ mod test { Ok(()) } + /// The symbol table is padded out to [`FSST_SYMBOL_TABLE_LEN`] so that `Decompressor` can + /// borrow it directly, but the logical accessors and the serialized buffers must still only + /// expose the populated prefix. + #[test] + fn symbol_table_padded_on_creation() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let strings = VarBinViewArray::from_iter_str(["abcabcab", "defghijk", "abcxyz"]); + let compressor = Compressor::rebuild_from( + [ + Symbol::from_slice(b"abc00000"), + Symbol::from_slice(b"defghijk"), + ], + [3u8, 8], + ); + let fsst_array = fsst_compress(&strings.into_array(), &compressor, &mut ctx)?; + + assert_eq!(fsst_array.padded_symbols().len(), FSST_SYMBOL_TABLE_LEN); + assert_eq!( + fsst_array.padded_symbol_lengths().len(), + FSST_SYMBOL_TABLE_LEN + ); + assert_eq!(fsst_array.padded_symbol_lengths().as_slice()[2..], [0; 253]); + + // Accessors and serialized buffers only see the two populated symbols. + assert_eq!(fsst_array.n_symbols(), 2); + assert_eq!(fsst_array.symbols().len(), 2); + assert_eq!(fsst_array.symbol_lengths(), &[3, 8]); + assert_eq!( + FSST::buffer(fsst_array.as_view(), 0).len(), + 2 * size_of::() + ); + assert_eq!(FSST::buffer(fsst_array.as_view(), 1).len(), 2); + + let decompressed = fsst_array + .into_array() + .execute::(&mut ctx)?; + assert_eq!(decompressed.bytes_at(0).as_slice(), b"abcabcab".as_ref()); + assert_eq!(decompressed.bytes_at(1).as_slice(), b"defghijk".as_ref()); + assert_eq!(decompressed.bytes_at(2).as_slice(), b"abcxyz".as_ref()); + Ok(()) + } + + /// Buffers arriving from a file hold the unpadded symbol table, so deserialization must pad + /// them before a `Decompressor` can borrow them. + #[test] + fn symbol_table_padded_on_deserialize() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let input = VarBinViewArray::from_iter_str(["abcabcab", "defghijk"]).into_array(); + let compressor = fsst_train_compressor(&input, &mut ctx)?; + let fsst_array = fsst_compress(&input, &compressor, &mut ctx)?; + + let buffers = [ + BufferHandle::new_host( + fsst_array + .padded_symbols() + .slice(0..fsst_array.n_symbols()) + .into_byte_buffer(), + ), + BufferHandle::new_host( + fsst_array + .padded_symbol_lengths() + .slice(0..fsst_array.n_symbols()) + .into_byte_buffer(), + ), + fsst_array.codes_bytes_handle().clone(), + ]; + assert!(buffers[1].len() < FSST_SYMBOL_TABLE_LEN); + + let children = vec![ + fsst_array.uncompressed_lengths().clone(), + fsst_array.codes_offsets().clone(), + ]; + + let deserialized = ArrayPlugin::deserialize( + &FSST, + ArrayDeserialization::new( + vortex_array::ArrayVTable::id(&FSST), + &DType::Utf8(Nullability::NonNullable), + 2, + &FSSTMetadata { + uncompressed_lengths_ptype: fsst_array + .uncompressed_lengths() + .dtype() + .as_ptype() + .into(), + codes_offsets_ptype: fsst_array.codes_offsets().dtype().as_ptype().into(), + } + .encode_to_vec(), + &buffers, + &children.as_slice(), + ), + &array_session(), + )?; + + let padded = deserialized + .clone() + .try_downcast::() + .map_err(|_| vortex_err!("deserialize must return an FSST array"))?; + assert_eq!(padded.padded_symbols().len(), FSST_SYMBOL_TABLE_LEN); + assert_eq!(padded.n_symbols(), fsst_array.symbols().len()); + + let decompressed = deserialized.execute::(&mut ctx)?; + assert_eq!(decompressed.bytes_at(0).as_slice(), b"abcabcab".as_ref()); + assert_eq!(decompressed.bytes_at(1).as_slice(), b"defghijk".as_ref()); + Ok(()) + } + + /// An already-padded table must be stored as-is rather than copied again. + #[test] + fn padded_constructor_does_not_repad() -> VortexResult<()> { + let symbols = padded_symbol_table(&[Symbol::from_slice(b"ab000000")], Symbol::ZERO); + let symbol_lengths = padded_symbol_table(&[2u8], 0); + let symbols_ptr = symbols.as_slice().as_ptr(); + let symbol_lengths_ptr = symbol_lengths.as_slice().as_ptr(); + + let data = FSSTSymbolTable::new_padded(symbols, symbol_lengths, 1)?; + + assert_eq!(data.padded_symbols().as_slice().as_ptr(), symbols_ptr); + assert_eq!( + data.padded_symbol_lengths().as_slice().as_ptr(), + symbol_lengths_ptr + ); + assert_eq!(data.symbols().len(), 1); + Ok(()) + } + + /// [`FSSTData::try_new_padded`] stores the buffers as given, so it must reject tables that are + /// not already padded. + #[test] + fn rejects_unpadded_input_to_padded_constructor() { + assert!( + FSSTSymbolTable::new_padded( + Buffer::::copy_from([Symbol::from_slice(b"ab000000")]), + Buffer::::copy_from([2]), + 1, + ) + .is_err() + ); + assert!( + FSSTSymbolTable::new_padded( + Buffer::::full(Symbol::ZERO, FSST_SYMBOL_TABLE_LEN), + Buffer::::full(0, FSST_SYMBOL_TABLE_LEN), + FSST_SYMBOL_TABLE_LEN + 1, + ) + .is_err() + ); + } + + #[test] + fn rejects_malformed_symbol_table() { + let codes_bytes = BufferHandle::new_host(Buffer::::empty()); + assert!( + FSSTData::try_new( + Buffer::::copy_from([Symbol::from_slice(b"ab000000")]), + Buffer::::copy_from([2, 2]), + codes_bytes.clone(), + 0, + ) + .is_err() + ); + assert!( + FSSTData::try_new( + Buffer::::full(Symbol::from_slice(b"ab000000"), FSST_SYMBOL_TABLE_LEN + 1,), + Buffer::::full(2, FSST_SYMBOL_TABLE_LEN + 1), + codes_bytes, + 0, + ) + .is_err() + ); + } + #[cfg_attr(miri, ignore)] #[test] fn test_fsst_metadata() { @@ -951,20 +1308,23 @@ mod test { let fsst = ArrayPlugin::deserialize( &FSST, - &DType::Utf8(Nullability::NonNullable), - 2, - &FSSTMetadata { - uncompressed_lengths_ptype: fsst_array - .uncompressed_lengths() - .dtype() - .as_ptype() - .into(), - // Legacy array did not store this field, use Protobuf default of 0. - codes_offsets_ptype: 0, - } - .encode_to_vec(), - &buffers, - &children.as_slice(), + ArrayDeserialization::new( + vortex_array::ArrayVTable::id(&FSST), + &DType::Utf8(Nullability::NonNullable), + 2, + &FSSTMetadata { + uncompressed_lengths_ptype: fsst_array + .uncompressed_lengths() + .dtype() + .as_ptype() + .into(), + // Legacy array did not store this field, use Protobuf default of 0. + codes_offsets_ptype: 0, + } + .encode_to_vec(), + &buffers, + &children.as_slice(), + ), &array_session(), )?; diff --git a/encodings/fsst/src/canonical.rs b/encodings/fsst/src/canonical.rs index 5c58dbbd501..004827dfdd2 100644 --- a/encodings/fsst/src/canonical.rs +++ b/encodings/fsst/src/canonical.rs @@ -1,8 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::mem::MaybeUninit; use std::sync::Arc; +use fsst::Decompressor; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; @@ -10,23 +12,31 @@ use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::varbin::VarBinArrayExt; -use vortex_array::arrays::varbinview::build_views::BinaryView; use vortex_array::arrays::varbinview::build_views::MAX_BUFFER_LEN; use vortex_array::arrays::varbinview::build_views::build_views; use vortex_array::match_each_integer_ptype; -use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use crate::FSST; use crate::FSSTArrayExt; +use crate::FSSTArraySlotsExt; pub(super) fn canonicalize_fsst( array: ArrayView<'_, FSST>, ctx: &mut ExecutionCtx, ) -> VortexResult { - let (buffers, views) = fsst_decode_views(array, 0, ctx)?; + let (uncompressed_bytes, uncompressed_lens) = fsst_decode_bytes(array, ctx)?; + let (buffers, views) = match_each_integer_ptype!(uncompressed_lens.ptype(), |P| { + build_views( + 0, + MAX_BUFFER_LEN, + uncompressed_bytes.freeze(), + uncompressed_lens.as_slice::

(), + ) + }); // SAFETY: FSST already validates the bytes for binary/UTF-8. We build views directly on // top of them, so the view pointers will all be valid. Ok(unsafe { @@ -40,50 +50,82 @@ pub(super) fn canonicalize_fsst( }) } -pub(crate) fn fsst_decode_views( +/// Extra headroom that keeps [`Decompressor::decompress_into`] on its fast path. +/// +/// It never writes past the slice it is given — every store is bounded by the end of the output — +/// but it emits whole 8-byte symbols, so it can only use the wide-store loop while at least 8 +/// bytes remain. Handing it 7 spare bytes lets that loop run through the final value instead of +/// finishing byte at a time. This is a performance knob, not a safety requirement. +/// +/// [`Decompressor::decompress_into`]: fsst::Decompressor::decompress_into +pub(crate) const FSST_DECODE_SLACK: usize = 7; + +/// Everything needed to decode an FSST array's values in one bulk `decompress_into` call. +pub(crate) struct FsstDecodePlan { + codes: ByteBuffer, + /// Per-row uncompressed lengths, zero for null rows. + pub(crate) lengths: PrimitiveArray, + /// Total decoded size, i.e. the sum of `lengths`. + pub(crate) total_size: usize, +} + +impl FsstDecodePlan { + pub(crate) fn new( + fsst_array: ArrayView<'_, FSST>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let codes = fsst_array.codes().sliced_bytes(); + let lengths = fsst_array + .uncompressed_lengths() + .clone() + .execute::(ctx)?; + + #[expect(clippy::cast_possible_truncation)] + let total_size: usize = match_each_integer_ptype!(lengths.ptype(), |P| { + lengths.as_slice::

().iter().map(|x| *x as usize).sum() + }); + + Ok(Self { + codes, + lengths, + total_size, + }) + } + + /// Bulk-decompresses the whole code stream into `out`, which must hold at least + /// `total_size + FSST_DECODE_SLACK` bytes. + /// + /// Kept inlinable so the decoder is not called from behind an extra frame in whichever + /// codegen unit the caller lands in; see OnPair's equivalent for why that matters. + #[inline] + pub(crate) fn decode_into( + &self, + decompressor: &Decompressor<'_>, + out: &mut [MaybeUninit], + ) -> VortexResult { + let len = decompressor.decompress_into(self.codes.as_slice(), out); + vortex_ensure!( + len == self.total_size, + "FSST decoded {len} bytes, expected {}", + self.total_size + ); + Ok(len) + } +} + +pub(crate) fn fsst_decode_bytes( fsst_array: ArrayView<'_, FSST>, - start_buf_index: u32, ctx: &mut ExecutionCtx, -) -> VortexResult<(Vec, Buffer)> { - // FSSTArray has two child arrays: - // 1. A VarBinArray, which holds the string heap of the compressed codes. - // 2. An uncompressed_lengths primitive array, storing the length of each original - // string element. - // To speed up canonicalization, we can decompress the entire string-heap in a single - // call. We then turn our uncompressed_lengths into an offsets buffer - // necessary for a VarBinViewArray and construct the canonical array. - let bytes = fsst_array.codes().sliced_bytes(); - - let uncompressed_lens_array = fsst_array - .uncompressed_lengths() - .clone() - .execute::(ctx)?; - - #[expect(clippy::cast_possible_truncation)] - let total_size: usize = match_each_integer_ptype!(uncompressed_lens_array.ptype(), |P| { - uncompressed_lens_array - .as_slice::

() - .iter() - .map(|x| *x as usize) - .sum() - }); - - // Bulk-decompress the entire array. - let decompressor = fsst_array.decompressor(); - let mut uncompressed_bytes = ByteBufferMut::with_capacity(total_size + 7); - let len = - decompressor.decompress_into(bytes.as_slice(), uncompressed_bytes.spare_capacity_mut()); +) -> VortexResult<(ByteBufferMut, PrimitiveArray)> { + let plan = FsstDecodePlan::new(fsst_array, ctx)?; + let mut uncompressed_bytes = ByteBufferMut::with_capacity(plan.total_size + FSST_DECODE_SLACK); + let len = plan.decode_into( + &fsst_array.decompressor(), + uncompressed_bytes.spare_capacity_mut(), + )?; + // SAFETY: `decode_into` initialized the first `len` bytes. unsafe { uncompressed_bytes.set_len(len) }; - - // Directly create the binary views. - match_each_integer_ptype!(uncompressed_lens_array.ptype(), |P| { - Ok(build_views( - start_buf_index, - MAX_BUFFER_LEN, - uncompressed_bytes, - uncompressed_lens_array.as_slice::

(), - )) - }) + Ok((uncompressed_bytes, plan.lengths)) } #[cfg(test)] @@ -97,15 +139,21 @@ mod tests { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; + use vortex_array::arrays::varbin::VarBinArrayExt; use vortex_array::builders::ArrayBuilder; + use vortex_array::builders::VarBinBuilder; use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_error::VortexResult; use vortex_session::VortexSession; + use super::fsst_decode_bytes; + use crate::FSST; + use crate::FSSTArrayExt; use crate::fsst_compress; use crate::fsst_train_compressor; @@ -172,8 +220,11 @@ mod tests { let mut ctx = SESSION.create_execution_ctx(); let (chunked_arr, data) = make_data_chunked(); - let mut builder = - VarBinViewBuilder::with_capacity(chunked_arr.dtype().clone(), chunked_arr.len()); + let mut builder = VarBinViewBuilder::with_capacity_in( + chunked_arr.dtype().clone(), + chunked_arr.len(), + ctx.allocator().clone(), + ); chunked_arr .clone() .into_array() @@ -199,13 +250,34 @@ mod tests { .collect::>(); assert_eq!(data, res2) }; + + { + let mut builder = VarBinBuilder::::with_capacity_in( + chunked_arr.dtype().clone(), + data.len(), + ctx.allocator(), + ); + chunked_arr + .into_array() + .append_to_builder(&mut builder, &mut ctx)?; + let arr = builder.finish_into_varbin(); + let mask = arr.validity()?.execute_mask(arr.len(), &mut ctx)?; + let actual = (0..arr.len()) + .map(|i| mask.value(i).then(|| arr.bytes_at(i).to_vec())) + .collect::>(); + assert_eq!(data, actual); + } Ok(()) } #[test] fn test_append_after_in_progress_buffer() -> VortexResult<()> { let dtype = DType::Binary(Nullability::NonNullable); - let mut builder = VarBinViewBuilder::with_capacity(dtype.clone(), 2); + let mut builder = VarBinViewBuilder::with_capacity_in( + dtype.clone(), + 2, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); builder.append_value(b"long enough!!!"); let varbin = VarBinArray::from_iter( @@ -225,4 +297,21 @@ mod tests { let _result = builder.finish_into_varbinview(); Ok(()) } + + #[test] + fn test_rejects_incorrect_uncompressed_lengths() -> VortexResult<()> { + let input = VarBinViewArray::from_iter_str(["hello"]).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let encoded = fsst_compress(&input, &fsst_train_compressor(&input, &mut ctx)?, &mut ctx)?; + let invalid = FSST::try_new_with_symbol_table( + encoded.dtype().clone(), + encoded.symbol_table(), + encoded.codes(), + PrimitiveArray::from_iter([4i32]).into_array(), + &mut ctx, + )?; + + assert!(fsst_decode_bytes(invalid.as_view(), &mut ctx).is_err()); + Ok(()) + } } diff --git a/encodings/fsst/src/compress.rs b/encodings/fsst/src/compress.rs index b9f815b474f..2cc28adcf11 100644 --- a/encodings/fsst/src/compress.rs +++ b/encodings/fsst/src/compress.rs @@ -10,6 +10,7 @@ use std::sync::Arc; use fsst::Compressor; +use fsst::Symbol; use num_traits::AsPrimitive; use vortex_array::ArrayRef; use vortex_array::ArrayView; @@ -18,14 +19,15 @@ use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBin; use vortex_array::arrays::VarBinView; -use vortex_array::arrays::varbin::VarBinArrayExt; +use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_array::arrays::varbin::builder::VarBinBuilder; use vortex_array::arrays::varbinview::BinaryView; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::IntegerPType; +use vortex_array::dtype::OffsetBuilderPType; use vortex_array::match_each_integer_ptype; -use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -35,6 +37,8 @@ use vortex_mask::Mask; use crate::FSST; use crate::FSSTArray; +use crate::array::FSSTSymbolTable; +use crate::array::padded_symbol_table; /// FSST worst case: every input byte expands to an escape + literal (2x). const FSST_PER_BYTE_OVERHEAD: usize = 2; @@ -197,9 +201,14 @@ fn compress_views( ctx: &mut ExecutionCtx, ) -> VortexResult where - O: IntegerPType + 'static, + O: OffsetBuilderPType + 'static, { - let mut sink = FsstSink::::with_capacity(strings.len(), compressor); + let mut sink = FsstSink::::with_capacity( + DType::Binary(strings.dtype().nullability()), + strings.len(), + compressor, + ctx.allocator(), + ); let views = strings.views(); let buffers = strings.data_buffers(); match mask.bit_buffer() { @@ -230,9 +239,14 @@ fn compress_varbin( ctx: &mut ExecutionCtx, ) -> VortexResult where - O: IntegerPType + 'static, + O: OffsetBuilderPType + 'static, { - let mut sink = FsstSink::::with_capacity(strings.len(), compressor); + let mut sink = FsstSink::::with_capacity( + DType::Binary(strings.dtype().nullability()), + strings.len(), + compressor, + ctx.allocator(), + ); let bytes = strings.bytes().as_slice(); match_each_integer_ptype!(offsets.ptype(), |I| { let off = offsets.as_slice::(); @@ -269,19 +283,24 @@ where } /// Per-row output state for an FSST compression pass. -struct FsstSink<'c, O: IntegerPType + 'static> { +struct FsstSink<'c, O: OffsetBuilderPType + 'static> { buffer: Vec, builder: VarBinBuilder, uncompressed_lengths: BufferMut, compressor: &'c Compressor, } -impl<'c, O: IntegerPType + 'static> FsstSink<'c, O> { - fn with_capacity(len: usize, compressor: &'c Compressor) -> Self { +impl<'c, O: OffsetBuilderPType + 'static> FsstSink<'c, O> { + fn with_capacity( + dtype: DType, + len: usize, + compressor: &'c Compressor, + allocator: &BufferAllocatorRef, + ) -> Self { Self { buffer: Vec::with_capacity(DEFAULT_BUFFER_LEN), - builder: VarBinBuilder::::with_capacity(len), - uncompressed_lengths: BufferMut::with_capacity(len), + builder: VarBinBuilder::::with_capacity_in(dtype, len, allocator), + uncompressed_lengths: BufferMut::with_capacity_in(len, allocator.clone()), compressor, } } @@ -289,7 +308,7 @@ impl<'c, O: IntegerPType + 'static> FsstSink<'c, O> { #[inline] fn emit(&mut self, row: Option<&[u8]>) { let Some(s) = row else { - self.builder.append_null(); + self.builder.push_null(); self.uncompressed_lengths.push(0); return; }; @@ -299,23 +318,32 @@ impl<'c, O: IntegerPType + 'static> FsstSink<'c, O> { i32::try_from(s.len()).vortex_expect("per-row uncompressed length must fit in i32"), ); - let target = FSST_PER_BYTE_OVERHEAD * s.len(); - if target > self.buffer.len() { - self.buffer.reserve(target - self.buffer.len()); - } + // `compress_into` writes into the spare capacity, so the buffer must be emptied first. + self.buffer.clear(); + self.buffer.reserve(FSST_PER_BYTE_OVERHEAD * s.len()); // SAFETY: `self.buffer` has capacity for the FSST worst-case output of `s`. - unsafe { self.compressor.compress_into(s, &mut self.buffer) }; + let written = unsafe { + self.compressor + .compress_into(s, self.buffer.spare_capacity_mut()) + }; + // SAFETY: `compress_into` initialized the first `written` bytes. + unsafe { self.buffer.set_len(written) }; self.builder.append_value(&self.buffer); } - fn finish(self, dtype: DType, ctx: &mut ExecutionCtx) -> VortexResult { - let codes = self.builder.finish(DType::Binary(dtype.nullability())); - FSST::try_new( + fn finish(mut self, dtype: DType, ctx: &mut ExecutionCtx) -> VortexResult { + let codes = self.builder.finish_into_varbin(); + // Pad the symbol table here so that the array can hand it straight to a `Decompressor` + // without copying. + FSST::try_new_with_symbol_table( dtype, - Buffer::copy_from(self.compressor.symbol_table()), - Buffer::::copy_from(self.compressor.symbol_lengths()), + Arc::new(FSSTSymbolTable::new_padded( + padded_symbol_table(self.compressor.symbol_table(), Symbol::ZERO), + padded_symbol_table(self.compressor.symbol_lengths(), 0), + self.compressor.n_symbols(), + )?), codes, self.uncompressed_lengths.into_array(), ctx, @@ -328,7 +356,7 @@ mod tests { use vortex_array::VortexSessionExecute; use vortex_array::array_session; use vortex_array::arrays::VarBinViewArray; - use vortex_array::arrays::varbin::VarBinArrayExt; + use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_array::dtype::PType; use vortex_error::VortexResult; diff --git a/encodings/fsst/src/compute/byte_length.rs b/encodings/fsst/src/compute/byte_length.rs index d78bd8be0ec..05505fe95e5 100644 --- a/encodings/fsst/src/compute/byte_length.rs +++ b/encodings/fsst/src/compute/byte_length.rs @@ -15,7 +15,7 @@ use vortex_array::validity::Validity; use vortex_error::VortexResult; use crate::FSST; -use crate::array::FSSTArrayExt; +use crate::array::FSSTArraySlotsExt; impl ByteLengthKernel for FSST { fn byte_length( diff --git a/encodings/fsst/src/compute/cast.rs b/encodings/fsst/src/compute/cast.rs index a1469681590..a3068adc437 100644 --- a/encodings/fsst/src/compute/cast.rs +++ b/encodings/fsst/src/compute/cast.rs @@ -6,7 +6,7 @@ use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::VarBinArray; -use vortex_array::arrays::varbin::VarBinArrayExt; +use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_array::dtype::DType; use vortex_array::scalar_fn::fns::cast::CastKernel; use vortex_array::scalar_fn::fns::cast::CastReduce; @@ -15,6 +15,7 @@ use vortex_error::VortexResult; use crate::FSST; use crate::FSSTArrayExt; +use crate::FSSTArraySlotsExt; fn build_with_codes_validity( array: ArrayView<'_, FSST>, diff --git a/encodings/fsst/src/compute/compare.rs b/encodings/fsst/src/compute/compare.rs index ec63af3735f..08b6681085b 100644 --- a/encodings/fsst/src/compute/compare.rs +++ b/encodings/fsst/src/compute/compare.rs @@ -21,6 +21,7 @@ use vortex_error::vortex_bail; use crate::FSST; use crate::FSSTArrayExt; +use crate::FSSTArraySlotsExt; impl CompareKernel for FSST { fn compare( lhs: ArrayView<'_, Self>, diff --git a/encodings/fsst/src/compute/filter.rs b/encodings/fsst/src/compute/filter.rs index eae72cac8dd..45a95520aca 100644 --- a/encodings/fsst/src/compute/filter.rs +++ b/encodings/fsst/src/compute/filter.rs @@ -13,6 +13,7 @@ use vortex_mask::Mask; use crate::FSST; use crate::FSSTArrayExt; +use crate::FSSTArraySlotsExt; impl FilterKernel for FSST { fn filter( diff --git a/encodings/fsst/src/compute/like.rs b/encodings/fsst/src/compute/like.rs index c922fba088e..388cdb342dc 100644 --- a/encodings/fsst/src/compute/like.rs +++ b/encodings/fsst/src/compute/like.rs @@ -7,7 +7,7 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; -use vortex_array::arrays::varbin::VarBinArrayExt; +use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_array::match_each_integer_ptype; use vortex_array::scalar_fn::fns::like::LikeKernel; use vortex_array::scalar_fn::fns::like::LikeOptions; @@ -47,11 +47,8 @@ impl LikeKernel for FSST { return Ok(None); }; - let symbols = array.symbols(); - let symbol_lengths = array.symbol_lengths(); - let Some(matcher) = - FsstMatcher::try_new(symbols.as_slice(), symbol_lengths.as_slice(), pattern_bytes)? + FsstMatcher::try_new(array.symbols(), array.symbol_lengths(), pattern_bytes)? else { return Ok(None); }; @@ -89,7 +86,6 @@ mod tests { use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::VarBinArray; - use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -122,8 +118,7 @@ mod tests { let len = array.len(); let arr = array.into_array(); let pattern = ConstantArray::new(pattern, len).into_array(); - let result = Like - .try_new_array(len, opts, [arr, pattern])? + let result = Like::try_new(arr, pattern, opts)? .into_array() .execute::(&mut SESSION.create_execution_ctx())?; Ok(result.into_bool()) diff --git a/encodings/fsst/src/compute/list_contains.rs b/encodings/fsst/src/compute/list_contains.rs index e40cb19b854..efb79bb04ad 100644 --- a/encodings/fsst/src/compute/list_contains.rs +++ b/encodings/fsst/src/compute/list_contains.rs @@ -7,14 +7,13 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; -use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; -use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::fns::list_contains::ListContains; use vortex_array::scalar_fn::fns::list_contains::ListContainsElementKernel; use vortex_buffer::ByteBuffer; @@ -105,14 +104,10 @@ impl ListContainsElementKernel for FSST { let needles = element.codes().into_array(); let len = needles.len(); - ListContains - .try_new_array( - len, - EmptyOptions, - [ConstantArray::new(code_list, len).into_array(), needles], - )? - .execute(ctx) - .map(Some) + ListContains::try_new(ConstantArray::new(code_list, len).into_array(), needles)? + .into_array() + .execute::(ctx) + .map(|array| Some(array.into_array())) } } @@ -127,11 +122,9 @@ mod tests { use vortex_array::arrays::ConstantArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::bool::BoolArrayExt; - use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; - use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::fns::list_contains::ListContains; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -223,13 +216,10 @@ mod tests { let mut results = Vec::new(); for needles in [plain, compressed] { let len = needles.len(); - let result = ListContains - .try_new_array( - len, - EmptyOptions, - [ConstantArray::new(list.clone(), len).into_array(), needles], - )? - .execute::(&mut ctx)?; + let result = + ListContains::try_new(ConstantArray::new(list.clone(), len).into_array(), needles)? + .into_array() + .execute::(&mut ctx)?; results.push(bool_answers(&result, &mut ctx)); } @@ -395,16 +385,14 @@ mod tests { let run = |needles: vortex_array::ArrayRef| { let len = needles.len(); - ListContains - .try_new_array( - len, - EmptyOptions, - [ - ConstantArray::new(binary_list.clone(), len).into_array(), - needles, - ], - ) - .and_then(|a| a.execute::(&mut SESSION.create_execution_ctx())) + ListContains::try_new( + ConstantArray::new(binary_list.clone(), len).into_array(), + needles, + ) + .and_then(|a| { + a.into_array() + .execute::(&mut SESSION.create_execution_ctx()) + }) }; let varbin = run(plain); diff --git a/encodings/fsst/src/compute/mod.rs b/encodings/fsst/src/compute/mod.rs index f7cc0582e27..6a49691a513 100644 --- a/encodings/fsst/src/compute/mod.rs +++ b/encodings/fsst/src/compute/mod.rs @@ -12,16 +12,15 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::arrays::VarBin; use vortex_array::arrays::dict::TakeExecute; +use vortex_array::arrays::varbin::take_varbin; use vortex_array::builtins::ArrayBuiltins; use vortex_array::scalar::Scalar; -use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_err; use crate::FSST; use crate::FSSTArrayExt; +use crate::FSSTArraySlotsExt; impl TakeExecute for FSST { fn take( @@ -36,14 +35,7 @@ impl TakeExecute for FSST { .clone() .union_nullability(indices.dtype().nullability()), array.symbol_table(), - { - let codes = array.codes(); - let codes = codes.as_view(); - ::take(codes, indices, ctx)? - .vortex_expect("VarBin take kernel always returns Some") - } - .try_downcast::() - .map_err(|_| vortex_err!("take for codes must return varbin array"))?, + take_varbin(array.codes().as_view(), indices, ctx)?, array .uncompressed_lengths() .take(indices.clone())? diff --git a/encodings/fsst/src/dfa/tests.rs b/encodings/fsst/src/dfa/tests.rs index 3ea9a507c20..88506206fcc 100644 --- a/encodings/fsst/src/dfa/tests.rs +++ b/encodings/fsst/src/dfa/tests.rs @@ -14,7 +14,6 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::VarBinArray; -use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -279,10 +278,8 @@ fn make_fsst_str(strings: &[Option<&str>]) -> FSSTArray { } fn run_like(array: FSSTArray, pattern_arr: ArrayRef) -> VortexResult { - let len = array.len(); let arr: ArrayRef = array.into_array(); - let result = Like - .try_new_array(len, LikeOptions::default(), [arr, pattern_arr])? + let result = Like::try_new(arr, pattern_arr, LikeOptions::default())? .into_array() .execute::(&mut SESSION.create_execution_ctx())?; Ok(result.into_bool()) diff --git a/encodings/fsst/src/kernel.rs b/encodings/fsst/src/kernel.rs index 47806fd0eac..199a1089903 100644 --- a/encodings/fsst/src/kernel.rs +++ b/encodings/fsst/src/kernel.rs @@ -68,7 +68,11 @@ mod tests { }); fn build_test_fsst_array() -> ArrayRef { - let mut builder = VarBinBuilder::::with_capacity(10); + let mut builder = VarBinBuilder::::with_capacity_in( + DType::Utf8(Nullability::NonNullable), + 10, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); builder.append_value(b"hello world"); builder.append_value(b"foo bar baz"); builder.append_value(b"testing fsst compression"); @@ -79,7 +83,7 @@ mod tests { builder.append_value(b"qrstuvwxyz"); builder.append_value(b"0123456789"); builder.append_value(b"final string"); - let input = builder.finish(DType::Utf8(Nullability::NonNullable)); + let input = builder.finish_into_varbin(); let mut ctx = SESSION.create_execution_ctx(); let arr = input.into_array(); @@ -138,7 +142,11 @@ mod tests { // Test case with special characters and nulls // Values: ["", "", "", "", "", "", "", "", "", "", "", ",", "A<<<<<<<", "", "", "", "", null, null, null, null, null, null] // Mask: only the last element is selected (true at index 22) - let mut builder = VarBinBuilder::::with_capacity(23); + let mut builder = VarBinBuilder::::with_capacity_in( + DType::Utf8(Nullability::Nullable), + 23, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); // 11 empty strings for _ in 0..11 { builder.append_value(b""); @@ -153,9 +161,9 @@ mod tests { } // 6 nulls for _ in 0..6 { - builder.append_null(); + builder.push_null(); } - let input = builder.finish(DType::Utf8(Nullability::Nullable)); + let input = builder.finish_into_varbin(); let array = input.clone().into_array(); let mut ctx = SESSION.create_execution_ctx(); @@ -179,12 +187,16 @@ mod tests { #[test] fn filter_only_null() -> VortexResult<()> { - let mut builder = VarBinBuilder::::with_capacity(3); - builder.append_null(); + let mut builder = VarBinBuilder::::with_capacity_in( + DType::Utf8(Nullability::Nullable), + 3, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + builder.push_null(); builder.append_value(b"A"); - builder.append_null(); + builder.push_null(); - let input = builder.finish(DType::Utf8(Nullability::Nullable)); + let input = builder.finish_into_varbin(); let array = input.clone().into_array(); let mut ctx = SESSION.create_execution_ctx(); @@ -220,15 +232,17 @@ mod tests { #[test] fn test_fsst_byte_length() -> VortexResult<()> { - let mut builder = VarBinBuilder::::with_capacity(3); + let mut builder = VarBinBuilder::::with_capacity_in( + DType::Utf8(Nullability::NonNullable), + 3, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); builder.append_value(b"hello"); builder.append_value(b"world!!"); builder.append_value("Пуховички"); // 9 characters, 18 bytes builder.append_value(b""); - let varbin = builder - .finish(DType::Utf8(Nullability::NonNullable)) - .into_array(); + let varbin = builder.finish_into_varbin().into_array(); let mut ctx = SESSION.create_execution_ctx(); let compressor = fsst_train_compressor(&varbin, &mut ctx)?; let fsst = fsst_compress(&varbin, &compressor, &mut ctx)?.into_array(); diff --git a/encodings/fsst/src/slice.rs b/encodings/fsst/src/slice.rs index b98f5fa03b3..c2b2e320d31 100644 --- a/encodings/fsst/src/slice.rs +++ b/encodings/fsst/src/slice.rs @@ -13,6 +13,7 @@ use vortex_error::vortex_err; use crate::FSST; use crate::FSSTArrayExt; +use crate::FSSTArraySlotsExt; impl SliceReduce for FSST { fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { diff --git a/encodings/fsst/src/test_utils.rs b/encodings/fsst/src/test_utils.rs index 7becb34ffbc..07a8bd429f0 100644 --- a/encodings/fsst/src/test_utils.rs +++ b/encodings/fsst/src/test_utils.rs @@ -75,7 +75,8 @@ pub fn gen_dict_fsst_test_data( // Benchmark dataset generators // --------------------------------------------------------------------------- -pub const NUM_STRINGS: usize = 100_000; +// Sized so per-iteration CodSpeed times of benches using it stay under ~1ms. +pub const NUM_STRINGS: usize = 2_000; // --------------------------------------------------------------------------- // URL generator (ClickBench-style weighted domains) diff --git a/encodings/fsst/src/tests.rs b/encodings/fsst/src/tests.rs index 43cdeb1a02c..eec1a3cff5f 100644 --- a/encodings/fsst/src/tests.rs +++ b/encodings/fsst/src/tests.rs @@ -30,15 +30,17 @@ static SESSION: LazyLock = LazyLock::new(|| { /// this function is VERY slow on miri, so we only want to run it once pub(crate) fn build_fsst_array(ctx: &mut ExecutionCtx) -> ArrayRef { - let mut input_array = VarBinBuilder::::with_capacity(3); + let mut input_array = VarBinBuilder::::with_capacity_in( + DType::Utf8(Nullability::NonNullable), + 3, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); input_array.append_value(b"The Greeks never said that the limit could not be overstepped"); input_array.append_value( b"They said it existed and that whoever dared to exceed it was mercilessly struck down", ); input_array.append_value(b"Nothing in present history can contradict them"); - let input_array = input_array - .finish(DType::Utf8(Nullability::NonNullable)) - .into_array(); + let input_array = input_array.finish_into_varbin().into_array(); let compressor = fsst_train_compressor(&input_array, ctx).unwrap(); fsst_compress(&input_array, &compressor, ctx) @@ -162,13 +164,15 @@ fn fsst_compress_offsets_overflow_i32() { println!("building large VarBinArray"); let string = vec![b'a'; STRING_LEN]; - let mut builder = VarBinBuilder::::with_capacity(N); + let mut builder = VarBinBuilder::::with_capacity_in( + DType::Utf8(Nullability::NonNullable), + N, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); for _ in 0..N { builder.append_value(&string); } - let array = builder - .finish(DType::Utf8(Nullability::NonNullable)) - .into_array(); + let array = builder.finish_into_varbin().into_array(); let compressor = CompressorBuilder::default().build(); let len = array.len(); diff --git a/encodings/experimental/onpair/Cargo.toml b/encodings/onpair/Cargo.toml similarity index 100% rename from encodings/experimental/onpair/Cargo.toml rename to encodings/onpair/Cargo.toml diff --git a/encodings/onpair/README.md b/encodings/onpair/README.md new file mode 100644 index 00000000000..f19697a412d --- /dev/null +++ b/encodings/onpair/README.md @@ -0,0 +1,38 @@ +# Vortex OnPair + +A Vortex Encoding for Binary and Utf8 data that uses the +[OnPair][onpair] short-string compression algorithm. OnPair is a +dictionary-based encoder with fast per-row random access. + +The trainer / encoder lives in the standalone [`onpair`][onpair-crate] +crate; this crate wraps the resulting column as a Vortex array with +cascading-compressor support on every integer child. + +## Compute + +Like the FSST encoding, this crate pushes down common operations over the +encoded representation. It supports `cast`, `filter`, byte length, and +constant equality / inequality. Unsupported operators fall back to ordinary +decompression. + +## Default Configuration + +The default training configuration uses OnPair's default dictionary budget and +a fixed seed. Vortex stores token codes as an integer child array; downstream +integer compression may narrow or bit-pack that child independently. + +## Layout + +- Buffer 0 — `dict_bytes`: dictionary blob built by the OnPair trainer, + including the read padding required by the decoder. +- Slot 0 — `dict_offsets`: integer child, len `dict_size + 1`. +- Slot 1 — `codes`: integer child, length `total_tokens`. +- Slot 2 — `codes_offsets`: integer child, length `num_rows + 1`. +- Slot 3 — `uncompressed_lengths`: integer child, length `num_rows`. +- Slot 4 — optional validity child. + +All four integer slot children flow through the standard cascading +compressor pipeline (FoR / BitPacking / RunEnd / etc.). + +[onpair]: https://arxiv.org/abs/2508.02280 +[onpair-crate]: https://github.com/spiraldb/onpair diff --git a/encodings/experimental/onpair/benches/decode.rs b/encodings/onpair/benches/decode.rs similarity index 81% rename from encodings/experimental/onpair/benches/decode.rs rename to encodings/onpair/benches/decode.rs index ff4158cda32..57debb32090 100644 --- a/encodings/experimental/onpair/benches/decode.rs +++ b/encodings/onpair/benches/decode.rs @@ -3,7 +3,7 @@ // //! Decode-path microbenchmarks for the OnPair Vortex array. //! -//! * `decompress_into` — the upstream `onpair::decompress_into` decoder hot +//! * `decode_into` — the upstream `onpair::decode_into` decoder hot //! loop, fed by pre-materialised [`DecodeInputs`]. Measures the inner loop //! only (no child `execute`, no allocation). //! * `canonicalize_to_varbinview` — the full Vortex @@ -28,7 +28,7 @@ use std::mem::MaybeUninit; use std::sync::LazyLock; use divan::Bencher; -use onpair::Parts; +use onpair::CompactDictionaryView; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -38,43 +38,45 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::filter::FilterKernel; +use vortex_array::buffer::BufferHandle; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_buffer::Buffer; -use vortex_buffer::ByteBuffer; use vortex_mask::Mask; -use vortex_onpair::DEFAULT_DICT12_CONFIG; +use vortex_onpair::DEFAULT_CONFIG; use vortex_onpair::OnPair; use vortex_onpair::OnPairArray; use vortex_onpair::OnPairArraySlotsExt; /// Host-resident decode inputs, materialised once so the decode-loop benchmark -/// measures only `onpair::decompress_into` (not child `execute`/allocation). +/// measures only `onpair::decode_into` (not child `execute`/allocation). struct DecodeInputs { - dict_bytes: ByteBuffer, + dict_bytes: BufferHandle, dict_offsets: Buffer, codes: Buffer, - bits: u32, } impl DecodeInputs { - fn as_parts(&self) -> Parts<'_> { - Parts { - dict_bytes: self.dict_bytes.as_slice(), - dict_offsets: self.dict_offsets.as_slice(), - bits: self.bits, - codes: self.codes.as_slice(), + fn dict(&self) -> CompactDictionaryView<'_> { + // SAFETY: `materialise` validates this borrowed dictionary once after + // widening offsets. The benchmark then keeps both buffers immutable. + unsafe { + CompactDictionaryView::new_unchecked( + self.dict_bytes.as_host().as_slice(), + self.dict_offsets.as_slice(), + ) } } - fn decompressed_len(&self) -> usize { - onpair::decompressed_len(self.as_parts()) + fn decoded_len(&self) -> usize { + onpair::decoded_len(self.codes.as_slice(), self.dict()) } - fn decompress_into(&self, out: &mut [MaybeUninit]) -> usize { - onpair::decompress_into(self.as_parts(), out) + fn decode_into(&self, out: &mut [MaybeUninit]) -> usize { + // SAFETY: callers allocate `decoded_len + DECODE_PADDING` bytes. + unsafe { onpair::decode_into(self.codes.as_slice(), self.dict(), out) } } } use vortex_onpair::onpair_compress; @@ -160,8 +162,10 @@ fn compress(n: usize, shape: Shape, ctx: &mut ExecutionCtx) -> OnPairArray { strings.iter().map(|s| Some(s.as_bytes())), DType::Utf8(Nullability::NonNullable), ); - onpair_compress(varbin.as_array(), DEFAULT_DICT12_CONFIG, ctx) + onpair_compress(varbin.as_array(), DEFAULT_CONFIG, ctx) .unwrap_or_else(|e| panic!("onpair_compress failed: {e}")) + .try_downcast::() + .unwrap_or_else(|array| panic!("expected OnPair array, got {}", array.encoding_id())) } /// Canonicalise a slot child to the decoder's native primitive width. @@ -175,13 +179,19 @@ fn widen(arr: &ArrayRef, ctx: &mut ExecutionCtx) -> Buffer { fn materialise(arr: &OnPairArray, ctx: &mut ExecutionCtx) -> (DecodeInputs, usize) { let view = arr.as_view(); + let dict_offsets = widen::(view.dict_offsets(), ctx); + let dict_bytes = view.dict_bytes_handle().clone(); + CompactDictionaryView::validate_safety( + dict_bytes.as_host().as_slice(), + dict_offsets.as_slice(), + ) + .expect("valid OnPair dictionary"); let inputs = DecodeInputs { - dict_bytes: view.dict_bytes().clone(), - dict_offsets: widen::(view.dict_offsets(), ctx), + dict_bytes, + dict_offsets, codes: widen::(view.codes(), ctx), - bits: view.bits(), }; - let total = inputs.decompressed_len(); + let total = inputs.decoded_len(); (inputs, total) } @@ -194,16 +204,16 @@ const CASES: &[(Shape, usize)] = &[ ]; /// Raw decode loop time, excluding child `execute` and the output allocation. -/// Hits `onpair::decompress_into` directly. +/// Hits `onpair::decode_into` directly. #[divan::bench(args = CASES)] -fn decompress_into_bench(bencher: Bencher, case: (Shape, usize)) { +fn decode_into_bench(bencher: Bencher, case: (Shape, usize)) { let mut ctx = SESSION.create_execution_ctx(); let (shape, n) = case; let arr = compress(n, shape, &mut ctx); let (inputs, total) = materialise(&arr, &mut ctx); bencher.bench_local(|| { - let mut out: Vec = Vec::with_capacity(total); - let written = inputs.decompress_into(out.spare_capacity_mut()); + let mut out: Vec = Vec::with_capacity(total + onpair::DECODE_PADDING); + let written = inputs.decode_into(out.spare_capacity_mut()); unsafe { out.set_len(written) }; divan::black_box(out); }); diff --git a/encodings/onpair/goldenfiles/onpair.metadata b/encodings/onpair/goldenfiles/onpair.metadata new file mode 100644 index 00000000000..80b7fcdef71 --- /dev/null +++ b/encodings/onpair/goldenfiles/onpair.metadata @@ -0,0 +1 @@ + (08 \ No newline at end of file diff --git a/encodings/experimental/onpair/src/array.rs b/encodings/onpair/src/array.rs similarity index 52% rename from encodings/experimental/onpair/src/array.rs rename to encodings/onpair/src/array.rs index 519243b100e..912fc27c278 100644 --- a/encodings/experimental/onpair/src/array.rs +++ b/encodings/onpair/src/array.rs @@ -5,7 +5,15 @@ use std::fmt::Debug; use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hasher; - +use std::mem::MaybeUninit; +use std::sync::Arc; +use std::sync::OnceLock; + +use num_traits::AsPrimitive; +use onpair::CompactDictionary; +use onpair::CompactDictionaryView; +use onpair::Dictionary; +use onpair::DictionaryStorage; use prost::Message as _; use vortex_array::Array; use vortex_array::ArrayEq; @@ -14,24 +22,27 @@ use vortex_array::ArrayId; use vortex_array::ArrayParts; use vortex_array::ArrayRef; use vortex_array::ArrayView; -use vortex_array::Canonical; use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; -use vortex_array::IntoArray; use vortex_array::array_slots; use vortex_array::buffer::BufferHandle; use vortex_array::builders::ArrayBuilder; +use vortex_array::builders::VarBinBuilder; use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; +use vortex_array::dtype::OffsetBuilderPType; use vortex_array::dtype::PType; +use vortex_array::match_each_integer_ptype; +use vortex_array::match_each_varbin_builder; use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTable; use vortex_array::vtable::child_to_validity; use vortex_array::vtable::validity_to_child; +use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -41,8 +52,10 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::canonical::OnPairDecodePlan; use crate::canonical::canonicalize_onpair; -use crate::canonical::onpair_decode_views; +use crate::canonical::onpair_decode_bytes; +use crate::decode::collect_widened; use crate::rules::RULES; /// An [`OnPair`]-encoded Vortex array. @@ -52,9 +65,8 @@ pub type OnPairArray = Array; /// /// On disk the layout is FSST-shape: /// -/// * Buffer 0 — `dict_bytes`: the dictionary blob built by the OnPair trainer, -/// padded with `onpair::MAX_TOKEN_SIZE` trailing zero -/// bytes so the over-copy decoder can read 16 bytes past the last token. +/// * Buffer 0 — `dict_bytes`: the read-padded dictionary blob built by the +/// OnPair trainer. /// * Slots — see [`OnPairSlots`]. /// /// The four integer slot children flow through the standard `compress_child` @@ -66,24 +78,18 @@ pub struct OnPairMetadata { /// Width of the per-row primitive `uncompressed_lengths` child. #[prost(enumeration = "PType", tag = "1")] pub uncompressed_lengths_ptype: i32, - /// Bits-per-token the column was compressed with (9..=16). Every value - /// in the `codes` child only uses its low `bits` bits. - #[prost(uint32, tag = "2")] - pub bits: u32, /// Number of dictionary tokens. `dict_offsets` has length `dict_size + 1`. - /// Bounded by `2^bits ≤ 2^16 = 65_536`, so `u32` is comfortably wide. #[prost(uint32, tag = "3")] pub dict_size: u32, - /// Total number of tokens across all rows. `codes` has this length; - /// `codes_offsets.last() == total_tokens`. + /// Length of the `codes` slot child. A sliced array may retain codes that + /// fall outside its visible row range. #[prost(uint64, tag = "4")] - pub total_tokens: u64, + pub codes_len: u64, /// PType of the `dict_offsets` slot child (defaults to U32, may be /// narrowed to U16/U8 by the cascading compressor when values fit). #[prost(enumeration = "PType", tag = "5")] pub dict_offsets_ptype: i32, - /// PType of the `codes` slot child (typically U16, may be narrowed to U8 - /// when `bits <= 8`). + /// PType of the `codes` slot child. #[prost(enumeration = "PType", tag = "6")] pub codes_ptype: i32, /// PType of the `codes_offsets` slot child. @@ -92,6 +98,7 @@ pub struct OnPairMetadata { } impl OnPairMetadata { + /// Decode the recorded [`PType`] of the `uncompressed_lengths` slot child. pub fn get_uncompressed_lengths_ptype(&self) -> VortexResult { PType::try_from(self.uncompressed_lengths_ptype) .map_err(|_| vortex_err!("Invalid PType {}", self.uncompressed_lengths_ptype)) @@ -100,95 +107,184 @@ impl OnPairMetadata { #[array_slots(OnPair)] pub struct OnPairSlots { - /// `PrimitiveArray`, length `dict_size + 1`. Cascading compressor may - /// narrow the ptype to U16/U8. + /// Dictionary-offset child, with length `dict_size + 1`. The cascading + /// compressor may re-encode this child independently; the materialised + /// offsets used by the runtime dictionary cache are derived from it on + /// first use. + #[slot(0)] pub dict_offsets: ArrayRef, - /// `PrimitiveArray`. Each value only uses its low `bits` bits; - /// downstream `FastLanes::BitPacking` losslessly shrinks the child to - /// exactly `bits`-bit codes on disk. + /// Primitive integer token codes. Downstream integer compression may + /// narrow or bit-pack this child independently of the OnPair metadata. + #[slot(1)] pub codes: ArrayRef, - /// `PrimitiveArray`, length `num_rows + 1`. FoR / RunEnd / etc. apply - /// naturally via the cascading compressor. + /// Primitive integer row offsets into `codes`, length `num_rows + 1`. The + /// cascading compressor may re-encode this child independently. + #[slot(2)] pub codes_offsets: ArrayRef, - /// Integer `PrimitiveArray`, length `num_rows`. Used to size the canonical - /// output buffer. + /// Integer decoded-length child, length `num_rows`. Used to size the + /// canonical output buffer. + #[slot(3)] pub uncompressed_lengths: ArrayRef, /// Optional validity child for the outer string column. + #[slot(4)] pub validity: Option, } -/// Inner data for an OnPair-encoded array. +/// Immutable storage for a materialised OnPair dictionary. +/// +/// The two buffers remain owned by Vortex. Implementing [`DictionaryStorage`] +/// this way lets `onpair::CompactDictionary` retain them without copying +/// either the dictionary bytes or the widened offsets. +#[derive(Clone, Debug)] +struct OnPairDictionaryStorage { + bytes: ByteBuffer, + offsets: Buffer, +} + +impl DictionaryStorage for OnPairDictionaryStorage { + #[inline] + fn bytes(&self) -> &[u8] { + self.bytes.as_slice() + } + + #[inline] + fn offsets(&self) -> &[u32] { + self.offsets.as_slice() + } +} + +/// Non-child data for an OnPair-encoded array. /// -/// Holds only the dictionary blob (buffer 0). Every other piece — -/// `dict_offsets`, the per-token `codes`, the per-row `codes_offsets`, the -/// per-row `uncompressed_lengths`, and the optional validity child — is a -/// Vortex slot child so it can be re-encoded by the cascading compressor. +/// The serialized dictionary bytes stay in buffer 0, while the dictionary +/// offsets remain a recursive child so the existing on-disk format is +/// unchanged. Once the offsets are materialised, `dictionary` retains an +/// upstream storage-backed [`CompactDictionary`] around the same immutable +/// buffers. Constructing an array from serialized parts does not materialise or +/// validate the dictionary: Vortex only reconstructs the recursive offset +/// child at that point. Converting that child to contiguous `u32` offsets and +/// validating the dictionary are deferred until the first decode/search +/// operation that needs dictionary token access. #[derive(Clone)] pub struct OnPairData { /// The dictionary blob (buffer 0). /// - /// INVARIANT: this buffer must be over-padded past its logical end - /// (`dict_offsets.last()`) by the decoder's fixed token read width, - /// `onpair::MAX_TOKEN_SIZE`. The over-copy decoder reads - /// every dictionary entry with one fixed-width load and then advances the - /// cursor by the token's true length, so the load for the final, shortest - /// token over-reads past the logical end of the dictionary. This is the - /// same over-read the decoder accounts for on the final few codes; the - /// trailing padding absorbs it so that any entry can be read in bounds. - /// `onpair_compress` establishes this padding (see `parts_to_children`); - /// the over-copy decoder lives in the `onpair` crate. + /// INVARIANT: this buffer is an OnPair compact dictionary byte buffer, + /// including its trailing read padding. dict_bytes: BufferHandle, - bits: u32, - len: usize, + /// The storage-backed dictionary, memoized after successful initialization. + /// Initialization decompresses the child and safety-validates the dictionary; + /// once cached, later operations do neither again. The dictionary owns + /// clones of the immutable Vortex buffer handles, so this does not copy the + /// bytes. A failed validation is not cached, and concurrent first users may + /// duplicate initialization work before one value wins the `OnceLock`. + /// + /// INVARIANT: once populated, the offsets passed + /// [`CompactDictionary::validate_safety`] against `dict_bytes`. + /// The `Arc` cell is shared only between arrays with identical dictionary + /// bytes and logically identical offsets (slice / filter / cast keep both). + dictionary: Arc>>, } impl OnPairData { - pub fn new(dict_bytes: BufferHandle, bits: u32, len: usize) -> Self { + /// Build [`OnPairData`] from the dictionary blob. + pub fn new(dict_bytes: BufferHandle) -> Self { Self { dict_bytes, - bits, - len, + dictionary: Arc::new(OnceLock::new()), } } - pub fn len(&self) -> usize { - self.len - } - - pub fn is_empty(&self) -> bool { - self.len == 0 - } - - pub fn bits(&self) -> u32 { - self.bits + /// Build [`OnPairData`] with the dictionary already materialised, so the + /// first decode skips both widening and validation. + /// + /// `offsets` must be the widened values of the `dict_offsets` child the + /// caller attaches to the array; validation proves only that they are + /// structurally safe against `dict_bytes`. + pub(crate) fn try_new_with_dictionary( + dict_bytes: BufferHandle, + offsets: Buffer, + ) -> VortexResult { + let dictionary = build_dictionary(dict_bytes.as_host().clone(), offsets)?; + Ok(Self { + dict_bytes, + dictionary: Arc::new(OnceLock::from(dictionary)), + }) } + /// The dictionary blob as a host byte buffer. pub fn dict_bytes(&self) -> &ByteBuffer { self.dict_bytes.as_host() } + /// The [`BufferHandle`] holding the dictionary blob (buffer 0). pub fn dict_bytes_handle(&self) -> &BufferHandle { &self.dict_bytes } } +/// Safety-validate `(bytes, offsets)` and seal them into a storage-backed +/// dictionary. +fn build_dictionary( + bytes: ByteBuffer, + offsets: Buffer, +) -> VortexResult> { + CompactDictionary::validate_safety(OnPairDictionaryStorage { bytes, offsets }) + .map_err(|e| vortex_err!(InvalidArgument: "Unsafe OnPair dictionary: {e}")) +} + +/// A safety-validated dictionary built from host-materialized parts, owned +/// independently of any array. +/// +/// [`dict_view`] reaches the same parts through host-only accessors, which panics when +/// the array's buffers are device-resident. Callers that have already copied the +/// dictionary blob and its widened offsets to the host — the CUDA executor, which stages +/// the dictionary on the host regardless — build through this instead. +pub struct OnPairDictionary(CompactDictionary); + +impl OnPairDictionary { + /// Validates `(bytes, offsets)` and seals them into a dictionary. + pub fn try_new(bytes: ByteBuffer, offsets: Buffer) -> VortexResult { + Ok(Self(build_dictionary(bytes, offsets)?)) + } + + /// Borrows the dictionary as a view. + pub fn as_view(&self) -> CompactDictionaryView<'_> { + self.0.as_view() + } +} + +/// A safety-validated [`CompactDictionaryView`] over `array`'s dictionary. +/// +/// The first successful initialization widens the `dict_offsets` child and +/// safety-validates the dictionary structure; the resulting storage-backed +/// dictionary is memoized in [`OnPairData`]. Once cached, subsequent calls — +/// including on arrays derived by slice / filter / cast, which share the cell — +/// pay neither cost again. +pub fn dict_view<'a>( + array: ArrayView<'a, OnPair>, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let data = array.data(); + let dictionary = match data.dictionary.get() { + Some(dictionary) => dictionary, + None => { + let widened = collect_widened::(array.dict_offsets(), ctx)?; + let dictionary = build_dictionary(data.dict_bytes().clone(), widened)?; + data.dictionary.get_or_init(|| dictionary) + } + }; + Ok(dictionary.as_view()) +} + impl Display for OnPairData { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!( - f, - "len: {}, bits: {}, dict_bytes_len: {}", - self.len, - self.bits, - self.dict_bytes.len() - ) + write!(f, "dict_bytes_len: {}", self.dict_bytes.len()) } } impl Debug for OnPairData { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("OnPairData") - .field("len", &self.len) - .field("bits", &self.bits) .field("dict_bytes_len", &self.dict_bytes.len()) .finish() } @@ -197,17 +293,14 @@ impl Debug for OnPairData { impl ArrayHash for OnPairData { fn array_hash(&self, state: &mut H, accuracy: EqMode) { self.dict_bytes.as_host().array_hash(state, accuracy); - state.write_u32(self.bits); } } impl ArrayEq for OnPairData { fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool { - self.bits == other.bits - && self - .dict_bytes - .as_host() - .array_eq(other.dict_bytes.as_host(), accuracy) + self.dict_bytes + .as_host() + .array_eq(other.dict_bytes.as_host(), accuracy) } } @@ -217,7 +310,6 @@ pub struct OnPair; impl OnPair { /// Build an [`OnPairArray`] from already-materialised parts. - #[expect(clippy::too_many_arguments, reason = "every child is a real input")] pub fn try_new( dtype: DType, dict_bytes: BufferHandle, @@ -226,7 +318,37 @@ impl OnPair { codes_offsets: ArrayRef, uncompressed_lengths: ArrayRef, validity: Validity, - bits: u32, + ) -> VortexResult { + Self::try_new_with_data( + dtype, + OnPairData::new(dict_bytes), + dict_offsets, + codes, + codes_offsets, + uncompressed_lengths, + validity, + ) + } + + /// Build an [`OnPairArray`] from already-materialised parts while reusing + /// an existing [`OnPairData`]. + /// + /// Reusing the data preserves the dictionary byte-buffer handle and the + /// shared lazy dictionary cache. This is useful when recursive compression + /// replaces the slot children without changing the logical dictionary. + /// + /// If `data` contains a memoized dictionary, `dict_offsets` must be + /// logically equivalent to the offsets used to create that dictionary. + /// The constructor intentionally does not validate the dictionary itself; + /// dictionary validation remains lazy and is performed on first use. + pub fn try_new_with_data( + dtype: DType, + data: OnPairData, + dict_offsets: ArrayRef, + codes: ArrayRef, + codes_offsets: ArrayRef, + uncompressed_lengths: ArrayRef, + validity: Validity, ) -> VortexResult { validate_parts( &dtype, @@ -234,36 +356,38 @@ impl OnPair { &codes, &codes_offsets, &uncompressed_lengths, - bits, )?; - let len = uncompressed_lengths.len(); - let data = OnPairData::new(dict_bytes, bits, len); - let slots = OnPairSlots { - dict_offsets, - codes, - codes_offsets, - uncompressed_lengths, - validity: validity_to_child(&validity, len), - } - .into_slots(); Ok(unsafe { - Array::from_parts_unchecked(ArrayParts::new(OnPair, dtype, len, data).with_slots(slots)) + Self::new_unchecked( + dtype, + data, + dict_offsets, + codes, + codes_offsets, + uncompressed_lengths, + validity, + ) }) } - #[expect(clippy::too_many_arguments, reason = "every child is a real input")] + /// Build an [`OnPairArray`] without validation, carrying `data` — and with + /// it the memoized dictionary — from an existing array. + /// + /// # Safety + /// The parts must satisfy the same invariants [`try_new`](Self::try_new) + /// checks. If `data`'s dictionary cell is populated (or shared with a live + /// array), `dict_offsets` must hold the same logical offsets the dictionary + /// was built from. pub(crate) unsafe fn new_unchecked( dtype: DType, - dict_bytes: BufferHandle, + data: OnPairData, dict_offsets: ArrayRef, codes: ArrayRef, codes_offsets: ArrayRef, uncompressed_lengths: ArrayRef, validity: Validity, - bits: u32, ) -> OnPairArray { let len = uncompressed_lengths.len(); - let data = OnPairData::new(dict_bytes, bits, len); let slots = OnPairSlots { dict_offsets, codes, @@ -284,13 +408,11 @@ fn validate_parts( codes: &ArrayRef, codes_offsets: &ArrayRef, uncompressed_lengths: &ArrayRef, - bits: u32, ) -> VortexResult<()> { vortex_ensure!( matches!(dtype, DType::Binary(_) | DType::Utf8(_)), "OnPair arrays must be Binary or Utf8, found {dtype}" ); - vortex_ensure!((9..=16).contains(&bits), "bits {bits} out of range [9, 16]"); if !dict_offsets.dtype().is_int() || dict_offsets.dtype().is_nullable() { vortex_bail!(InvalidArgument: "dict_offsets must be non-nullable integer"); @@ -326,7 +448,7 @@ impl VTable for OnPair { fn validate( &self, - data: &Self::TypedArrayData, + _data: &Self::TypedArrayData, dtype: &DType, len: usize, slots: &[Option], @@ -338,14 +460,10 @@ impl VTable for OnPair { s.codes, s.codes_offsets, s.uncompressed_lengths, - data.bits, )?; if s.uncompressed_lengths.len() != len { vortex_bail!(InvalidArgument: "uncompressed_lengths must have same len as outer array"); } - if data.len != len { - vortex_bail!(InvalidArgument: "OnPairData len {} != outer len {}", data.len, len); - } Ok(()) } @@ -379,6 +497,10 @@ impl VTable for OnPair { ); let mut data = array.data().clone(); data.dict_bytes = buffers[0].clone(); + // The replacement blob may differ from the one the memoized dictionary + // was validated against, so drop the (shared) cell rather than + // carry a claim we can no longer prove. + data.dictionary = Arc::new(OnceLock::new()); Ok( ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) .with_slots(array.slots().iter().cloned().collect()), @@ -391,13 +513,12 @@ impl VTable for OnPair { ) -> VortexResult>> { let dict_size = u32::try_from(array.dict_offsets().len().saturating_sub(1)) .map_err(|_| vortex_err!("OnPair dict_size exceeds u32"))?; - let total_tokens = array.codes().len() as u64; + let codes_len = array.codes().len() as u64; Ok(Some( OnPairMetadata { uncompressed_lengths_ptype: array.uncompressed_lengths().dtype().as_ptype().into(), - bits: array.bits(), dict_size, - total_tokens, + codes_len, dict_offsets_ptype: array.dict_offsets().dtype().as_ptype().into(), codes_ptype: array.codes().dtype().as_ptype().into(), codes_offsets_ptype: array.codes_offsets().dtype().as_ptype().into(), @@ -421,12 +542,11 @@ impl VTable for OnPair { let metadata = OnPairMetadata::decode(metadata)?; let uncompressed_ptype = metadata.get_uncompressed_lengths_ptype()?; - // Slot children. We pass `usize::MAX` for slots whose length we - // don't know up front (`dict_offsets` and `codes`). `codes_offsets` - // has known length `len + 1`. + // Slot children do not persist their own lengths, so metadata records + // the dictionary and code-stream sizes needed to deserialize them. let dict_offsets_len = metadata.dict_size as usize + 1; - let total_tokens = usize::try_from(metadata.total_tokens) - .map_err(|_| vortex_err!("total_tokens {} overflows usize", metadata.total_tokens))?; + let codes_len = usize::try_from(metadata.codes_len) + .map_err(|_| vortex_err!("codes_len {} overflows usize", metadata.codes_len))?; // The cascading compressor may have narrowed any of these integer // children to a tighter ptype; the recorded ptype tells the framework // exactly which dtype to materialise as. @@ -449,7 +569,7 @@ impl VTable for OnPair { let codes = children.get( 1, &DType::Primitive(codes_ptype, Nullability::NonNullable), - total_tokens, + codes_len, )?; let codes_offsets = children.get( 2, @@ -467,7 +587,7 @@ impl VTable for OnPair { other => vortex_bail!(InvalidArgument: "Expected 4 or 5 children, got {other}"), }; - let data = OnPairData::new(buffers[0].clone(), metadata.bits, len); + let data = OnPairData::new(buffers[0].clone()); let slots = OnPairSlots { dict_offsets, codes, @@ -492,25 +612,34 @@ impl VTable for OnPair { builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { + if let Some(result) = + match_each_varbin_builder!(builder, |builder| append_to_varbin(array, builder, ctx)) + { + return result; + } + + // The two arms here are every builder a `Utf8`/`Binary` dtype has: all four + // `VarBinBuilder` widths above, and `VarBinViewBuilder` below. There is deliberately no + // canonicalize-then-append fallback — it would decode to a `VarBinView` only for + // `VarBinView::append_to_builder` to reject the same remainder. let Some(builder) = builder.as_any_mut().downcast_mut::() else { - return array - .array() - .clone() - .execute::(ctx)? - .into_array() - .append_to_builder(builder, ctx); + vortex_bail!("append_to_builder for OnPair requires a variable-binary builder") }; - let next_buffer_index = builder.completed_block_count() + u32::from(builder.in_progress()); - let (buffers, views) = onpair_decode_views(array, next_buffer_index, ctx)?; - builder.push_buffer_and_adjusted_views( - &buffers, - &views, - array - .array() - .validity()? - .execute_mask(array.array().len(), ctx)?, - ); + // Decode the whole code stream into a new buffer, which the builder adopts as a data + // buffer with views built over it in place. + let validity = array + .array() + .validity()? + .execute_mask(array.array().len(), ctx)?; + let (out_bytes, lengths) = onpair_decode_bytes(array, ctx)?; + match_each_integer_ptype!(lengths.ptype(), |P| { + builder.append_buffer_with_lengths( + out_bytes.freeze(), + lengths.as_slice::

(), + &validity, + ) + }); Ok(()) } @@ -523,6 +652,41 @@ impl VTable for OnPair { } } +/// Decodes the code stream straight into `builder`'s byte storage. +/// +/// The offsets are the running sum of the uncompressed lengths the array already stores, so the +/// only work beyond the bulk `try_decode_into` is one prefix sum over them. +fn append_to_varbin( + array: ArrayView<'_, OnPair>, + builder: &mut VarBinBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> +where + usize: AsPrimitive, +{ + let plan = OnPairDecodePlan::new(array, ctx)?; + let validity = array + .array() + .validity()? + .execute_mask(array.array().len(), ctx)?; + // Built once, outside the ptype match: `append_decoded` takes it as `&mut dyn FnMut`, so + // creating the closure inside each arm would stamp out a shim per length type for no gain. + let mut decode = |out: &mut [MaybeUninit]| plan.decode_into(out); + match_each_integer_ptype!(plan.lengths.ptype(), |P| { + // SAFETY: `decode_into` initializes exactly the prefix whose length it returns. It needs + // no slack: it derives its bound from the slice it is handed and writes each value exactly. + unsafe { + builder.append_decoded( + plan.total_size, + 0, + plan.lengths.as_slice::

(), + &validity, + &mut decode, + ) + } + }) +} + impl ValidityVTable for OnPair { fn validity(array: ArrayView<'_, OnPair>) -> VortexResult { Ok(child_to_validity( @@ -534,6 +698,8 @@ impl ValidityVTable for OnPair { /// Convenience methods on top of the macro-generated [`OnPairArraySlotsExt`]. pub trait OnPairArrayExt: OnPairArraySlotsExt { + /// The array's [`Validity`], derived from the optional validity child and + /// the outer dtype's nullability. fn array_validity(&self) -> Validity { child_to_validity( self.as_ref().slots()[OnPairSlots::VALIDITY].as_ref(), diff --git a/encodings/onpair/src/canonical.rs b/encodings/onpair/src/canonical.rs new file mode 100644 index 00000000000..7ed5bfe6aa6 --- /dev/null +++ b/encodings/onpair/src/canonical.rs @@ -0,0 +1,152 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +// +//! Convert an [`OnPairArray`] to its canonical `VarBinViewArray` by handing +//! the materialised parts to `onpair::try_decode_into`. +//! +//! [`OnPairArray`]: crate::OnPairArray + +use std::mem::MaybeUninit; +use std::sync::Arc; + +use num_traits::AsPrimitive; +use onpair::CompactDictionaryView; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::varbinview::build_views::MAX_BUFFER_LEN; +use vortex_array::arrays::varbinview::build_views::build_views; +use vortex_array::match_each_integer_ptype; +use vortex_buffer::Buffer; +use vortex_buffer::ByteBufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; + +use crate::OnPair; +use crate::OnPairArraySlotsExt; +use crate::array::dict_view; +use crate::decode::code_boundary_at; +use crate::decode::collect_widened; + +pub(super) fn canonicalize_onpair( + array: ArrayView<'_, OnPair>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let (out_bytes, lengths) = onpair_decode_bytes(array, ctx)?; + let (buffers, views) = match_each_integer_ptype!(lengths.ptype(), |P| { + build_views( + 0, + MAX_BUFFER_LEN, + out_bytes.freeze(), + lengths.as_slice::

(), + ) + }); + let validity = array.array().validity()?; + Ok(unsafe { + VarBinViewArray::new_unchecked(views, Arc::from(buffers), array.dtype().clone(), validity) + .into_array() + }) +} + +pub(crate) struct OnPairDecodePlan<'a> { + codes: Buffer, + dict: CompactDictionaryView<'a>, + /// Per-row uncompressed lengths, zero for null rows. + pub(crate) lengths: PrimitiveArray, + /// Total decoded size, i.e. the sum of `lengths`. + pub(crate) total_size: usize, +} + +impl<'a> OnPairDecodePlan<'a> { + pub(crate) fn new(array: ArrayView<'a, OnPair>, ctx: &mut ExecutionCtx) -> VortexResult { + let lengths = array + .uncompressed_lengths() + .clone() + .execute::(ctx)?; + + let total_size: usize = match_each_integer_ptype!(lengths.ptype(), |P| { + lengths + .as_slice::

() + .iter() + .map(|&l| AsPrimitive::::as_(l)) + .sum() + }); + + // `codes_offsets` holds the per-row code boundaries and may itself be a + // sliced or filtered view of the original. Its first and last entries + // bound the contiguous run of `codes` belonging to the rows present in + // this array: `slice` keeps the full `codes` child and only narrows + // `codes_offsets` (so `code_start > 0` and/or `code_end < codes.len()`), + // while `filter` rebuilds both children so the window is the whole stream. + // OnPair has no `TakeExecute`, so a reordering take is served from the + // canonical `VarBinView` and never reaches this path. We only need those + // two boundaries, so point-look them up rather than decoding every offset. + let codes_offsets = array.codes_offsets(); + let code_start = code_boundary_at(codes_offsets, 0, ctx)?; + let code_end = code_boundary_at(codes_offsets, array.len(), ctx)?; + vortex_ensure!( + code_start <= code_end, + "OnPair codes_offsets must be nondecreasing" + ); + vortex_ensure!( + code_end <= array.codes().len(), + "OnPair codes_offsets end {} exceeds codes len {}", + code_end, + array.codes().len() + ); + + // Slice the `codes` child to that window *before* unpacking it, so a sliced + // array materialises only its own codes rather than the whole column's. The + // contiguous decoder walks `codes` in order and never reads the per-row + // boundaries, so an empty boundary slice is sound. + let codes = collect_widened::(&array.codes().slice(code_start..code_end)?, ctx)?; + let dict = dict_view(array, ctx)?; + + Ok(Self { + codes, + dict, + lengths, + total_size, + }) + } + + /// Bulk-decodes the whole code stream into `out`, which must hold at least `total_size` bytes. + /// + /// Do not reach for `#[inline(always)]` here. `try_decode_into` is monomorphized for + /// `CompactDictionaryView` whichever way this is annotated, and it is far too large to inline + /// either way — it stays an out-of-line call in both. All `always` buys is dropping this + /// wrapper's own frame, one `bl`/`ret` per decoded chunk against a whole-column decode, which + /// `benches/decode.rs::canonicalize_to_varbinview` cannot tell apart from noise. + #[inline] + pub(crate) fn decode_into(&self, out: &mut [MaybeUninit]) -> VortexResult { + let written = match onpair::try_decode_into(self.codes.as_slice(), self.dict, out) { + Ok(written) => written, + Err(_) => { + vortex_bail!("OnPair codes decode to more bytes than uncompressed_lengths records") + } + }; + + vortex_ensure!( + written == self.total_size, + "OnPair codes decoded to {written} bytes but uncompressed_lengths records {}", + self.total_size + ); + Ok(written) + } +} + +pub(crate) fn onpair_decode_bytes( + array: ArrayView<'_, OnPair>, + ctx: &mut ExecutionCtx, +) -> VortexResult<(ByteBufferMut, PrimitiveArray)> { + let plan = OnPairDecodePlan::new(array, ctx)?; + let mut out_bytes = ByteBufferMut::with_capacity(plan.total_size); + let written = plan.decode_into(out_bytes.spare_capacity_mut())?; + // SAFETY: `decode_into` initialised exactly `written` bytes. + unsafe { out_bytes.set_len(written) }; + Ok((out_bytes, plan.lengths)) +} diff --git a/encodings/onpair/src/compress.rs b/encodings/onpair/src/compress.rs new file mode 100644 index 00000000000..b46428445d8 --- /dev/null +++ b/encodings/onpair/src/compress.rs @@ -0,0 +1,171 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Train + compress entry points for the OnPair encoding. + +use onpair::Config; +use onpair::Rows; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::varbinview::BinaryView; +use vortex_array::buffer::BufferHandle; +use vortex_array::scalar::Scalar; +use vortex_buffer::Alignment; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_buffer::ByteBuffer; +use vortex_buffer::ByteBufferMut; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_mask::AllOr; + +use crate::OnPair; +use crate::OnPairData; + +/// Compress any [`ArrayRef`] whose canonical form is a string array. +/// +/// All-null inputs are returned as a [`ConstantArray`]. +pub fn onpair_compress( + array: &ArrayRef, + config: Config, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let array = array.clone().execute::(ctx)?; + let len = array.len(); + let validity = array.validity()?; + let mask = validity.execute_mask(len, ctx)?; + if matches!(mask.bit_buffer(), AllOr::None) { + // CascadingCompressor handles this earlier, but direct callers can reach it. + return Ok(ConstantArray::new(Scalar::null(array.dtype().clone()), len).into_array()); + } + + let views = array.views(); + let mut uncompressed_lengths: BufferMut = BufferMut::zeroed(len); + let buffers = array + .data_buffers() + .as_ref() + .iter() + .map(|b| b.as_host()) + .collect::>(); + + // Keep sums local to each arm: a shared callback accumulator can prevent + // vectorization of the all-valid loop. + let total_bytes = match mask.bit_buffer() { + AllOr::All => { + let mut total = 0; + for (view, length) in views.iter().zip(uncompressed_lengths.iter_mut()) { + *length = view.len(); + total += *length as usize; + } + total + } + AllOr::None => unreachable!("all-null input handled above"), + AllOr::Some(validity) => { + let lengths = uncompressed_lengths.as_mut_slice(); + let mut total = 0; + validity.for_each_set_index(|i| { + let length = views[i].len(); + lengths[i] = length; + total += length as usize; + }); + total + } + }; + + let rows = ViewRows { + views, + buffers: &buffers, + lengths: uncompressed_lengths.as_slice(), + total_bytes, + }; + let column = onpair::compress_rows::<_, u64>(&rows, config); + let (dict, codes, row_offsets) = column.into_raw(); + let (dict_bytes, dict_offsets) = dict.into_raw(); + let codes_offsets = codes_offsets_array(&row_offsets); + let codes = Buffer::from(codes).into_array(); + // The `dict_offsets` child and the memoized widened-offsets cell share + // this buffer, so seeding below costs no copy. + let dict_offsets = Buffer::from(dict_offsets); + + let uncompressed_lengths = uncompressed_lengths.into_array(); + + let data = OnPairData::try_new_with_dictionary( + dict_bytes_to_buffer(dict_bytes), + dict_offsets.clone(), + )?; + let encoded = OnPair::try_new_with_data( + array.dtype().clone(), + data, + dict_offsets.into_array(), + codes, + codes_offsets, + uncompressed_lengths, + validity, + )?; + Ok(encoded.into_array()) +} + +/// Reads inline and external values in place, treating null rows as empty. +struct ViewRows<'a> { + views: &'a [BinaryView], + buffers: &'a [&'a ByteBuffer], + lengths: &'a [u32], + total_bytes: usize, +} + +impl Rows for ViewRows<'_> { + fn num_rows(&self) -> usize { + self.views.len() + } + + fn total_bytes(&self) -> usize { + self.total_bytes + } + + #[inline] + fn row(&self, i: usize) -> &[u8] { + let view = &self.views[i]; + if self.lengths[i] == 0 { + &[] + } else if view.is_inlined() { + view.as_inlined().value() + } else { + let view_ref = view.as_view(); + &self.buffers[view_ref.buffer_index as usize][view_ref.as_range()] + } + } +} + +fn dict_bytes_to_buffer(dict_bytes: Vec) -> BufferHandle { + // Align dict_bytes to 8 bytes so the segment that ultimately holds the + // OnPair tree starts at an 8-aligned in-memory address. Without this anchor, + // downstream primitive children may deserialize from a misaligned segment. + let mut aligned = ByteBufferMut::with_capacity_aligned(dict_bytes.len(), Alignment::new(8)); + aligned.extend_from_slice(&dict_bytes); + BufferHandle::new_host(aligned.freeze()) +} + +/// Build the `codes_offsets` child from the library's per-row code boundaries, +/// storing the narrowest of `u32`/`u64` that holds the largest boundary. +/// `row_offsets` is non-decreasing, so its last entry is that maximum and one +/// bound check picks the width. `u32` covers the common case (the cascading +/// compressor narrows it further to `u16`/`u8`); `u64` engages only when a +/// single chunk carries more than `u32::MAX` tokens, matching the `u64` byte +/// offsets accepted at compression. +fn codes_offsets_array(row_offsets: &[u64]) -> ArrayRef { + let total_tokens = row_offsets.last().copied().unwrap_or(0); + if u32::try_from(total_tokens).is_ok() { + Buffer::from( + row_offsets + .iter() + .map(|&o| u32::try_from(o).vortex_expect("code boundary fits u32")) + .collect::>(), + ) + .into_array() + } else { + Buffer::from(row_offsets.to_vec()).into_array() + } +} diff --git a/encodings/experimental/onpair/src/compute/byte_length.rs b/encodings/onpair/src/compute/byte_length.rs similarity index 100% rename from encodings/experimental/onpair/src/compute/byte_length.rs rename to encodings/onpair/src/compute/byte_length.rs diff --git a/encodings/onpair/src/compute/cast.rs b/encodings/onpair/src/compute/cast.rs new file mode 100644 index 00000000000..a746d3dcc07 --- /dev/null +++ b/encodings/onpair/src/compute/cast.rs @@ -0,0 +1,114 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::dtype::DType; +use vortex_array::scalar_fn::fns::cast::CastKernel; +use vortex_array::scalar_fn::fns::cast::CastReduce; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; + +use crate::OnPair; +use crate::OnPairArraySlotsExt; + +/// Adjust nullability without touching any encoded payload — we only rewrap +/// into a new outer DType. +fn build_with_validity( + array: ArrayView<'_, OnPair>, + dtype: &DType, + validity: Validity, +) -> ArrayRef { + unsafe { + OnPair::new_unchecked( + dtype.clone(), + array.data().clone(), + array.dict_offsets().clone(), + array.codes().clone(), + array.codes_offsets().clone(), + array.uncompressed_lengths().clone(), + validity, + ) + } + .into_array() +} + +impl CastReduce for OnPair { + fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { + if !array.dtype().eq_ignore_nullability(dtype) { + return Ok(None); + } + let validity = array.array().validity()?; + let Some(new_validity) = + validity.trivially_cast_nullability(dtype.nullability(), array.array().len())? + else { + return Ok(None); + }; + Ok(Some(build_with_validity(array, dtype, new_validity))) + } +} + +impl CastKernel for OnPair { + fn cast( + array: ArrayView<'_, Self>, + dtype: &DType, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if !array.dtype().eq_ignore_nullability(dtype) { + return Ok(None); + } + let new_validity = array.array().validity()?.cast_nullability( + dtype.nullability(), + array.array().len(), + ctx, + )?; + Ok(Some(build_with_validity(array, dtype, new_validity))) + } +} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::VarBinArray; + use vortex_array::compute::conformance::cast::test_cast_conformance; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use crate::DEFAULT_CONFIG; + use crate::compress::onpair_compress; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + + #[rstest] + #[case(VarBinArray::from_iter( + vec![Some("hello"), Some("world"), Some("hello world")], + DType::Utf8(Nullability::NonNullable) + ))] + #[case(VarBinArray::from_iter( + vec![Some("foo"), None, Some("bar"), Some("foobar")], + DType::Utf8(Nullability::Nullable) + ))] + #[case(VarBinArray::from_iter( + vec![Some("test")], + DType::Utf8(Nullability::NonNullable) + ))] + fn test_cast_onpair_conformance(#[case] array: VarBinArray) -> VortexResult<()> { + let array = array.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let onpair = onpair_compress(&array, DEFAULT_CONFIG, &mut ctx)?; + test_cast_conformance(&onpair.into_array(), &mut ctx); + Ok(()) + } +} diff --git a/encodings/onpair/src/compute/compare.rs b/encodings/onpair/src/compute/compare.rs new file mode 100644 index 00000000000..3ed449731d9 --- /dev/null +++ b/encodings/onpair/src/compute/compare.rs @@ -0,0 +1,279 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use onpair::search; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::fns::binary::CompareKernel; +use vortex_array::scalar_fn::fns::operators::CompareOperator; +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; + +use crate::OnPair; +use crate::OnPairArraySlotsExt; +use crate::array::dict_view; +use crate::decode::collect_codes_window; + +impl CompareKernel for OnPair { + fn compare( + lhs: ArrayView<'_, Self>, + rhs: &ArrayRef, + operator: CompareOperator, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some(constant) = rhs.as_constant() else { + return Ok(None); + }; + let needle = match constant.dtype() { + DType::Utf8(_) => constant + .as_utf8() + .value() + .map(|value| value.as_bytes().to_vec()), + DType::Binary(_) => constant + .as_binary() + .value() + .map(|value| value.as_slice().to_vec()), + _ => return Ok(None), + }; + let Some(needle) = needle else { + return Ok(None); + }; + + let buffer = if needle.is_empty() { + let lengths = lhs.uncompressed_lengths(); + match operator { + // every value is greater than an empty string + CompareOperator::Gte => BitBuffer::new_set(lhs.len()), + // no value is less than an empty string + CompareOperator::Lt => BitBuffer::new_unset(lhs.len()), + _ => lengths + .binary( + ConstantArray::new(Scalar::zero_value(lengths.dtype()), lengths.len()) + .into_array(), + operator.into(), + )? + .execute(ctx)?, + } + } else { + if !matches!(operator, CompareOperator::Eq | CompareOperator::NotEq) { + return Ok(None); + } + + let dict = dict_view(lhs, ctx)?; + let query = search::tokenize(&needle, dict); + let window = collect_codes_window(lhs, ctx)?; + + let negated = operator == CompareOperator::NotEq; + BitBuffer::collect_bool(lhs.len(), |i| { + search::equals(window.row(i), &query) != negated + }) + }; + + Ok(Some( + BoolArray::new( + buffer, + lhs.validity()? + .union_nullability(constant.dtype().nullability()), + ) + .into_array(), + )) + } +} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::VarBinArray; + use vortex_array::assert_arrays_eq; + use vortex_array::builtins::ArrayBuiltins; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::scalar::Scalar; + use vortex_array::scalar_fn::fns::binary::CompareKernel; + use vortex_array::scalar_fn::fns::operators::CompareOperator; + use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + use vortex_session::VortexSession; + + use crate::DEFAULT_CONFIG; + use crate::OnPair; + use crate::compress::onpair_compress; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + + #[cfg_attr(miri, ignore)] + #[rstest] + #[case(Operator::Eq, [true, false, true, false])] + #[case(Operator::NotEq, [false, true, false, true])] + #[case(Operator::Gt, [false, true, false, true])] + #[case(Operator::Gte, [true, true, true, true])] + #[case(Operator::Lt, [false, false, false, false])] + #[case(Operator::Lte, [true, false, true, false])] + fn compare_empty_string(#[case] op: Operator, #[case] expected: [bool; 4]) -> VortexResult<()> { + let input = VarBinArray::from_iter( + [Some(""), Some("a"), Some(""), Some("bbb")], + DType::Utf8(Nullability::NonNullable), + ); + let mut ctx = SESSION.create_execution_ctx(); + let arr = onpair_compress(input.as_array(), DEFAULT_CONFIG, &mut ctx)?.into_array(); + + let result = arr + .binary(ConstantArray::new("", input.len()).into_array(), op)? + .execute::(&mut ctx)?; + assert_arrays_eq!(&result, &BoolArray::from_iter(expected), &mut ctx); + Ok(()) + } + + #[cfg_attr(miri, ignore)] + #[test] + fn compare_empty_string_nullable() -> VortexResult<()> { + let input = VarBinArray::from_iter( + [Some(""), None, Some("x")], + DType::Utf8(Nullability::Nullable), + ); + let mut ctx = SESSION.create_execution_ctx(); + let arr = onpair_compress(input.as_array(), DEFAULT_CONFIG, &mut ctx)?.into_array(); + + let eq_empty = arr + .clone() + .binary(ConstantArray::new("", arr.len()).into_array(), Operator::Eq)? + .execute::(&mut ctx)?; + assert_arrays_eq!( + &eq_empty, + &BoolArray::from_iter([Some(true), None, Some(false)]), + &mut ctx + ); + + let null_rhs = + ConstantArray::new(Scalar::null(DType::Utf8(Nullability::Nullable)), arr.len()); + let eq_null = arr + .binary(null_rhs.into_array(), Operator::Eq)? + .execute::(&mut ctx)?; + assert_arrays_eq!( + &eq_null, + &BoolArray::from_iter([None::, None, None]), + &mut ctx + ); + Ok(()) + } + + #[cfg_attr(miri, ignore)] + #[test] + fn compare_nonempty_string_nullable() -> VortexResult<()> { + let input = VarBinArray::from_iter( + [Some("hello"), None, Some("world"), Some("hello")], + DType::Utf8(Nullability::Nullable), + ); + let mut ctx = SESSION.create_execution_ctx(); + let arr = onpair_compress(input.as_array(), DEFAULT_CONFIG, &mut ctx)?.into_array(); + let rhs = ConstantArray::new("hello", arr.len()).into_array(); + + let eq = arr + .binary(rhs.clone(), Operator::Eq)? + .execute::(&mut ctx)?; + assert_arrays_eq!( + &eq, + &BoolArray::from_iter([Some(true), None, Some(false), Some(true)]), + &mut ctx + ); + + let neq = arr + .binary(rhs, Operator::NotEq)? + .execute::(&mut ctx)?; + assert_arrays_eq!( + &neq, + &BoolArray::from_iter([Some(false), None, Some(true), Some(false)]), + &mut ctx + ); + Ok(()) + } + + /// Call `CompareKernel::compare` directly and verify it returns `Some` + /// (i.e. the kernel handles the constant needle rather than silently + /// falling back to canonical decompression), and that ordering against a + /// non-empty needle declines — code sequences are not order-preserving. + #[cfg_attr(miri, ignore)] + #[test] + fn compare_kernel_handles_constant_needle() -> VortexResult<()> { + let input = VarBinArray::from_iter( + [Some("hello"), Some("world"), Some("hello"), Some("he")], + DType::Utf8(Nullability::NonNullable), + ); + let mut ctx = SESSION.create_execution_ctx(); + let arr = onpair_compress(input.as_array(), DEFAULT_CONFIG, &mut ctx)? + .try_downcast::() + .map_err(|array| vortex_err!("expected OnPair array, got {}", array.encoding_id()))?; + let rhs = ConstantArray::new("hello", arr.len()).into_array(); + + let eq = + ::compare(arr.as_view(), &rhs, CompareOperator::Eq, &mut ctx)? + .expect("OnPair CompareKernel should handle a constant needle"); + assert_arrays_eq!( + eq, + BoolArray::from_iter([true, false, true, false]), + &mut ctx + ); + + let lt = + ::compare(arr.as_view(), &rhs, CompareOperator::Lt, &mut ctx)?; + assert!( + lt.is_none(), + "ordering compare against a non-empty needle must fall back to canonical" + ); + Ok(()) + } + + /// The equality path resolves row windows relative to `codes_offsets[0]`, + /// which is nonzero for a sliced array. Exercise the kernel directly on a + /// slice so the windowed row arithmetic is covered. + #[cfg_attr(miri, ignore)] + #[test] + fn compare_kernel_on_sliced_array() -> VortexResult<()> { + let input = VarBinArray::from_iter( + [ + Some("aardvark"), + Some("hello"), + Some("world"), + Some("hello"), + Some("zebra"), + ], + DType::Utf8(Nullability::NonNullable), + ); + let mut ctx = SESSION.create_execution_ctx(); + let arr = onpair_compress(input.as_array(), DEFAULT_CONFIG, &mut ctx)?.into_array(); + let sliced = arr.slice(1..4)?; + assert!(sliced.is::(), "slice dropped OnPair encoding"); + let sliced = sliced + .try_downcast::() + .map_err(|_| vortex_err!("sliced array was not OnPair"))?; + + let rhs = ConstantArray::new("hello", sliced.len()).into_array(); + let eq = ::compare( + sliced.as_view(), + &rhs, + CompareOperator::Eq, + &mut ctx, + )? + .expect("OnPair CompareKernel should handle a constant needle"); + assert_arrays_eq!(eq, BoolArray::from_iter([true, false, true]), &mut ctx); + Ok(()) + } +} diff --git a/encodings/experimental/onpair/src/compute/filter.rs b/encodings/onpair/src/compute/filter.rs similarity index 95% rename from encodings/experimental/onpair/src/compute/filter.rs rename to encodings/onpair/src/compute/filter.rs index c26f3eeeacd..219a5a72907 100644 --- a/encodings/experimental/onpair/src/compute/filter.rs +++ b/encodings/onpair/src/compute/filter.rs @@ -20,7 +20,7 @@ use vortex_array::IntoArray; use vortex_array::arrays::List; use vortex_array::arrays::ListArray; use vortex_array::arrays::filter::FilterKernel; -use vortex_array::arrays::list::ListArrayExt; +use vortex_array::arrays::list::ListArraySlotsExt; use vortex_array::validity::Validity; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -63,13 +63,12 @@ impl FilterKernel for OnPair { unsafe { OnPair::new_unchecked( array.dtype().clone(), - array.dict_bytes_handle().clone(), + array.data().clone(), array.dict_offsets().clone(), filtered_codes.elements().clone(), filtered_codes.offsets().clone(), uncompressed_lengths, validity, - array.bits(), ) } .into_array(), diff --git a/encodings/experimental/onpair/src/compute/mod.rs b/encodings/onpair/src/compute/mod.rs similarity index 100% rename from encodings/experimental/onpair/src/compute/mod.rs rename to encodings/onpair/src/compute/mod.rs diff --git a/encodings/experimental/onpair/src/compute/slice.rs b/encodings/onpair/src/compute/slice.rs similarity index 94% rename from encodings/experimental/onpair/src/compute/slice.rs rename to encodings/onpair/src/compute/slice.rs index fcfebf413bf..49d840f2029 100644 --- a/encodings/experimental/onpair/src/compute/slice.rs +++ b/encodings/onpair/src/compute/slice.rs @@ -28,13 +28,12 @@ impl SliceReduce for OnPair { unsafe { OnPair::new_unchecked( array.dtype().clone(), - array.dict_bytes_handle().clone(), + array.data().clone(), array.dict_offsets().clone(), array.codes().clone(), codes_offsets, uncompressed_lengths, validity, - array.bits(), ) } .into_array(), diff --git a/encodings/onpair/src/decode.rs b/encodings/onpair/src/decode.rs new file mode 100644 index 00000000000..7aac8c3963f --- /dev/null +++ b/encodings/onpair/src/decode.rs @@ -0,0 +1,109 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +// +//! Helpers for turning [`OnPair`] slot children into the inputs the upstream +//! `onpair` decoder consumes. + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_buffer::Buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; + +use crate::OnPair; +use crate::OnPairArraySlotsExt; + +/// Canonicalise a slot child to the decoder's native primitive width. +pub(crate) fn collect_widened( + arr: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let dtype = DType::Primitive(T::PTYPE, arr.dtype().nullability()); + Ok(arr + .cast(dtype)? + .execute::(ctx)? + .into_buffer::()) +} + +pub(crate) fn code_boundary_at( + codes_offsets: &ArrayRef, + index: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { + codes_offsets + .execute_scalar(index, ctx)? + .as_primitive() + .as_::() + .ok_or_else(|| vortex_err!("OnPair codes_offsets[{index}] is null")) +} + +/// A validated, materialised window over an array's `codes`: the widened +/// per-row `codes_offsets` boundaries plus the codes they bound. +/// +/// `slice` keeps the full `codes` child and only narrows `codes_offsets`, so +/// for a sliced array the window starts at `offsets[0] > 0`; [`row`] resolves +/// row indices relative to that start. Built once per query by the +/// compressed-domain compare kernel. +/// +/// [`row`]: CodesWindow::row +pub(crate) struct CodesWindow { + offsets: Buffer, + codes: Buffer, + code_start: usize, +} + +impl CodesWindow { + /// The codes for row `i`. + pub(crate) fn row(&self, i: usize) -> &[u16] { + &self.codes[self.local(i)..self.local(i + 1)] + } + + /// Offset `i` rebased into the window's local `codes` slice. Offsets are + /// bounded by `codes.len()` (a `usize`), so the conversion never truncates. + fn local(&self, i: usize) -> usize { + usize::try_from(self.offsets[i]).vortex_expect("code offset fits usize") - self.code_start + } +} + +/// Materialise the [`CodesWindow`] for every row of `array`: offsets must be +/// nondecreasing and end within the `codes` child. Codes themselves are +/// trusted to the upstream decoder/search primitives, which bounds-check them +/// in-loop and panic on a malformed value. +pub(crate) fn collect_codes_window( + array: ArrayView<'_, OnPair>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = array.len(); + let offsets = collect_widened::(array.codes_offsets(), ctx)?; + vortex_ensure!( + offsets.len() == len + 1, + "OnPair codes_offsets has {} entries, expected len + 1 = {}", + offsets.len(), + len + 1 + ); + vortex_ensure!( + offsets.is_sorted(), + "OnPair codes_offsets must be nondecreasing" + ); + let code_start = usize::try_from(offsets[0]).vortex_expect("code offset fits usize"); + let code_end = usize::try_from(offsets[len]).vortex_expect("code offset fits usize"); + vortex_ensure!( + code_end <= array.codes().len(), + "OnPair codes_offsets end {} exceeds codes len {}", + code_end, + array.codes().len() + ); + let codes = collect_widened::(&array.codes().slice(code_start..code_end)?, ctx)?; + Ok(CodesWindow { + offsets, + codes, + code_start, + }) +} diff --git a/encodings/experimental/onpair/src/kernel.rs b/encodings/onpair/src/kernel.rs similarity index 84% rename from encodings/experimental/onpair/src/kernel.rs rename to encodings/onpair/src/kernel.rs index 8863d750a72..e6ae93248b5 100644 --- a/encodings/experimental/onpair/src/kernel.rs +++ b/encodings/onpair/src/kernel.rs @@ -10,13 +10,15 @@ use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; use vortex_array::scalar_fn::fns::byte_length::ByteLength; use vortex_array::scalar_fn::fns::byte_length::ByteLengthExecuteAdaptor; +use vortex_array::scalar_fn::fns::cast::Cast; +use vortex_array::scalar_fn::fns::cast::CastExecuteAdaptor; use vortex_session::VortexSession; use crate::OnPair; -// TODO: implement ListExecute & TakeExecute for OnPair pub(super) fn initialize(session: &VortexSession) { let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Cast.id(), OnPair, CastExecuteAdaptor(OnPair)); kernels.register_execute_parent_kernel(Filter.id(), OnPair, FilterExecuteAdaptor(OnPair)); kernels.register_execute_parent_kernel(Binary.id(), OnPair, CompareExecuteAdaptor(OnPair)); kernels.register_execute_parent_kernel( diff --git a/encodings/experimental/onpair/src/lib.rs b/encodings/onpair/src/lib.rs similarity index 66% rename from encodings/experimental/onpair/src/lib.rs rename to encodings/onpair/src/lib.rs index 11b22f63bc1..f9d8be5ff0d 100644 --- a/encodings/experimental/onpair/src/lib.rs +++ b/encodings/onpair/src/lib.rs @@ -2,11 +2,10 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors //! Vortex string array backed by the [OnPair][onpair] short-string -//! compression library, with `cast` and `filter` pushdown. +//! compression library, with pushdown for common string operations. //! -//! The default training preset is `dict-12` (12 bits per token, dictionary -//! capped at 4 096 entries). See [`onpair_compress`] for the entry point and -//! [`OnPairArray`] for the resulting array type. +//! See [`onpair_compress`] for the compression entry point and [`OnPairArray`] +//! for the encoded representation of non-null inputs. //! //! [onpair]: https://arxiv.org/abs/2508.02280 @@ -23,9 +22,13 @@ mod tests; pub use array::*; pub use compress::*; -pub use onpair::Bits; +pub use onpair::CompactDictionaryView; pub use onpair::Config; +pub use onpair::DEFAULT_CONFIG; +pub use onpair::DictionaryView; pub use onpair::Error as OnPairError; +pub use onpair::MAX_TOKEN_SIZE; +pub use onpair::MaxDictBits; pub use onpair::Threshold; use vortex_array::session::ArraySessionExt; use vortex_session::VortexSession; diff --git a/encodings/experimental/onpair/src/ops.rs b/encodings/onpair/src/ops.rs similarity index 73% rename from encodings/experimental/onpair/src/ops.rs rename to encodings/onpair/src/ops.rs index a6e097bbfd4..728e5a0e6f2 100644 --- a/encodings/experimental/onpair/src/ops.rs +++ b/encodings/onpair/src/ops.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use onpair::Parts; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::arrays::varbin::varbin_scalar; @@ -10,9 +9,11 @@ use vortex_array::vtable::OperationsVTable; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; use vortex_error::vortex_err; +use vortex_error::vortex_panic; use crate::OnPair; use crate::OnPairArraySlotsExt; +use crate::array::dict_view; use crate::decode::code_boundary_at; use crate::decode::collect_widened; @@ -32,13 +33,7 @@ impl OperationsVTable for OnPair { let row_end = code_boundary_at(codes_offsets, index + 1, ctx)?; let codes = collect_widened::(&array.codes().slice(row_start..row_end)?, ctx)?; - let dict_offsets = collect_widened::(array.dict_offsets(), ctx)?; - let parts = Parts { - dict_bytes: array.dict_bytes().as_slice(), - dict_offsets: dict_offsets.as_slice(), - bits: array.bits(), - codes: codes.as_slice(), - }; + let dict = dict_view(array, ctx)?; // The per-row decoded length is recorded in the `uncompressed_lengths` // child, so read it directly instead of asking the decoder to compute it. @@ -49,10 +44,19 @@ impl OperationsVTable for OnPair { .as_::() .ok_or_else(|| vortex_err!("OnPair uncompressed_lengths[{index}] is null"))?; let mut buf: Vec = Vec::with_capacity(len); - let written = onpair::decompress_into(parts, buf.spare_capacity_mut()); - debug_assert_eq!(written, len); - // SAFETY: `decompress_into` initialised `written` bytes of the spare - // capacity reserved above. + let written = + match onpair::try_decode_into(codes.as_slice(), dict, buf.spare_capacity_mut()) { + Ok(written) => written, + Err(_) => vortex_panic!( + "OnPair row {index} decodes to more bytes than uncompressed_lengths records" + ), + }; + if written != len { + vortex_panic!( + "OnPair row {index} decoded to {written} bytes but uncompressed_lengths records {len}" + ); + } + // SAFETY: `try_decode_into` initialised exactly `written` bytes. unsafe { buf.set_len(written) }; Ok(varbin_scalar(ByteBuffer::from(buf), array.dtype())) } diff --git a/encodings/experimental/onpair/src/rules.rs b/encodings/onpair/src/rules.rs similarity index 100% rename from encodings/experimental/onpair/src/rules.rs rename to encodings/onpair/src/rules.rs diff --git a/encodings/experimental/onpair/src/tests.rs b/encodings/onpair/src/tests.rs similarity index 70% rename from encodings/experimental/onpair/src/tests.rs rename to encodings/onpair/src/tests.rs index 4f2124025f2..7dac64befc2 100644 --- a/encodings/experimental/onpair/src/tests.rs +++ b/encodings/onpair/src/tests.rs @@ -3,31 +3,51 @@ use std::sync::LazyLock; +use onpair::CompactDictionaryView; use prost::Message; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::filter::FilterKernel; +use vortex_array::assert_arrays_eq; +use vortex_array::buffer::BufferHandle; +use vortex_array::builders::VarBinBuilder; +use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::match_each_integer_ptype; +use vortex_array::scalar_fn::fns::operators::Operator; use vortex_array::test_harness::check_metadata; use vortex_array::validity::Validity; use vortex_buffer::BufferMut; use vortex_session::VortexSession; +use crate::DEFAULT_CONFIG; use crate::OnPair; use crate::OnPairArrayExt; use crate::OnPairArraySlotsExt; use crate::OnPairMetadata; -use crate::compress::DEFAULT_DICT12_CONFIG; use crate::compress::onpair_compress; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +fn compress_onpair( + array: &vortex_array::ArrayRef, + ctx: &mut vortex_array::ExecutionCtx, +) -> vortex_error::VortexResult { + onpair_compress(array, DEFAULT_CONFIG, ctx)? + .try_downcast::() + .map_err(|array| { + vortex_error::vortex_err!("expected OnPair array, got {}", array.encoding_id()) + }) +} + fn sample_input() -> VarBinArray { VarBinArray::from_iter( [ @@ -41,6 +61,85 @@ fn sample_input() -> VarBinArray { ) } +#[test] +fn test_direct_offset_builder() -> vortex_error::VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = sample_input(); + let encoded = compress_onpair(input.as_ref(), &mut ctx)?; + let mut builder = VarBinBuilder::::with_capacity_in( + input.dtype().clone(), + input.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + encoded + .into_array() + .append_to_builder(&mut builder, &mut ctx)?; + assert_arrays_eq!(builder.finish_into_varbin(), input, &mut ctx); + Ok(()) +} + +#[test] +fn test_onpair_rejects_100k_token_dictionary() -> vortex_error::VortexResult<()> { + let num_tokens = 100_000usize; + let mut tokens = Vec::with_capacity(num_tokens); + tokens.extend((u8::MIN..=u8::MAX).map(|byte| vec![byte])); + let additional_tokens = u32::try_from(num_tokens - tokens.len())?; + tokens.extend((0..additional_tokens).map(|value| { + let bytes = value.to_be_bytes(); + bytes[1..].to_vec() + })); + tokens.sort_unstable(); + + let mut dict_bytes = Vec::new(); + let mut dict_offsets = Vec::with_capacity(num_tokens + 1); + dict_offsets.push(0u32); + for token in tokens { + dict_bytes.extend_from_slice(&token); + dict_offsets.push(u32::try_from(dict_bytes.len())?); + } + dict_bytes.resize(dict_bytes.len() + onpair::MAX_TOKEN_SIZE, 0); + + assert!(CompactDictionaryView::validate_safety(&dict_bytes, &dict_offsets).is_err()); + Ok(()) +} + +/// Dictionary safety is validated lazily — construction stays lightweight, and +/// a structurally corrupt dictionary is rejected by the first operation that +/// decodes or searches through it, including on derived (sliced) arrays. +#[cfg_attr(miri, ignore)] +#[test] +fn test_corrupt_dictionary_rejected_on_first_use() -> vortex_error::VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let arr = compress_onpair(&sample_input().into_array(), &mut ctx)?; + let view = arr.as_view(); + + // Chop one byte off the blob so the trailing read-padding invariant fails. + let truncated = view.dict_bytes().slice(0..view.dict_bytes().len() - 1); + let corrupt = OnPair::try_new( + view.dtype().clone(), + BufferHandle::new_host(truncated), + view.dict_offsets().clone(), + view.codes().clone(), + view.codes_offsets().clone(), + view.uncompressed_lengths().clone(), + view.array_validity(), + )?; + + let sliced = corrupt.clone().into_array().slice(1..3)?; + // Canonical decode, random access, and ops on a derived slice all pass + // through the same validation door. + assert!( + corrupt + .clone() + .into_array() + .execute::(&mut ctx) + .is_err() + ); + assert!(corrupt.into_array().execute_scalar(0, &mut ctx).is_err()); + assert!(sliced.execute::(&mut ctx).is_err()); + Ok(()) +} + #[cfg_attr(miri, ignore)] #[test] fn test_onpair_metadata_golden() { @@ -48,9 +147,8 @@ fn test_onpair_metadata_golden() { "onpair.metadata", &OnPairMetadata { uncompressed_lengths_ptype: PType::I32 as i32, - bits: 12, dict_size: 4096, - total_tokens: 128_000, + codes_len: 128_000, dict_offsets_ptype: PType::U32 as i32, codes_ptype: PType::U16 as i32, codes_offsets_ptype: PType::U32 as i32, @@ -65,7 +163,7 @@ fn test_onpair_roundtrip() -> vortex_error::VortexResult<()> { let input = sample_input(); let mut ctx = SESSION.create_execution_ctx(); - let compressed = onpair_compress(&input.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; + let compressed = onpair_compress(&input.into_array(), DEFAULT_CONFIG, &mut ctx)?; assert!(compressed.clone().into_array().is::()); let decoded = compressed @@ -88,6 +186,63 @@ fn test_onpair_roundtrip() -> vortex_error::VortexResult<()> { Ok(()) } +/// The `u64` `codes_offsets` branch only engages past `u32::MAX` tokens (a +/// multi-GiB chunk), so exercise the read path by widening a small array's +/// `codes_offsets` child to `u64` by hand and asserting the decode paths +/// (canonical and the compressed-domain equality compare, which builds a +/// `CodesWindow`) treat it identically to the default `u32` width. +#[cfg_attr(miri, ignore)] +#[test] +fn test_onpair_u64_codes_offsets() -> vortex_error::VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let narrow = compress_onpair(&sample_input().into_array(), &mut ctx)?; + + // Rebuild with only codes_offsets widened to u64; every other child is the + // input's, so the two arrays differ solely in codes_offsets width. + let wide = { + let view = narrow.as_view(); + let wide_offsets = view + .codes_offsets() + .cast(DType::Primitive(PType::U64, Nullability::NonNullable))? + .execute::(&mut ctx)? + .into_array(); + OnPair::try_new( + view.dtype().clone(), + view.dict_bytes_handle().clone(), + view.dict_offsets().clone(), + view.codes().clone(), + wide_offsets, + view.uncompressed_lengths().clone(), + view.array_validity(), + )? + }; + assert_eq!( + wide.as_view().codes_offsets().dtype(), + &DType::Primitive(PType::U64, Nullability::NonNullable) + ); + + // Canonical decode is byte-identical across the two widths. + let narrow_decoded = narrow.into_array().execute::(&mut ctx)?; + let wide_decoded = wide + .clone() + .into_array() + .execute::(&mut ctx)?; + assert_arrays_eq!(&wide_decoded, &narrow_decoded, &mut ctx); + + // Equality compare drives CodesWindow over the u64 codes_offsets. + let needle = ConstantArray::new("https://www.example.com/page", wide.len()).into_array(); + let eq = wide + .into_array() + .binary(needle, Operator::Eq)? + .execute::(&mut ctx)?; + assert_arrays_eq!( + &eq, + &BoolArray::from_iter([true, false, false, false, true]), + &mut ctx + ); + Ok(()) +} + #[cfg_attr(miri, ignore)] #[test] fn test_onpair_nullable_canonicalize() -> vortex_error::VortexResult<()> { @@ -96,7 +251,7 @@ fn test_onpair_nullable_canonicalize() -> vortex_error::VortexResult<()> { DType::Utf8(Nullability::Nullable), ); let mut ctx = SESSION.create_execution_ctx(); - let arr = onpair_compress(&input.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; + let arr = onpair_compress(&input.into_array(), DEFAULT_CONFIG, &mut ctx)?; let canonical = arr.into_array().execute::(&mut ctx)?; let mask = canonical .validity()? @@ -115,7 +270,7 @@ fn test_onpair_nullable_canonicalize() -> vortex_error::VortexResult<()> { fn test_onpair_scalar_at() -> vortex_error::VortexResult<()> { let input = sample_input(); let mut ctx = SESSION.create_execution_ctx(); - let arr = onpair_compress(&input.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; + let arr = onpair_compress(&input.into_array(), DEFAULT_CONFIG, &mut ctx)?; let s = arr.into_array().execute_scalar(2, &mut ctx)?; let v = s.as_utf8().value().unwrap(); assert_eq!(v.as_bytes(), b"https://www.test.org/page"); @@ -140,7 +295,7 @@ fn test_onpair_scalar_at_window() -> vortex_error::VortexResult<()> { DType::Utf8(Nullability::NonNullable), ); let mut ctx = SESSION.create_execution_ctx(); - let arr = onpair_compress(&varbin.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?.into_array(); + let arr = onpair_compress(&varbin.into_array(), DEFAULT_CONFIG, &mut ctx)?.into_array(); for &i in &[0usize, 1, 999, 1000, n - 1] { let got = arr.execute_scalar(i, &mut ctx)?; @@ -189,7 +344,7 @@ fn test_onpair_unroll_tail_boundaries(#[case] n: usize) -> vortex_error::VortexR DType::Utf8(Nullability::NonNullable), ); let mut ctx = SESSION.create_execution_ctx(); - let arr = onpair_compress(&input.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; + let arr = onpair_compress(&input.into_array(), DEFAULT_CONFIG, &mut ctx)?; let canonical = arr.into_array().execute::(&mut ctx)?; let mask = canonical .validity()? @@ -213,13 +368,30 @@ fn test_onpair_empty() -> vortex_error::VortexResult<()> { DType::Utf8(Nullability::NonNullable), ); let mut ctx = SESSION.create_execution_ctx(); - let arr = onpair_compress(&input.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; + let arr = onpair_compress(&input.into_array(), DEFAULT_CONFIG, &mut ctx)?; assert_eq!(arr.len(), 0); let canonical = arr.into_array().execute::(&mut ctx)?; assert_eq!(canonical.len(), 0); Ok(()) } +/// All-null input is already represented optimally by a constant null array. +#[cfg_attr(miri, ignore)] +#[test] +fn test_onpair_all_null() -> vortex_error::VortexResult<()> { + let input = VarBinArray::from_iter( + [None::<&str>, None, None], + DType::Utf8(Nullability::Nullable), + ) + .into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let arr = onpair_compress(&input, DEFAULT_CONFIG, &mut ctx)?; + + assert!(arr.is::()); + assert_arrays_eq!(arr, input, &mut ctx); + Ok(()) +} + /// Filter must share the dictionary — never recompress (this is the /// regression cause on TPC-H Q22 SF=10). Exercise both selectivities /// and check that the result is bit-exact and still an OnPairArray. @@ -235,7 +407,7 @@ fn test_onpair_filter_shares_dict() -> vortex_error::VortexResult<()> { DType::Utf8(Nullability::NonNullable), ); let mut ctx = SESSION.create_execution_ctx(); - let arr = onpair_compress(&varbin.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; + let arr = compress_onpair(&varbin.into_array(), &mut ctx)?; let dict_bytes_before = arr.dict_bytes().clone(); let dict_offsets_len_before = arr.dict_offsets().len(); @@ -309,13 +481,12 @@ fn narrow_codes_offsets(arr: &crate::OnPairArray, target: PType) -> crate::OnPai unsafe { OnPair::new_unchecked( view.dtype().clone(), - view.dict_bytes_handle().clone(), + view.data().clone(), view.dict_offsets().clone(), view.codes().clone(), narrowed_array, view.uncompressed_lengths().clone(), view.array_validity(), - view.bits(), ) } } @@ -339,7 +510,7 @@ fn test_onpair_filter_with_narrowed_codes_offsets_u16() -> vortex_error::VortexR DType::Utf8(Nullability::NonNullable), ); let mut ctx = SESSION.create_execution_ctx(); - let arr = onpair_compress(&varbin.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; + let arr = compress_onpair(&varbin.into_array(), &mut ctx)?; // Force `codes_offsets` to u16 so the panicking pre-fix // `as_slice::()` would fire. @@ -395,7 +566,7 @@ fn test_onpair_filter_with_narrowed_codes_offsets_u8() -> vortex_error::VortexRe DType::Utf8(Nullability::NonNullable), ); let mut ctx = SESSION.create_execution_ctx(); - let arr = onpair_compress(&varbin.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; + let arr = compress_onpair(&varbin.into_array(), &mut ctx)?; let arr = narrow_codes_offsets(&arr, PType::U8); assert_eq!(arr.as_view().codes_offsets().dtype().as_ptype(), PType::U8); @@ -428,7 +599,7 @@ fn test_onpair_slice_canonicalize() -> vortex_error::VortexResult<()> { DType::Utf8(Nullability::NonNullable), ); let mut ctx = SESSION.create_execution_ctx(); - let arr = onpair_compress(&varbin.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?.into_array(); + let arr = onpair_compress(&varbin.into_array(), DEFAULT_CONFIG, &mut ctx)?.into_array(); // interior (start>0, end0, // end=n), and a near-full window. diff --git a/encodings/experimental/onpair/tests/big_data.rs b/encodings/onpair/tests/big_data.rs similarity index 93% rename from encodings/experimental/onpair/tests/big_data.rs rename to encodings/onpair/tests/big_data.rs index e7f74fdd782..7feed9b1718 100644 --- a/encodings/experimental/onpair/tests/big_data.rs +++ b/encodings/onpair/tests/big_data.rs @@ -25,7 +25,7 @@ use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::scalar_fn::fns::operators::Operator; -use vortex_onpair::DEFAULT_DICT12_CONFIG; +use vortex_onpair::DEFAULT_CONFIG; use vortex_onpair::onpair_compress; use vortex_session::VortexSession; @@ -74,12 +74,11 @@ fn smoke_100k_rows() -> vortex_error::VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let t0 = Instant::now(); - let arr = onpair_compress(&varbin.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; + let arr = onpair_compress(&varbin.into_array(), DEFAULT_CONFIG, &mut ctx)?; let compress_elapsed = t0.elapsed(); - let bits = arr.bits(); eprintln!( - "compressed {} rows ({} raw bytes) in {:?}, bits={}", - n, raw_bytes, compress_elapsed, bits + "compressed {} rows ({} raw bytes) in {:?}", + n, raw_bytes, compress_elapsed ); let arr_ref = arr.into_array(); diff --git a/encodings/parquet-variant/Cargo.toml b/encodings/parquet-variant/Cargo.toml index 6cdc2b3fa8b..e90da6c8ebe 100644 --- a/encodings/parquet-variant/Cargo.toml +++ b/encodings/parquet-variant/Cargo.toml @@ -34,6 +34,7 @@ vortex-proto = { workspace = true, features = ["expr"] } vortex-session = { workspace = true } [dev-dependencies] +vortex-edition = { workspace = true } rstest = { workspace = true } tokio = { workspace = true, features = ["full"] } vortex-array = { workspace = true, features = ["_test-harness"] } diff --git a/encodings/parquet-variant/src/array.rs b/encodings/parquet-variant/src/array.rs index bc7fd2c22ec..5d364cae0ef 100644 --- a/encodings/parquet-variant/src/array.rs +++ b/encodings/parquet-variant/src/array.rs @@ -14,6 +14,7 @@ use vortex_array::EmptyArrayData; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::TypedArrayRef; +use vortex_array::array_slots; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::List; use vortex_array::arrays::ListArray; @@ -22,15 +23,15 @@ use vortex_array::arrays::StructArray; use vortex_array::arrays::Variant; use vortex_array::arrays::VariantArray; use vortex_array::arrays::list::ListArrayExt; +use vortex_array::arrays::list::ListArraySlotsExt; use vortex_array::arrays::struct_::StructArrayExt; -use vortex_array::arrays::variant::VariantArrayExt; +use vortex_array::arrays::variant::VariantArraySlotsExt; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::FieldName; use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; -use vortex_array::smallvec::smallvec; use vortex_array::validity::Validity; use vortex_array::vtable::child_to_validity; use vortex_array::vtable::validity_to_child; @@ -39,7 +40,7 @@ use vortex_array::vtable::validity_to_child; reason = "TODO(aduffy): figure out what to do with Parquet Variant" )] use vortex_arrow::ArrowArrayExecutor; -use vortex_arrow::FromArrowArray; +use vortex_arrow::ArrowSession; use vortex_arrow::to_arrow_null_buffer; use vortex_buffer::BitBuffer; use vortex_error::VortexExpect; @@ -49,16 +50,21 @@ use vortex_error::vortex_err; use crate::ParquetVariant; use crate::ParquetVariantArray; -/// The validity bitmap indicating which elements are non-null. -pub(crate) const VALIDITY_SLOT: usize = 0; -/// The metadata array for the Parquet variant values. -pub(crate) const METADATA_SLOT: usize = 1; -/// The value array containing the Parquet variant data. -pub(crate) const VALUE_SLOT: usize = 2; -/// The typed value array for strongly-typed Parquet variant data. -pub(crate) const TYPED_VALUE_SLOT: usize = 3; -pub(crate) const NUM_SLOTS: usize = 4; -pub(crate) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity", "metadata", "value", "typed_value"]; +#[array_slots(ParquetVariant)] +pub struct ParquetVariantSlots { + /// The validity bitmap indicating which elements are non-null. + #[slot(0)] + pub validity: Option, + /// The metadata array for the Parquet variant values. + #[slot(1)] + pub metadata: ArrayRef, + /// The value array containing the Parquet variant data. + #[slot(2)] + pub value: Option, + /// The typed value array for strongly-typed Parquet variant data. + #[slot(3)] + pub typed_value: Option, +} impl ParquetVariant { /// Creates a Parquet Variant array from canonical extension storage slots. @@ -73,31 +79,38 @@ impl ParquetVariant { ) -> VortexResult> { let len = metadata.len(); let dtype = DType::Variant(validity.nullability()); - let slots = smallvec![ - validity_to_child(&validity, len), - Some(metadata), + let slots = ParquetVariantSlots { + validity: validity_to_child(&validity, len), + metadata, value, typed_value, - ]; + } + .into_slots(); Array::try_from_parts( ArrayParts::new(ParquetVariant, dtype, len, EmptyArrayData).with_slots(slots), ) } - /// Converts an Arrow `parquet_variant_compute::VariantArray` into Parquet Variant storage. - pub fn from_arrow_variant(arrow_variant: &ArrowVariantArray) -> VortexResult { - Self::from_arrow_variant_impl(arrow_variant, false) + /// Converts an Arrow `parquet_variant_compute::VariantArray` into Parquet Variant storage, + /// converting the storage children through `session`. + pub fn from_arrow_variant( + arrow_variant: &ArrowVariantArray, + session: &ArrowSession, + ) -> VortexResult { + Self::from_arrow_variant_impl(arrow_variant, false, session) } pub(crate) fn from_arrow_variant_nullable( arrow_variant: &ArrowVariantArray, + session: &ArrowSession, ) -> VortexResult { - Self::from_arrow_variant_impl(arrow_variant, true) + Self::from_arrow_variant_impl(arrow_variant, true, session) } fn from_arrow_variant_impl( arrow_variant: &ArrowVariantArray, force_nullable: bool, + session: &ArrowSession, ) -> VortexResult { let storage = arrow_variant.inner(); let mut value_nullable = false; @@ -123,17 +136,17 @@ impl ParquetVariant { } else { Validity::NonNullable }); - let metadata = - ArrayRef::from_arrow(arrow_variant.metadata_field() as &dyn ArrowArray, false)?; + let metadata = session + .from_arrow_array(ArrowArrayRef::clone(arrow_variant.metadata_column()), false)?; - let value = arrow_variant - .value_field() - .map(|v| ArrayRef::from_arrow(v as &dyn ArrowArray, value_nullable)) - .transpose()?; + let value = Some(session.from_arrow_array( + ArrowArrayRef::clone(arrow_variant.value_column()), + value_nullable, + )?); let typed_value = arrow_variant - .typed_value_field() - .map(|tv| ArrayRef::from_arrow(tv.as_ref(), typed_value_nullable)) + .typed_value_column() + .map(|tv| session.from_arrow_array(ArrowArrayRef::clone(tv), typed_value_nullable)) .transpose()?; ParquetVariant::try_new(validity, metadata, value, typed_value).map(IntoArray::into_array) } @@ -147,8 +160,8 @@ pub(crate) fn core_storage_without_typed_value( // the outer `VariantArray::shredded` slot and the original had no `value`, synthesize an // all-null `value` so the remaining `ParquetVariant` still satisfies that invariant and // can round-trip back through `to_arrow`. - let value = array.value_array().cloned().or_else(|| { - array.typed_value_array().map(|_| { + let value = array.value().cloned().or_else(|| { + array.typed_value().map(|_| { ConstantArray::new( Scalar::null(DType::Binary(Nullability::Nullable)), array.len(), @@ -157,13 +170,8 @@ pub(crate) fn core_storage_without_typed_value( }) }); - ParquetVariant::try_new( - array.validity()?, - array.metadata_array().clone(), - value, - None, - ) - .map(IntoArray::into_array) + ParquetVariant::try_new(array.validity()?, array.metadata().clone(), value, None) + .map(IntoArray::into_array) } /// Converts a Parquet `typed_value` tree into the storage-agnostic canonical shredded tree. @@ -274,7 +282,7 @@ fn logical_shredded_from_parquet_field( let parquet_field = ParquetVariant::try_new(validity, metadata.clone(), Some(value), Some(typed_value))?; let shredded = parquet_field - .typed_value_array() + .typed_value() .cloned() .map(|typed_value| { logical_shredded_from_parquet_typed_value(metadata, typed_value, ctx) @@ -380,7 +388,7 @@ fn parquet_shredded_field_from_logical( "cannot rebuild Parquet shredded field: nested Variant lacks Parquet Variant core storage" ) })?; - let value = core.value_array().cloned().ok_or_else(|| { + let value = core.value().cloned().ok_or_else(|| { vortex_err!("cannot rebuild Parquet shredded field: partially shredded Variant has no residual value") })?; let typed_value = variant @@ -415,42 +423,27 @@ fn parquet_shredded_field_from_logical( .into_array()) } -/// Accessors and Arrow conversion for Parquet Variant storage arrays. -pub trait ParquetVariantArrayExt: TypedArrayRef { - /// Returns the non-nullable Parquet Variant metadata child. - fn metadata_array(&self) -> &ArrayRef { - self.as_ref().slots()[METADATA_SLOT] - .as_ref() - .vortex_expect("ParquetVariantArray metadata slot") - } - +/// Non-slot accessors and Arrow conversion for Parquet Variant storage arrays. +pub trait ParquetVariantArrayExt: + TypedArrayRef + ParquetVariantArraySlotsExt +{ /// Returns the outer row validity for the Variant values. - fn validity(&self) -> Validity { + fn parquet_variant_validity(&self) -> Validity { child_to_validity( - self.as_ref().slots()[VALIDITY_SLOT].as_ref(), + self.as_ref().slots()[ParquetVariantSlots::VALIDITY].as_ref(), self.as_ref().dtype().nullability(), ) } - /// Returns the optional raw Parquet Variant `value` child. - fn value_array(&self) -> Option<&ArrayRef> { - self.as_ref().slots()[VALUE_SLOT].as_ref() - } - - /// Returns the optional shredded Parquet Variant `typed_value` child. - fn typed_value_array(&self) -> Option<&ArrayRef> { - self.as_ref().slots()[TYPED_VALUE_SLOT].as_ref() - } - /// Converts this storage array to Arrow's canonical Parquet Variant extension storage. #[expect( deprecated, reason = "TODO(aduffy): figure out what to do with Parquet Variant" )] fn to_arrow(&self, ctx: &mut ExecutionCtx) -> VortexResult { - let metadata = self.metadata_array(); + let metadata = self.metadata(); let len = metadata.len(); - let nulls = to_arrow_null_buffer(self.validity(), len, ctx)?; + let nulls = to_arrow_null_buffer(self.parquet_variant_validity(), len, ctx)?; let mut fields = Vec::with_capacity(3); let mut arrays: Vec = Vec::with_capacity(3); @@ -463,7 +456,7 @@ pub trait ParquetVariantArrayExt: TypedArrayRef { ))); arrays.push(metadata_arrow); - if let Some(value) = self.value_array() { + if let Some(value) = self.value() { let value_arrow = value.clone().execute_arrow(None, ctx)?; fields.push(Arc::new(Field::new( "value", @@ -473,7 +466,7 @@ pub trait ParquetVariantArrayExt: TypedArrayRef { arrays.push(value_arrow); } - if let Some(typed_value) = self.typed_value_array() { + if let Some(typed_value) = self.typed_value() { let tv_arrow = typed_value.clone().execute_arrow(None, ctx)?; fields.push(Arc::new(Field::new( "typed_value", @@ -516,11 +509,12 @@ mod tests { use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; - use vortex_array::arrays::variant::VariantArrayExt; + use vortex_array::arrays::variant::VariantArraySlotsExt; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::validity::Validity; + use vortex_arrow::ArrowSessionExt; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -528,6 +522,7 @@ mod tests { use crate::ParquetVariant; use crate::array::ParquetVariantArrayExt; + use crate::array::ParquetVariantArraySlotsExt; use crate::array::parquet_typed_value_from_logical_shredded; static SESSION: LazyLock = LazyLock::new(|| { @@ -538,7 +533,7 @@ mod tests { fn assert_arrow_variant_storage_roundtrip(struct_array: StructArray) -> VortexResult<()> { let arrow_variant = ArrowVariantArray::try_new(&struct_array)?; - let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; let inner = vortex_arr .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant child"))?; @@ -589,7 +584,7 @@ mod tests { builder.append_variant(PqVariant::from(true)); let arrow_variant = builder.build(); - let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; assert_eq!(vortex_arr.len(), 3); assert_eq!( @@ -621,7 +616,7 @@ mod tests { let arrow_variant = ArrowVariantArray::try_new(&struct_array)?; - let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; assert_eq!(vortex_arr.len(), 3); assert_eq!( vortex_arr.dtype(), @@ -632,7 +627,7 @@ mod tests { .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant array"))?; let typed_value = parquet_array - .typed_value_array() + .typed_value() .ok_or_else(|| vortex_err!("expected typed_value child"))? .clone() .execute::(&mut SESSION.create_execution_ctx())?; @@ -697,7 +692,7 @@ mod tests { } #[test] - fn test_arrow_variant_import_typed_value_only_preserves_storage() -> VortexResult<()> { + fn test_arrow_variant_import_typed_value_only_normalizes_storage() -> VortexResult<()> { let metadata = binary_view_array([b"\x01\x00", b"\x01\x00", b"\x01\x00"]); let typed_value: ArrowArrayRef = Arc::new(Int32Array::from(vec![10, 20, 30])); @@ -712,21 +707,25 @@ mod tests { )?; let arrow_variant = ArrowVariantArray::try_new(&struct_array)?; - let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; let parquet_array = vortex_arr .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant array"))?; - assert!(parquet_array.value_array().is_none()); + let mut ctx = SESSION.create_execution_ctx(); + let value = parquet_array + .value() + .ok_or_else(|| vortex_err!("expected synthesized value child"))?; + assert!(value.all_invalid(&mut ctx)?); let typed_value = parquet_array - .typed_value_array() + .typed_value() .ok_or_else(|| vortex_err!("expected typed_value child"))? .clone() - .execute::(&mut SESSION.create_execution_ctx())?; + .execute::(&mut ctx)?; assert_arrays_eq!( typed_value, - PrimitiveArray::from_iter([10i32, 20, 30]), - &mut SESSION.create_execution_ctx() + PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30)]), + &mut ctx ); Ok(()) } @@ -749,12 +748,12 @@ mod tests { )?; let arrow_variant = ArrowVariantArray::try_new(&struct_array)?; - let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; let parquet_array = vortex_arr .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant array"))?; - assert!(parquet_array.value_array().is_some()); - assert!(parquet_array.typed_value_array().is_some()); + assert!(parquet_array.value().is_some()); + assert!(parquet_array.typed_value().is_some()); let mut ctx = SESSION.create_execution_ctx(); let roundtripped = parquet_array.to_arrow(&mut ctx)?; @@ -821,12 +820,12 @@ mod tests { .with_path("a", &DataType::Int32)? .build(); let shredded = shred_variant(&json_to_variant(&json)?, &shredding)?; - let original = ParquetVariant::from_arrow_variant(&shredded)?; + let original = ParquetVariant::from_arrow_variant(&shredded, &SESSION.arrow())?; assert!( original .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant"))? - .typed_value_array() + .typed_value() .is_some(), "fixture must be shredded" ); @@ -848,9 +847,9 @@ mod tests { // Inverse transform: rebuild a Parquet `typed_value` and reattach it. let rebuilt_typed_value = parquet_typed_value_from_logical_shredded(logical, &mut ctx)?; let rebuilt = ParquetVariant::try_new( - ParquetVariantArrayExt::validity(&core), - core.metadata_array().clone(), - core.value_array().cloned(), + ParquetVariantArrayExt::parquet_variant_validity(&core), + core.metadata().clone(), + core.value().cloned(), Some(rebuilt_typed_value), )? .into_array(); diff --git a/encodings/parquet-variant/src/arrow.rs b/encodings/parquet-variant/src/arrow.rs index 65efd6c571c..e1c03d704dc 100644 --- a/encodings/parquet-variant/src/arrow.rs +++ b/encodings/parquet-variant/src/arrow.rs @@ -19,7 +19,7 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::Variant; -use vortex_array::arrays::variant::VariantArrayExt; +use vortex_array::arrays::variant::VariantArraySlotsExt; use vortex_array::dtype::DType; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportVTable; @@ -36,6 +36,7 @@ use vortex_session::registry::Id; use crate::ParquetVariant; use crate::ParquetVariantArrayExt; +use crate::ParquetVariantArraySlotsExt; use crate::array::parquet_typed_value_from_logical_shredded; /// Arrow canonical extension name for Parquet Variant storage. @@ -68,9 +69,9 @@ pub(crate) fn export_storage_to_target( for field in target_fields { let child = match field.name().as_str() { - "metadata" => Some(parquet_array.metadata_array().clone()), - "value" => parquet_array.value_array().cloned(), - "typed_value" => parquet_array.typed_value_array().cloned(), + "metadata" => Some(parquet_array.metadata().clone()), + "value" => parquet_array.value().cloned(), + "typed_value" => parquet_array.typed_value().cloned(), _ => unreachable!("storage fields were validated before export"), }; let Some(child) = child else { @@ -89,7 +90,7 @@ pub(crate) fn export_storage_to_target( } let nulls = to_arrow_null_buffer( - ParquetVariantArrayExt::validity(parquet_array), + ParquetVariantArrayExt::parquet_variant_validity(parquet_array), parquet_array.as_ref().len(), ctx, )?; @@ -108,9 +109,9 @@ pub(crate) fn export_unshredded_storage_to_target( let arrow_variant = parquet_array.to_arrow(ctx)?; let unshredded = unshred_variant(&arrow_variant)?; let unshredded_array = if parquet_array.as_ref().dtype().is_nullable() { - ParquetVariant::from_arrow_variant_nullable(&unshredded)? + ParquetVariant::from_arrow_variant_nullable(&unshredded, &ctx.session().arrow())? } else { - ParquetVariant::from_arrow_variant(&unshredded)? + ParquetVariant::from_arrow_variant(&unshredded, &ctx.session().arrow())? }; let unshredded_parquet = unshredded_array.as_::(); export_storage_to_target(&unshredded_parquet, target_fields, ctx) @@ -145,9 +146,9 @@ pub(crate) fn parquet_variant_for_export( let typed_value = parquet_typed_value_from_logical_shredded(shredded.clone(), ctx)?; ParquetVariant::try_new( - ParquetVariantArrayExt::validity(&parquet_core), - parquet_core.metadata_array().clone(), - parquet_core.value_array().cloned(), + ParquetVariantArrayExt::parquet_variant_validity(&parquet_core), + parquet_core.metadata().clone(), + parquet_core.value().cloned(), Some(typed_value), ) .map(IntoArray::into_array) @@ -195,8 +196,8 @@ impl ArrowExportVTable for ParquetVariant { && let Some((request_has_value, request_has_typed_value)) = parquet_variant_storage_request(fields) { - let has_value = parquet_array.value_array().is_some(); - let has_typed_value = parquet_array.typed_value_array().is_some(); + let has_value = parquet_array.value().is_some(); + let has_typed_value = parquet_array.typed_value().is_some(); if request_has_value && !request_has_typed_value && has_typed_value { return Ok(ArrowExport::Exported(export_unshredded_storage_to_target( @@ -241,7 +242,11 @@ impl ArrowImportVTable for ParquetVariant { *ARROW_PARQUET_VARIANT } - fn from_arrow_field(&self, field: &Field) -> VortexResult> { + fn from_arrow_field( + &self, + field: &Field, + _session: &ArrowSession, + ) -> VortexResult> { if field .metadata() .get(EXTENSION_TYPE_NAME_KEY) @@ -258,6 +263,7 @@ impl ArrowImportVTable for ParquetVariant { array: ArrowArrayRef, field: &Field, dtype: &DType, + session: &ArrowSession, ) -> VortexResult { if !dtype.is_variant() || field @@ -271,9 +277,9 @@ impl ArrowImportVTable for ParquetVariant { let arrow_variant = ArrowVariantArray::try_new(array.as_struct())?; let imported = if dtype.is_nullable() { - ParquetVariant::from_arrow_variant_nullable(&arrow_variant)? + ParquetVariant::from_arrow_variant_nullable(&arrow_variant, session)? } else { - ParquetVariant::from_arrow_variant(&arrow_variant)? + ParquetVariant::from_arrow_variant(&arrow_variant, session)? }; Ok(ArrowImport::Imported(imported.into_array())) } @@ -709,7 +715,6 @@ mod tests { let actual = session.arrow().from_arrow_array(exported, &field)?; - assert_arrays_eq!(actual, expected, &mut ctx); - Ok(()) + assert_variant_scalars_eq(&actual, &expected, &session) } } diff --git a/encodings/parquet-variant/src/compute/allnondistinct.rs b/encodings/parquet-variant/src/compute/allnondistinct.rs index 4f5568fbcf6..7e6219bc542 100644 --- a/encodings/parquet-variant/src/compute/allnondistinct.rs +++ b/encodings/parquet-variant/src/compute/allnondistinct.rs @@ -16,7 +16,7 @@ use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use crate::ParquetVariant; -use crate::ParquetVariantArrayExt; +use crate::ParquetVariantArraySlotsExt; /// Lets `AllNonDistinct` compare two `ParquetVariant` arrays without canonicalizing them. /// @@ -51,7 +51,7 @@ impl DynAggregateKernel for AllNonDistinctParquetVariant { return Ok(None); }; - let typed_identical = match (lhs.typed_value_array(), rhs.typed_value_array()) { + let typed_identical = match (lhs.typed_value(), rhs.typed_value()) { (Some(lhs_typed), Some(rhs_typed)) => { if lhs_typed.dtype().eq_ignore_nullability(rhs_typed.dtype()) { all_non_distinct(lhs_typed, rhs_typed, ctx)? @@ -63,7 +63,7 @@ impl DynAggregateKernel for AllNonDistinctParquetVariant { }; if typed_identical { - let values_identical = match (lhs.value_array(), rhs.value_array()) { + let values_identical = match (lhs.value(), rhs.value()) { (Some(lhs_value), Some(rhs_value)) => all_non_distinct(lhs_value, rhs_value, ctx)?, (None, None) => true, // Mixed shredding layouts: let the generic canonical path handle it. diff --git a/encodings/parquet-variant/src/json_to_variant_tests.rs b/encodings/parquet-variant/src/json_to_variant_tests.rs index cd3b2f035ee..5bc09a426f7 100644 --- a/encodings/parquet-variant/src/json_to_variant_tests.rs +++ b/encodings/parquet-variant/src/json_to_variant_tests.rs @@ -20,7 +20,7 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::struct_::StructArrayExt; -use vortex_array::arrays::variant::VariantArrayExt; +use vortex_array::arrays::variant::VariantArraySlotsExt; use vortex_array::assert_arrays_eq; use vortex_array::assert_nth_scalar_is_null; use vortex_array::dtype::DType; @@ -36,7 +36,7 @@ use vortex_json::Json; use vortex_session::VortexSession; use crate::ParquetVariant; -use crate::ParquetVariantArrayExt; +use crate::ParquetVariantArraySlotsExt; use crate::ShreddingSpec; use crate::json_to_variant; @@ -145,7 +145,7 @@ fn converts_json_extension_input() -> VortexResult<()> { #[test] fn dict_encoded_input_converts_each_row() -> VortexResult<()> { // A dictionary-encoded JSON column exercises the dict-pushdown / canonicalization path: - // `json_to_variant` is not null-sensitive, so it pushes into the dict values (canonical + // `json_to_variant` is strict, so it pushes into the dict values (canonical // JSON extension values) where the kernel fires; either way every row must convert correctly. let values = json_input(VarBinViewArray::from_iter_str(["1", "2"]).into_array())?; let codes = PrimitiveArray::from_iter([0u8, 1, 0, 1, 0]).into_array(); @@ -200,7 +200,7 @@ fn shredding_produces_typed_value_child() -> VortexResult<()> { let result = execute_json_to_variant(input, shred_field_as_i64("a")?)?; assert!( - result.as_::().typed_value_array().is_some(), + result.as_::().typed_value().is_some(), "expected shredded typed_value child" ); @@ -283,7 +283,7 @@ fn shredding_root_path_shreds_top_level_values() -> VortexResult<()> { let result = execute_json_to_variant(input, spec)?; assert!( - result.as_::().typed_value_array().is_some(), + result.as_::().typed_value().is_some(), "expected shredded typed_value child" ); assert_variant_i64_rows(&result.slice(0..2)?, &[Some(1), Some(2)])?; diff --git a/encodings/parquet-variant/src/kernel.rs b/encodings/parquet-variant/src/kernel.rs index c783e6600ac..90aa9eba548 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/encodings/parquet-variant/src/kernel.rs @@ -6,7 +6,6 @@ use std::ops::Range; use std::sync::Arc; use arrow_array::ArrayRef as ArrowArrayRef; -use arrow_schema::Field; use arrow_schema::FieldRef; use parquet_variant::VariantPath as PqVariantPath; use parquet_variant::VariantPathElement as PqVariantPathElement; @@ -44,9 +43,8 @@ use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::variant_get::VariantGet; use vortex_array::scalar_fn::fns::variant_get::VariantPath; use vortex_array::scalar_fn::fns::variant_get::VariantPathElement; +use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; -use vortex_arrow::FromArrowArray; -use vortex_arrow::ToArrowType; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; @@ -57,6 +55,7 @@ use vortex_session::VortexSession; use crate::ParquetVariant; use crate::ParquetVariantArrayExt; +use crate::ParquetVariantArraySlotsExt; use crate::compute::AllNonDistinctParquetVariant; pub(crate) fn initialize(session: &VortexSession) { @@ -109,16 +108,26 @@ impl ExecuteParentKernel for VariantGetKernel { let arrow_variant = array.to_arrow(ctx)?; let arrow_input: ArrowArrayRef = Arc::new(arrow_variant.into_inner()); + let session = ctx.session().clone(); + let as_type = to_arrow_as_type(parent.options.dtype(), &session.arrow())?; let get_options = GetOptions::new_with_path(to_parquet_variant_path(parent.options.path())?) - .with_as_type(to_arrow_as_type(parent.options.dtype())?); + .with_as_type(as_type.clone()); let arrow_output = arrow_variant_get(&arrow_input, get_options)?; let output = if parent.options.dtype().is_none_or(DType::is_variant) { let arrow_variant_output = ArrowVariantArray::try_new(arrow_output.as_ref())?; - ParquetVariant::from_arrow_variant_nullable(&arrow_variant_output)? + ParquetVariant::from_arrow_variant_nullable( + &arrow_variant_output, + &ctx.session().arrow(), + )? } else { - ArrayRef::from_arrow(arrow_output.as_ref(), true)? + // Import through the same `as_type` field the cast targeted, so an extension target + // dtype comes back as that extension rather than as its bare storage type. + let as_type = as_type.ok_or_else(|| { + vortex_err!("a non-variant target dtype must produce an as_type field") + })?; + session.arrow().from_arrow_array(arrow_output, &as_type)? }; vortex_ensure_eq!( @@ -160,9 +169,9 @@ fn json_strings_to_variant( }; if nullable { - ParquetVariant::from_arrow_variant_nullable(&arrow_variant) + ParquetVariant::from_arrow_variant_nullable(&arrow_variant, &session.arrow()) } else { - ParquetVariant::from_arrow_variant(&arrow_variant) + ParquetVariant::from_arrow_variant(&arrow_variant, &session.arrow()) } } @@ -207,17 +216,20 @@ pub(crate) fn to_parquet_variant_path(path: &VariantPath) -> VortexResult) -> VortexResult> { +/// The Arrow field `variant_get` should cast to for a requested Vortex `dtype`. +/// +/// Built as a [`Field`](arrow_schema::Field) rather than a bare `DataType` so an extension target +/// keeps its `ARROW:extension:name`; the result is nullable because an absent path yields null. +fn to_arrow_as_type( + dtype: Option<&DType>, + session: &ArrowSession, +) -> VortexResult> { match dtype { - Some(dtype) if !dtype.is_variant() => Ok(Some(Arc::new(Field::new( - "variant_get", - dtype.to_arrow_dtype()?, - true, - )))), + Some(dtype) if !dtype.is_variant() => Ok(Some(Arc::new( + session + .to_arrow_field("variant_get", dtype)? + .with_nullable(true), + ))), Some(_) | None => Ok(None), } } @@ -229,15 +241,9 @@ impl SliceKernel for ParquetVariant { _ctx: &mut ExecutionCtx, ) -> VortexResult> { let validity = array.validity()?.slice(range.clone())?; - let metadata = array.metadata_array().slice(range.clone())?; - let value = array - .value_array() - .map(|v| v.slice(range.clone())) - .transpose()?; - let typed_value = array - .typed_value_array() - .map(|tv| tv.slice(range)) - .transpose()?; + let metadata = array.metadata().slice(range.clone())?; + let value = array.value().map(|v| v.slice(range.clone())).transpose()?; + let typed_value = array.typed_value().map(|tv| tv.slice(range)).transpose()?; Ok(Some( ParquetVariant::try_new(validity, metadata, value, typed_value)?.into_array(), )) @@ -251,13 +257,10 @@ impl FilterKernel for ParquetVariant { _ctx: &mut ExecutionCtx, ) -> VortexResult> { let validity = array.validity()?.filter(mask)?; - let metadata = array.metadata_array().filter(mask.clone())?; - let value = array - .value_array() - .map(|v| v.filter(mask.clone())) - .transpose()?; + let metadata = array.metadata().filter(mask.clone())?; + let value = array.value().map(|v| v.filter(mask.clone())).transpose()?; let typed_value = array - .typed_value_array() + .typed_value() .map(|tv| tv.filter(mask.clone())) .transpose()?; Ok(Some( @@ -273,13 +276,10 @@ impl TakeExecute for ParquetVariant { _ctx: &mut ExecutionCtx, ) -> VortexResult> { let validity = array.validity()?.take(indices)?; - let metadata = array.metadata_array().take(indices.clone())?; - let value = array - .value_array() - .map(|v| v.take(indices.clone())) - .transpose()?; + let metadata = array.metadata().take(indices.clone())?; + let value = array.value().map(|v| v.take(indices.clone())).transpose()?; let typed_value = array - .typed_value_array() + .typed_value() .map(|tv| tv.take(indices.clone())) .transpose()?; Ok(Some( @@ -318,7 +318,7 @@ mod tests { use vortex_array::arrays::Variant; use vortex_array::arrays::VariantArray; use vortex_array::arrays::struct_::StructArrayExt; - use vortex_array::arrays::variant::VariantArrayExt; + use vortex_array::arrays::variant::VariantArraySlotsExt; use vortex_array::assert_arrays_eq; use vortex_array::assert_nth_scalar_is_null; use vortex_array::dtype::DType as VortexDType; @@ -329,7 +329,7 @@ mod tests { use vortex_array::scalar_fn::fns::variant_get::VariantPath; use vortex_array::scalar_fn::fns::variant_get::VariantPathElement; use vortex_array::validity::Validity; - use vortex_arrow::FromArrowArray; + use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -337,8 +337,9 @@ mod tests { use vortex_mask::Mask; use vortex_session::VortexSession; + use super::to_arrow_as_type; use crate::ParquetVariant; - use crate::ParquetVariantArrayExt; + use crate::ParquetVariantArraySlotsExt; static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session(); @@ -346,13 +347,76 @@ mod tests { session }); + /// `variant_get`'s cast target must be a `Field` carrying `ARROW:extension:name`, so an + /// extension target dtype survives both the cast and the import back into Vortex. A bare + /// `DataType` would silently degrade the target to its storage type. + #[test] + fn to_arrow_as_type_keeps_extension_metadata() -> VortexResult<()> { + use arrow_schema::extension::ExtensionType; + use arrow_schema::extension::Uuid as ArrowUuid; + use vortex_array::dtype::extension::ExtDType; + use vortex_array::extension::uuid::Uuid; + use vortex_array::extension::uuid::UuidMetadata; + + let storage = VortexDType::FixedSizeList( + Arc::new(VortexDType::Primitive(PType::U8, Nullability::NonNullable)), + 16, + Nullability::NonNullable, + ); + let uuid = VortexDType::Extension( + ExtDType::try_with_vtable(Uuid, UuidMetadata::default(), storage)?.erased(), + ); + + let arrow = SESSION.arrow(); + let as_type = to_arrow_as_type(Some(&uuid), &arrow)? + .ok_or_else(|| vortex_err!("expected an as_type field"))?; + + assert_eq!(as_type.data_type(), &DataType::FixedSizeBinary(16)); + assert_eq!(as_type.extension_type_name(), Some(ArrowUuid::NAME)); + // An absent path yields null, so the cast target must stay nullable. + assert!(as_type.is_nullable()); + + // And the field round-trips back to the extension dtype the caller asked for. + assert_eq!(arrow.from_arrow_field(&as_type)?, uuid.as_nullable()); + Ok(()) + } + + /// A non-extension target stays a plain nullable field. + #[test] + fn to_arrow_as_type_plain_dtype() -> VortexResult<()> { + let arrow = SESSION.arrow(); + let as_type = to_arrow_as_type( + Some(&VortexDType::Primitive( + PType::I32, + Nullability::NonNullable, + )), + &arrow, + )? + .ok_or_else(|| vortex_err!("expected an as_type field"))?; + assert_eq!(as_type.data_type(), &DataType::Int32); + assert_eq!(as_type.extension_type_name(), None); + assert!(as_type.is_nullable()); + Ok(()) + } + + /// Variant (and absent) targets ask for no cast at all. + #[test] + fn to_arrow_as_type_variant_targets_are_none() -> VortexResult<()> { + let arrow = SESSION.arrow(); + assert!( + to_arrow_as_type(Some(&VortexDType::Variant(Nullability::Nullable)), &arrow)?.is_none() + ); + assert!(to_arrow_as_type(None, &arrow)?.is_none()); + Ok(()) + } + fn make_unshredded_array() -> VortexResult { let mut builder = VariantArrayBuilder::new(4); builder.append_variant(PqVariant::from(42i32)); builder.append_variant(PqVariant::from("hello")); builder.append_variant(PqVariant::from(true)); builder.append_variant(PqVariant::from(99i64)); - ParquetVariant::from_arrow_variant(&builder.build()) + ParquetVariant::from_arrow_variant(&builder.build(), &SESSION.arrow()) } fn make_nullable_array() -> VortexResult { @@ -369,13 +433,13 @@ mod tests { Some(NullBuffer::from(vec![true, false, true, false])), )?; let arrow_variant = ArrowVariantArray::try_new(&null_struct)?; - ParquetVariant::from_arrow_variant(&arrow_variant) + ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow()) } fn make_unshredded_json_array(values: Vec>) -> VortexResult { let json: ArrowArrayRef = Arc::new(StringArray::from(values)); let arrow_variant = json_to_variant(&json)?; - ParquetVariant::from_arrow_variant(&arrow_variant) + ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow()) } fn parse_path(path: &str) -> VortexResult { @@ -775,15 +839,20 @@ mod tests { .map(|field| field.is_nullable()) .unwrap_or(false); - let metadata = - ArrayRef::from_arrow(arrow_variant.metadata_field() as &dyn ArrowArray, false)?; - let value = arrow_variant - .value_field() - .map(|value| ArrayRef::from_arrow(value as &dyn ArrowArray, value_nullable)) - .transpose()?; + let metadata = SESSION + .arrow() + .from_arrow_array(ArrowArrayRef::clone(arrow_variant.metadata_column()), false)?; + let value = Some(SESSION.arrow().from_arrow_array( + ArrowArrayRef::clone(arrow_variant.value_column()), + value_nullable, + )?); let typed_value = arrow_variant - .typed_value_field() - .map(|typed_value| ArrayRef::from_arrow(typed_value.as_ref(), typed_value_nullable)) + .typed_value_column() + .map(|typed_value| { + SESSION + .arrow() + .from_arrow_array(ArrowArrayRef::clone(typed_value), typed_value_nullable) + }) .transpose()?; Ok( @@ -794,7 +863,7 @@ mod tests { fn make_partially_shredded_object_array() -> VortexResult { let arrow_variant = make_partially_shredded_arrow_variant()?; - let parquet_array = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let parquet_array = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; let mut ctx = SESSION.create_execution_ctx(); let Canonical::Variant(canonical) = parquet_array.execute::(&mut ctx)? else { return Err(vortex_err!("expected canonical variant array")); @@ -911,7 +980,7 @@ mod tests { None, )?; let arrow_variant = ArrowVariantArray::try_new(&struct_array)?; - ParquetVariant::from_arrow_variant(&arrow_variant) + ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow()) } fn assert_typed_value_i32( @@ -922,7 +991,7 @@ mod tests { let executed = array.clone().execute::(&mut ctx)?; let typed_value = executed .as_::() - .typed_value_array() + .typed_value() .ok_or_else(|| vortex_err!("expected typed_value child"))? .clone() .execute::(&mut ctx)?; @@ -1004,8 +1073,8 @@ mod tests { .core_storage() .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant core storage"))?; - assert!(core_storage.typed_value_array().is_none()); - assert!(core_storage.value_array().is_some()); + assert!(core_storage.typed_value().is_none()); + assert!(core_storage.value().is_some()); let shredded = canonical_variant .shredded() @@ -1103,7 +1172,7 @@ mod tests { assert!( parquet_array .as_::() - .typed_value_array() + .typed_value() .is_some() ); @@ -1117,8 +1186,8 @@ mod tests { .core_storage() .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant core storage"))?; - assert!(core_storage.typed_value_array().is_none()); - assert!(core_storage.value_array().is_some()); + assert!(core_storage.typed_value().is_none()); + assert!(core_storage.value().is_some()); Ok(()) } diff --git a/encodings/parquet-variant/src/lib.rs b/encodings/parquet-variant/src/lib.rs index 88125af04fd..5a82702b8f2 100644 --- a/encodings/parquet-variant/src/lib.rs +++ b/encodings/parquet-variant/src/lib.rs @@ -37,6 +37,9 @@ mod vtable; use std::sync::Arc; pub use array::ParquetVariantArrayExt; +pub use array::ParquetVariantArraySlotsExt; +pub use array::ParquetVariantSlots; +pub use array::ParquetVariantSlotsView; use vortex_array::session::ArraySessionExt; use vortex_arrow::ArrowSessionExt; pub use vortex_json::JsonToVariant; diff --git a/encodings/parquet-variant/src/operations.rs b/encodings/parquet-variant/src/operations.rs index 3cc3d2e5b4a..8517df27f22 100644 --- a/encodings/parquet-variant/src/operations.rs +++ b/encodings/parquet-variant/src/operations.rs @@ -27,7 +27,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; -use crate::ParquetVariantArrayExt; +use crate::ParquetVariantArraySlotsExt; use crate::vtable::ParquetVariant; impl OperationsVTable for ParquetVariant { @@ -46,7 +46,7 @@ impl OperationsVTable for ParquetVariant { } let metadata = array - .metadata_array() + .metadata() .execute_scalar(index, ctx)? .as_binary() .value() @@ -54,8 +54,8 @@ impl OperationsVTable for ParquetVariant { .vortex_expect("non-null metadata row must have binary value"); let inner = scalar_from_variant_storage( metadata.as_ref(), - array.value_array(), - array.typed_value_array(), + array.value(), + array.typed_value(), index, ctx, )?; @@ -372,6 +372,7 @@ fn parquet_variant_to_scalar(variant: PqVariant<'_, '_>) -> VortexResult #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::LazyLock; use arrow_array::Array as _; use arrow_array::ArrayRef as ArrowArrayRef; @@ -393,12 +394,20 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; + use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::ParquetVariant; use crate::ParquetVariantArrayExt; use crate::operations::parquet_variant_to_scalar; + static SESSION: LazyLock = LazyLock::new(|| { + let session = array_session(); + crate::initialize(&session); + session + }); + fn binary_view_array(values: &[&[u8]]) -> ArrowArrayRef { let mut builder = BinaryViewBuilder::new(); for value in values { @@ -411,7 +420,7 @@ mod tests { arrow_variant: &ArrowVariantArray, rows: impl IntoIterator, ) -> VortexResult<()> { - let vortex_arr = ParquetVariant::from_arrow_variant(arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(arrow_variant, &SESSION.arrow())?; for index in rows { let expected_inner = parquet_variant_to_scalar(arrow_variant.try_value(index)?)?; @@ -443,7 +452,7 @@ mod tests { )?; let arrow_variant = ArrowVariantArray::try_new(&null_struct)?; - let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; assert_eq!(vortex_arr.dtype(), &DType::Variant(Nullability::Nullable)); @@ -484,7 +493,7 @@ mod tests { )?; let arrow_variant = ArrowVariantArray::try_new(&null_struct)?; - let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; let present_variant_null = vortex_arr.execute_scalar(0, &mut array_session().create_execution_ctx())?; @@ -521,7 +530,7 @@ mod tests { )?; let arrow_variant = ArrowVariantArray::try_new(&null_struct)?; - let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; assert_eq!(vortex_arr.dtype(), &DType::Variant(Nullability::Nullable)); assert!( @@ -550,7 +559,7 @@ mod tests { builder.append_variant(PqVariant::from(2i32)); let arrow_variant = builder.build(); - let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; assert_eq!( vortex_arr.dtype(), @@ -671,7 +680,7 @@ mod tests { )?; let arrow_variant = ArrowVariantArray::try_new(&struct_array)?; - let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; let row0 = vortex_arr.execute_scalar(0, &mut array_session().create_execution_ctx())?; let row0 = row0.as_variant().value().unwrap().as_list(); @@ -763,7 +772,7 @@ mod tests { )?; let arrow_variant = ArrowVariantArray::try_new(&struct_array)?; - let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; + let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant, &SESSION.arrow())?; let object = vortex_arr.execute_scalar(0, &mut array_session().create_execution_ctx())?; let object = object.as_variant().value().unwrap().as_struct(); diff --git a/encodings/parquet-variant/src/validity.rs b/encodings/parquet-variant/src/validity.rs index b28c7102e22..c3b62c028f8 100644 --- a/encodings/parquet-variant/src/validity.rs +++ b/encodings/parquet-variant/src/validity.rs @@ -11,6 +11,6 @@ use crate::vtable::ParquetVariant; impl ValidityVTable for ParquetVariant { fn validity(array: ArrayView<'_, ParquetVariant>) -> VortexResult { - Ok(ParquetVariantArrayExt::validity(&array)) + Ok(ParquetVariantArrayExt::parquet_variant_validity(&array)) } } diff --git a/encodings/parquet-variant/src/vtable.rs b/encodings/parquet-variant/src/vtable.rs index 4fa3ea10690..297198df438 100644 --- a/encodings/parquet-variant/src/vtable.rs +++ b/encodings/parquet-variant/src/vtable.rs @@ -16,7 +16,6 @@ use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::serde::ArrayChildren; -use vortex_array::smallvec::smallvec; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_array::vtable::child_to_validity; @@ -30,13 +29,8 @@ use vortex_proto::dtype as pb; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::array::METADATA_SLOT; -use crate::array::NUM_SLOTS; -use crate::array::ParquetVariantArrayExt; -use crate::array::SLOT_NAMES; -use crate::array::TYPED_VALUE_SLOT; -use crate::array::VALIDITY_SLOT; -use crate::array::VALUE_SLOT; +use crate::array::ParquetVariantArraySlotsExt; +use crate::array::ParquetVariantSlots; use crate::array::core_storage_without_typed_value; use crate::array::logical_shredded_from_parquet_typed_value; @@ -103,16 +97,20 @@ impl VTable for ParquetVariant { slots: &[Option], ) -> VortexResult<()> { vortex_ensure!( - slots.len() == NUM_SLOTS, - "ParquetVariantArray expects {NUM_SLOTS} slots, got {}", + slots.len() == ParquetVariantSlots::COUNT, + "ParquetVariantArray expects {} slots, got {}", + ParquetVariantSlots::COUNT, slots.len() ); - let validity = child_to_validity(slots[VALIDITY_SLOT].as_ref(), dtype.nullability()); - let metadata = slots[METADATA_SLOT] + let validity = child_to_validity( + slots[ParquetVariantSlots::VALIDITY].as_ref(), + dtype.nullability(), + ); + let metadata = slots[ParquetVariantSlots::METADATA] .as_ref() .ok_or_else(|| vortex_err!("ParquetVariantArray metadata slot"))?; - let value = slots[VALUE_SLOT].as_ref(); - let typed_value = slots[TYPED_VALUE_SLOT].as_ref(); + let value = slots[ParquetVariantSlots::VALUE].as_ref(); + let typed_value = slots[ParquetVariantSlots::TYPED_VALUE].as_ref(); vortex_ensure!( matches!(dtype, DType::Variant(_)), @@ -180,7 +178,7 @@ impl VTable for ParquetVariant { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + ParquetVariantSlots::NAMES[idx].to_string() } fn serialize( @@ -188,14 +186,14 @@ impl VTable for ParquetVariant { _session: &VortexSession, ) -> VortexResult>> { let typed_value_dtype = array - .typed_value_array() + .typed_value() .map(|tv| tv.dtype().try_into()) .transpose()?; Ok(Some( ParquetVariantMetadataProto { - has_value: array.value_array().is_some(), + has_value: array.value().is_some(), typed_value_dtype, - value_nullable: array.value_array().is_some_and(|v| v.dtype().is_nullable()), + value_nullable: array.value().is_some_and(|v| v.dtype().is_nullable()), } .encode_to_vec(), )) @@ -265,21 +263,22 @@ impl VTable for ParquetVariant { None }; - let slots = smallvec![ - validity_to_child(&validity, len), - Some(variant_metadata), + let slots = ParquetVariantSlots { + validity: validity_to_child(&validity, len), + metadata: variant_metadata, value, typed_value, - ]; + } + .into_slots(); Ok(ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData).with_slots(slots)) } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { let shredded = array - .typed_value_array() + .typed_value() .cloned() .map(|typed_value| { - logical_shredded_from_parquet_typed_value(array.metadata_array(), typed_value, ctx) + logical_shredded_from_parquet_typed_value(array.metadata(), typed_value, ctx) }) .transpose()?; let core_storage = core_storage_without_typed_value(&array)?; @@ -309,24 +308,30 @@ mod tests { use vortex_array::Canonical; use vortex_array::EqMode; use vortex_array::IntoArray; - use vortex_array::VTable; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::VariantArray; - use vortex_array::arrays::variant::VariantArrayExt; + use vortex_array::arrays::variant::VariantArraySlotsExt; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; + use vortex_array::dtype::session::DTypeSessionExt; use vortex_array::serde::SerializeOptions; use vortex_array::serde::SerializedArray; use vortex_array::session::ArraySessionExt; use vortex_array::stream::ArrayStreamExt; use vortex_array::validity::Validity; + use vortex_arrow::ArrowSessionExt; use vortex_buffer::BitBuffer; use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; + use vortex_edition::ComponentKind; + use vortex_edition::Edition; + use vortex_edition::EditionId; + use vortex_edition::EditionInclusion; + use vortex_edition::EditionSessionExt; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_file::OpenOptionsSessionExt; @@ -335,11 +340,12 @@ mod tests { use vortex_layout::LayoutStrategy; use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; use vortex_layout::session::LayoutSession; + use vortex_layout::session::LayoutSessionExt; use vortex_session::VortexSession; use vortex_session::registry::ReadContext; use crate::ParquetVariant; - use crate::array::ParquetVariantArrayExt; + use crate::array::ParquetVariantArraySlotsExt; static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session(); @@ -385,26 +391,83 @@ mod tests { None, )?; - ParquetVariant::from_arrow_variant(&ArrowVariantArray::try_new(&arrow_storage)?) + ParquetVariant::from_arrow_variant( + &ArrowVariantArray::try_new(&arrow_storage)?, + &SESSION.arrow(), + ) } #[fixture] - fn parquet_variant_file_session() -> VortexSession { + fn parquet_variant_file_session() -> VortexResult { + const TEST_EDITION: EditionId = EditionId::new("test", 2026, 7, 0); + let session = vortex_array::array_session() .with::() .with::(); vortex_file::register_default_encodings(&session); session.arrays().register(ParquetVariant); + let editions = session.editions(); + editions + .declare_edition(Edition { + id: TEST_EDITION, + min_library_version: None, + }) + .map_err(|error| vortex_err!("{error}"))?; + let component_ids = [ + ( + ComponentKind::Array, + session + .arrays() + .registry() + .read(|map| map.keys().copied().collect::>()), + ), + ( + ComponentKind::Layout, + session + .layouts() + .registry() + .read(|map| map.keys().copied().collect::>()), + ), + ( + ComponentKind::DType, + session + .dtypes() + .registry() + .read(|map| map.keys().copied().collect::>()), + ), + ]; + for (kind, ids) in component_ids { + for id in ids { + editions + .declare_inclusion(EditionInclusion::new(kind, &id, TEST_EDITION)) + .map_err(|error| vortex_err!("{error}"))?; + } + } + for id in [ + "vortex.bounded_max", + "vortex.bounded_min", + "vortex.max", + "vortex.min", + "vortex.nan_count", + "vortex.null_count", + ] { + editions + .declare_inclusion(EditionInclusion::new( + ComponentKind::Aggregate, + id, + TEST_EDITION, + )) + .map_err(|error| vortex_err!("{error}"))?; + } session + .enable_edition(TEST_EDITION) + .map_err(|error| vortex_err!("{error}"))?; + Ok(session) } #[fixture] fn write_strategy() -> Arc { - let mut allowed = vortex_file::ALLOWED_ENCODINGS.clone(); - allowed.insert(ParquetVariant.id()); - vortex_file::WriteStrategyBuilder::default() - .with_allow_encodings(allowed) - .build() + vortex_file::WriteStrategyBuilder::default().build() } #[test] @@ -416,7 +479,7 @@ mod tests { let parquet_variant = ParquetVariant::try_new(Validity::NonNullable, metadata, None, Some(typed_value))?; - assert!(parquet_variant.typed_value_array().is_some()); + assert!(parquet_variant.typed_value().is_some()); let mut ctx = SESSION.create_execution_ctx(); let Canonical::Variant(variant) = parquet_variant @@ -430,7 +493,7 @@ mod tests { .core_storage() .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant core storage"))?; - assert!(core_storage.typed_value_array().is_none()); + assert!(core_storage.typed_value().is_none()); let shredded = variant .shredded() .ok_or_else(|| vortex_err!("expected canonical shredded child"))?; @@ -452,9 +515,10 @@ mod tests { #[tokio::test] async fn test_file_roundtrip_typed_value_variant_with_statistics( #[from(typed_value_variant_array)] expected: VortexResult, - parquet_variant_file_session: VortexSession, + parquet_variant_file_session: VortexResult, ) -> VortexResult<()> { let expected = expected?; + let parquet_variant_file_session = parquet_variant_file_session?; let mut bytes = ByteBufferMut::empty(); parquet_variant_file_session @@ -479,10 +543,11 @@ mod tests { #[tokio::test] async fn test_file_roundtrip_typed_value_variant_with_zoned_strategy( #[from(typed_value_variant_array)] expected: VortexResult, - parquet_variant_file_session: VortexSession, + parquet_variant_file_session: VortexResult, write_strategy: Arc, ) -> VortexResult<()> { let expected = expected?; + let parquet_variant_file_session = parquet_variant_file_session?; let mut bytes = ByteBufferMut::empty(); parquet_variant_file_session @@ -533,7 +598,7 @@ mod tests { .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant array"))?; let typed = decoded_pv - .typed_value_array() + .typed_value() .ok_or_else(|| vortex_err!("expected typed_value child"))?; assert_eq!(typed.dtype(), &DType::Variant(Nullability::NonNullable)); Ok(()) @@ -555,8 +620,8 @@ mod tests { let decoded_pv = decoded .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant array"))?; - assert!(decoded_pv.value_array().is_some()); - assert!(decoded_pv.typed_value_array().is_none()); + assert!(decoded_pv.value().is_some()); + assert!(decoded_pv.typed_value().is_none()); Ok(()) } @@ -580,7 +645,7 @@ mod tests { .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant array"))?; let typed = decoded_pv - .typed_value_array() + .typed_value() .ok_or_else(|| vortex_err!("expected typed_value child"))?; assert_eq!( typed.dtype(), diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index d47a935d845..44a6a8e4045 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -29,6 +29,8 @@ use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::IntoArray; +use vortex_array::TypedArrayRef; +use vortex_array::array_slots; use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::buffer::BufferHandle; @@ -37,7 +39,6 @@ use vortex_array::dtype::PType; use vortex_array::dtype::half; use vortex_array::scalar::Scalar; use vortex_array::serde::ArrayChildren; -use vortex_array::smallvec::smallvec; use vortex_array::validity::Validity; use vortex_array::vtable::OperationsVTable; use vortex_array::vtable::VTable; @@ -46,7 +47,6 @@ use vortex_array::vtable::child_to_validity; use vortex_array::vtable::validity_to_child; use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; -use vortex_buffer::ByteBufferMut; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -139,7 +139,10 @@ impl VTable for Pco { len: usize, slots: &[Option], ) -> VortexResult<()> { - let validity = child_to_validity(slots[0].as_ref(), dtype.nullability()); + let validity = child_to_validity( + PcoSlotsView::from_slots(slots).validity, + dtype.nullability(), + ); data.validate(dtype, len, &validity) } @@ -229,20 +232,23 @@ impl VTable for Pco { .sum::(); vortex_ensure!(pages.len() == expected_n_pages); - let slots = smallvec![validity_to_child(&validity, len)]; - let data = PcoData::new(chunk_metas, pages, dtype.as_ptype(), metadata, len); + let slots = PcoSlots { + validity: validity_to_child(&validity, len), + } + .into_slots(); + // SAFETY: `Array::try_from_parts`, which consumes these parts, validates the data before + // publishing the array. + let data = + unsafe { PcoData::new_unchecked(chunk_metas, pages, dtype.as_ptype(), metadata, len) }; Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + PcoSlots::NAMES[idx].to_string() } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - let unsliced_validity = child_to_validity( - array.as_ref().slots()[0].as_ref(), - array.dtype().nullability(), - ); + let unsliced_validity = array.unsliced_validity(); Ok(ExecutionResult::done( array .data() @@ -308,17 +314,29 @@ pub(crate) fn vortex_err_from_pco(err: PcoError) -> VortexError { pub struct Pco; impl Pco { - pub(crate) fn try_new( - dtype: DType, - data: PcoData, - validity: Validity, - ) -> VortexResult { + /// Constructs a Pco array after validating its data and validity invariants. + pub fn try_new(dtype: DType, data: PcoData, validity: Validity) -> VortexResult { let len = data.len(); data.validate(&dtype, len, &validity)?; - let slots = smallvec![validity_to_child(&validity, data.unsliced_n_rows())]; - Ok(unsafe { + // SAFETY: `validate` checked the dtype, length, validity, metadata, and buffer invariants. + Ok(unsafe { Self::new_unchecked(dtype, data, validity) }) + } + + /// Constructs a Pco array without validating its data and validity invariants. + /// + /// # Safety + /// + /// The caller must ensure that [`PcoData::validate`] would succeed for `dtype`, `data.len()`, + /// and `validity`. + pub unsafe fn new_unchecked(dtype: DType, data: PcoData, validity: Validity) -> PcoArray { + let len = data.len(); + let slots = PcoSlots { + validity: validity_to_child(&validity, data.unsliced_n_rows()), + } + .into_slots(); + unsafe { Array::from_parts_unchecked(ArrayParts::new(Pco, dtype, len, data).with_slots(slots)) - }) + } } /// Compress a primitive array using pcodec. @@ -335,9 +353,24 @@ impl Pco { } } -/// The validity bitmap indicating which elements are non-null. -pub(super) const NUM_SLOTS: usize = 1; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"]; +#[array_slots(Pco)] +pub struct PcoSlots { + /// The validity bitmap indicating which elements are non-null. + #[slot(0)] + pub validity: Option, +} + +/// Additional typed accessors for Pco arrays. +pub trait PcoArrayExt: PcoArraySlotsExt { + /// Reconstruct the unsliced [`Validity`] from the validity slot. + fn unsliced_validity(&self) -> Validity { + child_to_validity( + self.as_ref().slots()[PcoSlots::VALIDITY].as_ref(), + self.as_ref().dtype().nullability(), + ) + } +} +impl> PcoArrayExt for T {} #[derive(Clone, Debug)] /// Encoding-specific data for a [`PcoArray`]. @@ -413,11 +446,52 @@ impl PcoData { .sum::(), "page count does not match metadata" ); + + let mut n_values = 0usize; + for (chunk_idx, chunk) in self.metadata.chunks.iter().enumerate() { + let mut chunk_n_values = 0usize; + for page in &chunk.pages { + let page_n_values = page.n_values as usize; + vortex_ensure!( + page_n_values != 0, + "Pco chunk {chunk_idx} contains an empty page" + ); + chunk_n_values = chunk_n_values.checked_add(page_n_values).ok_or_else(|| { + vortex_err!("Pco chunk {chunk_idx} value count overflows usize") + })?; + } + vortex_ensure!( + chunk_n_values <= VALUES_PER_CHUNK, + "Pco chunk {chunk_idx} contains {chunk_n_values} values, exceeding the maximum of {VALUES_PER_CHUNK}" + ); + n_values = n_values + .checked_add(chunk_n_values) + .ok_or_else(|| vortex_err!("Pco value count overflows usize"))?; + } + vortex_ensure!( + n_values <= self.unsliced_n_rows, + "Pco contains {n_values} values for only {} rows", + self.unsliced_n_rows + ); + if validity.definitely_no_nulls() { + vortex_ensure!( + n_values == self.unsliced_n_rows, + "Pco contains {n_values} values for {} non-null rows", + self.unsliced_n_rows + ); + } else if validity.definitely_all_null() { + vortex_ensure!(n_values == 0, "Pco contains values for an all-null array"); + } Ok(()) } - /// Construct unsliced Pco data from chunk metadata, pages, and serialized metadata. - pub fn new( + /// Constructs unsliced Pco data without validating its metadata and buffers. + /// + /// # Safety + /// + /// The returned data must pass [`Self::validate`] before it is decompressed or published as + /// part of an array. + pub unsafe fn new_unchecked( chunk_metas: Vec, pages: Vec, ptype: PType, @@ -493,17 +567,17 @@ impl PcoData { } ); - let mut chunk_meta_buffer = ByteBufferMut::with_capacity(cc.meta_size_hint()); + let mut chunk_meta_buffer = Vec::with_capacity(cc.meta_size_hint()); cc.write_meta(&mut chunk_meta_buffer) .map_err(vortex_err_from_pco)?; - chunk_meta_buffers.push(chunk_meta_buffer.freeze()); + chunk_meta_buffers.push(ByteBuffer::from(chunk_meta_buffer)); let mut page_infos = vec![]; for (page_idx, page_n_values) in cc.n_per_page().into_iter().enumerate() { - let mut page = ByteBufferMut::with_capacity(cc.page_size_hint(page_idx)); + let mut page = Vec::with_capacity(cc.page_size_hint(page_idx)); cc.write_page(page_idx, &mut page) .map_err(vortex_err_from_pco)?; - page_buffers.push(page.freeze()); + page_buffers.push(ByteBuffer::from(page)); page_infos.push(PcoPageInfo { n_values: u32::try_from(page_n_values)?, }); @@ -515,13 +589,17 @@ impl PcoData { header, chunks: chunk_infos, }; - Ok(PcoData::new( - chunk_meta_buffers, - page_buffers, - parray.dtype().as_ptype(), - metadata, - parray.len(), - )) + // SAFETY: The compressor produced matching chunk metadata and pages from `parray`, with + // the same ptype, logical length, and valid-value count. + Ok(unsafe { + PcoData::new_unchecked( + chunk_meta_buffers, + page_buffers, + parray.dtype().as_ptype(), + metadata, + parray.len(), + ) + }) } /// Downcast and compress an array into Pco data. @@ -545,7 +623,7 @@ impl PcoData { } /// Decompress this Pco data into a primitive array. - pub fn decompress( + pub(crate) fn decompress( &self, unsliced_validity: &Validity, ctx: &mut ExecutionCtx, @@ -586,9 +664,12 @@ impl PcoData { // may exceed the bounds of the slice, so we need to slice later. let (fd, _) = FileDecompressor::new(self.metadata.header.as_slice()).map_err(vortex_err_from_pco)?; - let mut decompressed_values = BufferMut::::with_capacity(slice_n_values); + let mut decompressed_values = + BufferMut::::with_capacity(slice_n_values.min(VALUES_PER_CHUNK)); + // Validation bounds every page-count prefix and selected-page subtotal by + // `unsliced_n_rows`, so the arithmetic below cannot overflow. let mut page_idx = 0; - let mut page_value_start = 0; + let mut page_value_start = 0usize; let mut n_skipped_values = 0; for (chunk_info, chunk_meta) in self.metadata.chunks.iter().zip(&self.chunk_metas) { // lazily initialize chunk decompressor @@ -609,7 +690,11 @@ impl PcoData { unsafe { decompressed_values.set_len(new_len); } - let page: &[u8] = self.pages[page_idx].as_ref(); + let page: &[u8] = self + .pages + .get(page_idx) + .ok_or_else(|| vortex_err!("Missing Pco page {page_idx}"))? + .as_ref(); let mut cd = match chunk_decompressor.take() { Some(d) => d, @@ -638,10 +723,18 @@ impl PcoData { } // Slice only the values requested. + // Skipped pages end before `slice_value_start`, and the resulting stop is at most + // `slice_value_stop`, which is bounded by `unsliced_n_rows`. let value_offset = slice_value_start - n_skipped_values; + let value_stop = value_offset + slice_n_values; + vortex_ensure!( + value_stop <= decompressed_values.len(), + "Pco contains {} decompressed values, but the requested range ends at {value_stop}", + decompressed_values.len() + ); Ok(decompressed_values .freeze() - .slice(value_offset..value_offset + slice_n_values) + .slice(value_offset..value_stop) .into_byte_buffer()) } @@ -678,9 +771,9 @@ impl PcoData { impl ValidityVTable for Pco { fn validity(array: ArrayView<'_, Pco>) -> VortexResult { - let unsliced_validity = - child_to_validity(array.slots()[0].as_ref(), array.dtype().nullability()); - unsliced_validity.slice(array.slice_start()..array.slice_stop()) + array + .unsliced_validity() + .slice(array.slice_start()..array.slice_stop()) } } @@ -690,8 +783,7 @@ impl OperationsVTable for Pco { index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let unsliced_validity = - child_to_validity(array.slots()[0].as_ref(), array.dtype().nullability()); + let unsliced_validity = array.unsliced_validity(); array ._slice(index, index + 1) .decompress(&unsliced_validity, ctx)? @@ -709,7 +801,9 @@ mod tests { use vortex_array::assert_arrays_eq; use vortex_array::validity::Validity; use vortex_buffer::buffer; + use vortex_error::VortexResult; + use super::VALUES_PER_CHUNK; use crate::Pco; #[test] @@ -741,4 +835,36 @@ mod tests { .into_array(); assert_arrays_eq!(sliced, expected, &mut ctx); } + + #[test] + fn test_decompress_bounds_initial_allocation() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let values = PrimitiveArray::from_iter([42u32]); + let pco = Pco::from_primitive(values.as_view(), 0, 128, &mut ctx)?; + let mut data = pco.data().clone(); + + // Simulate a corrupt logical length whose validity claims far more values than the + // encoded pages contain. This must not be used directly as an allocation size. + data.unsliced_n_rows = usize::MAX; + data.slice_stop = usize::MAX; + + let result = data.decompress_values_typed::(&Validity::NonNullable, &mut ctx); + assert!(result.is_err()); + Ok(()) + } + + #[test] + fn test_validate_rejects_oversized_chunk() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let values = PrimitiveArray::from_iter([42u32]); + let pco = Pco::from_primitive(values.as_view(), 0, 128, &mut ctx)?; + let mut data = pco.data().clone(); + data.metadata.chunks[0].pages[0].n_values = u32::try_from(VALUES_PER_CHUNK + 1)?; + + assert!( + data.validate(pco.dtype(), pco.len(), &Validity::NonNullable) + .is_err() + ); + Ok(()) + } } diff --git a/encodings/pco/src/compute/cast.rs b/encodings/pco/src/compute/cast.rs index c59e205e2e1..6ac0d7de3b9 100644 --- a/encodings/pco/src/compute/cast.rs +++ b/encodings/pco/src/compute/cast.rs @@ -6,10 +6,10 @@ use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::dtype::DType; use vortex_array::scalar_fn::fns::cast::CastReduce; -use vortex_array::vtable::child_to_validity; use vortex_error::VortexResult; use crate::Pco; +use crate::PcoArrayExt; use crate::PcoData; impl CastReduce for Pco { @@ -28,21 +28,24 @@ impl CastReduce for Pco { return Ok(None); } - let unsliced_validity = - child_to_validity(array.slots()[0].as_ref(), array.dtype().nullability()); + let unsliced_validity = array.unsliced_validity(); let Some(new_validity) = unsliced_validity.trivially_cast_nullability(dtype.nullability(), array.len())? else { return Ok(None); }; - let data = PcoData::new( - array.chunk_metas.clone(), - array.pages.clone(), - dtype.as_ptype(), - array.metadata.clone(), - array.unsliced_n_rows(), - ) + // SAFETY: The buffers and metadata come from a validated Pco array, the primitive type is + // unchanged, and `Pco::try_new` validates the adjusted nullability and slice below. + let data = unsafe { + PcoData::new_unchecked( + array.chunk_metas.clone(), + array.pages.clone(), + dtype.as_ptype(), + array.metadata.clone(), + array.unsliced_n_rows(), + ) + } ._slice(array.slice_start(), array.slice_stop()); Ok(Some( diff --git a/encodings/pco/src/slice.rs b/encodings/pco/src/slice.rs index 30dd1ee3e93..316ea835aa1 100644 --- a/encodings/pco/src/slice.rs +++ b/encodings/pco/src/slice.rs @@ -7,15 +7,14 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::arrays::slice::SliceReduce; -use vortex_array::vtable::child_to_validity; use vortex_error::VortexResult; use crate::Pco; +use crate::PcoArrayExt; impl SliceReduce for Pco { fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { - let unsliced_validity = - child_to_validity(array.slots()[0].as_ref(), array.dtype().nullability()); + let unsliced_validity = array.unsliced_validity(); Ok(Some( Pco::try_new( array.dtype().clone(), diff --git a/encodings/pco/src/tests.rs b/encodings/pco/src/tests.rs index 5cc771601f0..8a0e1083b16 100644 --- a/encodings/pco/src/tests.rs +++ b/encodings/pco/src/tests.rs @@ -18,7 +18,6 @@ use vortex_array::serde::SerializeOptions; use vortex_array::serde::SerializedArray; use vortex_array::session::ArraySessionExt; use vortex_array::validity::Validity; -use vortex_array::vtable::child_to_validity; use vortex_arrow::ArrowSessionExt; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -28,6 +27,7 @@ use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::ReadContext; +use crate::PcoArrayExt; use crate::PcoData; static SESSION: LazyLock = LazyLock::new(|| { @@ -47,10 +47,7 @@ fn test_compress_decompress() { assert!(compressed.pages.len() < array.into_array().nbytes() as usize); // check full decompression works - let unsliced_validity = child_to_validity( - compressed.as_ref().slots()[0].as_ref(), - compressed.dtype().nullability(), - ); + let unsliced_validity = compressed.unsliced_validity(); let decompressed = compressed.decompress(&unsliced_validity, &mut ctx).unwrap(); assert_arrays_eq!(decompressed, PrimitiveArray::from_iter(data), &mut ctx); @@ -82,10 +79,7 @@ fn test_compress_decompress_small() { let expected = array.into_array(); assert_arrays_eq!(compressed, expected, &mut ctx); - let unsliced_validity = child_to_validity( - compressed.as_ref().slots()[0].as_ref(), - compressed.dtype().nullability(), - ); + let unsliced_validity = compressed.unsliced_validity(); let decompressed = compressed.decompress(&unsliced_validity, &mut ctx).unwrap(); assert_arrays_eq!(decompressed, expected, &mut ctx); } @@ -96,10 +90,7 @@ fn test_empty() { let data: Vec = vec![]; let array = PrimitiveArray::from_iter(data.clone()); let compressed = Pco::from_primitive(array.as_view(), 3, 100, &mut ctx).unwrap(); - let unsliced_validity = child_to_validity( - compressed.as_ref().slots()[0].as_ref(), - compressed.dtype().nullability(), - ); + let unsliced_validity = compressed.unsliced_validity(); let primitive = compressed.decompress(&unsliced_validity, &mut ctx).unwrap(); assert_arrays_eq!(primitive, PrimitiveArray::from_iter(data), &mut ctx); } diff --git a/encodings/runend/Cargo.toml b/encodings/runend/Cargo.toml index 4d36d0c1548..5e607b78cc6 100644 --- a/encodings/runend/Cargo.toml +++ b/encodings/runend/Cargo.toml @@ -29,7 +29,9 @@ workspace = true [dev-dependencies] divan = { workspace = true } +insta = { workspace = true } itertools = { workspace = true } +mimalloc = { workspace = true } rand = { workspace = true } rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } diff --git a/encodings/runend/benches/run_end_compress.rs b/encodings/runend/benches/run_end_compress.rs index cd241ad42c6..6e77dd76efc 100644 --- a/encodings/runend/benches/run_end_compress.rs +++ b/encodings/runend/benches/run_end_compress.rs @@ -7,6 +7,7 @@ use std::sync::LazyLock; use divan::Bencher; use itertools::repeat_n; +use mimalloc::MiMalloc; use vortex_array::IntoArray; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; @@ -19,6 +20,12 @@ use vortex_runend::RunEnd; use vortex_runend::compress::runend_encode; use vortex_session::VortexSession; +// `runend_encode` allocates its output buffers inside the timed region, so route allocation +// through vendored mimalloc to keep glibc malloc (which varies across runner images) out of +// the measured trace. +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + fn main() { divan::main(); } @@ -33,15 +40,15 @@ const BENCH_ARGS: &[(usize, usize)] = &[ (1000, 4), (1000, 16), (1000, 256), + (4_000, 4), + (4_000, 16), + (4_000, 256), + (4_000, 1024), (10_000, 4), (10_000, 16), (10_000, 256), (10_000, 1024), - (100_000, 4), - (100_000, 16), - (100_000, 256), - (100_000, 1024), - (100_000, 4096), + (10_000, 4096), ]; #[divan::bench(args = BENCH_ARGS)] diff --git a/encodings/runend/benches/run_end_null_count.rs b/encodings/runend/benches/run_end_null_count.rs index f214889daae..46d8710c38c 100644 --- a/encodings/runend/benches/run_end_null_count.rs +++ b/encodings/runend/benches/run_end_null_count.rs @@ -35,18 +35,18 @@ const BENCH_ARGS: &[(usize, usize, f64)] = &[ (10_000, 1024, 0.01), (10_000, 1024, 0.1), (10_000, 1024, 0.5), - (100_000, 4, 0.01), - (100_000, 4, 0.1), - (100_000, 4, 0.5), - (100_000, 16, 0.01), - (100_000, 16, 0.1), - (100_000, 16, 0.5), - (100_000, 256, 0.01), - (100_000, 256, 0.1), - (100_000, 256, 0.5), - (100_000, 1024, 0.01), - (100_000, 1024, 0.1), - (100_000, 1024, 0.5), + (32_000, 4, 0.01), + (32_000, 4, 0.1), + (32_000, 4, 0.5), + (32_000, 16, 0.01), + (32_000, 16, 0.1), + (32_000, 16, 0.5), + (32_000, 256, 0.01), + (32_000, 256, 0.1), + (32_000, 256, 0.5), + (32_000, 1024, 0.01), + (32_000, 1024, 0.1), + (32_000, 1024, 0.5), ]; static SESSION: LazyLock = LazyLock::new(|| { diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs index 2a9659efdd8..3a0b10c5c4a 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -21,15 +21,19 @@ use vortex_array::ExecutionResult; use vortex_array::IntoArray; use vortex_array::TypedArrayRef; use vortex_array::VortexSessionExecute; +use vortex_array::array_slots; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::ListViewArray; use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::legacy_session; use vortex_array::serde::ArrayChildren; -use vortex_array::smallvec::smallvec; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTable; @@ -41,6 +45,7 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::compress::runend_decode_decimal; use crate::compress::runend_decode_primitive; use crate::compress::runend_decode_varbinview; use crate::compress::runend_encode; @@ -93,12 +98,9 @@ impl VTable for RunEnd { len: usize, slots: &[Option], ) -> VortexResult<()> { - let ends = slots[ENDS_SLOT] - .as_ref() - .vortex_expect("RunEndArray ends slot"); - let values = slots[VALUES_SLOT] - .as_ref() - .vortex_expect("RunEndArray values slot"); + let run_end_slots = RunEndSlotsView::from_slots(slots); + let ends = run_end_slots.ends; + let values = run_end_slots.values; // TODO(ctx): trait fixes - VTable::validate has a fixed signature. let mut ctx = legacy_session().create_execution_ctx(); RunEndData::validate_parts(ends, values, data.offset, len, &mut ctx)?; @@ -162,13 +164,13 @@ impl VTable for RunEnd { let values = children.get(1, dtype, runs)?; let offset = usize::try_from(metadata.offset).vortex_expect("Offset must be a valid usize"); - let slots = smallvec![Some(ends), Some(values)]; + let slots = RunEndSlots { ends, values }.into_slots(); let data = RunEndData::new(offset); Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + RunEndSlots::NAMES[idx].to_string() } fn reduce_parent( @@ -184,12 +186,15 @@ impl VTable for RunEnd { } } -/// The run-end positions marking where each run terminates. -pub(super) const ENDS_SLOT: usize = 0; -/// The values for each run. -pub(super) const VALUES_SLOT: usize = 1; -pub(super) const NUM_SLOTS: usize = 2; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["ends", "values"]; +#[array_slots(RunEnd)] +pub struct RunEndSlots { + /// The run-end positions marking where each run terminates. + #[slot(0)] + pub ends: ArrayRef, + /// The values for each run. + #[slot(1)] + pub values: ArrayRef, +} #[derive(Clone, Debug)] pub struct RunEndData { @@ -208,23 +213,11 @@ pub struct RunEndDataParts { pub offset: usize, } -pub trait RunEndArrayExt: TypedArrayRef { +pub trait RunEndArrayExt: RunEndArraySlotsExt { fn offset(&self) -> usize { self.offset } - fn ends(&self) -> &ArrayRef { - self.as_ref().slots()[ENDS_SLOT] - .as_ref() - .vortex_expect("RunEndArray ends slot") - } - - fn values(&self) -> &ArrayRef { - self.as_ref().slots()[VALUES_SLOT] - .as_ref() - .vortex_expect("RunEndArray values slot") - } - fn dtype(&self) -> &DType { self.values().dtype() } @@ -255,7 +248,7 @@ impl RunEnd { length: usize, ) -> RunEndArray { let dtype = values.dtype().clone(); - let slots = smallvec![Some(ends), Some(values)]; + let slots = RunEndSlots { ends, values }.into_slots(); let data = unsafe { RunEndData::new_unchecked(offset) }; unsafe { Array::from_parts_unchecked( @@ -273,7 +266,7 @@ impl RunEnd { let len = RunEndData::logical_len_from_ends(&ends, ctx)?; RunEndData::validate_parts(&ends, &values, 0, len, ctx)?; let dtype = values.dtype().clone(); - let slots = smallvec![Some(ends), Some(values)]; + let slots = RunEndSlots { ends, values }.into_slots(); let data = RunEndData::new(0); Array::try_from_parts(ArrayParts::new(RunEnd, dtype, len, data).with_slots(slots)) } @@ -288,7 +281,7 @@ impl RunEnd { ) -> VortexResult { RunEndData::validate_parts(&ends, &values, offset, length, ctx)?; let dtype = values.dtype().clone(); - let slots = smallvec![Some(ends), Some(values)]; + let slots = RunEndSlots { ends, values }.into_slots(); let data = RunEndData::new(offset); Array::try_from_parts(ArrayParts::new(RunEnd, dtype, length, data).with_slots(slots)) } @@ -305,7 +298,7 @@ impl RunEnd { let ends = ends.into_array(); let len = array.len(); let dtype = values.dtype().clone(); - let slots = smallvec![Some(ends), Some(values)]; + let slots = RunEndSlots { ends, values }.into_slots(); let data = unsafe { RunEndData::new_unchecked(0) }; Array::try_from_parts(ArrayParts::new(RunEnd, dtype, len, data).with_slots(slots)) } else { @@ -498,6 +491,13 @@ pub(super) fn run_end_canonicalize( let pvalues = array.values().clone().execute_as("values", ctx)?; runend_decode_primitive(pends, pvalues, array.offset(), array.len(), ctx)?.into_array() } + DType::Decimal(..) => { + let values = array + .values() + .clone() + .execute_as::("values", ctx)?; + runend_decode_decimal(pends, values, array.offset(), array.len(), ctx)?.into_array() + } DType::Utf8(_) | DType::Binary(_) => { let values = array .values() @@ -505,23 +505,74 @@ pub(super) fn run_end_canonicalize( .execute_as::("values", ctx)?; runend_decode_varbinview(pends, values, array.offset(), array.len(), ctx)?.into_array() } + DType::List(..) => { + let values = array + .values() + .clone() + .execute_as::("values", ctx)?; + runend_decode_listview(pends, values, array.offset(), array.len())?.into_array() + } _ => vortex_bail!("Unsupported RunEnd value type: {}", array.dtype()), }) } +fn runend_decode_listview( + ends: PrimitiveArray, + values: ListViewArray, + offset: usize, + length: usize, +) -> VortexResult { + let validity = match values.validity()? { + Validity::NonNullable => Validity::NonNullable, + Validity::AllValid => Validity::AllValid, + Validity::AllInvalid => Validity::AllInvalid, + Validity::Array(validity) => Validity::Array(unsafe { + RunEnd::new_unchecked(ends.clone().into_array(), validity, offset, length).into_array() + }), + }; + + // SAFETY: the `RunEndArray`s re-express valid per-run ListView metadata over the logical output + // length. The original `elements` child is reused, so every view still points at a valid range. + Ok(unsafe { + ListViewArray::new_unchecked( + values.elements().clone(), + RunEnd::new_unchecked( + ends.clone().into_array(), + values.offsets().clone(), + offset, + length, + ) + .into_array(), + RunEnd::new_unchecked(ends.into_array(), values.sizes().clone(), offset, length) + .into_array(), + validity, + ) + }) +} + #[cfg(test)] mod tests { + use std::sync::Arc; use std::sync::LazyLock; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; + use vortex_array::arrays::DecimalArray; use vortex_array::arrays::DictArray; + use vortex_array::arrays::ListArray; + use vortex_array::arrays::ListViewArray; use vortex_array::arrays::VarBinViewArray; + use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::assert_arrays_eq; + use vortex_array::builders::VarBinBuilder; use vortex_array::dtype::DType; + use vortex_array::dtype::DecimalDType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; + use vortex_array::dtype::i256; + use vortex_array::validity::Validity; use vortex_buffer::buffer; + use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::RunEnd; @@ -556,17 +607,210 @@ mod tests { #[test] fn test_runend_utf8() { let mut ctx = SESSION.create_execution_ctx(); - let values = VarBinViewArray::from_iter_str(["a", "b", "c"]).into_array(); + let values = + VarBinViewArray::from_iter_nullable_str([Some("a"), None, Some("c")]).into_array(); let arr = RunEnd::new(buffer![2u32, 5, 10].into_array(), values, &mut ctx); assert_eq!(arr.len(), 10); - assert_eq!(arr.dtype(), &DType::Utf8(Nullability::NonNullable)); + assert_eq!(arr.dtype(), &DType::Utf8(Nullability::Nullable)); + + let expected = VarBinViewArray::from_iter_nullable_str([ + Some("a"), + Some("a"), + None, + None, + None, + Some("c"), + Some("c"), + Some("c"), + Some("c"), + Some("c"), + ]) + .into_array(); + let mut builder = VarBinBuilder::::with_capacity_in( + arr.dtype().clone(), + arr.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + arr.append_to_builder(&mut builder, &mut ctx).unwrap(); + assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx); + assert_arrays_eq!(arr.into_array(), expected, &mut ctx); + } - let expected = - VarBinViewArray::from_iter_str(["a", "a", "b", "b", "b", "c", "c", "c", "c", "c"]) + #[test] + fn test_runend_decimal() { + let mut ctx = SESSION.create_execution_ctx(); + let decimal_dtype = DecimalDType::new(10, 2); + let values = DecimalArray::from_iter([12345i64, 67890, -12300], decimal_dtype).into_array(); + let arr = RunEnd::new(buffer![2u32, 5, 10].into_array(), values, &mut ctx); + assert_eq!(arr.len(), 10); + assert_eq!( + arr.dtype(), + &DType::Decimal(decimal_dtype, Nullability::NonNullable) + ); + + let expected = DecimalArray::from_iter( + [ + 12345i64, 12345, 67890, 67890, 67890, -12300, -12300, -12300, -12300, -12300, + ], + decimal_dtype, + ) + .into_array(); + assert_arrays_eq!(arr.into_array(), expected, &mut ctx); + } + + #[test] + fn test_runend_list_i64() { + let mut ctx = SESSION.create_execution_ctx(); + let values = ListArray::from_iter_slow::( + vec![vec![1i64, 2], vec![3], vec![4, 5, 6]], + Arc::new(DType::Primitive(PType::I64, Nullability::NonNullable)), + ) + .unwrap() + .into_array(); + let arr = RunEnd::new(buffer![2u32, 5, 10].into_array(), values, &mut ctx); + + let expected = ListArray::from_iter_slow::( + vec![ + vec![1i64, 2], + vec![1, 2], + vec![3], + vec![3], + vec![3], + vec![4, 5, 6], + vec![4, 5, 6], + vec![4, 5, 6], + vec![4, 5, 6], + vec![4, 5, 6], + ], + Arc::new(DType::Primitive(PType::I64, Nullability::NonNullable)), + ) + .unwrap() + .into_array(); + assert_arrays_eq!(arr.into_array(), expected, &mut ctx); + } + + #[test] + fn test_runend_nullable_decimal() { + let mut ctx = SESSION.create_execution_ctx(); + let decimal_dtype = DecimalDType::new(10, 2); + let values = + DecimalArray::from_option_iter([Some(12345i64), None, Some(-12300)], decimal_dtype) .into_array(); + let arr = RunEnd::new(buffer![2u32, 5, 10].into_array(), values, &mut ctx); + assert_eq!(arr.len(), 10); + assert_eq!( + arr.dtype(), + &DType::Decimal(decimal_dtype, Nullability::Nullable) + ); + + let expected = DecimalArray::from_option_iter( + [ + Some(12345i64), + Some(12345), + None, + None, + None, + Some(-12300), + Some(-12300), + Some(-12300), + Some(-12300), + Some(-12300), + ], + decimal_dtype, + ) + .into_array(); assert_arrays_eq!(arr.into_array(), expected, &mut ctx); } + #[test] + fn test_runend_list_bool() { + let mut ctx = SESSION.create_execution_ctx(); + let values = ListArray::from_iter_slow::( + vec![vec![true, false], vec![false], vec![true, true, false]], + Arc::new(DType::Bool(Nullability::NonNullable)), + ) + .unwrap() + .into_array(); + let arr = RunEnd::new(buffer![2u32, 5, 10].into_array(), values, &mut ctx); + + let expected = ListArray::from_iter_slow::( + vec![ + vec![true, false], + vec![true, false], + vec![false], + vec![false], + vec![false], + vec![true, true, false], + vec![true, true, false], + vec![true, true, false], + vec![true, true, false], + vec![true, true, false], + ], + Arc::new(DType::Bool(Nullability::NonNullable)), + ) + .unwrap() + .into_array(); + assert_arrays_eq!(arr.into_array(), expected, &mut ctx); + } + + #[test] + fn test_runend_list_utf8() { + let mut ctx = SESSION.create_execution_ctx(); + let values = ListArray::try_new( + VarBinViewArray::from_iter_str(["a", "b", "c", "d", "e", "f"]).into_array(), + buffer![0u32, 2, 3, 6].into_array(), + Validity::NonNullable, + ) + .unwrap() + .into_array(); + let arr = RunEnd::new(buffer![2u32, 5, 10].into_array(), values, &mut ctx); + + let expected = ListArray::try_new( + VarBinViewArray::from_iter_str([ + "a", "b", "a", "b", "c", "c", "c", "d", "e", "f", "d", "e", "f", "d", "e", "f", + "d", "e", "f", "d", "e", "f", + ]) + .into_array(), + buffer![0u32, 2, 4, 5, 6, 7, 10, 13, 16, 19, 22].into_array(), + Validity::NonNullable, + ) + .unwrap() + .into_array(); + assert_arrays_eq!(arr.into_array(), expected, &mut ctx); + } + + #[test] + fn test_runend_list_canonicalizes_to_runend_listview_slots() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = ListArray::try_new( + buffer![1i64, 2, 3, 4, 5, 6].into_array(), + buffer![0u32, 2, 3, 6].into_array(), + Validity::from_iter([true, false, true]), + )? + .into_array(); + let arr = RunEnd::try_new(buffer![2u32, 5, 6].into_array(), values, &mut ctx)?; + + let listview = arr + .clone() + .into_array() + .execute::(&mut ctx)?; + assert!(listview.offsets().is::()); + assert!(listview.sizes().is::()); + match listview.validity()? { + Validity::Array(validity) => assert!(validity.is::()), + validity => panic!("expected array-backed validity, got {validity:?}"), + } + + let expected = ListArray::try_new( + buffer![1i64, 2, 1, 2, 3, 3, 3, 4, 5, 6].into_array(), + buffer![0u32, 2, 4, 5, 6, 7, 10].into_array(), + Validity::from_iter([true, true, false, false, false, true]), + )? + .into_array(); + assert_arrays_eq!(arr.into_array(), expected, &mut ctx); + Ok(()) + } + #[test] fn test_runend_dict() { let mut ctx = SESSION.create_execution_ctx(); @@ -587,4 +831,87 @@ mod tests { .into_array(); assert_arrays_eq!(arr.into_array(), expected, &mut ctx); } + + #[test] + fn test_runend_decimal_i128() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let decimal_dtype = DecimalDType::new(20, 2); + let values = DecimalArray::from_iter([12_345i128, -67_890, 100], decimal_dtype); + let arr = RunEnd::try_new( + buffer![2u32, 5, 6].into_array(), + values.into_array(), + &mut ctx, + )?; + + let decoded = arr.into_array().execute::(&mut ctx)?; + let expected = DecimalArray::from_iter( + [12_345i128, 12_345, -67_890, -67_890, -67_890, 100], + decimal_dtype, + ); + assert_arrays_eq!(decoded, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_runend_decimal_nullable() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let decimal_dtype = DecimalDType::new(20, 2); + let values = + DecimalArray::from_option_iter([Some(12_345i128), None, Some(-67_890)], decimal_dtype); + let arr = RunEnd::try_new( + buffer![2u32, 5, 7].into_array(), + values.into_array(), + &mut ctx, + )?; + + let decoded = arr.into_array().execute::(&mut ctx)?; + let expected = DecimalArray::from_option_iter( + [ + Some(12_345i128), + Some(12_345), + None, + None, + None, + Some(-67_890), + Some(-67_890), + ], + decimal_dtype, + ); + assert_arrays_eq!(decoded, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_runend_decimal_slice() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let decimal_dtype = DecimalDType::new(20, 2); + let values = DecimalArray::from_iter([100i128, 200, 300], decimal_dtype); + let arr = RunEnd::try_new( + buffer![3u32, 5, 10].into_array(), + values.into_array(), + &mut ctx, + )?; + + let sliced = arr.slice(2..8)?; + let decoded = sliced.execute::(&mut ctx)?; + let expected = DecimalArray::from_iter([100i128, 200, 200, 300, 300, 300], decimal_dtype); + assert_arrays_eq!(decoded, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_runend_decimal_i256() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let decimal_dtype = DecimalDType::new(40, 4); + let first = i256::from_i128(123_456); + let second = i256::from_i128(-789_012); + let values = DecimalArray::from_iter([first, second], decimal_dtype); + let arr = RunEnd::try_new(buffer![2u32, 5].into_array(), values.into_array(), &mut ctx)?; + + let decoded = arr.into_array().execute::(&mut ctx)?; + let expected = + DecimalArray::from_iter([first, first, second, second, second], decimal_dtype); + assert_arrays_eq!(decoded, expected, &mut ctx); + Ok(()) + } } diff --git a/encodings/runend/src/compress.rs b/encodings/runend/src/compress.rs index 544b677dfa3..05de5fb438f 100644 --- a/encodings/runend/src/compress.rs +++ b/encodings/runend/src/compress.rs @@ -8,16 +8,19 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::DecimalArray; use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::bool::BoolArrayExt; +use vortex_array::arrays::decimal::DecimalArrayExt; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::expr::stats::Precision; use vortex_array::expr::stats::Stat; +use vortex_array::match_each_decimal_value_type; use vortex_array::match_each_native_ptype; use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::scalar::Scalar; @@ -284,6 +287,33 @@ pub fn runend_decode_typed_primitive( PrimitiveArray::new(decoded, validity) } +/// Decode a run-end encoded decimal array by expanding its native values. +pub fn runend_decode_decimal( + ends: PrimitiveArray, + values: DecimalArray, + offset: usize, + length: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let validity_mask = values + .as_ref() + .validity()? + .execute_mask(values.as_ref().len(), ctx)?; + + Ok(match_each_decimal_value_type!(values.values_type(), |D| { + let (decoded, validity) = match_each_unsigned_integer_ptype!(ends.ptype(), |E| { + runend_decode_slice( + trimmed_ends_iter(ends.as_slice::(), offset, length), + values.buffer::().as_slice(), + validity_mask, + values.dtype().nullability(), + length, + ) + }); + DecimalArray::new(decoded, values.decimal_dtype(), validity) + })) +} + /// Decode a run-end encoded VarBinView array by expanding views directly. pub fn runend_decode_varbinview( ends: PrimitiveArray, diff --git a/encodings/runend/src/compute/cast.rs b/encodings/runend/src/compute/cast.rs index 4722f9e5549..aa2dc5eb406 100644 --- a/encodings/runend/src/compute/cast.rs +++ b/encodings/runend/src/compute/cast.rs @@ -11,6 +11,7 @@ use vortex_error::VortexResult; use crate::RunEnd; use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; impl CastReduce for RunEnd { fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { // Cast the values array to the target type diff --git a/encodings/runend/src/compute/compare.rs b/encodings/runend/src/compute/compare.rs index ecdd03e3686..f3facb1167a 100644 --- a/encodings/runend/src/compute/compare.rs +++ b/encodings/runend/src/compute/compare.rs @@ -16,6 +16,7 @@ use vortex_error::VortexResult; use crate::RunEnd; use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; use crate::decompress_bool::runend_decode_bools; impl CompareKernel for RunEnd { diff --git a/encodings/runend/src/compute/fill_null.rs b/encodings/runend/src/compute/fill_null.rs index ddce31184eb..90284be3e51 100644 --- a/encodings/runend/src/compute/fill_null.rs +++ b/encodings/runend/src/compute/fill_null.rs @@ -11,6 +11,7 @@ use vortex_error::VortexResult; use crate::RunEnd; use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; impl FillNullReduce for RunEnd { fn fill_null( diff --git a/encodings/runend/src/compute/filter.rs b/encodings/runend/src/compute/filter.rs index 6d8d4aab6ea..4dc00ea2aba 100644 --- a/encodings/runend/src/compute/filter.rs +++ b/encodings/runend/src/compute/filter.rs @@ -23,8 +23,22 @@ use vortex_mask::Mask; use crate::RunEnd; use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; use crate::compute::take::take_indices_unchecked; -const FILTER_TAKE_THRESHOLD: f64 = 0.1; + +/// Takes directly below this average number of selected rows per source run. +/// +/// Take scales with the selection size; run filtering scans every run. [#1969] introduced this +/// heuristic without a focused threshold benchmark. A larger value favors take. +/// +/// [#1969]: https://github.com/vortex-data/vortex/pull/1969 +const TAKE_SELECTED_ROWS_PER_RUN_THRESHOLD: f64 = 0.1; + +/// Takes directly below this row count to avoid fixed run and value filtering costs. +/// [#1969] introduced the cutoff. A larger value favors take. +/// +/// [#1969]: https://github.com/vortex-data/vortex/pull/1969 +const MIN_RUN_FILTER_SELECTED_ROWS: usize = 25; impl FilterKernel for RunEnd { fn filter( @@ -35,87 +49,98 @@ impl FilterKernel for RunEnd { let mask_values = mask .values() .vortex_expect("FilterKernel precondition: mask is Mask::Values"); - - let runs_ratio = mask_values.true_count() as f64 / array.ends().len() as f64; - - if runs_ratio < FILTER_TAKE_THRESHOLD || mask_values.true_count() < 25 { - Ok(Some(take_indices_unchecked( + let selected_rows = mask_values.true_count(); + let source_run_count = array.ends().len(); + let selected_rows_per_run = selected_rows as f64 / source_run_count as f64; + let use_direct_take = selected_rows_per_run < TAKE_SELECTED_ROWS_PER_RUN_THRESHOLD + || selected_rows < MIN_RUN_FILTER_SELECTED_ROWS; + + if use_direct_take { + return Ok(Some(take_indices_unchecked( array, mask_values.indices(), &Validity::NonNullable, ctx, - )?)) - } else { - let primitive_run_ends = array.ends().clone().execute::(ctx)?; - let (run_ends, values_mask) = - match_each_unsigned_integer_ptype!(primitive_run_ends.ptype(), |P| { - filter_run_end_primitive( - primitive_run_ends.as_slice::

(), - array.offset() as u64, - array.len() as u64, - mask_values.bit_buffer(), - )? - }); - let values = array.values().filter(values_mask)?; - - // SAFETY: guaranteed by implementation of filter_run_end_primitive - unsafe { - Ok(Some( - RunEnd::new_unchecked( - run_ends.into_array(), - values, - 0, - mask_values.true_count(), - ) - .into_array(), - )) - } + )?)); } + + let primitive_run_ends = array.ends().clone().execute::(ctx)?; + let (filtered_run_ends, values_mask) = + match_each_unsigned_integer_ptype!(primitive_run_ends.ptype(), |P| { + filter_run_end_primitive( + primitive_run_ends.as_slice::

(), + array.offset() as u64, + array.len() as u64, + mask_values.bit_buffer(), + )? + }); + let filtered_values = array.values().filter(values_mask)?; + + // SAFETY: `filter_run_end_primitive` returns one strictly increasing end for each retained + // run value, with the final end equal to `selected_rows`. + let filtered = unsafe { + RunEnd::new_unchecked( + filtered_run_ends.into_array(), + filtered_values, + 0, + selected_rows, + ) + }; + + Ok(Some(filtered.into_array())) } } -// Code adapted from apache arrow-rs https://github.com/apache/arrow-rs/blob/b1f5c250ebb6c1252b4e7c51d15b8e77f4c361fa/arrow-select/src/filter.rs#L425 +/// Recomputes cumulative run ends and selects the run values retained by `mask`. +/// +/// The caller supplies validated RunEnd metadata for `offset..offset + length` and a mask containing +/// exactly `length` bits. The returned ends are strictly increasing and contain one entry for each +/// selected run value. +/// +/// Adapted from the [Apache Arrow Rust implementation](https://github.com/apache/arrow-rs/blob/b1f5c250ebb6c1252b4e7c51d15b8e77f4c361fa/arrow-select/src/filter.rs#L425). pub fn filter_run_end_primitive + AsPrimitive>( run_ends: &[R], offset: u64, length: u64, mask: &BitBuffer, ) -> VortexResult<(PrimitiveArray, Mask)> { - let mut new_run_ends = buffer_mut![R::zero(); run_ends.len()]; + let mut filtered_run_ends = buffer_mut![R::zero(); run_ends.len()]; - let mut start = 0u64; - let mut j = 0; - let mut count = R::zero(); + let mut run_start = 0u64; + let mut retained_run_count = 0; + let mut filtered_end = R::zero(); - let new_mask: Mask = BitBuffer::collect_bool(run_ends.len(), |i| { - let end = min(run_ends[i].as_() - offset, length); + let values_mask: Mask = BitBuffer::collect_bool(run_ends.len(), |run_idx| { + let run_end = min(run_ends[run_idx].as_() - offset, length); - // SIMD popcount of the predicate bits in this run. The range matches the - // bit-by-bit `value_unchecked` read it replaces, so `end <= mask.len()`. - let start_usize = start + // Bulk popcount is SIMD-capable and avoids per-bit reads. The input contract and clamp prove + // `run_start_idx <= run_end_idx <= mask.len()`. + let run_start_idx = run_start .try_into() .vortex_expect("run start index must fit in usize"); - let end_usize = end + let run_end_idx = run_end .try_into() .vortex_expect("run end index must fit in usize"); - let run_trues = mask.count_range(start_usize, end_usize); - count += ::from(run_trues) + let selected_in_run = mask.count_range(run_start_idx, run_end_idx); + filtered_end += ::from(selected_in_run) .vortex_expect("run popcount must fit in run-end native type"); - let keep = run_trues > 0; + let retain_run = selected_in_run > 0; - // this is to avoid branching - new_run_ends[j] = count; - j += keep as usize; + // Always write the current end, then advance only for a retained run. This keeps the loop + // branchless. + filtered_run_ends[retained_run_count] = filtered_end; + retained_run_count += retain_run as usize; - start = end; - keep + run_start = run_end; + retain_run }) .into(); - new_run_ends.truncate(j); + filtered_run_ends.truncate(retained_run_count); + Ok(( - PrimitiveArray::new(new_run_ends, Validity::NonNullable), - new_mask, + PrimitiveArray::new(filtered_run_ends, Validity::NonNullable), + values_mask, )) } diff --git a/encodings/runend/src/compute/is_constant.rs b/encodings/runend/src/compute/is_constant.rs index d21f2334018..710fb7a3712 100644 --- a/encodings/runend/src/compute/is_constant.rs +++ b/encodings/runend/src/compute/is_constant.rs @@ -11,7 +11,7 @@ use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use crate::RunEnd; -use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; /// RunEnd-specific is_constant kernel. /// diff --git a/encodings/runend/src/compute/is_sorted.rs b/encodings/runend/src/compute/is_sorted.rs index 65de2a174f2..17ae94272cb 100644 --- a/encodings/runend/src/compute/is_sorted.rs +++ b/encodings/runend/src/compute/is_sorted.rs @@ -14,7 +14,7 @@ use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use crate::RunEnd; -use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; /// RunEnd-specific is_sorted kernel. /// diff --git a/encodings/runend/src/compute/min_max.rs b/encodings/runend/src/compute/min_max.rs index 4dd39edc5f3..d17c5b3c65c 100644 --- a/encodings/runend/src/compute/min_max.rs +++ b/encodings/runend/src/compute/min_max.rs @@ -12,7 +12,7 @@ use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use crate::RunEnd; -use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; /// RunEnd-specific min/max kernel. /// diff --git a/encodings/runend/src/compute/take.rs b/encodings/runend/src/compute/take.rs index 4ec601e71a6..067bebef742 100644 --- a/encodings/runend/src/compute/take.rs +++ b/encodings/runend/src/compute/take.rs @@ -24,6 +24,7 @@ use vortex_mask::Mask; use crate::RunEnd; use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; use crate::iter::trimmed_ends_iter; const SORTED_LINEAR_RUNS_PER_INDEX_THRESHOLD: usize = 16; diff --git a/encodings/runend/src/compute/take_from.rs b/encodings/runend/src/compute/take_from.rs index 0987cd3fc55..e1bd33a3f97 100644 --- a/encodings/runend/src/compute/take_from.rs +++ b/encodings/runend/src/compute/take_from.rs @@ -13,6 +13,7 @@ use vortex_error::VortexResult; use crate::RunEnd; use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; #[derive(Debug)] pub(crate) struct RunEndTakeFrom; @@ -68,7 +69,7 @@ mod tests { use crate::RunEnd; use crate::RunEndArray; - use crate::array::RunEndArrayExt; + use crate::array::RunEndArraySlotsExt; use crate::compute::take_from::RunEndTakeFrom; static SESSION: LazyLock = LazyLock::new(|| { diff --git a/encodings/runend/src/decompress_bool.rs b/encodings/runend/src/decompress_bool.rs index d2b778fdbbf..658138a905e 100644 --- a/encodings/runend/src/decompress_bool.rs +++ b/encodings/runend/src/decompress_bool.rs @@ -101,6 +101,7 @@ pub fn runend_decode_typed_bool( /// Fast path for few runs with no offset. Uses direct slice access to minimize overhead. /// This avoids the `trimmed_ends_iter` iterator chain which adds significant overhead /// for small numbers of runs. +#[allow(clippy::inline_always)] #[inline(always)] fn decode_few_runs_no_offset( ends: &[E], @@ -220,6 +221,7 @@ fn decode_bool_nullable( } /// Sequential decode for few runs - avoids prefill overhead. +#[allow(clippy::inline_always)] #[inline(always)] fn decode_nullable_sequential( run_ends: impl Iterator, diff --git a/encodings/runend/src/kernel.rs b/encodings/runend/src/kernel.rs index 13f32df07a1..2eed964e3a9 100644 --- a/encodings/runend/src/kernel.rs +++ b/encodings/runend/src/kernel.rs @@ -24,6 +24,7 @@ use vortex_session::VortexSession; use crate::RunEnd; use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; use crate::compute::take_from::RunEndTakeFrom; pub(super) fn initialize(session: &VortexSession) { diff --git a/encodings/runend/src/lib.rs b/encodings/runend/src/lib.rs index ed3ff946d31..b991609c19c 100644 --- a/encodings/runend/src/lib.rs +++ b/encodings/runend/src/lib.rs @@ -16,6 +16,9 @@ mod iter; mod kernel; pub mod ops; mod rules; +#[cfg(test)] +#[cfg(not(codspeed))] +mod trace_tests; #[doc(hidden)] pub mod _benchmarking { diff --git a/encodings/runend/src/ops.rs b/encodings/runend/src/ops.rs index 39fb004c3e7..e2c2e3fc99b 100644 --- a/encodings/runend/src/ops.rs +++ b/encodings/runend/src/ops.rs @@ -15,6 +15,7 @@ use vortex_error::VortexResult; use crate::RunEnd; use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; impl OperationsVTable for RunEnd { fn scalar_at( diff --git a/encodings/runend/src/rules.rs b/encodings/runend/src/rules.rs index ac7d6038146..a6f3d6a4ae3 100644 --- a/encodings/runend/src/rules.rs +++ b/encodings/runend/src/rules.rs @@ -19,6 +19,7 @@ use vortex_error::VortexResult; use crate::RunEnd; use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; pub(super) const RULES: ParentRuleSet = ParentRuleSet::new(&[ // CastReduceAdaptor must come before RunEndScalarFnRule so that cast operations are executed diff --git a/encodings/runend/src/trace_tests.rs b/encodings/runend/src/trace_tests.rs new file mode 100644 index 00000000000..a9347d95089 --- /dev/null +++ b/encodings/runend/src/trace_tests.rs @@ -0,0 +1,147 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Snapshot tests tracing how reductions and executions flow through run-end arrays. + +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::FilterArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::optimizer::ArrayOptimizer; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::fns::binary::Binary; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::session::ArraySession; +use vortex_array::test_harness::trace::trace_op; +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +use crate::RunEnd; + +fn execution_ctx() -> ExecutionCtx { + let session = VortexSession::empty().with::(); + // Execute-parent kernels live in the session registry; without registering RunEnd's own + // kernels the traces below would show plain canonicalization instead of the pushdowns. + crate::initialize(&session); + ExecutionCtx::new(session) +} + +/// Run-end encoding of `[1, 1, 1, 2, 2, 3, 3, 3, 3]`: ends `[3, 5, 9]`, values `[1, 2, 3]`. +fn runend_array() -> VortexResult { + Ok(RunEnd::encode( + PrimitiveArray::from_iter([1i32, 1, 1, 2, 2, 3, 3, 3, 3]).into_array(), + &mut execution_ctx(), + )? + .into_array()) +} + +#[test] +fn trace_compare_on_runend() -> VortexResult<()> { + let runend = runend_array()?; + let rhs = ConstantArray::new(Scalar::from(2i32), runend.len()).into_array(); + let compared = Binary::try_new(runend, rhs, Operator::Eq)?.into_array(); + + let traced = trace_op(|| compared.optimize())?; + insta::assert_snapshot!(traced.trace.to_string(), @" + optimize root=vortex.binary(bool, len=9) session=false + reduce_parent static:RunEndScalarFnRule slot=0 parent=vortex.binary(bool, len=9) child=vortex.runend(i32, len=9) -> vortex.runend(bool, len=9) + done output=vortex.runend(bool, len=9) + "); + + let optimized = traced.output; + let traced = trace_op(|| { + optimized + .execute::(&mut execution_ctx()) + .map(IntoArray::into_array) + })?; + + assert_arrays_eq!( + traced.output, + BoolArray::from_iter([false, false, false, true, true, false, false, false, false]), + &mut execution_ctx() + ); + insta::assert_snapshot!(traced.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.runend(bool, len=9) + iter 0 current=vortex.runend(bool, len=9) builder_active=false + execute_until target=AnyCanonical root=vortex.binary(bool, len=3) + iter 0 current=vortex.binary(bool, len=3) builder_active=false + Done array=vortex.bool(bool, len=3) + iter 1 current=vortex.bool(bool, len=3) builder_active=false + return output=vortex.bool(bool, len=3) + Done array=vortex.bool(bool, len=9) + iter 1 current=vortex.bool(bool, len=9) builder_active=false + return output=vortex.bool(bool, len=9) + "); + + Ok(()) +} + +#[test] +fn trace_filter_on_runend() -> VortexResult<()> { + let runend = runend_array()?; + let filtered = FilterArray::try_new( + runend, + Mask::from_iter([true, false, false, true, true, false, false, true, false]), + )? + .into_array(); + + let traced = trace_op(|| { + filtered + .execute::(&mut execution_ctx()) + .map(IntoArray::into_array) + })?; + + assert_arrays_eq!( + traced.output, + PrimitiveArray::from_iter([1i32, 2, 2, 3]), + &mut execution_ctx() + ); + insta::assert_snapshot!(traced.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.filter(i32, len=4) + iter 0 current=vortex.filter(i32, len=4) builder_active=false + child_execute_parent session[0]:execute_parent_fn slot=0 parent=vortex.filter(i32, len=4) child=vortex.runend(i32, len=9) -> vortex.dict(i32, len=4) + iter 1 current=vortex.dict(i32, len=4) builder_active=false + child_execute_parent session[0]:execute_parent_fn slot=1 parent=vortex.dict(i32, len=4) child=vortex.primitive(i32, len=3) -> vortex.primitive(i32, len=4) + iter 2 current=vortex.primitive(i32, len=4) builder_active=false + return output=vortex.primitive(i32, len=4) + "); + + Ok(()) +} + +#[test] +fn trace_take_on_runend() -> VortexResult<()> { + let runend = runend_array()?; + // A take is expressed as a `DictArray` whose codes are the take indices. + let indices = PrimitiveArray::from_iter([8u64, 0, 4, 4]).into_array(); + let take = DictArray::try_new(indices, runend)?.into_array(); + + let traced = trace_op(|| { + take.execute::(&mut execution_ctx()) + .map(IntoArray::into_array) + })?; + + assert_arrays_eq!( + traced.output, + PrimitiveArray::from_iter([3i32, 1, 2, 2]), + &mut execution_ctx() + ); + insta::assert_snapshot!(traced.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.dict(i32, len=4) + iter 0 current=vortex.dict(i32, len=4) builder_active=false + child_execute_parent session[0]:execute_parent_fn slot=1 parent=vortex.dict(i32, len=4) child=vortex.runend(i32, len=9) -> vortex.dict(i32, len=4) + iter 1 current=vortex.dict(i32, len=4) builder_active=false + child_execute_parent session[0]:execute_parent_fn slot=1 parent=vortex.dict(i32, len=4) child=vortex.primitive(i32, len=3) -> vortex.primitive(i32, len=4) + iter 2 current=vortex.primitive(i32, len=4) builder_active=false + return output=vortex.primitive(i32, len=4) + "); + + Ok(()) +} diff --git a/encodings/sequence/src/array.rs b/encodings/sequence/src/array.rs index 3b8e0065f06..10a222544b6 100644 --- a/encodings/sequence/src/array.rs +++ b/encodings/sequence/src/array.rs @@ -6,7 +6,7 @@ use std::fmt::Formatter; use std::hash::Hash; use std::hash::Hasher; -use num_traits::cast::FromPrimitive; +use num_traits::AsPrimitive; use prost::Message; use smallvec::smallvec; use vortex_array::Array; @@ -28,7 +28,6 @@ use vortex_array::dtype::PType; use vortex_array::expr::stats::Precision as StatPrecision; use vortex_array::expr::stats::Stat; use vortex_array::match_each_integer_ptype; -use vortex_array::match_each_native_ptype; use vortex_array::match_each_pvalue; use vortex_array::scalar::PValue; use vortex_array::scalar::Scalar; @@ -49,6 +48,8 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::compress::sequence_decompress; +use crate::eval; +use crate::eval::SequenceValue; use crate::rules::RULES; /// A [`Sequence`]-encoded Vortex array. @@ -64,8 +65,11 @@ pub struct SequenceMetadata { pub(super) const SLOT_NAMES: [&str; 0] = []; -#[derive(Clone, Debug)] /// An array representing the equation `A[i] = base + i * multiplier`. +/// +/// The base uses the output ptype, while the step is normalized to `i64` or `u64` and may fall +/// outside that ptype. Construction ensures every sequence value fits the output ptype. +#[derive(Clone, Debug)] pub struct SequenceData { base: PValue, multiplier: PValue, @@ -99,7 +103,7 @@ impl SequenceData { ) } - /// Constructs a sequence array using two integer values (with the same ptype). + /// Constructs a sequence array using two integer values, validated against output `ptype`. pub(crate) fn try_new( base: PValue, multiplier: PValue, @@ -125,26 +129,91 @@ impl SequenceData { }; if !ptype.is_int() { - vortex_bail!("only integer ptype are supported in SequenceArray currently") + vortex_bail!("only integer ptypes are supported in SequenceArray currently") } vortex_ensure!(length > 0, "SequenceArray length must be greater than zero"); - Self::try_last(base, multiplier, *ptype, length).map_err(|e| { - e.with_context(format!( - "final value not expressible, base = {base:?}, multiplier = {multiplier:?}, len = {length} ", - )) - })?; + Self::narrowed_base(base, *ptype)?; + Self::ensure_last_expressible(base, multiplier, *ptype, length) + } + + /// Ensures the final value fits `ptype` without overflowing intermediate arithmetic. + fn ensure_last_expressible( + base: PValue, + multiplier: PValue, + ptype: PType, + length: usize, + ) -> VortexResult<()> { + let steps = (length - 1) as u64; + let (ascending, magnitude) = eval::step_parts(multiplier) + .ok_or_else(|| vortex_err!("step {multiplier} must be an integer"))?; + if steps == 0 || magnitude == 0 { + return Ok(()); + } + + // Measure room in the ptype's signedness so large `u64` bases remain exact. + let room = if ptype.is_signed_int() { + let base = base.cast::()?; + let max = i64::try_from(ptype.max_value_as_u64()) + .vortex_expect("a signed ptype's max fits i64"); + base.abs_diff(if ascending { max } else { -max - 1 }) + } else { + let base = base.cast::()?; + if ascending { + ptype.max_value_as_u64() - base + } else { + base + } + }; + + vortex_ensure!( + steps <= room / magnitude, + "final value not expressible, base = {base:?}, multiplier = {multiplier:?}, len = {length}" + ); Ok(()) } + /// The step's ptype: the serialized form preserves its signedness but not its width. + fn multiplier_ptype_from_proto( + multiplier: &vortex_proto::scalar::ScalarValue, + ) -> VortexResult { + use vortex_proto::scalar::scalar_value::Kind; + match multiplier + .kind + .as_ref() + .ok_or_else(|| vortex_err!("multiplier value missing kind"))? + { + Kind::Int64Value(_) => Ok(PType::I64), + Kind::Uint64Value(_) => Ok(PType::U64), + _ => vortex_bail!("only integer ptypes are supported in SequenceArray currently"), + } + } + + fn narrowed_base(base: PValue, ptype: PType) -> VortexResult { + vortex_ensure!(base.ptype().is_int(), "base {base} must be an integer"); + match_each_integer_ptype!(ptype, |P| { Ok(PValue::from(base.cast::

()?)) }) + } + + /// Puts `base` into the output ptype and the step into its canonical ptype. fn normalize(base: PValue, multiplier: PValue, ptype: PType) -> VortexResult<(PValue, PValue)> { - match_each_integer_ptype!(ptype, |P| { - Ok(( - PValue::from(base.cast::

()?), - PValue::from(multiplier.cast::

()?), - )) - }) + let base = Self::narrowed_base(base, ptype)?; + + // Give equivalent steps the same representation. + let multiplier = match_each_pvalue!( + multiplier, + uint: |v| { + let v: u64 = v.as_(); + i64::try_from(v).map(PValue::from).unwrap_or(PValue::U64(v)) + }, + int: |v| { + let v: i64 = v.as_(); + PValue::from(v) + }, + float: |v| { vortex_bail!("step {v} must be an integer") } + ); + + Ok((base, multiplier)) } /// Constructs a [`SequenceArray`] payload without validation. @@ -152,12 +221,14 @@ impl SequenceData { /// # Safety /// /// The caller must ensure that: - /// - `base` and `multiplier` are both normalized to the same integer `ptype`. - /// - they are logically compatible with the outer dtype and len. + /// - `base` uses the outer dtype's integer ptype. + /// - `multiplier` is a canonical `i64` or `u64`. + /// - every sequence value fits the outer dtype's ptype. pub(crate) unsafe fn new_unchecked(base: PValue, multiplier: PValue) -> Self { Self { base, multiplier } } + /// The array's output ptype. pub fn ptype(&self) -> PType { self.base.ptype() } @@ -170,48 +241,42 @@ impl SequenceData { self.multiplier } + /// `base` and `multiplier` reduced into `O`, the type the values are computed in. + pub(crate) fn wrapping_parts(&self) -> VortexResult<(O, O)> { + eval::wrapping_parts(self.base, self.multiplier).ok_or_else(|| { + vortex_err!( + "SequenceArray values must be integers, got base {:?} and step {:?}", + self.base, + self.multiplier + ) + }) + } + + /// The two's-complement bits of `base` and `multiplier`, widened to 64 bits. + pub fn wrapping_bits(&self) -> VortexResult<(u64, u64)> { + self.wrapping_parts::() + } + pub fn into_parts(self) -> SequenceDataParts { SequenceDataParts { base: self.base, multiplier: self.multiplier, - ptype: self.base.ptype(), + ptype: self.ptype(), } } - pub(crate) fn try_last( - base: PValue, - multiplier: PValue, - ptype: PType, - length: usize, - ) -> VortexResult { - match_each_integer_ptype!(ptype, |P| { - let len_t =

::from_usize(length - 1) - .ok_or_else(|| vortex_err!("cannot convert length {} into {}", length, ptype))?; - - let base = base.cast::

()?; - let multiplier = multiplier.cast::

()?; - let last = len_t - .checked_mul(multiplier) - .and_then(|offset| offset.checked_add(base)) - .ok_or_else(|| vortex_err!("last value computation overflows"))?; - Ok(PValue::from(last)) - }) - } - pub(crate) fn index_value(&self, idx: usize) -> PValue { - match_each_native_ptype!(self.ptype(), |P| { - let base = self.base.cast::

().vortex_expect("must be able to cast"); - let multiplier = self - .multiplier - .cast::

() - .vortex_expect("must be able to cast"); - let value = base + (multiplier *

::from_usize(idx).vortex_expect("must fit")); - - PValue::from(value) + match_each_integer_ptype!(self.ptype(), |O| { + let (base, multiplier) = self + .wrapping_parts::() + .vortex_expect("sequence values are integers"); + PValue::from(eval::wrapping_value(base, multiplier, idx)) }) } } +// Normalization gives equal sequences the same value tags. Compare tags first because `PValue` +// equality can panic for mixed signedness outside the `i64` range. impl ArrayHash for SequenceData { fn array_hash(&self, state: &mut H, _accuracy: EqMode) { self.base.hash(state); @@ -221,7 +286,10 @@ impl ArrayHash for SequenceData { impl ArrayEq for SequenceData { fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool { - self.base == other.base && self.multiplier == other.multiplier + self.base.ptype() == other.base.ptype() + && self.multiplier.ptype() == other.multiplier.ptype() + && self.base == other.base + && self.multiplier == other.multiplier } } @@ -297,36 +365,46 @@ impl VTable for Sequence { "SequenceArray expects 0 children, got {}", children.len() ); + let DType::Primitive(output_ptype, _) = dtype else { + vortex_bail!( + "only primitive dtypes are supported in SequenceArray currently, got {dtype}" + ); + }; let metadata = SequenceMetadata::decode(metadata)?; - let ptype = dtype.as_ptype(); + let base_metadata = metadata + .base + .as_ref() + .ok_or_else(|| vortex_err!("base required"))?; + + let multiplier_metadata = metadata + .multiplier + .as_ref() + .ok_or_else(|| vortex_err!("multiplier required"))?; // We go via Scalar to validate that the value is valid for the ptype. let base = Scalar::from_proto_value( - metadata - .base - .as_ref() - .ok_or_else(|| vortex_err!("base required"))?, - &DType::Primitive(ptype, NonNullable), + base_metadata, + &DType::Primitive(*output_ptype, NonNullable), session, )? .as_primitive() .pvalue() .vortex_expect("sequence array base should be a non-nullable primitive"); + // The serialized step preserves signedness independently of the output ptype. + let multiplier_ptype = SequenceData::multiplier_ptype_from_proto(multiplier_metadata)?; let multiplier = Scalar::from_proto_value( - metadata - .multiplier - .as_ref() - .ok_or_else(|| vortex_err!("multiplier required"))?, - &DType::Primitive(ptype, NonNullable), + multiplier_metadata, + &DType::Primitive(multiplier_ptype, NonNullable), session, )? .as_primitive() .pvalue() .vortex_expect("sequence array multiplier should be a non-nullable primitive"); - let data = SequenceData::try_new(base, multiplier, ptype, dtype.nullability(), len)?; + let data = + SequenceData::try_new(base, multiplier, *output_ptype, dtype.nullability(), len)?; Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data)) } @@ -394,6 +472,8 @@ impl Sequence { /// Construct a new [`SequenceArray`] from pre-validated parts. /// + /// Arguments are normalized before constructing the array. + /// /// # Safety /// /// Caller must ensure the sequence is logically compatible with the provided dtype and len. @@ -406,7 +486,7 @@ impl Sequence { ) -> SequenceArray { let dtype = DType::Primitive(ptype, nullability); let (base, multiplier) = SequenceData::normalize(base, multiplier, ptype) - .vortex_expect("SequenceArray parts must be normalized to the target ptype"); + .vortex_expect("SequenceArray parts must be representable in the output ptype"); let stats = Self::stats(multiplier); let data = unsafe { SequenceData::new_unchecked(base, multiplier) }; unsafe { Array::from_parts_unchecked(ArrayParts::new(Sequence, dtype, length, data)) } @@ -452,17 +532,30 @@ impl Sequence { mod tests { use std::sync::LazyLock; + use rstest::rstest; + use vortex_array::ArrayContext; + use vortex_array::ArrayEq; + use vortex_array::EqMode; + use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; use vortex_array::expr::stats::Precision as StatPrecision; use vortex_array::expr::stats::Stat; use vortex_array::expr::stats::StatsProviderExt; + use vortex_array::scalar::PValue; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; + use vortex_array::serde::SerializeOptions; + use vortex_array::serde::SerializedArray; + use vortex_buffer::ByteBufferMut; use vortex_error::VortexResult; + use vortex_error::vortex_err; use vortex_session::VortexSession; + use vortex_session::registry::ReadContext; use crate::Sequence; @@ -586,4 +679,194 @@ mod tests { Ok(()) } + + #[rstest] + #[case::descending_step_unsigned_output(PValue::from(100i32), PValue::from(-10i32), PType::U8)] + #[case::narrow_unsigned(PValue::from(1000u32), PValue::from(100u32), PType::U16)] + #[case::signed_output(PValue::from(0i16), PValue::from(1i16), PType::I32)] + #[case::signed_step_past_i64_max(PValue::from(0i64), PValue::from(1i64 << 62), PType::U64)] + #[case::unsigned_step_past_i64_max(PValue::from(0u64), PValue::from(u64::MAX / 4), PType::U64)] + fn serde_roundtrip_preserves_values( + #[case] base: PValue, + #[case] multiplier: PValue, + #[case] output_ptype: PType, + ) -> VortexResult<()> { + let array = Sequence::try_new(base, multiplier, output_ptype, Nullability::NonNullable, 4)?; + assert_eq!(array.ptype(), output_ptype); + + let dtype = array.dtype().clone(); + let len = array.len(); + let ctx = ArrayContext::empty(); + let serialized = + array + .clone() + .into_array() + .serialize(&ctx, &SESSION, &SerializeOptions::default())?; + + let mut concat = ByteBufferMut::empty(); + for buf in serialized { + concat.extend_from_slice(buf.as_ref()); + } + + let decoded = SerializedArray::try_from(concat.freeze())?.decode( + &dtype, + len, + &ReadContext::new(ctx.to_ids()), + &SESSION, + )?; + + let decoded_sequence = decoded + .as_opt::() + .ok_or_else(|| vortex_err!("decoded array should still be a SequenceArray"))?; + assert_eq!(decoded_sequence.ptype(), output_ptype); + assert_eq!(decoded_sequence.multiplier(), array.multiplier()); + assert_eq!(decoded.dtype(), &dtype); + assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx()); + + Ok(()) + } + + #[test] + fn descending_step_unsigned_output() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = Sequence::try_new( + PValue::from(100i32), + PValue::from(-10i32), + PType::U8, + Nullability::NonNullable, + 5, + )?; + + assert_arrays_eq!( + array, + PrimitiveArray::from_iter([100u8, 90, 80, 70, 60]), + &mut ctx + ); + assert_eq!( + array.clone().into_array().execute_scalar(1, &mut ctx)?, + Scalar::from(90u8) + ); + assert_arrays_eq!( + array.slice(3..5)?, + PrimitiveArray::from_iter([70u8, 60]), + &mut ctx + ); + + Ok(()) + } + + #[test] + fn step_spanning_output_range() -> VortexResult<()> { + let array = Sequence::try_new( + PValue::from(255u8), + PValue::from(-255i32), + PType::U8, + Nullability::NonNullable, + 2, + )?; + + assert_arrays_eq!( + array, + PrimitiveArray::from_iter([255u8, 0]), + &mut SESSION.create_execution_ctx() + ); + + Ok(()) + } + + #[test] + fn values_past_i64_max() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let step = 1u64 << 62; + let array = Sequence::try_new( + PValue::from(0i64), + PValue::from(1i64 << 62), + PType::U64, + Nullability::NonNullable, + 4, + )?; + + assert_arrays_eq!( + array, + PrimitiveArray::from_iter([0, step, 2 * step, 3 * step]), + &mut ctx + ); + assert_eq!( + array.into_array().execute_scalar(3, &mut ctx)?, + Scalar::from(3 * step) + ); + + Ok(()) + } + + #[test] + fn eq_across_step_signedness() -> VortexResult<()> { + let base = PValue::from((1u64 << 63) - 1); + let descending = Sequence::try_new( + base, + PValue::from(-1i64), + PType::U64, + Nullability::NonNullable, + 2, + )? + .into_array(); + let ascending = Sequence::try_new( + base, + PValue::from(1u64 << 63), + PType::U64, + Nullability::NonNullable, + 2, + )? + .into_array(); + + assert!(descending.array_eq(&descending.clone(), EqMode::Value)); + assert!(!descending.array_eq(&ascending, EqMode::Value)); + + Ok(()) + } + + #[test] + fn constant_sequence_longer_than_output_range() -> VortexResult<()> { + let array = Sequence::try_new( + PValue::from(7u8), + PValue::from(0i32), + PType::U8, + Nullability::NonNullable, + 300, + )?; + + assert_arrays_eq!( + array, + PrimitiveArray::from_iter([7u8; 300]), + &mut SESSION.create_execution_ctx() + ); + + Ok(()) + } + + #[test] + fn deserialize_rejects_values_outside_output_ptype() -> VortexResult<()> { + let array = Sequence::try_new_typed(-5i32, 1i32, Nullability::NonNullable, 5)?; + let len = array.len(); + let ctx = ArrayContext::empty(); + let serialized = + array + .into_array() + .serialize(&ctx, &SESSION, &SerializeOptions::default())?; + + let mut concat = ByteBufferMut::empty(); + for buf in serialized { + concat.extend_from_slice(buf.as_ref()); + } + + let decoded = SerializedArray::try_from(concat.freeze())?.decode( + &DType::Primitive(PType::U8, Nullability::NonNullable), + len, + &ReadContext::new(ctx.to_ids()), + &SESSION, + ); + assert!(decoded.is_err()); + + Ok(()) + } } diff --git a/encodings/sequence/src/compress.rs b/encodings/sequence/src/compress.rs index b20980cf132..fb912e6259b 100644 --- a/encodings/sequence/src/compress.rs +++ b/encodings/sequence/src/compress.rs @@ -1,8 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::ops::Add; - use num_traits::CheckedAdd; use num_traits::CheckedSub; use vortex_array::ArrayRef; @@ -11,10 +9,10 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::match_each_integer_ptype; -use vortex_array::match_each_native_ptype; use vortex_array::scalar::PValue; use vortex_array::validity::Validity; use vortex_buffer::BufferMut; @@ -24,14 +22,16 @@ use vortex_error::VortexResult; use crate::Sequence; use crate::SequenceArray; use crate::SequenceData; -/// An iterator that yields `base, base + step, base + 2*step, ...` via repeated addition. +use crate::eval::SequenceValue; + +/// Iterates a sequence using wrapping addition. struct SequenceIter { acc: T, step: T, remaining: usize, } -impl> Iterator for SequenceIter { +impl Iterator for SequenceIter { type Item = T; #[inline] @@ -41,9 +41,7 @@ impl> Iterator for SequenceIter { } let val = self.acc; self.remaining -= 1; - if self.remaining > 0 { - self.acc = self.acc + self.step; - } + self.acc = self.acc.wrapping_add(&self.step); Some(val) } @@ -54,14 +52,14 @@ impl> Iterator for SequenceIter { } // SAFETY: `size_hint` returns an exact count and `next` yields exactly that many items. -unsafe impl> TrustedLen for SequenceIter {} +unsafe impl TrustedLen for SequenceIter {} /// Decompresses a [`SequenceArray`] into a [`PrimitiveArray`]. #[inline] pub fn sequence_decompress(array: &SequenceArray) -> VortexResult { - fn decompress_inner( - base: P, - multiplier: P, + fn decompress_inner( + base: O, + multiplier: O, len: usize, nullability: Nullability, ) -> PrimitiveArray { @@ -73,9 +71,8 @@ pub fn sequence_decompress(array: &SequenceArray) -> VortexResult { PrimitiveArray::new(values, Validity::from(nullability)) } - let prim = match_each_native_ptype!(array.ptype(), |P| { - let base = array.base().cast::

()?; - let multiplier = array.multiplier().cast::

()?; + let prim = match_each_integer_ptype!(array.dtype().as_ptype(), |O| { + let (base, multiplier) = array.wrapping_parts::()?; decompress_inner(base, multiplier, array.len(), array.dtype().nullability()) }); Ok(prim.into_array()) @@ -131,8 +128,15 @@ fn encode_primitive_array + CheckedAdd + CheckedSu return Ok(None); } - if SequenceData::try_last(base.into(), multiplier.into(), P::PTYPE, slice.len()).is_err() { - // If the last value is out of range, we cannot encode + // Reject an out-of-range sequence before scanning the slice. + if SequenceData::validate( + base.into(), + multiplier.into(), + &DType::Primitive(P::PTYPE, nullability), + slice.len(), + ) + .is_err() + { return Ok(None); } diff --git a/encodings/sequence/src/compute/cast.rs b/encodings/sequence/src/compute/cast.rs index fb18cf11181..fa5d3a7661d 100644 --- a/encodings/sequence/src/compute/cast.rs +++ b/encodings/sequence/src/compute/cast.rs @@ -5,12 +5,8 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::scalar::Scalar; -use vortex_array::scalar::ScalarValue; use vortex_array::scalar_fn::fns::cast::CastReduce; use vortex_error::VortexResult; -use vortex_error::vortex_err; use crate::Sequence; impl CastReduce for Sequence { @@ -26,62 +22,21 @@ impl CastReduce for Sequence { return Ok(None); } - // Check if this is just a nullability change - if array.ptype() == *target_ptype && array.dtype().nullability() != *target_nullability { - // For SequenceArray, we can just create a new one with the same parameters - // but different nullability - return Ok(Some( - Sequence::try_new( - array.base(), - array.multiplier(), - *target_ptype, - *target_nullability, - array.len(), - )? - .into_array(), - )); - } - - // For type changes, we need to cast the base and multiplier - if array.ptype() != *target_ptype { - // Create scalars from PValues and cast them - let base_scalar = Scalar::try_new( - DType::Primitive(array.ptype(), Nullability::NonNullable), - Some(ScalarValue::Primitive(array.base())), - )?; - let multiplier_scalar = Scalar::try_new( - DType::Primitive(array.ptype(), Nullability::NonNullable), - Some(ScalarValue::Primitive(array.multiplier())), - )?; - - let new_base_scalar = - base_scalar.cast(&DType::Primitive(*target_ptype, Nullability::NonNullable))?; - let new_multiplier_scalar = multiplier_scalar - .cast(&DType::Primitive(*target_ptype, Nullability::NonNullable))?; - - // Extract PValues from the casted scalars - let new_base = new_base_scalar - .as_primitive() - .pvalue() - .ok_or_else(|| vortex_err!("Cast resulted in null base value"))?; - let new_multiplier = new_multiplier_scalar - .as_primitive() - .pvalue() - .ok_or_else(|| vortex_err!("Cast resulted in null multiplier value"))?; - - return Ok(Some( - Sequence::try_new( - new_base, - new_multiplier, - *target_ptype, - *target_nullability, - array.len(), - )? - .into_array(), - )); + if array.dtype() == dtype { + return Ok(None); } - Ok(None) + // try_new also validates that the produced values fit the target ptype. + Ok(Some( + Sequence::try_new( + array.base(), + array.multiplier(), + *target_ptype, + *target_nullability, + array.len(), + )? + .into_array(), + )) } } @@ -99,6 +54,9 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; + use vortex_array::scalar::Scalar; + use vortex_array::scalar::ScalarValue; + use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::Sequence; @@ -197,14 +155,40 @@ mod tests { ); } + #[test] + fn test_cast_sequence_narrows_to_output_dtype() -> VortexResult<()> { + let casted = Sequence::try_new_typed(100i32, -10i32, Nullability::NonNullable, 5)? + .into_array() + .cast(DType::Primitive(PType::U8, Nullability::NonNullable))?; + + let sequence = casted + .as_typed::() + .expect("integer sequence cast should preserve SequenceArray"); + assert_eq!(sequence.ptype(), PType::U8); + assert_eq!(sequence.multiplier(), (-10i64).into()); + assert_eq!( + sequence.dtype(), + &DType::Primitive(PType::U8, Nullability::NonNullable) + ); + + let scalar = casted.execute_scalar(1, &mut SESSION.create_execution_ctx())?; + assert_eq!( + scalar, + Scalar::try_new( + DType::Primitive(PType::U8, Nullability::NonNullable), + Some(ScalarValue::from(90u8)), + )? + ); + + Ok(()) + } + #[rstest] #[case::i32(Sequence::try_new_typed(0i32, 1i32, Nullability::NonNullable, 5).unwrap())] #[case::u64(Sequence::try_new_typed(1000u64, 100u64, Nullability::NonNullable, 4).unwrap())] - // TODO(DK): SequenceArray does not actually conform. You cannot cast this array to u8 even - // though all its values are representable therein. - // - // #[case::negative_step(Sequence::try_new_typed(100i32, -10i32, Nullability::NonNullable, - // 5).unwrap())] + #[case::negative_step( + Sequence::try_new_typed(100i32, -10i32, Nullability::NonNullable, 5).unwrap() + )] #[case::single(Sequence::try_new_typed(42i64, 0i64, Nullability::NonNullable, 1).unwrap())] #[case::constant(Sequence::try_new_typed( 100i32, diff --git a/encodings/sequence/src/compute/compare.rs b/encodings/sequence/src/compute/compare.rs index 5895330ea81..aa53cc44e00 100644 --- a/encodings/sequence/src/compute/compare.rs +++ b/encodings/sequence/src/compute/compare.rs @@ -7,20 +7,17 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; -use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; -use vortex_array::match_each_integer_ptype; use vortex_array::scalar::PValue; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::binary::CompareKernel; use vortex_array::scalar_fn::fns::operators::CompareOperator; -use vortex_buffer::BitBuffer; +use vortex_buffer::BitBufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_error::vortex_err; use crate::array::Sequence; +use crate::eval; impl CompareKernel for Sequence { fn compare( @@ -38,16 +35,16 @@ impl CompareKernel for Sequence { return Ok(None); }; - // Check if there exists an integer solution to const = base + (0..len) * multiplier. - let set_idx = find_intersection_scalar( - lhs.base(), - lhs.multiplier(), - lhs.len(), - constant - .as_primitive() - .pvalue() - .vortex_expect("null constant handled in adaptor"), - ); + let value = constant + .as_primitive() + .pvalue() + .vortex_expect("null constant handled in adaptor"); + + // Fall back for non-integer constants. + let Some(intersection) = find_intersection(lhs.base(), lhs.multiplier(), lhs.len(), value) + else { + return Ok(None); + }; let nullability = lhs.dtype().nullability() | rhs.dtype().nullability(); let validity = match nullability { @@ -55,82 +52,73 @@ impl CompareKernel for Sequence { Nullability::Nullable => vortex_array::validity::Validity::AllValid, }; - if let Ok(set_idx) = set_idx { - let buffer = BitBuffer::from_iter((0..lhs.len()).map(|idx| idx == set_idx)); - Ok(Some(BoolArray::new(buffer, validity).into_array())) - } else { - Ok(Some( - ConstantArray::new(Scalar::bool(false, nullability), lhs.len()).into_array(), - )) - } + let array = match intersection { + Intersection::None => { + ConstantArray::new(Scalar::bool(false, nullability), lhs.len()).into_array() + } + Intersection::All => { + ConstantArray::new(Scalar::bool(true, nullability), lhs.len()).into_array() + } + Intersection::At(idx) => { + let mut buffer = BitBufferMut::new_unset(lhs.len()); + buffer.set(idx); + BoolArray::new(buffer.freeze(), validity).into_array() + } + }; + + Ok(Some(array)) } } -/// Find the index where `base + idx * multiplier == intercept`, if one exists. -/// -/// # Errors -/// Return `VortexError` if: -/// - `len` is 0 -/// - `intercept` or `multiplier` can't be cast to `base`'s PType -/// - `intercept` is outside the range of the sequence -/// - `intercept` doesn't fall exactly on a sequence value -pub(crate) fn find_intersection_scalar( - base: PValue, - multiplier: PValue, - len: usize, - intercept: PValue, -) -> VortexResult { - match_each_integer_ptype!(base.ptype(), |P| { - let intercept = intercept.cast::

()?; - let base = base.cast::

()?; - let multiplier = multiplier.cast::

()?; - find_intersection(base, multiplier, len, intercept) - }) +/// The intersection of a sequence with a value. +pub(crate) enum Intersection { + None, + At(usize), + All, } -fn find_intersection( - base: P, - multiplier: P, +/// Finds where a sequence equals `value`, or returns `None` for non-integer inputs. +pub(crate) fn find_intersection( + base: PValue, + multiplier: PValue, len: usize, - intercept: P, -) -> VortexResult { - if len == 0 { - vortex_bail!("len == 0") + value: PValue, +) -> Option { + if !value.ptype().is_int() || len == 0 { + return (len == 0).then_some(Intersection::None); } + let (ascending, magnitude) = eval::step_parts(multiplier)?; - let count = P::from_usize(len - 1).vortex_expect("idx must fit into type"); - let end_element = base + (multiplier * count); - - // Handle ascending vs descending sequences - let (min_val, max_val) = if multiplier.is_ge(P::zero()) { - (base, end_element) + // Use the base's signedness to preserve values above `i64::MAX`. + let (towards, offset) = if base.ptype().is_signed_int() { + let base = base.cast::().vortex_expect("base fits its ptype"); + let Ok(value) = value.cast::() else { + return Some(Intersection::None); + }; + (value >= base, base.abs_diff(value)) } else { - (end_element, base) + let base = base.cast::().vortex_expect("base fits its ptype"); + let Ok(value) = value.cast::() else { + return Some(Intersection::None); + }; + (value >= base, base.abs_diff(value)) }; - // Check if intercept is in range - if !intercept.is_ge(min_val) || !intercept.is_le(max_val) { - vortex_bail!("{intercept} is outside of ({min_val}, {max_val}) range") - } - - // Handle zero multiplier (constant sequence) - if multiplier == P::zero() { - if intercept == base { - return Ok(0); + if offset == 0 { + return Some(if magnitude == 0 { + Intersection::All } else { - vortex_bail!("{intercept} != {base} with zero multiplier") - } + Intersection::At(0) + }); } - - // Check if (intercept - base) is evenly divisible by multiplier - let diff = intercept - base; - if diff % multiplier != P::zero() { - vortex_bail!("{diff} % {multiplier} != 0") + if magnitude == 0 || towards != ascending || offset % magnitude != 0 { + return Some(Intersection::None); } - let idx = diff / multiplier; - idx.to_usize() - .ok_or_else(|| vortex_err!("Cannot represent {idx} as usize")) + Some(match usize::try_from(offset / magnitude) { + Ok(idx) if idx < len => Intersection::At(idx), + _ => Intersection::None, + }) } #[cfg(test)] @@ -145,7 +133,10 @@ mod tests { use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::Nullability::NonNullable; use vortex_array::dtype::Nullability::Nullable; + use vortex_array::dtype::PType; + use vortex_array::scalar::PValue; use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::Sequence; @@ -191,4 +182,66 @@ mod tests { let expected = BoolArray::from_iter([false, false, false, false]); assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); } + + #[test] + fn test_compare_descending_unsigned() -> VortexResult<()> { + let lhs = Sequence::try_new( + PValue::from(100i32), + PValue::from(-10i32), + PType::U8, + NonNullable, + 5, + )?; + let rhs = ConstantArray::new(80u8, lhs.len()); + let result = lhs.into_array().binary(rhs.into_array(), Operator::Eq)?; + let expected = BoolArray::from_iter([false, false, true, false, false]); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); + + Ok(()) + } + + #[test] + fn test_compare_past_i64_max() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let step = (1u64 << 63) + 1; + let lhs = Sequence::try_new( + PValue::from(1u64 << 62), + PValue::from(step), + PType::U64, + NonNullable, + 2, + )?; + + let hit = lhs.clone().into_array().binary( + ConstantArray::new((1u64 << 62) + step, lhs.len()).into_array(), + Operator::Eq, + )?; + assert_arrays_eq!(hit, BoolArray::from_iter([false, true]), &mut ctx); + + let miss = lhs + .into_array() + .binary(ConstantArray::new(u64::MAX, 2).into_array(), Operator::Eq)?; + assert_arrays_eq!(miss, BoolArray::from_iter([false, false]), &mut ctx); + + Ok(()) + } + + #[test] + fn test_compare_constant_sequence() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let lhs = Sequence::try_new_typed(100i32, 0i32, NonNullable, 5)?; + + let matches = lhs.clone().into_array().binary( + ConstantArray::new(100i32, lhs.len()).into_array(), + Operator::Eq, + )?; + assert_arrays_eq!(matches, BoolArray::from_iter([true; 5]), &mut ctx); + + let misses = lhs + .into_array() + .binary(ConstantArray::new(7i32, 5).into_array(), Operator::Eq)?; + assert_arrays_eq!(misses, BoolArray::from_iter([false; 5]), &mut ctx); + + Ok(()) + } } diff --git a/encodings/sequence/src/compute/filter.rs b/encodings/sequence/src/compute/filter.rs index aae06610d82..52da57654f5 100644 --- a/encodings/sequence/src/compute/filter.rs +++ b/encodings/sequence/src/compute/filter.rs @@ -7,8 +7,7 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::filter::FilterKernel; -use vortex_array::dtype::NativePType; -use vortex_array::match_each_native_ptype; +use vortex_array::match_each_integer_ptype; use vortex_array::validity::Validity; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; @@ -16,6 +15,8 @@ use vortex_error::VortexResult; use vortex_mask::Mask; use crate::Sequence; +use crate::eval; +use crate::eval::SequenceValue; impl FilterKernel for Sequence { fn filter( @@ -24,23 +25,29 @@ impl FilterKernel for Sequence { _ctx: &mut ExecutionCtx, ) -> VortexResult> { let validity = Validity::from(array.dtype().nullability()); - match_each_native_ptype!(array.ptype(), |P| { - let mul = array.multiplier().cast::

()?; - let base = array.base().cast::

()?; - Ok(Some(filter_impl(mul, base, mask, validity))) + match_each_integer_ptype!(array.dtype().as_ptype(), |O| { + let (base, multiplier) = array.wrapping_parts::()?; + Ok(Some(filter_impl(base, multiplier, mask, validity))) }) } } -fn filter_impl(mul: T, base: T, mask: &Mask, validity: Validity) -> ArrayRef { +fn filter_impl( + base: O, + multiplier: O, + mask: &Mask, + validity: Validity, +) -> ArrayRef { let mask_values = mask .values() .vortex_expect("FilterKernel precondition: mask is Mask::Values"); - let mut buffer = BufferMut::::with_capacity(mask_values.true_count()); - buffer.extend(mask_values.indices().iter().map(|&idx| { - let i = T::from_usize(idx).vortex_expect("all valid indices fit"); - base + i * mul - })); + let mut buffer = BufferMut::::with_capacity(mask_values.true_count()); + buffer.extend( + mask_values + .indices() + .iter() + .map(|&idx| eval::wrapping_value(base, multiplier, idx)), + ); PrimitiveArray::new(buffer.freeze(), validity).into_array() } @@ -54,6 +61,8 @@ mod tests { use vortex_array::compute::conformance::filter::MEDIUM_SIZE; use vortex_array::compute::conformance::filter::test_filter_conformance; use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::scalar::PValue; use crate::Sequence; use crate::SequenceArray; @@ -71,6 +80,18 @@ mod tests { #[case(Sequence::try_new_typed(0u32, 1, Nullability::NonNullable, 5).unwrap())] #[case(Sequence::try_new_typed(0u32, 5, Nullability::NonNullable, MEDIUM_SIZE).unwrap())] #[case(Sequence::try_new_typed(0u64, 1, Nullability::NonNullable, LARGE_SIZE).unwrap())] + #[case::descending_u8( + Sequence::try_new(PValue::from(200i32), PValue::from(-3i32), PType::U8, + Nullability::NonNullable, 60).unwrap() + )] + #[case::past_i64_max( + Sequence::try_new(PValue::from(0i64), PValue::from(1i64 << 62), PType::U64, + Nullability::NonNullable, 4).unwrap() + )] + #[case::constant_longer_than_u8( + Sequence::try_new(PValue::from(7u8), PValue::from(0i32), PType::U8, + Nullability::NonNullable, MEDIUM_SIZE).unwrap() + )] fn test_filter_sequence_conformance(#[case] array: SequenceArray) { test_filter_conformance( &array.into_array(), diff --git a/encodings/sequence/src/compute/list_contains.rs b/encodings/sequence/src/compute/list_contains.rs index 93bd823a5f2..80ffcad24cd 100644 --- a/encodings/sequence/src/compute/list_contains.rs +++ b/encodings/sequence/src/compute/list_contains.rs @@ -5,12 +5,15 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::list_contains::ListContainsElementReduce; use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::array::Sequence; -use crate::compute::compare::find_intersection_scalar; +use crate::compute::compare::Intersection; +use crate::compute::compare::find_intersection; impl ListContainsElementReduce for Sequence { fn list_contains( @@ -26,23 +29,31 @@ impl ListContainsElementReduce for Sequence { .elements() .vortex_expect("non-null element (checked in entry)"); + let nullability = list.dtype().nullability() | element.dtype().nullability(); + let mut set_indices: Vec = Vec::new(); for intercept in list_elements.iter() { let Some(intercept) = intercept.as_primitive().pvalue() else { continue; }; - if let Ok(intersection) = find_intersection_scalar( + match find_intersection( element.base(), element.multiplier(), element.len(), intercept, ) { - set_indices.push(intersection) + // Non-integer elements do not match the sequence. + None | Some(Intersection::None) => {} + Some(Intersection::At(idx)) => set_indices.push(idx), + Some(Intersection::All) => { + return Ok(Some( + ConstantArray::new(Scalar::bool(true, nullability), element.len()) + .into_array(), + )); + } } } - let nullability = list.dtype().nullability() | element.dtype().nullability(); - Ok(Some( BoolArray::from_indices(element.len(), set_indices, nullability.into()).into_array(), )) @@ -110,4 +121,22 @@ mod tests { assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); } } + + #[test] + fn test_list_contains_constant_sequence() { + let list_scalar = Scalar::list( + Arc::new(I32.into()), + vec![7.into(), 42.into()], + Nullability::Nullable, + ); + + let array = Sequence::try_new_typed(42i32, 0, Nullability::NonNullable, 3) + .unwrap() + .into_array(); + + let expr = list_contains(lit(list_scalar), root()); + let result = array.apply(&expr).unwrap(); + let expected = BoolArray::from_iter([Some(true), Some(true), Some(true)]); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); + } } diff --git a/encodings/sequence/src/compute/min_max.rs b/encodings/sequence/src/compute/min_max.rs index 2ff8b2abd58..ab6d4faf3ab 100644 --- a/encodings/sequence/src/compute/min_max.rs +++ b/encodings/sequence/src/compute/min_max.rs @@ -9,14 +9,13 @@ use vortex_array::aggregate_fn::fns::min_max::make_minmax_dtype; use vortex_array::aggregate_fn::kernels::DynAggregateKernel; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; -use vortex_array::match_each_pvalue; -use vortex_array::scalar::PValue; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; use vortex_error::VortexResult; +use vortex_error::vortex_err; use crate::Sequence; -use crate::SequenceData; +use crate::eval; /// Sequence-specific min/max kernel. /// @@ -47,25 +46,20 @@ impl DynAggregateKernel for SequenceMinMaxKernel { return Ok(Some(Scalar::null(struct_dtype))); } - let base = seq.base(); - let last = SequenceData::try_last(base, seq.multiplier(), seq.ptype(), seq.len())?; - - // Determine min and max based on multiplier direction. - // For unsigned types, multiplier is always >= 0. - let (min_pvalue, max_pvalue) = match_each_pvalue!( - seq.multiplier(), - uint: |_v| { (base, last) }, - int: |v| { - if v >= 0 { - (base, last) - } else { - (last, base) - } - }, - float: |_v| { unreachable!("float multiplier not supported for SequenceArray") } - ); - - let non_nullable_dtype = DType::Primitive(seq.ptype(), Nullability::NonNullable); + let output_ptype = seq.dtype().as_ptype(); + + // A sequence's extrema are its first and last values. + let last = seq.index_value(seq.len() - 1); + let (ascending, _) = eval::step_parts(seq.multiplier()) + .ok_or_else(|| vortex_err!("step {} must be an integer", seq.multiplier()))?; + + let (min_pvalue, max_pvalue) = if ascending { + (seq.base(), last) + } else { + (last, seq.base()) + }; + + let non_nullable_dtype = DType::Primitive(output_ptype, Nullability::NonNullable); let min_scalar = Scalar::try_new( non_nullable_dtype.clone(), Some(ScalarValue::Primitive(min_pvalue)), @@ -79,3 +73,49 @@ impl DynAggregateKernel for SequenceMinMaxKernel { ))) } } + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::NumericalAggregateOpts; + use vortex_array::aggregate_fn::fns::min_max::MinMaxResult; + use vortex_array::aggregate_fn::fns::min_max::min_max; + use vortex_array::builtins::ArrayBuiltins; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::scalar::Scalar; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + use vortex_session::VortexSession; + + use crate::Sequence; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + + #[test] + fn min_max_uses_output_dtype() -> VortexResult<()> { + let array = Sequence::try_new_typed(100i32, -10i32, Nullability::NonNullable, 5)? + .into_array() + .cast(DType::Primitive(PType::U8, Nullability::NonNullable))?; + + let MinMaxResult { min, max } = min_max( + &array, + &mut SESSION.create_execution_ctx(), + NumericalAggregateOpts::default(), + )? + .ok_or_else(|| vortex_err!("min_max of a non-empty sequence should not be null"))?; + + assert_eq!(min, Scalar::from(60u8)); + assert_eq!(max, Scalar::from(100u8)); + + Ok(()) + } +} diff --git a/encodings/sequence/src/compute/mod.rs b/encodings/sequence/src/compute/mod.rs index 22b4e3ac05b..74539f302d5 100644 --- a/encodings/sequence/src/compute/mod.rs +++ b/encodings/sequence/src/compute/mod.rs @@ -18,6 +18,8 @@ mod tests { use vortex_array::array_session; use vortex_array::compute::conformance::consistency::test_array_consistency; use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::scalar::PValue; use crate::Sequence; use crate::SequenceArray; @@ -71,6 +73,28 @@ mod tests { 5 ).unwrap())] // Results in [100, 90, 80, 70, 60] + // Steps that are not representable in the output ptype + #[case::sequence_descending_u8(Sequence::try_new( + PValue::from(200i32), + PValue::from(-3i32), // Every value fits u8, the step does not + PType::U8, + Nullability::NonNullable, + 60 + ).unwrap())] + #[case::sequence_past_i64_max(Sequence::try_new( + PValue::from(0i64), + PValue::from(1i64 << 62), // Values run past i64::MAX + PType::U64, + Nullability::NonNullable, + 4 + ).unwrap())] + #[case::sequence_constant_longer_than_u8(Sequence::try_new( + PValue::from(7u8), + PValue::from(0i32), + PType::U8, + Nullability::NonNullable, + 500 + ).unwrap())] // Large arrays #[case::sequence_large(Sequence::try_new_typed( 0i64, diff --git a/encodings/sequence/src/compute/slice.rs b/encodings/sequence/src/compute/slice.rs index c2b64b68ef4..8d38c4076b8 100644 --- a/encodings/sequence/src/compute/slice.rs +++ b/encodings/sequence/src/compute/slice.rs @@ -19,7 +19,7 @@ impl SliceReduce for Sequence { Sequence::new_unchecked( array.index_value(range.start), array.multiplier(), - array.ptype(), + array.dtype().as_ptype(), array.dtype().nullability(), range.len(), ) diff --git a/encodings/sequence/src/compute/take.rs b/encodings/sequence/src/compute/take.rs index c59842224e2..2d3e9a01d68 100644 --- a/encodings/sequence/src/compute/take.rs +++ b/encodings/sequence/src/compute/take.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use num_traits::cast::NumCast; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; @@ -11,24 +10,23 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::dict::TakeExecute; use vortex_array::dtype::DType; use vortex_array::dtype::IntegerPType; -use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::match_each_integer_ptype; -use vortex_array::match_each_native_ptype; use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::Buffer; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_panic; use vortex_mask::AllOr; use vortex_mask::Mask; use crate::Sequence; +use crate::eval; +use crate::eval::SequenceValue; -fn take_inner( - mul: S, - base: S, +fn take_inner( + base: O, + multiplier: O, indices: &[T], indices_mask: Mask, result_nullability: Nullability, @@ -40,14 +38,13 @@ fn take_inner( if i.as_() >= len { vortex_panic!(OutOfBounds: i.as_(), 0, len); } - let i = ::from::(*i).vortex_expect("all indices fit"); - base + i * mul + eval::wrapping_value(base, multiplier, i.as_()) })), Validity::from(result_nullability), ) .into_array(), AllOr::None => ConstantArray::new( - Scalar::null(DType::Primitive(S::PTYPE, Nullability::Nullable)), + Scalar::null(DType::Primitive(O::PTYPE, Nullability::Nullable)), indices.len(), ) .into_array(), @@ -59,11 +56,9 @@ fn take_inner( vortex_panic!(OutOfBounds: i.as_(), 0, len); } - let i = - ::from::(*i).vortex_expect("all valid indices fit"); - base + i * mul + eval::wrapping_value(base, multiplier, i.as_()) } else { - S::zero() + O::zero() } })); PrimitiveArray::new(buffer, Validity::from(b.clone())).into_array() @@ -71,6 +66,41 @@ fn take_inner( } } +fn take_with_typed_indices( + array: ArrayView<'_, Sequence>, + indices: &[T], + indices_mask: Mask, + result_nullability: Nullability, +) -> VortexResult { + match_each_integer_ptype!(array.dtype().as_ptype(), |O| { + let (base, multiplier) = array.wrapping_parts::()?; + Ok(take_inner::( + base, + multiplier, + indices, + indices_mask, + result_nullability, + array.len(), + )) + }) +} + +fn take_sequence( + array: ArrayView<'_, Sequence>, + indices: &PrimitiveArray, + indices_mask: Mask, + result_nullability: Nullability, +) -> VortexResult { + match_each_integer_ptype!(indices.ptype(), |T| { + take_with_typed_indices::( + array, + indices.as_slice::(), + indices_mask, + result_nullability, + ) + }) +} + impl TakeExecute for Sequence { fn take( array: ArrayView<'_, Self>, @@ -81,21 +111,7 @@ impl TakeExecute for Sequence { let indices = indices.clone().execute::(ctx)?; let result_nullability = array.dtype().nullability() | indices.dtype().nullability(); - match_each_integer_ptype!(indices.ptype(), |T| { - let indices = indices.as_slice::(); - match_each_native_ptype!(array.ptype(), |S| { - let mul = array.multiplier().cast::()?; - let base = array.base().cast::()?; - Ok(Some(take_inner( - mul, - base, - indices, - mask, - result_nullability, - array.len(), - ))) - }) - }) + take_sequence(array, &indices, mask, result_nullability).map(Some) } } @@ -109,6 +125,8 @@ mod test { use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::take::test_take_conformance; use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::scalar::PValue; use crate::Sequence; use crate::SequenceArray; @@ -162,6 +180,27 @@ mod test { Nullability::Nullable, 1000 ).unwrap())] + #[case::sequence_descending_u8(Sequence::try_new( + PValue::from(200i32), + PValue::from(-3i32), + PType::U8, + Nullability::NonNullable, + 60 + ).unwrap())] + #[case::sequence_past_i64_max(Sequence::try_new( + PValue::from(0i64), + PValue::from(1i64 << 62), + PType::U64, + Nullability::NonNullable, + 4 + ).unwrap())] + #[case::sequence_constant_longer_than_u8(Sequence::try_new( + PValue::from(7u8), + PValue::from(0i32), + PType::U8, + Nullability::NonNullable, + 500 + ).unwrap())] fn sequence_take_conformance(#[case] sequence: SequenceArray) { test_take_conformance( &sequence.into_array(), diff --git a/encodings/sequence/src/eval.rs b/encodings/sequence/src/eval.rs new file mode 100644 index 00000000000..96f59bddac3 --- /dev/null +++ b/encodings/sequence/src/eval.rs @@ -0,0 +1,83 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Evaluation of the sequence equation `A[i] = base + i * multiplier`. +//! +//! Arithmetic remains exact within 64 bits by keeping unsigned values out of signed types. +//! Materialization uses wrapping arithmetic in the output ptype. + +use num_traits::AsPrimitive; +use num_traits::WrappingAdd; +use num_traits::WrappingMul; +use vortex_array::dtype::IntegerPType; +use vortex_array::match_each_pvalue; +use vortex_array::scalar::PValue; + +/// An integer type that sequence values can be materialized in. +pub(crate) trait SequenceValue: IntegerPType + WrappingAdd + WrappingMul { + /// The low bits of a sign-extended `i64`, reinterpreted as `Self`. + fn wrapping_from_i64(value: i64) -> Self; + + /// The low bits of a `u64`, reinterpreted as `Self`. + fn wrapping_from_u64(value: u64) -> Self; + + /// The low bits of an index, reinterpreted as `Self`. + fn wrapping_from_usize(index: usize) -> Self; +} + +impl SequenceValue for T +where + T: IntegerPType + WrappingAdd + WrappingMul, + i64: AsPrimitive, + u64: AsPrimitive, + usize: AsPrimitive, +{ + #[inline] + fn wrapping_from_i64(value: i64) -> Self { + value.as_() + } + + #[inline] + fn wrapping_from_u64(value: u64) -> Self { + value.as_() + } + + #[inline] + fn wrapping_from_usize(index: usize) -> Self { + index.as_() + } +} + +/// Reduces `base` and `multiplier` into the arithmetic type `O`. +pub(crate) fn wrapping_parts(base: PValue, multiplier: PValue) -> Option<(O, O)> { + Some((wrapping_from(base)?, wrapping_from(multiplier)?)) +} + +/// The two's-complement bits of an integer [`PValue`], reduced into `O`. `None` for a float. +fn wrapping_from(value: PValue) -> Option { + match_each_pvalue!( + value, + uint: |v| { Some(O::wrapping_from_u64(v.as_())) }, + int: |v| { Some(O::wrapping_from_i64(v.as_())) }, + float: |_v| { None } + ) +} + +/// Returns an integer step's direction and magnitude. +pub(crate) fn step_parts(step: PValue) -> Option<(bool, u64)> { + match_each_pvalue!( + step, + uint: |v| { Some((true, v.as_())) }, + int: |v| { + let v: i64 = v.as_(); + Some((v >= 0, v.unsigned_abs())) + }, + float: |_v| { None } + ) +} + +/// The sequence value at `index`, computed modulo `2^bits`. +#[inline] +pub(crate) fn wrapping_value(base: O, multiplier: O, index: usize) -> O { + base.wrapping_add(&multiplier.wrapping_mul(&O::wrapping_from_usize(index))) +} diff --git a/encodings/sequence/src/lib.rs b/encodings/sequence/src/lib.rs index b898963c346..6c1a4acc5e8 100644 --- a/encodings/sequence/src/lib.rs +++ b/encodings/sequence/src/lib.rs @@ -4,7 +4,10 @@ mod array; mod compress; mod compute; +mod eval; mod kernel; +#[cfg(test)] +mod model_tests; mod rules; /// Represents the equation A\[i\] = a * i + b. diff --git a/encodings/sequence/src/model_tests.rs b/encodings/sequence/src/model_tests.rs new file mode 100644 index 00000000000..eb93c4199c3 --- /dev/null +++ b/encodings/sequence/src/model_tests.rs @@ -0,0 +1,221 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Differential tests of sequence construction and its kernels against an exact `i128` model. + +use num_traits::Bounded; +use num_traits::ToPrimitive; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::fns::min_max::MinMaxResult; +use vortex_array::aggregate_fn::fns::min_max::min_max; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::bool::BoolArrayExt; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability::NonNullable; +use vortex_array::dtype::PType; +use vortex_array::match_each_integer_ptype; +use vortex_array::scalar::PValue; +use vortex_array::scalar::Scalar; +use vortex_array::scalar::ScalarValue; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_mask::Mask; + +use crate::Sequence; + +// Keep the `i128` oracle independent of the production code's 64-bit arithmetic. +fn widen(value: PValue) -> Option { + vortex_array::match_each_pvalue!( + value, + uint: |v| { Some(i128::from(v)) }, + int: |v| { Some(i128::from(v)) }, + float: |_v| { None } + ) +} + +fn narrow(value: i128, ptype: PType) -> Option { + match_each_integer_ptype!(ptype, |O| { + num_traits::cast::(value).map(PValue::from) + }) +} + +fn exact_value(base: PValue, multiplier: PValue, index: usize) -> Option { + let base = widen(base)?; + let multiplier = widen(multiplier)?; + i128::try_from(index) + .ok() + .and_then(|index| multiplier.checked_mul(index)) + .and_then(|offset| base.checked_add(offset)) +} + +const VALUES: [PValue; 15] = [ + PValue::I8(i8::MIN), + PValue::I8(i8::MAX), + PValue::I16(-3), + PValue::I32(-10), + PValue::I32(0), + PValue::I32(100), + PValue::I64(i64::MIN), + PValue::I64(i64::MAX), + PValue::I64(1 << 62), + PValue::I64(-255), + PValue::U8(0), + PValue::U8(u8::MAX), + PValue::U32(1000), + PValue::U64(u64::MAX), + PValue::U64(1 << 63), +]; + +const PTYPES: [PType; 8] = [ + PType::I8, + PType::I16, + PType::I32, + PType::I64, + PType::U8, + PType::U16, + PType::U32, + PType::U64, +]; + +const LENGTHS: [usize; 4] = [1, 2, 5, 300]; + +fn model(base: PValue, multiplier: PValue, ptype: PType, length: usize) -> Option> { + let (min, max) = match_each_integer_ptype!(ptype, |P| { + ( +

::min_value().to_i128()?, +

::max_value().to_i128()?, + ) + }); + + (0..length) + .map(|idx| exact_value(base, multiplier, idx).filter(|value| (min..=max).contains(value))) + .collect() +} + +fn values(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { + let primitive = array.clone().execute::(ctx)?; + Ok(match_each_integer_ptype!(primitive.ptype(), |P| { + primitive + .as_slice::

() + .iter() + .filter_map(|value| value.to_i128()) + .collect() + })) +} + +fn constant(value: i128, ptype: PType, len: usize) -> VortexResult { + let value = narrow(value, ptype) + .ok_or_else(|| vortex_err!("{value} is not representable in {ptype}"))?; + let scalar = Scalar::try_new( + DType::Primitive(ptype, NonNullable), + Some(ScalarValue::Primitive(value)), + )?; + Ok(ConstantArray::new(scalar, len).into_array()) +} + +#[test] +fn sequence_kernels_match_exact_model() -> VortexResult<()> { + let session = vortex_array::array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + let mut built = 0; + for base in VALUES { + for multiplier in VALUES { + for ptype in PTYPES { + for length in LENGTHS { + let expected = model(base, multiplier, ptype, length); + let array = Sequence::try_new(base, multiplier, ptype, NonNullable, length); + let case = format!("base {base:?}, step {multiplier:?}, {ptype}, len {length}"); + + let Some(expected) = expected else { + assert!(array.is_err(), "built an unrepresentable sequence: {case}"); + continue; + }; + let array = array + .map_err(|e| e.with_context(format!("representable sequence: {case}")))? + .into_array(); + built += 1; + + assert_eq!(values(&array, &mut ctx)?, expected, "canonical: {case}"); + + for idx in [0, length / 2, length - 1] { + let scalar = array.clone().execute_scalar(idx, &mut ctx)?; + assert_eq!(scalar.dtype(), array.dtype(), "scalar dtype: {case}"); + assert_eq!( + widen(scalar.as_primitive().pvalue().unwrap()).unwrap(), + expected[idx], + "scalar at {idx}: {case}" + ); + } + + let indices = PrimitiveArray::from_iter((0..length as u64).rev()).into_array(); + let taken = array.clone().take(indices)?; + let reversed = expected.iter().copied().rev().collect::>(); + assert_eq!(values(&taken, &mut ctx)?, reversed, "take: {case}"); + + let mask = Mask::from_iter((0..length).map(|idx| idx % 2 == 0)); + let filtered = array.clone().filter(mask)?; + let every_other = expected.iter().copied().step_by(2).collect::>(); + assert_eq!(values(&filtered, &mut ctx)?, every_other, "filter: {case}"); + + if length > 2 { + let sliced = array.clone().slice(1..length - 1)?; + assert_eq!( + values(&sliced, &mut ctx)?, + expected[1..length - 1], + "slice: {case}" + ); + } + + let MinMaxResult { min, max } = + min_max(&array, &mut ctx, NumericalAggregateOpts::default())? + .expect("min_max of a non-empty sequence is not null"); + assert_eq!( + widen(min.as_primitive().pvalue().unwrap()).unwrap(), + *expected.iter().min().unwrap(), + "min: {case}" + ); + assert_eq!( + widen(max.as_primitive().pvalue().unwrap()).unwrap(), + *expected.iter().max().unwrap(), + "max: {case}" + ); + + let held = expected[length / 2]; + for value in [Some(held), held.checked_add(1).filter(|v| *v != held)] { + let Some(value) = value.filter(|value| model_holds(ptype, *value)) else { + continue; + }; + let compared = array + .clone() + .binary(constant(value, ptype, length)?, Operator::Eq)? + .execute::(&mut ctx)?; + let bits = compared.to_bit_buffer(); + let matches = (0..length).map(|idx| bits.value(idx)).collect::>(); + let expected_matches = expected + .iter() + .map(|element| *element == value) + .collect::>(); + assert_eq!(matches, expected_matches, "compare {value}: {case}"); + } + } + } + } + } + + assert!(built > 1000, "only built {built} sequences"); + Ok(()) +} + +fn model_holds(ptype: PType, value: i128) -> bool { + narrow(value, ptype).is_some() +} diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index 0be6d0a4162..d2317dbae2c 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -4,8 +4,11 @@ use std::sync::Arc; use itertools::Itertools; +use num_traits::AsPrimitive; use num_traits::NumCast; use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::BoolArray; @@ -22,6 +25,8 @@ use vortex_array::arrays::VarBinView; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; +use vortex_array::arrays::listview::ListViewRebuildMode; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::arrays::varbinview::build_views::BinaryView; @@ -30,7 +35,9 @@ use vortex_array::builders::ArrayBuilder; use vortex_array::builders::DecimalBuilder; use vortex_array::builders::FixedSizeListBuilder; use vortex_array::builders::ListViewBuilder; -use vortex_array::builders::builder_with_capacity; +use vortex_array::builders::VarBinBuilder; +use vortex_array::builders::VarBinViewBuilder; +use vortex_array::builders::builder_with_capacity_in; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::DecimalType; @@ -38,12 +45,13 @@ use vortex_array::dtype::IntegerPType; use vortex_array::dtype::NativeDecimalType; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; +use vortex_array::dtype::OffsetBuilderPType; use vortex_array::dtype::StructFields; use vortex_array::match_each_decimal_value_type; use vortex_array::match_each_integer_ptype; use vortex_array::match_each_native_ptype; use vortex_array::match_each_unsigned_integer_ptype; -use vortex_array::match_smallest_offset_type; +use vortex_array::match_smallest_list_offset_type; use vortex_array::patches::Patches; use vortex_array::scalar::DecimalScalar; use vortex_array::scalar::ListScalar; @@ -60,11 +68,198 @@ use vortex_error::VortexError; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_mask::Mask; use crate::ConstantArray; use crate::Sparse; +use crate::SparseExt; use crate::SparseParts; +/// The pieces every string decode of a sparse array starts from: the offset-resolved patch +/// indices and canonical patch values, plus the fill value's bytes (`None` for a null fill). +struct SparseStringParts { + len: usize, + patches: Patches, + indices: PrimitiveArray, + values: VarBinViewArray, + fill: Option, +} + +impl SparseStringParts { + fn new(array: ArrayView<'_, Sparse>, ctx: &mut ExecutionCtx) -> VortexResult { + let patches = array.resolved_patches()?; + let indices = patches.indices().clone().execute::(ctx)?; + let values = patches + .values() + .clone() + .execute::(ctx)? + .into_varbinview(); + let fill_scalar = array.data().fill_scalar(); + let fill = match array.dtype() { + DType::Utf8(_) => fill_scalar + .as_utf8() + .value() + .cloned() + .map(BufferString::into_inner), + DType::Binary(_) => fill_scalar.as_binary().value().cloned(), + dtype => vortex_bail!("Sparse string decode of non-string dtype {dtype}"), + }; + Ok(Self { + len: array.len(), + patches, + indices, + values, + fill, + }) + } +} + +/// Scatters the patch views over the fill value straight into `builder`. +/// +/// The canonical route materializes the same scattered views buffer, wraps it in an array, and +/// then `append_varbinview_array` walks all `len` views a second time to rebase their buffer +/// indices onto the builder's. Scattering into the builder instead adopts the patch buffers once +/// and costs one bulk view fill plus one view write per patch, with no byte copy. +pub(super) fn append_sparse_to_varbinview( + array: ArrayView<'_, Sparse>, + builder: &mut VarBinViewBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let parts = SparseStringParts::new(array, ctx)?; + let validity = sparse_validity( + &parts.patches, + array.data().fill_scalar(), + array.dtype().nullability(), + parts.len, + ctx, + )? + .execute_mask(parts.len, ctx)?; + + let mut buffers = parts + .values + .data_buffers() + .iter() + .map(|buffer| buffer.as_host().clone()) + .collect::>(); + // A fill value short enough to inline lives in its view; only a longer one needs its bytes + // adopted as a (tiny) data buffer. + let fill_view = match &parts.fill { + Some(bytes) if bytes.len() > BinaryView::MAX_INLINED_SIZE => { + buffers.push(bytes.clone()); + BinaryView::make_view( + bytes.as_slice(), + u32::try_from(buffers.len() - 1).vortex_expect("too many buffers"), + 0, + ) + } + Some(bytes) => BinaryView::make_view(bytes.as_slice(), 0, 0), + None => BinaryView::make_view(&[], 0, 0), + }; + + let views = parts.values.views(); + match_each_integer_ptype!(parts.indices.ptype(), |I| { + let indices = parts.indices.as_slice::(); + builder.append_views_scattered( + buffers, + parts.len, + fill_view, + indices + .iter() + .zip(views.iter()) + .map(|(row, view)| (AsPrimitive::::as_(*row), *view)), + &validity, + ); + }); + Ok(()) +} + +/// Walks the rows in order, appending fill runs between the patches, straight into `builder`. +/// +/// The canonical route materializes the scattered views array first; the fill-value copies per +/// row are inherent to the offsets layout either way, so appending directly just skips that +/// intermediate. +pub(super) fn append_sparse_to_varbin( + array: ArrayView<'_, Sparse>, + builder: &mut VarBinBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> +where + usize: AsPrimitive, +{ + let parts = SparseStringParts::new(array, ctx)?; + let patch_validity = parts + .values + .validity()? + .execute_mask(parts.values.len(), ctx)?; + + let views = parts.values.views(); + let buffers = parts + .values + .data_buffers() + .iter() + .map(|buffer| buffer.as_host().as_slice()) + .collect::>(); + + match_each_integer_ptype!(parts.indices.ptype(), |I| { + append_patch_walk_to_varbin( + builder, + parts.indices.as_slice::(), + parts.len, + views, + &buffers, + &patch_validity, + parts.fill.as_ref(), + ) + }) +} + +/// The typed patch walk behind [`append_sparse_to_varbin`]. +fn append_patch_walk_to_varbin>( + builder: &mut VarBinBuilder, + indices: &[I], + len: usize, + views: &[BinaryView], + buffers: &[&[u8]], + patch_validity: &Mask, + fill: Option<&ByteBuffer>, +) -> VortexResult<()> +where + usize: AsPrimitive, +{ + let fill_run = |builder: &mut VarBinBuilder, n: usize| -> VortexResult<()> { + match fill { + Some(bytes) => builder.append_n_values(bytes, n), + None => { + builder.push_nulls(n); + Ok(()) + } + } + }; + + let mut previous = 0usize; + for (patch, row) in indices.iter().enumerate() { + let row = AsPrimitive::::as_(*row); + // The canonical decode scatters the patches in order, so of duplicate indices the + // last one wins; skip a patch that the next one lands on top of. + if patch + 1 < indices.len() && AsPrimitive::::as_(indices[patch + 1]) == row { + continue; + } + vortex_ensure!( + previous <= row && row < len, + "Sparse patch indices must be ascending within the array length {len}" + ); + fill_run(builder, row - previous)?; + if patch_validity.value(patch) { + builder.append_value(views[patch].bytes(buffers)); + } else { + builder.push_null(); + } + previous = row + 1; + } + fill_run(builder, len - previous) +} + fn sparse_validity( patches: &Patches, fill_value: &Scalar, @@ -151,6 +346,7 @@ pub(super) fn execute_sparse(parts: SparseParts, ctx: &mut ExecutionCtx) -> Vort DType::FixedSizeList(.., nullability) => { execute_sparse_fixed_size_list(&patches, &fill_value, len, *nullability, ctx)? } + DType::Map(..) => vortex_bail!("Sparse canonicalization does not support Map arrays yet"), DType::Struct(struct_fields, ..) => execute_sparse_struct( struct_fields, fill_value.as_struct(), @@ -174,21 +370,34 @@ fn execute_sparse_lists( ctx: &mut ExecutionCtx, ) -> VortexResult { // Patch indices are non-negative; reinterpret to unsigned so this dispatches over 4 widths - // instead of 8. `O` is already unsigned (from `match_smallest_offset_type`). + // instead of 8. `O` is already unsigned (from `match_smallest_list_offset_type`). let indices = resolved.indices().as_::().into_owned(); let indices = indices.reinterpret_cast(indices.ptype().to_unsigned()); - let values = resolved.values().as_::().into_owned(); let fill_list = fill_value.as_list(); + // Flatten the patch views up front so that runs of them can be appended as slices below. An + // exact layout also trims `elements` to exactly what the patches reference, so the patch half + // of the count below is what the builder will actually hold rather than an over-estimate. + let values = resolved + .values() + .clone() + .downcast::() + .rebuild(ListViewRebuildMode::MakeExact, ctx)?; + + // Each gap between patches is appended as one constant array, whose canonical form points + // every view at a single copy of the fill value, so a gap costs the fill value's elements once + // however many rows it covers. Bound the number of gaps: there is at most one on either side of + // each patch, and each one covers at least one row. let n_filled = len - resolved.num_patches(); - let total_canonical_values = values.elements().len() + fill_list.len() * n_filled; + let n_fill_runs = (resolved.num_patches() + 1).min(n_filled); + let total_canonical_values = values.elements().len() + fill_list.len() * n_fill_runs; Ok(match_each_unsigned_integer_ptype!(indices.ptype(), |I| { - match_smallest_offset_type!(total_canonical_values, |O| { + match_smallest_list_offset_type!(total_canonical_values, |O| { execute_sparse_lists_inner::( indices.as_slice(), values, - fill_list, + fill_value, values_dtype, len, total_canonical_values, @@ -200,10 +409,10 @@ fn execute_sparse_lists( } #[expect(clippy::too_many_arguments)] -fn execute_sparse_lists_inner( +fn execute_sparse_lists_inner( patch_indices: &[I], patch_values: ListViewArray, - fill_scalar: ListScalar, + fill_value: &Scalar, values_dtype: Arc, len: usize, total_canonical_values: usize, @@ -212,24 +421,27 @@ fn execute_sparse_lists_inner( ) -> ArrayRef { // Create the builder with appropriate types. It is easy to just use the same type for both // `offsets` and `sizes` since we have no other constraints. - let mut builder = ListViewBuilder::::with_capacity( + let mut builder = ListViewBuilder::::with_capacity_in( values_dtype, nullability, total_canonical_values, len, + ctx.allocator(), ); - let fill_elements = list_scalar_elements_array(fill_scalar); - let patch_values_validity = patch_values + // The fill's elements become an array once, up front. Every gap then appends that same array, + // so the fill's elements are stored once for the whole result however many gaps reference them. + let fill_elements = list_scalar_elements_array(fill_value.as_list(), ctx.allocator()); + + // One mask for the whole patch array rather than a validity lookup per patch. + let patch_validity = patch_values .listview_validity() .execute_mask(patch_values.len(), ctx) .vortex_expect("sparse list validity mask failed to execute"); let mut next_index = 0; - for ((patch_idx, sparse_idx), patch_valid) in patch_indices - .iter() - .enumerate() - .zip(patch_values_validity.iter()) + for ((patch_idx, sparse_idx), patch_valid) in + patch_indices.iter().enumerate().zip(patch_validity.iter()) { let sparse_idx = sparse_idx .to_usize() @@ -242,6 +454,9 @@ fn execute_sparse_lists_inner( ctx, ); + // Take each patch's elements rather than slicing the patch array itself: slicing a + // `ListView` slices its offsets, its sizes and its elements, and every one of those slices + // pays an optimizer pass, where this pays one for the elements alone. if patch_valid { let patch_list = patch_values .list_elements_at(patch_idx) @@ -261,6 +476,46 @@ fn execute_sparse_lists_inner( builder.finish() } +/// Materializes a list scalar's elements into an array, or `None` if the scalar is null. +fn list_scalar_elements_array( + list: ListScalar, + allocator: &vortex_buffer::BufferAllocatorRef, +) -> Option { + list.elements().map(|elements| { + let mut builder = builder_with_capacity_in(list.element_dtype(), elements.len(), allocator); + for element in elements { + builder + .append_scalar(&element) + .vortex_expect("list element scalar was invalid"); + } + builder.finish() + }) +} + +/// Appends the run of `count` fill lists that covers the gap before the next patch. +/// +/// The whole run goes in as one append that points `count` views at a single copy of +/// `fill_elements`, so a gap costs nothing per row it covers. +fn append_list_fill( + builder: &mut ListViewBuilder, + fill_elements: Option<&ArrayRef>, + count: usize, + ctx: &mut ExecutionCtx, +) { + if count == 0 { + return; + } + + match fill_elements { + Some(fill_elements) => builder + .append_array_as_repeated_list(fill_elements, count, ctx) + .vortex_expect("Failed to append sparse fill value"), + // A null fill has no elements to share, and the builder can record the nulls without + // going through an array at all. + None => builder.append_nulls(count), + } +} + /// Canonicalize a sparse [`FixedSizeListArray`] by expanding it into a dense representation. fn execute_sparse_fixed_size_list( resolved: &Patches, @@ -271,13 +526,12 @@ fn execute_sparse_fixed_size_list( ) -> VortexResult { let indices = resolved.indices().as_::().into_owned(); let values = resolved.values().as_::().into_owned(); - let fill_scalar = fill_value.as_list(); Ok(match_each_integer_ptype!(indices.ptype(), |I| { execute_sparse_fixed_size_list_inner::( indices.as_slice(), values, - fill_scalar, + fill_value, len, nullability, ctx, @@ -293,9 +547,9 @@ fn execute_sparse_fixed_size_list( /// elements (or defaults if null). Since all lists have the same size, we can directly append /// elements without tracking offsets. fn execute_sparse_fixed_size_list_inner( - indices: &[I], + patch_indices: &[I], values: FixedSizeListArray, - fill_scalar: ListScalar, + fill_value: &Scalar, array_len: usize, nullability: Nullability, ctx: &mut ExecutionCtx, @@ -305,26 +559,35 @@ fn execute_sparse_fixed_size_list_inner( .dtype() .as_fixed_size_list_element_opt() .vortex_expect("sparse fixed-size-list values must have fixed-size-list dtype"); - let mut builder = FixedSizeListBuilder::with_capacity( + let mut builder = FixedSizeListBuilder::with_capacity_in( Arc::clone(element_dtype), list_size, nullability, array_len, + ctx.allocator(), ); - let fill_elements = list_scalar_elements_array(fill_scalar); - let values_validity = values + // The fill's elements become an array once, up front, so that a gap does not rebuild them. + // They are tiled per row rather than shared - a fixed-size list holds its elements back to + // back - unless they are all the same scalar, in which case the tile stays constant-encoded + // and the tiling costs nothing. + let fill_elements = fixed_size_list_fill_tile(fill_value.as_list(), list_size, ctx.allocator()); + + // One mask for the whole patch array rather than a validity lookup per patch. + let patch_validity = values .validity() .vortex_expect("sparse fixed-size-list validity should be derivable") .execute_mask(values.len(), ctx) .vortex_expect("sparse fixed-size-list validity mask failed to execute"); let mut next_index = 0; - let indices = indices - .iter() - .map(|x| (*x).to_usize().vortex_expect("index must fit in usize")); - for ((patch_idx, sparse_idx), patch_valid) in indices.enumerate().zip(values_validity.iter()) { + for ((patch_idx, sparse_idx), patch_valid) in + patch_indices.iter().enumerate().zip(patch_validity.iter()) + { // Fill gap before this patch with fill values. + let sparse_idx = sparse_idx + .to_usize() + .vortex_expect("patch index must fit in usize"); append_fixed_size_list_fill( &mut builder, fill_elements.as_ref(), @@ -332,7 +595,9 @@ fn execute_sparse_fixed_size_list_inner( ctx, ); - // Append the patch value, handling null patches by appending defaults. + // Take each patch's elements rather than slicing the patch array itself: slicing a + // `FixedSizeList` slices its elements and its validity, and every one of those slices pays + // an optimizer pass, where this pays one for the elements alone. if patch_valid { let patch_list = values .fixed_size_list_elements_at(patch_idx) @@ -358,49 +623,50 @@ fn execute_sparse_fixed_size_list_inner( builder.finish_into_fixed_size_list() } -fn list_scalar_elements_array(list: ListScalar) -> Option { - list.elements().map(|elements| { - let mut builder = builder_with_capacity(list.element_dtype(), elements.len()); - for element in elements { - builder - .append_scalar(&element) - .vortex_expect("list element scalar was invalid"); +/// Materializes the elements a fixed-size-list fill value covers each of its rows with, or `None` +/// if the fill is null. +/// +/// Elements that are all the same scalar stay a constant array, so tiling them over a gap costs +/// nothing however many rows it covers. +fn fixed_size_list_fill_tile( + fill: ListScalar, + list_size: u32, + allocator: &vortex_buffer::BufferAllocatorRef, +) -> Option { + let elements = fill.elements()?; + + Some(match elements.iter().all_equal_value() { + Ok(uniform) => ConstantArray::new(uniform.clone(), list_size as usize).into_array(), + Err(_) => { + let mut builder = + builder_with_capacity_in(fill.element_dtype(), elements.len(), allocator); + for element in &elements { + builder + .append_scalar(element) + .vortex_expect("fixed-size-list element scalar was invalid"); + } + builder.finish() } - builder.finish() }) } -fn append_list_fill( - builder: &mut ListViewBuilder, - fill_elements: Option<&ArrayRef>, - count: usize, - ctx: &mut ExecutionCtx, -) { - if let Some(fill_elements) = fill_elements { - for _ in 0..count { - builder - .append_array_as_list(fill_elements, ctx) - .vortex_expect("Failed to append sparse fill value"); - } - } else { - builder.append_nulls(count); - } -} - +/// Appends the run of `count` fill lists that covers the gap before the next patch. fn append_fixed_size_list_fill( builder: &mut FixedSizeListBuilder, fill_elements: Option<&ArrayRef>, count: usize, ctx: &mut ExecutionCtx, ) { - if let Some(fill_elements) = fill_elements { - for _ in 0..count { - builder - .append_array_as_list(fill_elements, ctx) - .vortex_expect("Failed to append sparse fixed-size-list fill value"); - } - } else { - builder.append_nulls(count); + if count == 0 { + return; + } + + match fill_elements { + Some(fill_elements) => builder + .append_array_as_repeated_list(fill_elements, count, ctx) + .vortex_expect("Failed to append sparse fixed-size-list fill value"), + // A null fill has no elements of its own, only the placeholders the builder writes. + None => builder.append_nulls(count), } } @@ -474,7 +740,6 @@ fn execute_sparse_struct( ), }; let patch_values_as_struct = unresolved_patches.values().as_::().into_owned(); - let columns_patch_values = patch_values_as_struct.unmasked_fields(); let names = patch_values_as_struct.names(); let validity = top_level_fill_validity.patch( len, @@ -493,8 +758,8 @@ fn execute_sparse_struct( Ok(StructArray::try_from_iter_with_validity( names.iter().zip_eq( - columns_patch_values - .iter() + patch_values_as_struct + .iter_unmasked_fields() .cloned() .zip_eq(fill_values) .map(|(patch_values, fill_value)| unsafe { @@ -520,7 +785,8 @@ fn execute_sparse_decimal( len: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut builder = DecimalBuilder::with_capacity::(len, decimal_dtype, nullability); + let mut builder = + DecimalBuilder::with_capacity_in::(len, decimal_dtype, nullability, ctx.allocator()); match fill_value.decimal_value() { Some(fill_value) => { let fill_value = fill_value @@ -569,16 +835,20 @@ fn execute_varbin_inner( let n_patch_buffers = values.data_buffers().len(); let mut buffers = values.data_buffers().to_vec(); - let fill = if let Some(buffer) = &fill_value { - buffers.push(BufferHandle::new_host(buffer.clone())); - BinaryView::make_view( - buffer.as_ref(), - u32::try_from(n_patch_buffers).vortex_expect("too many buffers"), - 0, - ) - } else { + let fill = match &fill_value { + // A fill value short enough to inline lives in its view; pushing its buffer would leave + // that buffer entirely unreferenced. + Some(buffer) if buffer.len() > BinaryView::MAX_INLINED_SIZE => { + buffers.push(BufferHandle::new_host(buffer.clone())); + BinaryView::make_view( + buffer.as_ref(), + u32::try_from(n_patch_buffers).vortex_expect("too many buffers"), + 0, + ) + } + Some(buffer) => BinaryView::make_view(buffer.as_ref(), 0, 0), // any <=12 character value will do - BinaryView::make_view(&[], 0, 0) + None => BinaryView::make_view(&[], 0, 0), }; let mut views = buffer_mut![fill; len]; @@ -604,6 +874,7 @@ mod test { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; + use vortex_array::arrays::Chunked; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::ListArray; @@ -612,8 +883,12 @@ mod test { use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; + use vortex_array::arrays::chunked::ChunkedArrayExt; use vortex_array::arrays::listview::ListViewArrayExt; + use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::assert_arrays_eq; + use vortex_array::builders::VarBinBuilder; + use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::FieldNames; @@ -1302,6 +1577,149 @@ mod test { Ok(()) } + /// Each gap between patches is appended as a single constant array, so the fill value's + /// elements are stored once per gap however many rows the gap covers. + #[test] + fn test_sparse_list_fill_stores_one_copy_per_gap() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + + // Two single-element patch lists: [1] and [2]. + let lists = unsafe { + ListViewArray::new_unchecked( + buffer![1i32, 2].into_array(), + buffer![0u32, 1].into_array(), + buffer![1u32, 1].into_array(), + Validity::AllValid, + ) + .with_zero_copy_to_list(true) + } + .into_array(); + + // Patches at 10 and 20 of 10,000 rows, so the fill covers three gaps. + let indices = buffer![10u32, 20].into_array(); + let fill = vec![7i32, 8, 9]; + let sparse = + Sparse::try_new(indices, lists, 10_000, Scalar::from(Some(fill.clone())))?.into_array(); + + let actual = sparse.execute::(&mut ctx)?; + assert_eq!(actual.len(), 10_000); + assert_eq!( + actual.elements().len(), + 2 + 3 * fill.len(), + "the fill value should be stored once per gap, not once per row", + ); + + let fill_elements = PrimitiveArray::from_iter(fill); + for index in [0, 9, 11, 19, 21, 9_999] { + assert_arrays_eq!( + actual.list_elements_at(index).vortex_expect("fill list"), + fill_elements, + &mut ctx + ); + } + assert_arrays_eq!( + actual.list_elements_at(10).vortex_expect("patch list"), + PrimitiveArray::from_iter([1i32]), + &mut ctx + ); + assert_arrays_eq!( + actual.list_elements_at(20).vortex_expect("patch list"), + PrimitiveArray::from_iter([2i32]), + &mut ctx + ); + + Ok(()) + } + + /// Nested builders chunk a child on the boundaries it is appended on, so the number of appends + /// canonicalization makes is visible in the elements child. Patches go in one at a time, so + /// they cost a chunk each; a gap covers all its rows with a single append, so it costs one + /// chunk however many rows it fills. + #[test] + fn test_sparse_list_chunks_elements_per_patch_and_once_per_gap() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + + const PATCHES: usize = 100; + let patches_u32 = u32::try_from(PATCHES).vortex_expect("fits in u32"); + let patches_i32 = i32::try_from(PATCHES).vortex_expect("fits in i32"); + + // `PATCHES` single-element lists, patched onto rows 0..PATCHES of a 2 * PATCHES-row array, + // so there is exactly one patch run followed by exactly one gap. + let patch_values = ListViewArray::new( + PrimitiveArray::from_iter(0..patches_i32).into_array(), + PrimitiveArray::from_iter(0..patches_u32).into_array(), + PrimitiveArray::from_iter(std::iter::repeat_n(1u32, PATCHES)).into_array(), + Validity::AllValid, + ) + .into_array(); + + let indices = PrimitiveArray::from_iter(0..patches_u32).into_array(); + let fill = Scalar::from(Some(vec![-1i32])); + let sparse = Sparse::try_new(indices, patch_values, 2 * PATCHES, fill)?.into_array(); + + let actual = sparse.execute::(&mut ctx)?; + assert_eq!( + actual.elements().as_::().nchunks(), + PATCHES + 1, + "expected one chunk per patch and a single chunk for the whole gap", + ); + + let expected_lists = (0..patches_i32) + .map(|i| Some(vec![i])) + .chain(std::iter::repeat_n(Some(vec![-1i32]), PATCHES)); + let expected = ListArray::from_iter_opt_slow::( + expected_lists, + Arc::new(PType::I32.into()), + )?; + assert_arrays_eq!(actual, expected, &mut ctx); + + Ok(()) + } + + /// Patches on consecutive rows are appended as one slice of the patch array, so this covers the + /// run arithmetic together with everything that has to survive it: a null patch inside a run, + /// patch views that overlap and are out of order, runs separated by gaps, and a trailing gap. + #[test] + fn test_sparse_list_appends_consecutive_patches_as_one_run() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + + // Overlapping, out-of-order patch views over three elements: + // - patch 0: [20, 30] + // - patch 1: [10] (null, so its elements are never read) + // - patch 2: [10, 20, 30] + // - patch 3: [10, 20] + let patches = unsafe { + ListViewArray::new_unchecked( + buffer![10i32, 20, 30].into_array(), + buffer![1u32, 0, 0, 0].into_array(), + buffer![2u32, 1, 3, 2].into_array(), + Validity::from_iter([true, false, true, true]), + ) + }; + assert!(!patches.is_zero_copy_to_list()); + + // Rows 0..3 are one run of patches, row 5 is another, and rows 3-4 and 6 are gaps. + let indices = buffer![0u8, 1, 2, 5].into_array(); + let fill = Scalar::from(Some(vec![7i32, 8])); + let sparse = Sparse::try_new(indices, patches.into_array(), 7, fill)?.into_array(); + + let actual = sparse.execute::(&mut ctx)?; + + // The seven elements the patches reference, plus one copy of the two-element fill for each + // of the two gaps. The run of patches on rows 0..3 has no gap inside it to pay for. + assert_eq!(actual.elements().len(), 7 + 2 * 2); + + let expected = ListViewArray::new( + buffer![20i32, 30, 10, 20, 30, 7, 8, 7, 8, 10, 20, 7, 8].into_array(), + buffer![0u8, 2, 2, 5, 7, 9, 11].into_array(), + buffer![2u8, 0, 3, 2, 2, 2, 2].into_array(), + Validity::from_iter([true, false, true, true, true, true, true]), + ); + assert_arrays_eq!(actual, expected, &mut ctx); + + Ok(()) + } + #[test] fn test_sparse_binary_varbin_null_fill() { let mut ctx = SESSION.create_execution_ctx(); @@ -1565,15 +1983,17 @@ mod test { expected_elements.extend(buffer![42i32; 252]); let expected = ListArray::try_new( expected_elements.freeze().into_array(), - buffer![0u16, 1, 2, 3, 4, 256].into_array(), + buffer![0u32, 1, 2, 3, 4, 256].into_array(), Validity::AllValid, )? .into_array(); + // The canonical offsets outgrow the source's u8 offsets; u32 is the smallest offset type + // a list builder can produce (see `OffsetBuilderPType`). let actual_listview = actual.clone().execute::(&mut ctx)?; assert_eq!( actual_listview.offsets().dtype(), - &DType::Primitive(PType::U16, NonNullable) + &DType::Primitive(PType::U32, NonNullable) ); assert_arrays_eq!(&actual, &expected, &mut ctx); @@ -1756,4 +2176,75 @@ mod test { // when dealing with sliced arrays that have non-zero starting offsets. Ok(()) } + + /// Appending a sparse string array to either variable-binary builder must match its + /// canonical decode exactly — across long, short (inlinable), and null fill values, null + /// patch values, and a sliced array whose patch indices carry an offset. + #[rstest] + #[case::long_fill(Some("a fill value that is too long to inline"))] + #[case::short_fill(Some("123"))] + #[case::null_fill(None)] + fn test_sparse_append_to_string_builders(#[case] fill: Option<&str>) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let long = "a patch value that is far too long to inline"; + let strings = + >::from_iter([Some(long), None, Some("tiny")]) + .into_array(); + let fill_scalar = match fill { + Some(value) => Scalar::from(value.to_owned()).into_nullable(), + None => Scalar::null(DType::Utf8(Nullable)), + }; + let array = Sparse::try_new(buffer![1u16, 4, 8].into_array(), strings, 10, fill_scalar)? + .into_array(); + + for candidate in [array.clone(), array.slice(1..9)?] { + let expected = candidate.clone().execute::(&mut ctx)?; + + let mut view_builder = VarBinViewBuilder::with_capacity_in( + candidate.dtype().clone(), + 4, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); + candidate.append_to_builder(&mut view_builder, &mut ctx)?; + assert_arrays_eq!(view_builder.finish_into_varbinview(), expected, &mut ctx); + + let mut varbin_builder = VarBinBuilder::::new_in( + candidate.dtype().clone(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + candidate.append_to_builder(&mut varbin_builder, &mut ctx)?; + assert_arrays_eq!(varbin_builder.finish_into_varbin(), expected, &mut ctx); + } + Ok(()) + } + + /// A fill value short enough to inline lives in its views; the canonical decode must not + /// retain a data buffer nothing references. + #[test] + fn test_sparse_short_fill_pushes_no_buffer() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let strings = VarBinViewArray::from_iter_str(["hello", "goodbye"]); + assert!(strings.data_buffers().is_empty()); + + let array = Sparse::try_new( + buffer![1u16, 3].into_array(), + strings.into_array(), + 6, + Scalar::from("123".to_owned()), + )?; + + let actual = array + .as_array() + .clone() + .execute::(&mut ctx)?; + assert!( + actual.data_buffers().is_empty(), + "an inlinable fill value must not push an unreferenced buffer" + ); + + let expected = + VarBinViewArray::from_iter_str(["123", "hello", "123", "goodbye", "123", "123"]); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } } diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs index 4d4ec2623de..70f993c10ff 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -28,9 +28,12 @@ use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::bool::BoolArrayExt; use vortex_array::buffer::BufferHandle; +use vortex_array::builders::ArrayBuilder; +use vortex_array::builders::VarBinViewBuilder; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; +use vortex_array::match_each_varbin_builder; use vortex_array::patches::PatchSlotIndices; use vortex_array::patches::Patches; use vortex_array::patches::PatchesData; @@ -45,7 +48,6 @@ use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTable; use vortex_buffer::Buffer; -use vortex_buffer::ByteBufferMut; use vortex_error::VortexExpect as _; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -57,6 +59,8 @@ use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::canonical::append_sparse_to_varbin; +use crate::canonical::append_sparse_to_varbinview; use crate::canonical::execute_sparse; use crate::rules::RULES; @@ -117,8 +121,11 @@ pub type SparseArray = Array; #[vortex_array::array_slots(Sparse)] pub struct SparseSlots { + #[slot(0)] pub patch_indices: ArrayRef, + #[slot(1)] pub patch_values: ArrayRef, + #[slot(2)] pub patch_chunk_offsets: Option, } @@ -209,7 +216,7 @@ impl VTable for Sparse { match idx { 0 => { let fill_value_buffer = - ScalarValue::to_proto_bytes::(array.fill_value.value()).freeze(); + ScalarValue::to_proto_bytes::>(array.fill_value.value()).into(); BufferHandle::new_host(fill_value_buffer) } _ => vortex_panic!("SparseArray buffer index {idx} out of bounds"), @@ -350,6 +357,33 @@ impl VTable for Sparse { // TODO(joe): remove ctx from execute_sparse since all slots should be canonical. execute_sparse(parts, ctx).map(ExecutionResult::done) } + + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + // Strings decode as a scatter over the fill value; appending that scatter straight into + // a variable-binary builder skips materializing the canonical intermediate. + if matches!(array.dtype(), DType::Utf8(_) | DType::Binary(_)) { + if let Some(builder) = builder.as_any_mut().downcast_mut::() { + return append_sparse_to_varbinview(array, builder, ctx); + } + if let Some(result) = match_each_varbin_builder!(builder, |builder| { + append_sparse_to_varbin(array, builder, ctx) + }) { + return result; + } + } + + // Everything else decodes through the canonical array, like the default implementation. + array + .array() + .clone() + .execute::(ctx)? + .into_array() + .append_to_builder(builder, ctx) + } } const PATCH_SLOTS: PatchSlotIndices = PatchSlotIndices { @@ -697,7 +731,9 @@ mod test { use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; + use vortex_array::builders::VarBinBuilder; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -905,6 +941,52 @@ mod test { ); } + #[test] + fn test_append_utf8_to_dyn_varbin_builder() { + let mut ctx = SESSION.create_execution_ctx(); + let values = VarBinViewArray::from_iter_nullable_str([Some("patched"), None, Some("last")]) + .into_array(); + let fill = Scalar::null(values.dtype().clone()); + let array = Sparse::try_new(buffer![1u8, 3, 5].into_array(), values, 6, fill).unwrap(); + let expected = VarBinViewArray::from_iter_nullable_str([ + None, + Some("patched"), + None, + None, + None, + Some("last"), + ]) + .into_array(); + let mut builder = VarBinBuilder::::with_capacity_in( + array.dtype().clone(), + array.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + array.append_to_builder(&mut builder, &mut ctx).unwrap(); + assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx); + } + + #[test] + fn test_append_utf8_with_duplicate_patch_indices() { + let mut ctx = SESSION.create_execution_ctx(); + let values = VarBinViewArray::from_iter_str(["first", "second"]).into_array(); + let array = Sparse::try_new( + buffer![1u8, 1].into_array(), + values, + 2, + Scalar::utf8("fill".to_string(), Nullability::NonNullable), + ) + .unwrap(); + let expected = VarBinViewArray::from_iter_str(["fill", "second"]).into_array(); + let mut builder = VarBinBuilder::::with_capacity_in( + array.dtype().clone(), + array.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + array.append_to_builder(&mut builder, &mut ctx).unwrap(); + assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx); + } + #[test] #[should_panic] fn test_invalid_length() { diff --git a/encodings/zigzag/src/array.rs b/encodings/zigzag/src/array.rs index 05fb8a261e0..ef3165132f4 100644 --- a/encodings/zigzag/src/array.rs +++ b/encodings/zigzag/src/array.rs @@ -17,6 +17,7 @@ use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::IntoArray; use vortex_array::TypedArrayRef; +use vortex_array::array_slots; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::PType; @@ -62,9 +63,7 @@ impl VTable for ZigZag { len: usize, slots: &[Option], ) -> VortexResult<()> { - let encoded = slots[ENCODED_SLOT] - .as_ref() - .vortex_expect("ZigZagArray encoded slot"); + let encoded = ZigZagSlotsView::from_slots(slots).encoded; let expected_dtype = ZigZagData::dtype_from_encoded_dtype(encoded.dtype())?; vortex_ensure!( dtype == &expected_dtype, @@ -134,7 +133,7 @@ impl VTable for ZigZag { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + ZigZagSlots::NAMES[idx].to_string() } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { @@ -162,10 +161,12 @@ impl ArrayEq for ZigZagData { } } -/// The zigzag-encoded values (signed integers mapped to unsigned). -pub(super) const ENCODED_SLOT: usize = 0; -pub(super) const NUM_SLOTS: usize = 1; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["encoded"]; +#[array_slots(ZigZag)] +pub struct ZigZagSlots { + /// The zigzag-encoded values (signed integers mapped to unsigned). + #[slot(0)] + pub encoded: ArrayRef, +} #[derive(Clone, Debug)] pub struct ZigZagData {} @@ -176,13 +177,7 @@ impl Display for ZigZagData { } } -pub trait ZigZagArrayExt: TypedArrayRef { - fn encoded(&self) -> &ArrayRef { - self.as_ref().slots()[ENCODED_SLOT] - .as_ref() - .vortex_expect("ZigZagArray encoded slot") - } - +pub trait ZigZagArrayExt: ZigZagArraySlotsExt { fn ptype(&self) -> PType { PType::try_from(self.encoded().dtype()) .vortex_expect("ZigZagArray encoded dtype") diff --git a/encodings/zigzag/src/compute/cast.rs b/encodings/zigzag/src/compute/cast.rs index 94123e4498d..12064f432e8 100644 --- a/encodings/zigzag/src/compute/cast.rs +++ b/encodings/zigzag/src/compute/cast.rs @@ -10,7 +10,7 @@ use vortex_array::scalar_fn::fns::cast::CastReduce; use vortex_error::VortexResult; use crate::ZigZag; -use crate::array::ZigZagArrayExt; +use crate::array::ZigZagArraySlotsExt; impl CastReduce for ZigZag { fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { if !dtype.is_signed_int() { diff --git a/encodings/zigzag/src/compute/mod.rs b/encodings/zigzag/src/compute/mod.rs index e1c2f27f5d7..354593e284a 100644 --- a/encodings/zigzag/src/compute/mod.rs +++ b/encodings/zigzag/src/compute/mod.rs @@ -9,15 +9,13 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::dict::TakeExecute; use vortex_array::arrays::filter::FilterReduce; -use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; -use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::fns::mask::Mask as MaskExpr; use vortex_array::scalar_fn::fns::mask::MaskReduce; use vortex_error::VortexResult; use vortex_mask::Mask; use crate::ZigZag; -use crate::array::ZigZagArrayExt; +use crate::array::ZigZagArraySlotsExt; impl FilterReduce for ZigZag { fn filter(array: ArrayView<'_, Self>, mask: &Mask) -> VortexResult> { @@ -39,11 +37,7 @@ impl TakeExecute for ZigZag { impl MaskReduce for ZigZag { fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult> { - let masked_encoded = MaskExpr.try_new_array( - array.encoded().len(), - EmptyOptions, - [array.encoded().clone(), mask.clone()], - )?; + let masked_encoded = MaskExpr::try_new(array.encoded().clone(), mask.clone())?.into_array(); Ok(Some(ZigZag::try_new(masked_encoded)?.into_array())) } } diff --git a/encodings/zigzag/src/slice.rs b/encodings/zigzag/src/slice.rs index 6a7128d5111..d21e17302c0 100644 --- a/encodings/zigzag/src/slice.rs +++ b/encodings/zigzag/src/slice.rs @@ -10,7 +10,7 @@ use vortex_array::arrays::slice::SliceReduce; use vortex_error::VortexResult; use crate::ZigZag; -use crate::array::ZigZagArrayExt; +use crate::array::ZigZagArraySlotsExt; impl SliceReduce for ZigZag { fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { diff --git a/encodings/zstd/Cargo.toml b/encodings/zstd/Cargo.toml index 1b2cb4f2dd0..6ab7b2ab31b 100644 --- a/encodings/zstd/Cargo.toml +++ b/encodings/zstd/Cargo.toml @@ -19,15 +19,13 @@ all-features = true [lints] workspace = true -[features] -# This feature enables unstable encodings for which we don't guarantee stability. -unstable_encodings = [] - [dependencies] itertools = { workspace = true } +num-traits = { workspace = true } prost = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } +vortex-edition = { workspace = true } vortex-error = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index dc9f075e44e..fb3551e539b 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -6,9 +6,12 @@ use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hash; use std::hash::Hasher; +use std::mem::MaybeUninit; +use std::ops::Range; use std::sync::Arc; use itertools::Itertools as _; +use num_traits::AsPrimitive; use prost::Message as _; use vortex_array::Array; use vortex_array::ArrayEq; @@ -22,13 +25,19 @@ use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::IntoArray; +use vortex_array::array_slots; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::varbinview::build_views::BinaryView; use vortex_array::arrays::varbinview::build_views::MAX_BUFFER_LEN; use vortex_array::buffer::BufferHandle; +use vortex_array::builders::ArrayBuilder; +use vortex_array::builders::VarBinBuilder; +use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; +use vortex_array::dtype::OffsetBuilderPType; +use vortex_array::match_each_varbin_builder; use vortex_array::scalar::Scalar; use vortex_array::serde::ArrayChildren; use vortex_array::smallvec::smallvec; @@ -49,13 +58,15 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; -use vortex_error::vortex_panic; use vortex_mask::AllOr; +use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use zstd::zstd_safe::WriteBuf; use crate::ZstdFrameMetadata; use crate::ZstdMetadata; +use crate::validate_frame_content_size; // Zstd doesn't support training dictionaries on very few samples. const MIN_SAMPLES_FOR_DICTIONARY: usize = 8; @@ -143,7 +154,7 @@ impl VTable for Zstd { len: usize, slots: &[Option], ) -> VortexResult<()> { - let validity = child_to_validity(slots[0].as_ref(), dtype.nullability()); + let validity = child_to_validity(slots[ZstdSlots::VALIDITY].as_ref(), dtype.nullability()); data.validate(dtype, len, &validity) } @@ -253,12 +264,12 @@ impl VTable for Zstd { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + ZstdSlots::NAMES[idx].to_string() } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { let unsliced_validity = child_to_validity( - array.as_ref().slots()[0].as_ref(), + array.as_ref().slots()[ZstdSlots::VALIDITY].as_ref(), array.dtype().nullability(), ); array @@ -268,6 +279,26 @@ impl VTable for Zstd { .map(ExecutionResult::done) } + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if let Some(result) = + match_each_varbin_builder!(builder, |builder| append_to_varbin(array, builder, ctx)) + { + return result; + } + // The two arms here are every builder a `Utf8`/`Binary` dtype has: all four + // `VarBinBuilder` widths above, and `VarBinViewBuilder` below. There is deliberately no + // canonicalize-then-append fallback — it would decompress to a `VarBinView` only for + // `VarBinView::append_to_builder` to reject the same remainder. + let Some(builder) = builder.as_any_mut().downcast_mut::() else { + vortex_bail!("append_to_builder for Zstd requires a variable-binary builder") + }; + append_to_varbinview(array, builder, ctx) + } + fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, @@ -277,6 +308,111 @@ impl VTable for Zstd { } } +fn unsliced_validity(array: ArrayView<'_, Zstd>) -> Validity { + child_to_validity( + array.slots()[ZstdSlots::VALIDITY].as_ref(), + array.dtype().nullability(), + ) +} + +/// Copies the decompressed values straight into `builder`'s byte storage. +/// +/// The decompressed frames interleave a length prefix with each value, so the bytes have to be +/// compacted; sizing the offsets, byte storage and validity from the slice's own metadata keeps +/// that down to one offset store plus one `memcpy` per value. +fn append_to_varbin( + array: ArrayView<'_, Zstd>, + builder: &mut VarBinBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> +where + usize: AsPrimitive, +{ + let slice = array + .data() + .decompress_slice(array.dtype(), &unsliced_validity(array), ctx)?; + let mask = slice.validity.execute_mask(slice.n_rows, ctx)?; + append_slice_to_varbin(&slice, &mask, builder) +} + +/// Copies the values `mask` marks as valid out of `slice` and into `builder`. +fn append_slice_to_varbin( + slice: &DecompressedSlice, + mask: &Mask, + builder: &mut VarBinBuilder, +) -> VortexResult<()> +where + usize: AsPrimitive, +{ + let (values, num_bytes) = slice.value_bytes()?; + // Each value is length-prefixed, so the frames can only be walked in order — which is the + // order `append_valid_slices` visits the valid rows in. A walk that runs out early hands back + // its remainder, which the builder rejects as a byte-count mismatch. + let mut values = ZstdValues::new(values); + builder.append_valid_slices(num_bytes, mask, |_| values.next_value()) +} + +/// Hands the decompressed frames to `builder` as data buffers with views built over them. +/// +/// The frames already hold the values contiguously, so the views can reference them in place and +/// the only per-row work is building one view; going through the canonical array instead would +/// rewrite every view a second time to rebase its buffer index. +fn append_to_varbinview( + array: ArrayView<'_, Zstd>, + builder: &mut VarBinViewBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let slice = array + .data() + .decompress_slice(array.dtype(), &unsliced_validity(array), ctx)?; + let mask = slice.validity.execute_mask(slice.n_rows, ctx)?; + + // No values were stored, so there is nothing to reference and the frames can be dropped. + if mask.all_false() { + builder.append_nulls(slice.n_rows); + return Ok(()); + } + + // The decompressed frames cover whole frames and so can extend past the requested values on + // either side. Reconstructing over just the requested region keeps the pushed buffers fully + // utilized, which is what `append_views_built_at` requires: it hands them to the finished + // array as they are, without compacting them. + let value_bytes = slice.bytes.slice(slice.value_byte_range()?); + // The values only reveal themselves while walking the length-prefixed frames, so the views + // are built inside the builder's numbering callback rather than from a lengths slice. + builder.append_views_built_at(&mask, |next_buffer_index| { + let (buffers, valid_views) = + try_reconstruct_views(&value_bytes, next_buffer_index, MAX_BUFFER_LEN)?; + vortex_ensure!( + valid_views.len() == mask.true_count(), + "Corrupt zstd metadata: the decompressed frames hold {} values for the {} valid rows \ + of the slice", + valid_views.len(), + mask.true_count() + ); + + let views = match mask.bit_buffer() { + AllOr::All => valid_views, + AllOr::None => unreachable!("handled above"), + AllOr::Some(bits) => { + // Null rows carry an empty view, so scatter the stored values into their rows. + // Walking the set bits a word at a time avoids materializing the mask's indices, + // which the views are the only consumer of. + let mut views = BufferMut::::zeroed(slice.n_rows); + let mut valid_row = 0; + bits.for_each_set_index(|index| { + // In bounds: `valid_views.len() == mask.true_count()` was checked above, and + // `index < slice.n_rows` because `bits` is the mask over those rows. + views[index] = valid_views[valid_row]; + valid_row += 1; + }); + views.freeze() + } + }; + Ok((buffers, views)) + }) +} + #[derive(Clone, Debug)] /// Zstd array encoding marker. pub struct Zstd; @@ -340,7 +476,7 @@ impl Zstd { /// Decompress a [`ZstdArray`] into its canonical Vortex representation. pub fn decompress(array: &ZstdArray, ctx: &mut ExecutionCtx) -> VortexResult { let unsliced_validity = child_to_validity( - array.as_ref().slots()[0].as_ref(), + array.as_ref().slots()[ZstdSlots::VALIDITY].as_ref(), array.dtype().nullability(), ); array @@ -349,9 +485,12 @@ impl Zstd { } } -/// The validity bitmap indicating which elements are non-null. -pub(super) const NUM_SLOTS: usize = 1; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"]; +#[array_slots(Zstd)] +pub struct ZstdSlots { + /// The validity bitmap indicating which elements are non-null. + #[slot(0)] + pub validity: Option, +} #[derive(Clone, Debug)] /// Encoding-specific data for a [`ZstdArray`]. @@ -474,22 +613,48 @@ fn collect_valid_vbv( /// tests to exercise the splitting path without allocating >2 GiB. pub fn reconstruct_views( buffer: &ByteBuffer, + start_buf_index: u32, max_buffer_len: usize, ) -> (Vec, Buffer) { + let (buffers, views, _) = walk_views(buffer, start_buf_index, max_buffer_len); + (buffers, views) +} + +/// [`reconstruct_views`], but rejecting a buffer that the length prefixes do not tile exactly. +fn try_reconstruct_views( + buffer: &ByteBuffer, + start_buf_index: u32, + max_buffer_len: usize, +) -> VortexResult<(Vec, Buffer)> { + match walk_views(buffer, start_buf_index, max_buffer_len) { + (buffers, views, None) => Ok((buffers, views)), + (_, _, Some(error)) => Err(error), + } +} + +/// Walks `buffer` until it is exhausted or a length prefix leaves it, returning what was decoded +/// along with the error that stopped the walk. +fn walk_views( + buffer: &ByteBuffer, + start_buf_index: u32, + max_buffer_len: usize, +) -> (Vec, Buffer, Option) { let mut views = BufferMut::::empty(); let mut buffers = Vec::new(); let mut segment_start: usize = 0; let mut offset = 0; + // Only a new segment changes the buffer index, so it is tracked instead of recomputed per view. + let mut buf_index = start_buf_index; + let mut error = None; while offset < buffer.len() { - let str_len = ViewLen::from_le_bytes( - buffer - .get(offset..offset + size_of::()) - .vortex_expect("corrupted zstd length") - .try_into() - .ok() - .vortex_expect("must fit ViewLen size"), - ) as usize; + let str_len = match zstd_value_len(buffer.as_slice(), offset) { + Ok(str_len) => str_len, + Err(err) => { + error = Some(err); + break; + } + }; let value_data_offset = offset + size_of::(); let local_offset = value_data_offset - segment_start; @@ -497,12 +662,28 @@ pub fn reconstruct_views( if local_offset + str_len > max_buffer_len && offset > segment_start { buffers.push(buffer.slice(segment_start..offset)); segment_start = offset; + let Some(next_index) = buf_index.checked_add(1) else { + error = Some(vortex_err!("Zstd values need more than u32::MAX buffers")); + break; + }; + buf_index = next_index; } - let local_offset = u32::try_from(value_data_offset - segment_start) - .vortex_expect("local offset within segment must fit in u32"); - let buf_index = u32::try_from(buffers.len()).vortex_expect("buffer index must fit in u32"); - let value = &buffer[value_data_offset..value_data_offset + str_len]; + let Ok(local_offset) = u32::try_from(value_data_offset - segment_start) else { + error = Some(vortex_err!( + "Zstd value offset {} does not fit in u32; max_buffer_len {max_buffer_len} is too large", + value_data_offset - segment_start + )); + break; + }; + let Some(value) = buffer.get(value_data_offset..value_data_offset + str_len) else { + error = Some(vortex_err!( + "Corrupt zstd value: {str_len} bytes at offset {value_data_offset} run past the \ + end of the {} byte frame buffer", + buffer.len() + )); + break; + }; views.push(BinaryView::make_view(value, buf_index, local_offset)); offset = value_data_offset + str_len; } @@ -511,7 +692,233 @@ pub fn reconstruct_views( buffers.push(buffer.slice(segment_start..buffer.len())); } - (buffers, views.freeze()) + (buffers, views.freeze(), error) +} + +/// Narrows the views decoded from the frames down to the values a slice requests. +fn slice_views( + views: &Buffer, + range: Range, +) -> VortexResult> { + vortex_ensure!( + range.end <= views.len(), + "Corrupt zstd metadata: values {}..{} are out of bounds of the {} values held by the \ + decompressed frames", + range.start, + range.end, + views.len() + ); + Ok(views.slice(range)) +} + +/// A zstd output buffer over uninitialized spare capacity. +/// +/// `decompress_to_buffer` writes through a raw pointer and reports how many bytes it produced, so +/// it never reads its destination — but handing it a `&mut [u8]` covering uninitialized memory +/// would be undefined behaviour regardless of what it does with it. [`WriteBuf`] is the interface +/// zstd provides for exactly this case, and it keeps the alternative (zeroing the whole buffer +/// before every decompression) off the hot path. +struct UninitDestination<'a> { + spare: &'a mut [MaybeUninit], + filled: usize, +} + +impl<'a> UninitDestination<'a> { + fn new(spare: &'a mut [MaybeUninit]) -> Self { + Self { spare, filled: 0 } + } +} + +// SAFETY: `as_mut_ptr` and `capacity` describe the whole spare region, so zstd only ever writes +// within it, and `filled_until` merely records the count it reports. `as_slice` is bounded by that +// count, so it never exposes a byte zstd did not write. +unsafe impl WriteBuf for UninitDestination<'_> { + fn as_slice(&self) -> &[u8] { + // SAFETY: zstd reported writing `filled` bytes from the start of `spare`. + unsafe { std::slice::from_raw_parts(self.spare.as_ptr().cast::(), self.filled) } + } + + fn capacity(&self) -> usize { + self.spare.len() + } + + fn as_mut_ptr(&mut self) -> *mut u8 { + self.spare.as_mut_ptr().cast::() + } + + unsafe fn filled_until(&mut self, n: usize) { + self.filled = n; + } +} + +struct DecompressedSlice { + bytes: ByteBuffer, + validity: Validity, + byte_width: usize, + n_rows: usize, + value_idx_start: usize, + value_idx_stop: usize, + n_skipped_values: usize, + /// Number of stored values held by `bytes`, which covers whole frames and so may extend past + /// the requested slice on either side. + n_buffered_values: usize, +} + +impl DecompressedSlice { + /// The range of stored values this slice requests, as an index into `bytes`. + /// + /// The frame metadata that drives `n_skipped_values` is untrusted, so a frame claiming to hold + /// values that precede the ones we decompressed is rejected instead of wrapping. + fn value_range(&self) -> VortexResult> { + let start = self + .value_idx_start + .checked_sub(self.n_skipped_values) + .ok_or_else(|| { + vortex_err!( + "Corrupt zstd metadata: skipped frames hold {} values, past the first \ + requested value {}", + self.n_skipped_values, + self.value_idx_start + ) + })?; + let end = self + .value_idx_stop + .checked_sub(self.n_skipped_values) + .ok_or_else(|| { + vortex_err!( + "Corrupt zstd metadata: skipped frames hold {} values, past the last \ + requested value {}", + self.n_skipped_values, + self.value_idx_stop + ) + })?; + vortex_ensure!( + start <= end, + "Corrupt zstd metadata: value range {start}..{end} is not ascending" + ); + Ok(start..end) + } + + /// The bounds within `bytes` of the length-prefixed region that should hold exactly this + /// slice's values. + /// + /// Walking the length prefixes is a dependent load chain, so both ends are derived from the + /// slice metadata where possible: an unsliced array skips both walks. The far end is then a + /// claim rather than a checked fact, so a caller must hold the values it reads to the count + /// from [`Self::value_range`] — with [`ZstdValues`], whose shortfall shows up in the byte total + /// from [`Self::value_bytes`], or by counting decoded values as [`try_reconstruct_views`] does. + /// Otherwise a region ending part way through a value passes its trailing bytes off as values. + fn value_byte_range(&self) -> VortexResult> { + let Range { start, end } = self.value_range()?; + let buffer = self.bytes.as_slice(); + let from = zstd_value_offset(buffer, 0, start)?; + let to = if end == self.n_buffered_values { + buffer.len() + } else { + zstd_value_offset(buffer, from, end - start)? + }; + vortex_ensure!( + from <= to && to <= buffer.len(), + "Corrupt zstd metadata: values {from}..{to} are out of bounds of the {} byte frame \ + buffer", + buffer.len() + ); + Ok(from..to) + } + + /// The length-prefixed region of `bytes` holding exactly this slice's values, along with the + /// total size of those values once the length prefixes are dropped. + fn value_bytes(&self) -> VortexResult<(&[u8], usize)> { + let range = self.value_byte_range()?; + let n_values = self.value_range()?.len(); + let buffer = self.bytes.as_slice(); + let bytes = buffer.get(range.clone()).ok_or_else(|| { + vortex_err!( + "Corrupt zstd metadata: values {}..{} are out of bounds of the {} byte frame \ + buffer", + range.start, + range.end, + buffer.len() + ) + })?; + // Every value carries a length prefix, so the region must be at least that large. + let prefix_bytes = n_values.checked_mul(size_of::()).ok_or_else(|| { + vortex_err!("Corrupt zstd metadata: value count {n_values} overflows a byte count") + })?; + let num_bytes = bytes.len().checked_sub(prefix_bytes).ok_or_else(|| { + vortex_err!( + "Corrupt zstd metadata: {n_values} values do not fit in the {} bytes holding them", + bytes.len() + ) + })?; + Ok((bytes, num_bytes)) + } +} + +/// Returns the byte offset `count` length-prefixed values past `offset`. +/// +/// Each step is bounds-checked by the next prefix read, so only the offset the walk lands on needs +/// a check of its own. +fn zstd_value_offset(buffer: &[u8], mut offset: usize, count: usize) -> VortexResult { + for _ in 0..count { + offset += size_of::() + zstd_value_len(buffer, offset)?; + } + vortex_ensure!( + offset <= buffer.len(), + "Corrupt zstd values: walking {count} values ended at offset {offset}, past the end of \ + the {} byte frame buffer", + buffer.len() + ); + Ok(offset) +} + +/// Reads the length prefix of the value starting at `offset`. +#[inline] +fn zstd_value_len(buffer: &[u8], offset: usize) -> VortexResult { + let prefix = buffer + .get(offset..) + .and_then(|rest| rest.first_chunk::<{ size_of::() }>()) + .ok_or_else(|| { + vortex_err!( + "Corrupt zstd values: length prefix at offset {offset} runs past the end of the \ + {} byte frame buffer", + buffer.len() + ) + })?; + Ok(ViewLen::from_le_bytes(*prefix) as usize) +} + +/// A forward walk over the values of a length-prefixed region. +struct ZstdValues<'a> { + buffer: &'a [u8], + offset: usize, +} + +impl<'a> ZstdValues<'a> { + fn new(buffer: &'a [u8]) -> Self { + Self { buffer, offset: 0 } + } + + /// The next value, or every byte the walk could not decode once a prefix leaves the region. + /// + /// The remainder is what makes a shortfall visible in the byte total alone, without a second + /// walk to validate the region up front. A caller sizes that total as the region minus one + /// prefix per value, so `k` of `n` values decoded leaves it still expecting the `n - k` + /// prefixes the walk abandoned; the remainder covers those and more, or is empty only because + /// the region ended exactly and the decoded bytes already fall short. Neither can add up. + fn next_value(&mut self) -> &'a [u8] { + let value_start = self.offset + size_of::(); + let value = zstd_value_len(self.buffer, self.offset) + .ok() + .and_then(|len| self.buffer.get(value_start..value_start.checked_add(len)?)); + match value { + Some(value) => { + self.offset = value_start + value.len(); + value + } + None => &self.buffer[self.offset..], + } + } } impl ZstdData { @@ -586,6 +993,10 @@ impl ZstdData { self.frames.len(), self.metadata.frames.len() ); + for (index, (frame, metadata)) in self.frames.iter().zip(&self.metadata.frames).enumerate() + { + validate_frame_content_size(frame.as_slice(), metadata.uncompressed_size, index)?; + } Ok(()) } @@ -658,9 +1069,10 @@ impl ZstdData { .unwrap_or(value_bytes.len()); let uncompressed = &value_bytes.slice(frame_byte_starts[i]..frame_byte_end); - let compressed = compressor + let mut compressed = compressor .compress(uncompressed) .map_err(|err| VortexError::from(err).with_context("while compressing"))?; + compressed.shrink_to_fit(); frame_metas.push(ZstdFrameMetadata { uncompressed_size: uncompressed.len() as u64, n_values: values_per_frame.min(n_values - i * values_per_frame) as u64, @@ -732,7 +1144,7 @@ impl ZstdData { let value_bytes = values.buffer_handle().try_to_host_sync()?; // Align frames to buffer alignment. This is necessary for overaligned buffers. - let alignment = *value_bytes.alignment(); + let alignment = value_bytes.alignment().as_usize(); let step_width = (values_per_frame * byte_width).div_ceil(alignment) * alignment; let frame_byte_starts = (0..n_values * byte_width) @@ -897,46 +1309,76 @@ impl ZstdData { } } - fn decompress( + fn decompress_slice( &self, dtype: &DType, unsliced_validity: &Validity, ctx: &mut ExecutionCtx, - ) -> VortexResult { + ) -> VortexResult { // To start, we figure out which frames we need to decompress, and with // what row offset into the first such frame. let byte_width = Self::byte_width(dtype); let slice_n_rows = self.slice_stop - self.slice_start; - let slice_value_indices = unsliced_validity - .execute_mask(self.unsliced_n_rows, ctx)? - .valid_counts_for_indices(&[self.slice_start, self.slice_stop]); + let unsliced_mask = unsliced_validity.execute_mask(self.unsliced_n_rows, ctx)?; + let slice_value_indices = + unsliced_mask.valid_counts_for_indices(&[self.slice_start, self.slice_stop]); let slice_value_idx_start = slice_value_indices[0]; let slice_value_idx_stop = slice_value_indices[1]; let mut frames_to_decompress = vec![]; let mut value_idx_start = 0; - let mut uncompressed_size_to_decompress = 0; + let mut uncompressed_size_to_decompress = 0usize; let mut n_skipped_values = 0; + let mut n_buffered_values = 0; for (frame, frame_meta) in self.frames.iter().zip(&self.metadata.frames) { if value_idx_start >= slice_value_idx_stop { break; } - let frame_uncompressed_size = usize::try_from(frame_meta.uncompressed_size) - .vortex_expect("Uncompressed size must fit in usize"); - let frame_n_values = if frame_meta.n_values == 0 { - // possibly older primitive-only metadata that just didn't store this + let frame_uncompressed_size = + usize::try_from(frame_meta.uncompressed_size).map_err(|_| { + vortex_err!( + "Zstd frame uncompressed size {} does not fit in a usize", + frame_meta.uncompressed_size + ) + })?; + let frame_n_values = if frame_meta.n_values != 0 { + usize::try_from(frame_meta.n_values).map_err(|_| { + vortex_err!( + "Zstd frame value count {} does not fit in a usize", + frame_meta.n_values + ) + })? + } else if dtype.is_primitive() { + // Possibly older primitive-only metadata that just didn't store this. Fixed-width + // values make the byte count an exact value count. frame_uncompressed_size / byte_width } else { - usize::try_from(frame_meta.n_values).vortex_expect("frame size must fit usize") + // The same fallback would read a byte count as a value count for variable-width + // values, which misattributes values to frames. A single frame holds every stored + // value, so that case is still recoverable; anything else is not. + vortex_ensure!( + self.frames.len() == 1, + "Zstd frame metadata for a variable-width array is missing its value count" + ); + unsliced_mask.true_count() }; - let value_idx_stop = value_idx_start + frame_n_values; + // Bounding the running total also bounds the two accumulators below, which partition + // it between the frames we keep and the ones we skip. + let value_idx_stop = value_idx_start.checked_add(frame_n_values).ok_or_else(|| { + vortex_err!("Corrupt zstd metadata: frame value counts overflow a usize") + })?; if value_idx_stop > slice_value_idx_start { // we need this frame frames_to_decompress.push(frame); - uncompressed_size_to_decompress += frame_uncompressed_size; + uncompressed_size_to_decompress = uncompressed_size_to_decompress + .checked_add(frame_uncompressed_size) + .ok_or_else(|| { + vortex_err!("Corrupt zstd metadata: frame sizes overflow a usize") + })?; + n_buffered_values += frame_n_values; } else { n_skipped_values += frame_n_values; } @@ -953,24 +1395,28 @@ impl ZstdData { uncompressed_size_to_decompress, Alignment::new(byte_width), ); - unsafe { - // safety: we immediately fill all bytes in the following loop, - // assuming our metadata's uncompressed size is correct - decompressed.set_len(uncompressed_size_to_decompress); - } let mut uncompressed_start = 0; for frame in frames_to_decompress { - let uncompressed_written = decompressor - .decompress_to_buffer(frame.as_slice(), &mut decompressed[uncompressed_start..])?; - uncompressed_start += uncompressed_written; + // Decompress straight into the spare capacity. Each frame gets only the region after + // the ones before it, bounded by the size the metadata declared, so a frame that + // expands further than advertised is refused by zstd rather than overrunning. + let mut destination = UninitDestination::new( + &mut decompressed.spare_capacity_mut() + [uncompressed_start..uncompressed_size_to_decompress], + ); + uncompressed_start += + decompressor.decompress_to_buffer(frame.as_slice(), &mut destination)?; } if uncompressed_start != uncompressed_size_to_decompress { - vortex_panic!( + vortex_bail!( "Zstd metadata or frames were corrupt; expected {} bytes but decompressed {}", uncompressed_size_to_decompress, uncompressed_start ); } + // SAFETY: the loop above decompressed exactly `uncompressed_start` bytes into the front of + // the spare capacity, and the check above pins that to the requested length. + unsafe { decompressed.set_len(uncompressed_start) }; let decompressed = decompressed.freeze(); // Last, we slice the exact values requested out of the decompressed data. @@ -985,7 +1431,7 @@ impl ZstdData { // We ensure that the validity of the decompressed array ALWAYS matches the validity // implied by the DType. if !dtype.is_nullable() && !matches!(slice_validity, Validity::NonNullable) { - assert!( + vortex_ensure!( matches!(slice_validity, Validity::AllValid), "ZSTD array expects to be non-nullable but there are nulls after decompression" ); @@ -997,30 +1443,58 @@ impl ZstdData { // END OF IMPORTANT BLOCK // + Ok(DecompressedSlice { + bytes: decompressed, + validity: slice_validity, + byte_width, + n_rows: slice_n_rows, + value_idx_start: slice_value_idx_start, + value_idx_stop: slice_value_idx_stop, + n_skipped_values, + n_buffered_values, + }) + } + + fn decompress( + &self, + dtype: &DType, + unsliced_validity: &Validity, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let slice = self.decompress_slice(dtype, unsliced_validity, ctx)?; match dtype { DType::Primitive(..) => { - let slice_values_buffer = decompressed.slice( - (slice_value_idx_start - n_skipped_values) * byte_width - ..(slice_value_idx_stop - n_skipped_values) * byte_width, - ); + let Range { start, end } = slice.value_range()?; + let byte_range = start + .checked_mul(slice.byte_width) + .zip(end.checked_mul(slice.byte_width)) + .filter(|(_, byte_stop)| *byte_stop <= slice.bytes.len()) + .map(|(byte_start, byte_stop)| byte_start..byte_stop) + .ok_or_else(|| { + vortex_err!( + "Corrupt zstd metadata: values {start}..{end} of {} bytes each are \ + out of bounds of the {} byte frame buffer", + slice.byte_width, + slice.bytes.len() + ) + })?; + let slice_values_buffer = slice.bytes.slice(byte_range); let primitive = PrimitiveArray::from_values_byte_buffer( slice_values_buffer, dtype.as_ptype(), - slice_validity, - slice_n_rows, + slice.validity, + slice.n_rows, ctx, ); Ok(primitive.into_array()) } DType::Binary(_) | DType::Utf8(_) => { - match slice_validity.execute_mask(slice_n_rows, ctx)?.indices() { + match slice.validity.execute_mask(slice.n_rows, ctx)?.indices() { AllOr::All => { - let (buffers, all_views) = reconstruct_views(&decompressed, MAX_BUFFER_LEN); - let valid_views = all_views.slice( - slice_value_idx_start - n_skipped_values - ..slice_value_idx_stop - n_skipped_values, - ); + let (buffers, all_views) = + try_reconstruct_views(&slice.bytes, 0, MAX_BUFFER_LEN)?; + let valid_views = slice_views(&all_views, slice.value_range()?)?; // SAFETY: we properly construct the views inside `reconstruct_views` Ok(unsafe { @@ -1028,24 +1502,22 @@ impl ZstdData { valid_views, Arc::from(buffers), dtype.clone(), - slice_validity, + slice.validity, ) } .into_array()) } AllOr::None => Ok(ConstantArray::new( Scalar::null(dtype.clone()), - slice_n_rows, + slice.n_rows, ) .into_array()), AllOr::Some(valid_indices) => { - let (buffers, all_views) = reconstruct_views(&decompressed, MAX_BUFFER_LEN); - let valid_views = all_views.slice( - slice_value_idx_start - n_skipped_values - ..slice_value_idx_stop - n_skipped_values, - ); + let (buffers, all_views) = + try_reconstruct_views(&slice.bytes, 0, MAX_BUFFER_LEN)?; + let valid_views = slice_views(&all_views, slice.value_range()?)?; - let mut views = BufferMut::::zeroed(slice_n_rows); + let mut views = BufferMut::::zeroed(slice.n_rows); for (view, index) in valid_views.into_iter().zip_eq(valid_indices) { views[*index] = view } @@ -1056,14 +1528,14 @@ impl ZstdData { views.freeze(), Arc::from(buffers), dtype.clone(), - slice_validity, + slice.validity, ) } .into_array()) } } } - _ => vortex_panic!("Unsupported dtype for Zstd array: {}", dtype), + _ => vortex_bail!("Unsupported dtype for Zstd array: {}", dtype), } } @@ -1107,8 +1579,10 @@ impl ZstdData { impl ValidityVTable for Zstd { fn validity(array: ArrayView<'_, Zstd>) -> VortexResult { - let unsliced_validity = - child_to_validity(array.slots()[0].as_ref(), array.dtype().nullability()); + let unsliced_validity = child_to_validity( + array.slots()[ZstdSlots::VALIDITY].as_ref(), + array.dtype().nullability(), + ); unsliced_validity.slice(array.slice_start()..array.slice_stop()) } } @@ -1119,8 +1593,10 @@ impl OperationsVTable for Zstd { index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let unsliced_validity = - child_to_validity(array.slots()[0].as_ref(), array.dtype().nullability()); + let unsliced_validity = child_to_validity( + array.slots()[ZstdSlots::VALIDITY].as_ref(), + array.dtype().nullability(), + ); let sliced = array.data().with_slice(index, index + 1); sliced .decompress(array.dtype(), &unsliced_validity, ctx)? @@ -1131,9 +1607,23 @@ impl OperationsVTable for Zstd { #[cfg(test)] #[expect(clippy::cast_possible_truncation)] mod tests { + use rstest::rstest; + use vortex_array::arrays::varbin::VarBinArrayExt as _; + use vortex_array::builders::VarBinBuilder; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability::NonNullable; + use vortex_array::validity::Validity; use vortex_buffer::ByteBuffer; + use vortex_error::VortexResult; + use vortex_mask::Mask; + use super::DecompressedSlice; + use super::ViewLen; + use super::append_slice_to_varbin; use super::reconstruct_views; + use super::try_reconstruct_views; + use super::zstd_value_len; + use super::zstd_value_offset; use crate::array::BinaryView; /// Build a Zstd-style interleaved buffer: [u32-LE length][string bytes] repeated. @@ -1147,11 +1637,31 @@ mod tests { ByteBuffer::copy_from(buf.as_slice()) } + /// A slice over `bytes` that requests `value_idx_start..value_idx_stop`. + fn decompressed_slice( + bytes: ByteBuffer, + value_idx_start: usize, + value_idx_stop: usize, + n_skipped_values: usize, + n_buffered_values: usize, + ) -> DecompressedSlice { + DecompressedSlice { + bytes, + validity: Validity::NonNullable, + byte_width: 1, + n_rows: value_idx_stop - value_idx_start, + value_idx_start, + value_idx_stop, + n_skipped_values, + n_buffered_values, + } + } + #[test] fn test_reconstruct_views_no_split() { let strings: &[&[u8]] = &[b"hello", b"world"]; let buf = make_interleaved(strings); - let (buffers, views) = reconstruct_views(&buf, 1024); + let (buffers, views) = reconstruct_views(&buf, 0, 1024); assert_eq!(buffers.len(), 1); assert_eq!(views.len(), 2); @@ -1168,7 +1678,7 @@ mod tests { // so it rolls into a second segment. let strings: &[&[u8]] = &[b"aaaaaaaaaaaaa", b"bbbbbbbbbbbbb"]; let buf = make_interleaved(strings); - let (buffers, views) = reconstruct_views(&buf, 20); + let (buffers, views) = reconstruct_views(&buf, 0, 20); assert_eq!(buffers.len(), 2); assert_eq!(views.len(), 2); @@ -1176,4 +1686,119 @@ mod tests { // Second entry starts a new segment at byte 17 (the length prefix), so local offset = 4. assert_eq!(views[1], BinaryView::make_view(b"bbbbbbbbbbbbb", 1, 4)); } + + /// A buffer whose last entry claims more bytes than remain, as corrupt frame data would. + fn make_overrunning() -> ByteBuffer { + let mut buf = Vec::new(); + buf.extend_from_slice(&5u32.to_le_bytes()); + buf.extend_from_slice(b"hello"); + buf.extend_from_slice(&9u32.to_le_bytes()); + buf.extend_from_slice(b"ab"); + ByteBuffer::copy_from(buf.as_slice()) + } + + #[test] + fn test_reconstruct_views_rejects_overrunning_value() { + let buf = make_overrunning(); + assert!(try_reconstruct_views(&buf, 0, 1024).is_err()); + + // The lenient walk keeps the decodable prefix instead of panicking. + let (buffers, views) = reconstruct_views(&buf, 0, 1024); + assert_eq!(buffers.len(), 1); + assert_eq!(views.len(), 1); + assert_eq!(views[0], BinaryView::make_view(b"hello", 0, 4)); + } + + #[test] + fn test_reconstruct_views_rejects_truncated_length_prefix() { + // A trailing partial length prefix cannot start a value. + let buf = + ByteBuffer::copy_from([5u8, 0, 0, 0, b'h', b'e', b'l', b'l', b'o', 1, 0].as_ref()); + assert!(try_reconstruct_views(&buf, 0, 1024).is_err()); + assert_eq!(reconstruct_views(&buf, 0, 1024).1.len(), 1); + } + + #[rstest] + #[case::truncated_buffer(&[0u8, 0, 0], 0)] + #[case::truncated_tail(&[4u8, 0, 0, 0], 2)] + #[case::offset_at_end(&[4u8, 0, 0, 0], 4)] + #[case::offset_past_end(&[4u8, 0, 0, 0], 64)] + fn test_zstd_value_len_rejects_out_of_bounds(#[case] buffer: &[u8], #[case] offset: usize) { + assert!(zstd_value_len(buffer, offset).is_err()); + } + + #[test] + fn test_zstd_value_offset_rejects_walking_past_the_end() -> VortexResult<()> { + let buf = make_interleaved(&[b"hello", b"world"]); + assert_eq!(zstd_value_offset(buf.as_slice(), 0, 2)?, buf.len()); + // Only two values are stored, so the third step leaves the buffer. + assert!(zstd_value_offset(buf.as_slice(), 0, 3).is_err()); + Ok(()) + } + + #[test] + fn test_value_range_rejects_skipping_past_the_requested_values() { + // Frame metadata claiming more skipped values than the slice starts at would wrap. + let slice = decompressed_slice(make_interleaved(&[b"hello"]), 2, 3, 4, 1); + assert!(slice.value_range().is_err()); + assert!(slice.value_bytes().is_err()); + } + + #[rstest] + // The buffered value count matches, so both ends come from the metadata. + #[case::exact_metadata(2)] + // It does not, so the far end is walked instead. + #[case::walked_end(9)] + fn test_value_bytes_totals_the_stored_values( + #[case] n_buffered_values: usize, + ) -> VortexResult<()> { + let buf = make_interleaved(&[b"hello", b"world"]); + let slice = decompressed_slice(buf.clone(), 0, 2, 0, n_buffered_values); + let (bytes, num_bytes) = slice.value_bytes()?; + assert_eq!(bytes, buf.as_slice()); + assert_eq!(num_bytes, buf.len() - 2 * size_of::()); + Ok(()) + } + + #[test] + fn test_value_bytes_rejects_more_values_than_the_buffer_holds() { + // Frame metadata claims nine values but only two are stored. + let slice = decompressed_slice(make_interleaved(&[b"hello", b"world"]), 0, 5, 0, 9); + assert!(slice.value_bytes().is_err()); + } + + #[test] + fn test_append_to_varbin_copies_the_stored_values() -> VortexResult<()> { + let slice = decompressed_slice(make_interleaved(&[b"hello", b"world"]), 0, 2, 0, 2); + let mut builder = VarBinBuilder::::new_in( + DType::Utf8(NonNullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + append_slice_to_varbin(&slice, &Mask::new_true(2), &mut builder)?; + + let appended = builder.finish_into_varbin(); + assert_eq!(appended.bytes_at(0).as_slice(), b"hello"); + assert_eq!(appended.bytes_at(1).as_slice(), b"world"); + Ok(()) + } + + #[test] + fn test_append_to_varbin_rejects_a_dangling_length_prefix() { + let mut buffer = Vec::new(); + buffer.extend_from_slice(&3u32.to_le_bytes()); + buffer.extend_from_slice(b"cat"); + // A prefix with no value after it. It takes up exactly the four bytes the missing value's + // own prefix would have, so treating that value as empty still totals the byte count the + // metadata implies and appends ["cat", ""]. + buffer.extend_from_slice(&1u32.to_le_bytes()); + + let slice = decompressed_slice(ByteBuffer::copy_from(buffer.as_slice()), 0, 2, 0, 2); + let mut builder = VarBinBuilder::::new_in( + DType::Utf8(NonNullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + assert!(append_slice_to_varbin(&slice, &Mask::new_true(2), &mut builder).is_err()); + // The builder rejected the values before committing any of them. + assert_eq!(builder.finish_into_varbin().len(), 0); + } } diff --git a/encodings/zstd/src/compute/cast.rs b/encodings/zstd/src/compute/cast.rs index dba59198c88..598a5d7195e 100644 --- a/encodings/zstd/src/compute/cast.rs +++ b/encodings/zstd/src/compute/cast.rs @@ -12,6 +12,7 @@ use vortex_error::VortexResult; use crate::Zstd; use crate::ZstdData; +use crate::ZstdSlots; impl CastReduce for Zstd { fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { @@ -33,12 +34,17 @@ impl CastReduce for Zstd { } (Nullability::NonNullable, Nullability::Nullable) => { // nonnull => null, trivial cast by altering the validity - child_to_validity(array.slots()[0].as_ref(), array.dtype().nullability()) + child_to_validity( + array.slots()[ZstdSlots::VALIDITY].as_ref(), + array.dtype().nullability(), + ) } (Nullability::Nullable, Nullability::NonNullable) => { // null => non-null works if there are no nulls in the sliced range - let unsliced_validity = - child_to_validity(array.slots()[0].as_ref(), array.dtype().nullability()); + let unsliced_validity = child_to_validity( + array.slots()[ZstdSlots::VALIDITY].as_ref(), + array.dtype().nullability(), + ); let has_nulls = !unsliced_validity .slice(array.slice_start()..array.slice_stop())? .definitely_no_nulls(); diff --git a/encodings/zstd/src/editions.rs b/encodings/zstd/src/editions.rs new file mode 100644 index 00000000000..43f695a42e2 --- /dev/null +++ b/encodings/zstd/src/editions.rs @@ -0,0 +1,61 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The `zstd` edition family. +//! +//! Zstd buffer wrapping is opt-in. This module declares its persisted array encoding. + +use vortex_edition::Edition; +use vortex_edition::EditionDeclaration; +use vortex_edition::EditionFamily; +use vortex_edition::EditionId; +use vortex_edition::EditionMember; + +/// The `zstd` family: optional Zstd-backed serialized array representations. +pub static FAMILY: EditionFamily = EditionFamily { + name: "zstd", + origin: "vortex-zstd", + doc: "Optional Zstd-backed serialized array representations. A reader built without \ +`vortex-zstd` cannot resolve these members, so they are versioned independently of `core` and \ +must be enabled explicitly by the caller.", +}; + +/// The February 2026 draft edition of the `zstd` family. +pub const ZSTD_2026_02: EditionId = EditionId::new("zstd", 2026, 2, 0); + +/// The declaration of [`ZSTD_2026_02`] and the components that join the family at it. +pub static DECLARATION: EditionDeclaration = EditionDeclaration { + edition: Edition { + id: ZSTD_2026_02, + min_library_version: None, + }, + added: &[EditionMember::array(&"vortex.zstd_buffers")], +}; + +#[cfg(test)] +mod tests { + use vortex_edition::EditionSessionExt; + use vortex_edition::test_harness::validate_edition; + use vortex_error::VortexResult; + + use super::*; + + #[test] + fn zstd_edition_is_valid() -> VortexResult<()> { + let session = vortex_array::array_session(); + crate::initialize(&session); + validate_edition(&session.editions(), &ZSTD_2026_02) + } + + #[test] + fn initialize_does_not_enable_zstd_edition() { + let session = vortex_array::array_session(); + crate::initialize(&session); + assert!( + !session + .enabled_editions() + .editions() + .contains(&ZSTD_2026_02) + ); + } +} diff --git a/encodings/zstd/src/lib.rs b/encodings/zstd/src/lib.rs index a178fdec142..9280d99958c 100644 --- a/encodings/zstd/src/lib.rs +++ b/encodings/zstd/src/lib.rs @@ -7,8 +7,8 @@ //! trained dictionary across frames. Frame metadata lets slices decompress only the frames that can //! contribute values to the requested row range. //! -//! With the `unstable_encodings` feature, `ZstdBuffers` stores the buffers of another encoding as -//! independently compressed zstd buffers while preserving the inner encoding metadata. +//! [`ZstdBuffers`] stores the buffers of another encoding as independently compressed zstd +//! buffers while preserving the inner encoding metadata. //! //! This crate exposes array encodings only. Compression scheme selection is wired through //! `vortex-btrblocks` and file writing. To deserialize arrays manually, register the encoding in the @@ -23,19 +23,58 @@ pub use array::*; use vortex_array::dtype::proto::dtype as pb; -#[cfg(feature = "unstable_encodings")] +use vortex_array::session::ArraySessionExt; +use vortex_edition::EditionSessionExt; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::VortexSession; pub use zstd_buffers::*; mod array; mod compute; +pub mod editions; mod rules; mod slice; -#[cfg(feature = "unstable_encodings")] mod zstd_buffers; #[cfg(test)] mod test; +/// Register the Zstd encodings and their optional edition declaration with a Vortex session. +pub fn initialize(session: &VortexSession) { + session.arrays().register(Zstd); + session.arrays().register(ZstdBuffers); + if session.editions().find(&editions::ZSTD_2026_02).is_none() { + session + .editions() + .declare_family(&editions::FAMILY) + .map_err(|error| vortex_err!("{error}")) + .vortex_expect("Zstd edition family is valid"); + session + .register_edition(&editions::DECLARATION) + .map_err(|error| vortex_err!("{error}")) + .vortex_expect("Zstd edition declaration is valid"); + } +} + +/// Ensure Vortex metadata agrees with the content size declared by a zstd frame. +pub(crate) fn validate_frame_content_size( + frame: &[u8], + metadata_size: u64, + index: usize, +) -> VortexResult<()> { + let frame_content_size = zstd::zstd_safe::get_frame_content_size(frame) + .map_err(|error| vortex_err!("Invalid zstd frame {index}: {error}"))? + .ok_or_else(|| vortex_err!("Zstd frame {index} does not declare a content size"))?; + vortex_ensure!( + metadata_size == frame_content_size, + "Zstd frame {index} metadata declares {metadata_size} uncompressed bytes, but its header declares {frame_content_size}" + ); + Ok(()) +} + #[derive(Clone, prost::Message)] /// Metadata for one zstd frame. pub struct ZstdFrameMetadata { diff --git a/encodings/zstd/src/slice.rs b/encodings/zstd/src/slice.rs index 54f4441635d..d1ccf6f565d 100644 --- a/encodings/zstd/src/slice.rs +++ b/encodings/zstd/src/slice.rs @@ -11,11 +11,14 @@ use vortex_array::vtable::child_to_validity; use vortex_error::VortexResult; use crate::Zstd; +use crate::ZstdSlots; impl SliceReduce for Zstd { fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { - let unsliced_validity = - child_to_validity(array.slots()[0].as_ref(), array.dtype().nullability()); + let unsliced_validity = child_to_validity( + array.slots()[ZstdSlots::VALIDITY].as_ref(), + array.dtype().nullability(), + ); Ok(Some( Zstd::try_new( array.dtype().clone(), diff --git a/encodings/zstd/src/test.rs b/encodings/zstd/src/test.rs index cf45eb97c04..fafbfcaf159 100644 --- a/encodings/zstd/src/test.rs +++ b/encodings/zstd/src/test.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors #![expect(clippy::cast_possible_truncation)] +use rstest::rstest; +use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; @@ -10,14 +12,23 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; use vortex_array::assert_nth_scalar; +use vortex_array::builders::VarBinBuilder; +use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; use vortex_array::validity::Validity; use vortex_buffer::Alignment; use vortex_buffer::Buffer; +use vortex_buffer::ByteBuffer; +use vortex_error::VortexResult; use vortex_mask::Mask; use crate::Zstd; +use crate::ZstdArray; +use crate::ZstdData; +use crate::ZstdFrameMetadata; +use crate::ZstdMetadata; #[test] fn test_zstd_compress_decompress() { @@ -205,6 +216,85 @@ fn test_zstd_var_bin_view() { assert_nth_scalar!(sliced, 2, "Lorem ipsum dolor sit amet", &mut ctx); } +#[test] +fn test_zstd_append_to_offset_builder() { + let mut ctx = array_session().create_execution_ctx(); + let array = VarBinViewArray::from_iter( + [ + Some(b"foo".as_slice()), + Some(b"bar".as_slice()), + None, + Some(b"Lorem ipsum dolor sit amet".as_slice()), + Some(b"baz".as_slice()), + ], + DType::Utf8(Nullability::Nullable), + ); + let compressed = Zstd::from_var_bin_view(&array, 0, 3, &mut ctx) + .unwrap() + .slice(1..4) + .unwrap(); + let mut builder = VarBinBuilder::::with_capacity_in( + compressed.dtype().clone(), + compressed.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + compressed + .append_to_builder(&mut builder, &mut ctx) + .unwrap(); + assert_arrays_eq!( + builder.finish_into_varbin(), + array.into_array().slice(1..4).unwrap(), + &mut ctx + ); +} + +/// A slice decompresses whole frames, so the frames hold values on either side of the ones it +/// requests. `append_views_built_at` publishes the buffers it is handed as they are, so only the +/// requested region may reach it — otherwise the finished array retains the whole frames. +#[test] +fn test_zstd_append_to_view_builder_keeps_only_the_sliced_bytes() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // Long enough that the values live in the data buffers rather than inline in the views. + let values = (0..12) + .map(|i| format!("value number {i} padded well past the inline limit")) + .collect::>(); + let array = VarBinViewArray::from_iter_str(&values); + // Three values per frame, so the slice below starts and ends inside a frame holding more. + let compressed = Zstd::from_var_bin_view(&array, 0, 3, &mut ctx)?.slice(4..8)?; + + // Seeded with a value of its own so the pushed buffers land after an in-progress buffer, and + // appended to twice so the second push has to rebase past the first. + let mut builder = VarBinViewBuilder::with_capacity_in( + compressed.dtype().clone(), + 9, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); + builder.append_value(&values[0]); + compressed.append_to_builder(&mut builder, &mut ctx)?; + compressed.append_to_builder(&mut builder, &mut ctx)?; + let appended = builder.finish_into_varbinview(); + + let expected = VarBinViewArray::from_iter_str( + std::iter::once(&values[0]) + .chain(&values[4..8]) + .chain(&values[4..8]), + ); + assert_arrays_eq!(appended, expected, &mut ctx); + + // Each stored value costs its bytes plus the u32 length prefix zstd interleaves. + let sliced_bytes: usize = values[4..8] + .iter() + .map(|value| value.len() + size_of::()) + .sum(); + let buffered: usize = appended + .data_buffers() + .iter() + .map(|buffer| buffer.as_host().len()) + .sum(); + assert_eq!(buffered, values[0].len() + 2 * sliced_bytes); + Ok(()) +} + #[test] fn test_zstd_decompress_var_bin_view() { let mut ctx = array_session().create_execution_ctx(); @@ -246,6 +336,147 @@ fn test_sliced_array_children() { sliced.children(); } +/// Six rows, five of them stored, compressed into `values_per_frame`-sized frames with the frame +/// metadata a reader would have deserialized rewritten by `corrupt`. +fn corrupt_var_bin_view_metadata( + values_per_frame: usize, + corrupt: impl FnOnce(&mut ZstdMetadata), + ctx: &mut ExecutionCtx, +) -> VortexResult<(VarBinViewArray, ZstdArray)> { + let array = VarBinViewArray::from_iter( + [ + Some(b"foo".as_slice()), + Some(b"bar".as_slice()), + None, + Some(b"Lorem ipsum dolor sit amet".as_slice()), + Some(b"baz".as_slice()), + Some(b"quux".as_slice()), + ], + DType::Utf8(Nullability::Nullable), + ); + let mut data = ZstdData::from_var_bin_view(&array, 0, values_per_frame, ctx)?; + corrupt(&mut data.metadata); + let compressed = Zstd::try_new(array.dtype().clone(), data, array.validity()?)?; + Ok((array, compressed)) +} + +/// Frame metadata comes straight off disk, so an inconsistent value count has to surface as an +/// error from both read paths rather than a panic or a wrapped length. +#[rstest] +#[case::frame_holds_fewer_values_than_claimed(|metadata: &mut ZstdMetadata| { + metadata.frames[0].n_values = 1000; +})] +#[case::frame_value_counts_overflow(|metadata: &mut ZstdMetadata| { + metadata.frames[1].n_values = u64::MAX; +})] +#[case::missing_value_count_across_frames(|metadata: &mut ZstdMetadata| { + metadata.frames[0].n_values = 0; +})] +fn test_zstd_rejects_corrupt_frame_metadata( + #[case] corrupt: fn(&mut ZstdMetadata), +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let (_, compressed) = corrupt_var_bin_view_metadata(3, corrupt, &mut ctx)?; + + assert!(Zstd::decompress(&compressed, &mut ctx).is_err()); + + let mut builder = VarBinBuilder::::with_capacity_in( + compressed.dtype().clone(), + compressed.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + assert!( + compressed + .append_to_builder(&mut builder, &mut ctx) + .is_err() + ); + Ok(()) +} + +/// Frame bytes are as untrusted as the metadata describing them: a frame whose last value is +/// nothing but a length prefix has to error from every read path, not read back as an empty value. +#[test] +fn test_zstd_rejects_a_frame_ending_in_a_dangling_length_prefix() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let mut values = Vec::new(); + values.extend_from_slice(&3u32.to_le_bytes()); + values.extend_from_slice(b"cat"); + // A prefix with no value after it. Treating the missing value as empty still totals the three + // bytes the metadata implies for two values. + values.extend_from_slice(&1u32.to_le_bytes()); + + let dtype = DType::Utf8(Nullability::NonNullable); + let compressed = Zstd::try_new( + dtype.clone(), + ZstdData::new( + None, + vec![ByteBuffer::from(zstd::bulk::compress(&values, 3)?)], + ZstdMetadata { + dictionary_size: 0, + frames: vec![ZstdFrameMetadata { + uncompressed_size: values.len() as u64, + n_values: 2, + }], + }, + 2, + ), + Validity::NonNullable, + )?; + + assert!(Zstd::decompress(&compressed, &mut ctx).is_err()); + let mut varbin = VarBinBuilder::::with_capacity_in( + dtype.clone(), + 2, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + assert!(compressed.append_to_builder(&mut varbin, &mut ctx).is_err()); + let mut views = VarBinViewBuilder::with_capacity_in( + dtype, + 2, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); + assert!(compressed.append_to_builder(&mut views, &mut ctx).is_err()); + Ok(()) +} + +/// Metadata written before frames recorded their value count leaves it at zero. A single frame +/// holds every stored value, so those arrays still read back. +#[test] +fn test_zstd_reads_legacy_single_frame_var_bin_metadata() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let (array, compressed) = + corrupt_var_bin_view_metadata(0, |metadata| metadata.frames[0].n_values = 0, &mut ctx)?; + + assert_arrays_eq!(compressed, array.clone().into_array(), &mut ctx); + assert_arrays_eq!( + compressed.slice(2..5)?, + array.into_array().slice(2..5)?, + &mut ctx + ); + Ok(()) +} + +/// The same legacy metadata for fixed-width values recovers the count from the frame size, which +/// stays exact across frames. +#[test] +fn test_zstd_reads_legacy_primitive_frame_metadata() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = PrimitiveArray::from_iter(0..200_i32); + let mut data = ZstdData::from_primitive(&array, 3, 30, &mut ctx)?; + for frame in &mut data.metadata.frames { + frame.n_values = 0; + } + let compressed = Zstd::try_new(array.dtype().clone(), data, array.validity()?)?; + + assert_arrays_eq!(compressed, PrimitiveArray::from_iter(0..200_i32), &mut ctx); + assert_arrays_eq!( + compressed.slice(100..105)?, + PrimitiveArray::from_iter(100..105_i32), + &mut ctx + ); + Ok(()) +} + /// Tests that each beginning of a frame in ZSTD matches /// the buffer alignment when compressing primitive arrays. #[test] @@ -259,3 +490,20 @@ fn test_zstd_frame_start_buffer_alignment() { assert!(compressed.is_ok()); } + +#[test] +fn test_zstd_rejects_mismatched_frame_content_size() { + let mut ctx = array_session().create_execution_ctx(); + let compressed = + Zstd::from_primitive(&PrimitiveArray::from_iter([1_i32, 2, 3]), 0, 0, &mut ctx).unwrap(); + let mut data = compressed.data().clone(); + data.metadata.frames[0].uncompressed_size = 16 * 1024 * 1024 * 1024; + + let error = Zstd::try_new( + DType::Primitive(PType::I32, Nullability::NonNullable), + data, + Validity::NonNullable, + ) + .unwrap_err(); + assert!(error.to_string().contains("metadata declares")); +} diff --git a/encodings/zstd/src/zstd_buffers.rs b/encodings/zstd/src/zstd_buffers.rs index 73361a593a1..f21deee6f64 100644 --- a/encodings/zstd/src/zstd_buffers.rs +++ b/encodings/zstd/src/zstd_buffers.rs @@ -10,6 +10,7 @@ use std::sync::Arc; use prost::Message as _; use vortex_array::Array; +use vortex_array::ArrayDeserialization; use vortex_array::ArrayEq; use vortex_array::ArrayHash; use vortex_array::ArrayId; @@ -39,6 +40,7 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::ZstdBuffersMetadata; +use crate::validate_frame_content_size; /// A [`ZstdBuffers`]-encoded Vortex array. pub type ZstdBuffersArray = Array; @@ -59,42 +61,38 @@ impl ZstdBuffers { /// Compress every top-level buffer of `array` independently with zstd. /// - /// Children are preserved as slots and the wrapped array's serialized metadata is stored so the - /// original array can be rebuilt after decompression. + /// The wrapped array's serialized representation is captured so it can be rebuilt after + /// decompression, including any buffers or children selected by its serializer. pub fn compress( array: &ArrayRef, level: i32, session: &VortexSession, ) -> VortexResult { - let encoding_id = array.encoding_id(); - let metadata = session + let serialization = session .array_serialize(array)? .ok_or_else(|| vortex_err!("[ZstdBuffers]: Array does not support serialization"))?; - let buffer_handles = array.buffer_handles(); - let children = array.children(); - let mut compressed_buffers = Vec::with_capacity(buffer_handles.len()); - let mut uncompressed_sizes = Vec::with_capacity(buffer_handles.len()); - let mut buffer_alignments = Vec::with_capacity(buffer_handles.len()); + let mut compressed_buffers = Vec::with_capacity(serialization.buffers.len()); + let mut uncompressed_sizes = Vec::with_capacity(serialization.buffers.len()); + let mut buffer_alignments = Vec::with_capacity(serialization.buffers.len()); let mut compressor = zstd::bulk::Compressor::new(level)?; - // Compression is currently CPU-only, so we gather all buffers on the host. - for handle in &buffer_handles { - buffer_alignments.push(u32::from(handle.alignment())); - let host_buf = handle.clone().try_to_host_sync()?; - uncompressed_sizes.push(host_buf.len() as u64); - let compressed = compressor.compress(&host_buf)?; + for buffer in &serialization.buffers { + buffer_alignments.push(u32::from(buffer.alignment())); + uncompressed_sizes.push(buffer.len() as u64); + let mut compressed = compressor.compress(buffer)?; + compressed.shrink_to_fit(); compressed_buffers.push(BufferHandle::new_host(ByteBuffer::from(compressed))); } let data = ZstdBuffersData { - inner_encoding_id: encoding_id, - inner_metadata: metadata, + inner_encoding_id: serialization.serialized_id, + inner_metadata: serialization.metadata, compressed_buffers, uncompressed_sizes, buffer_alignments, }; - let slots: ArraySlots = children.into_iter().map(Some).collect(); + let slots: ArraySlots = serialization.children.into_iter().map(Some).collect(); let compressed = Array::try_from_parts( ArrayParts::new(ZstdBuffers, array.dtype().clone(), array.len(), data) .with_slots(slots), @@ -109,20 +107,24 @@ impl ZstdBuffers { buffer_handles: &[BufferHandle], session: &VortexSession, ) -> VortexResult { - let registry = session.arrays().registry().clone(); - let inner_vtable = registry - .find(&array.data().inner_encoding_id) + let inner_vtable = session + .arrays() + .registry() + .get(&array.data().inner_encoding_id) .ok_or_else(|| { vortex_err!("Unknown inner encoding: {}", array.data().inner_encoding_id) })?; let children: Vec = array.slots().iter().flatten().cloned().collect(); inner_vtable.deserialize( - array.dtype(), - array.len(), - &array.data().inner_metadata, - buffer_handles, - &children.as_slice(), + ArrayDeserialization::new( + array.data().inner_encoding_id, + array.dtype(), + array.len(), + &array.data().inner_metadata, + buffer_handles, + &children.as_slice(), + ), session, ) } @@ -257,6 +259,8 @@ impl ZstdBuffersData { let alignment = self.buffer_alignments.get(i).copied().unwrap_or(1); let aligned = Alignment::try_from(alignment)?; + let compressed = buf.clone().try_to_host_sync()?; + validate_frame_content_size(compressed.as_slice(), uncompressed_size, i)?; let mut output = ByteBufferMut::with_capacity_aligned(size, aligned); let spare = output.spare_capacity_mut(); @@ -273,7 +277,6 @@ impl ZstdBuffersData { // `set_len(size)` after zstd reports how many bytes were written. let dst = unsafe { std::slice::from_raw_parts_mut(spare.as_mut_ptr().cast::(), size) }; - let compressed = buf.clone().try_to_host_sync()?; let written = decompressor.decompress_to_buffer(compressed.as_slice(), dst)?; if written != size { return Err(vortex_err!( @@ -294,6 +297,18 @@ impl ZstdBuffersData { // If invariants are somehow broken, device decompression could have UB, so ensure // they still hold. self.validate()?; + // Only host-resident frames are checked: reading a device frame's header would cost a D2H + // copy per frame, and the device output allocation is fallible so it cannot abort. + for (index, (buffer, &metadata_size)) in self + .compressed_buffers + .iter() + .zip(&self.uncompressed_sizes) + .enumerate() + { + if let Some(frame) = buffer.as_host_opt() { + validate_frame_content_size(frame.as_slice(), metadata_size, index)?; + } + } let output_sizes = self .uncompressed_sizes @@ -341,7 +356,7 @@ fn compute_output_layout( let mut total_size = 0usize; for (&size, &alignment) in output_sizes.iter().zip(output_alignments.iter()) { - total_size = total_size.next_multiple_of(*alignment); + total_size = total_size.next_multiple_of(alignment.as_usize()); offsets.push(total_size); total_size += size; } @@ -652,6 +667,24 @@ mod tests { Ok(()) } + #[test] + fn test_rejects_mismatched_frame_content_size_before_output_allocation() -> VortexResult<()> { + let compressed = ZstdBuffers::compress(&make_primitive_array(), 3, &array_session())?; + let mut data = compressed.data().clone(); + data.uncompressed_sizes[0] = 16 * 1024 * 1024 * 1024; + + for error in [ + data.decompress_buffers().unwrap_err(), + data.decode_plan().unwrap_err(), + ] { + assert!( + error.to_string().contains("metadata declares"), + "unexpected error: {error}" + ); + } + Ok(()) + } + #[test] fn test_validity_delegates_for_nullable_input() -> VortexResult<()> { let input = make_nullable_primitive_array(); diff --git a/fuzz/Cargo.toml b/fuzz/Cargo.toml index 58c728b23a5..13f307e8b14 100644 --- a/fuzz/Cargo.toml +++ b/fuzz/Cargo.toml @@ -37,8 +37,9 @@ vortex-btrblocks = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } vortex-fsst = { workspace = true, features = ["_test-harness"] } -vortex-io = { workspace = true } +vortex-io = { workspace = true, features = ["wasm-bindgen"] } vortex-mask = { workspace = true } +vortex-row = { workspace = true } vortex-runend = { workspace = true, features = ["arbitrary"] } vortex-session = { workspace = true } vortex-utils = { workspace = true } @@ -104,3 +105,11 @@ name = "compress_gpu" path = "fuzz_targets/compress_gpu.rs" test = false required-features = ["native", "cuda"] + +[[bin]] +bench = false +doc = false +name = "row_encode" +path = "fuzz_targets/row_encode.rs" +test = false +required-features = ["native"] diff --git a/fuzz/fuzz_targets/file_io.rs b/fuzz/fuzz_targets/file_io.rs index f62693ccc49..6d9c8906fc9 100644 --- a/fuzz/fuzz_targets/file_io.rs +++ b/fuzz/fuzz_targets/file_io.rs @@ -19,7 +19,6 @@ use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_btrblocks::BtrBlocksCompressorBuilder; -use vortex_buffer::ByteBufferMut; use vortex_error::VortexExpect; use vortex_error::vortex_panic; use vortex_file::OpenOptionsSessionExt; @@ -72,20 +71,34 @@ fuzz_target!(|fuzz: FuzzFileAction| -> Corpus { ), }; - let mut full_buff = ByteBufferMut::empty(); + let mut full_buff = Vec::new(); let _footer = write_options .blocking(&*RUNTIME) .write(&mut full_buff, array_data.to_array_iterator()) .vortex_expect("file write should succeed in fuzz test"); - let mut output = SESSION + let file = SESSION .open_options() .open_buffer(full_buff) - .vortex_expect("open_buffer should succeed in fuzz test") + .vortex_expect("open_buffer should succeed in fuzz test"); + let projection = projection_expr + .unwrap_or_else(root) + .optimize_recursive(file.dtype()) + .and_then(|expr| expr.bind(file.dtype())) + .vortex_expect("projection should bind in fuzz test"); + let filter = filter_expr + .map(|filter| { + filter + .optimize_recursive(file.dtype()) + .and_then(|expr| expr.bind(file.dtype())) + }) + .transpose() + .vortex_expect("filter should bind in fuzz test"); + let mut output = file .scan() .vortex_expect("scan should succeed in fuzz test") - .with_projection(projection_expr.unwrap_or_else(root)) - .with_some_filter(filter_expr) + .with_projection(projection) + .with_some_filter(filter) .into_array_iter(&*RUNTIME) .vortex_expect("into_array_iter should succeed in fuzz test") .try_collect::<_, Vec<_>, _>() diff --git a/fuzz/fuzz_targets/row_encode.rs b/fuzz/fuzz_targets/row_encode.rs new file mode 100644 index 00000000000..07d63cc322c --- /dev/null +++ b/fuzz/fuzz_targets/row_encode.rs @@ -0,0 +1,35 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![no_main] +#![expect(clippy::unwrap_used)] + +use std::str::FromStr; + +use libfuzzer_sys::Corpus; +use libfuzzer_sys::fuzz_target; +use tracing::level_filters::LevelFilter; +use vortex_error::vortex_panic; +use vortex_fuzz::FuzzRowEncode; +use vortex_fuzz::run_row_encode; + +fuzz_target!( + init: { + let fmt = tracing_subscriber::fmt::format() + .with_ansi(false) // Colour output is messed up in raw logs + .without_time() // We run fuzzer in CI which prepends timestamps + .compact(); + let level = std::env::var("RUST_LOG").map( + |v| LevelFilter::from_str(v.as_str()).unwrap()).unwrap_or(LevelFilter::INFO); + tracing_subscriber::fmt() + .event_format(fmt) + .with_max_level(level) + .init(); + }, + |input: FuzzRowEncode| -> Corpus { + match run_row_encode(input) { + Ok(true) => Corpus::Keep, + Ok(false) => Corpus::Reject, + Err(e) => vortex_panic!("{e}"), + } +}); diff --git a/fuzz/src/array/compare.rs b/fuzz/src/array/compare.rs index f594a9b4102..03ce50a7170 100644 --- a/fuzz/src/array/compare.rs +++ b/fuzz/src/array/compare.rs @@ -186,7 +186,11 @@ pub fn compare_canonical_array( })) .into_array() } - d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => { + d @ (DType::Null + | DType::Map(..) + | DType::Union(..) + | DType::Variant(_) + | DType::Extension(_)) => { unreachable!("DType {d} not supported for fuzzing") } } diff --git a/fuzz/src/array/fill_null.rs b/fuzz/src/array/fill_null.rs index f3a49b25637..fd8d8ec88c3 100644 --- a/fuzz/src/array/fill_null.rs +++ b/fuzz/src/array/fill_null.rs @@ -47,8 +47,12 @@ pub fn fill_null_canonical_array( } Canonical::Struct(_) | Canonical::List(_) + | Canonical::Map(_) | Canonical::FixedSizeList(_) | Canonical::Extension(_) => canonical.into_array().fill_null(fill_value.clone())?, + Canonical::Union(_) => { + todo!("TODO(connor)[Union]: support Union arrays in the fill_null fuzzer") + } Canonical::Variant(_) => unreachable!("Variant arrays are not fuzzed"), }) } diff --git a/fuzz/src/array/filter.rs b/fuzz/src/array/filter.rs index aa54b0d9488..5b90d1cecc8 100644 --- a/fuzz/src/array/filter.rs +++ b/fuzz/src/array/filter.rs @@ -11,6 +11,7 @@ use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::bool::BoolArrayExt; use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::builders::builder_with_capacity_in; use vortex_array::dtype::DType; use vortex_array::match_each_decimal_value_type; use vortex_array::match_each_native_ptype; @@ -121,6 +122,19 @@ pub fn filter_canonical_array( ) .map(|a| a.into_array()) } + DType::Map(..) => { + let mut builder = builder_with_capacity_in( + array.dtype(), + filter.iter().filter(|b| **b).count(), + ctx.allocator(), + ); + for (idx, keep) in filter.iter().enumerate() { + if *keep { + builder.append_scalar(&array.execute_scalar(idx, ctx)?)?; + } + } + Ok(builder.finish()) + } d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => { unreachable!("DType {d} not supported for fuzzing") } diff --git a/fuzz/src/array/mask.rs b/fuzz/src/array/mask.rs index f3cecec2b10..e59c2f4fed9 100644 --- a/fuzz/src/array/mask.rs +++ b/fuzz/src/array/mask.rs @@ -18,8 +18,10 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::bool::BoolArrayExt; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; -use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::builders::builder_with_capacity_in; use vortex_array::dtype::Nullability; use vortex_array::match_each_decimal_value_type; use vortex_array::validity::Validity; @@ -130,7 +132,7 @@ pub fn mask_canonical_array( Canonical::Struct(array) => { let new_validity = mask_validity(&array.validity()?, mask, ctx); StructArray::try_new_with_dtype( - array.unmasked_fields(), + array.iter_unmasked_fields().cloned(), array.struct_fields().clone(), array.len(), new_validity, @@ -138,6 +140,18 @@ pub fn mask_canonical_array( .vortex_expect("StructArray creation should succeed in fuzz test") .into_array() } + Canonical::Map(array) => { + let result_dtype = array.dtype().as_nullable(); + let mut builder = builder_with_capacity_in(&result_dtype, array.len(), ctx.allocator()); + for idx in 0..array.len() { + if mask.value(idx) { + builder.append_scalar(&array.execute_scalar(idx, ctx)?.cast(&result_dtype)?)?; + } else { + builder.append_null(); + } + } + builder.finish() + } Canonical::Extension(array) => { // Recursively mask the storage array let storage_canonical = array.storage_array().clone().execute::(ctx)?; @@ -149,6 +163,9 @@ pub fn mask_canonical_array( .with_nullability(masked_storage.dtype().nullability()); ExtensionArray::new(ext_dtype, masked_storage).into_array() } + Canonical::Union(_) => { + todo!("TODO(connor)[Union]: support Union arrays in the mask fuzzer") + } Canonical::Variant(_) => unreachable!("Variant arrays are not fuzzed"), }) } diff --git a/fuzz/src/array/mod.rs b/fuzz/src/array/mod.rs index b6402cb1e1d..e513c5daf81 100644 --- a/fuzz/src/array/mod.rs +++ b/fuzz/src/array/mod.rs @@ -518,6 +518,7 @@ fn actions_for_dtype(dtype: &DType) -> HashSet { acc.intersection(&actions).copied().collect() }) } + DType::Map(..) => [Compress, Slice, Take, Filter, Mask, ScalarAt].into(), DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), // Currently, no support at all DType::Variant(_) => unreachable!("Variant dtype shouldn't be fuzzed"), diff --git a/fuzz/src/array/scalar_at.rs b/fuzz/src/array/scalar_at.rs index e13ad75343d..a391822f67e 100644 --- a/fuzz/src/array/scalar_at.rs +++ b/fuzz/src/array/scalar_at.rs @@ -6,10 +6,12 @@ use std::sync::Arc; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::arrays::StructArray; use vortex_array::arrays::bool::BoolArrayExt; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::map::MapArrayExt; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::arrays::varbin::varbin_scalar; use vortex_array::dtype::DType; @@ -104,6 +106,29 @@ pub fn scalar_at_canonical_array( let storage_scalar = scalar_at_canonical_array(storage_canonical, index, ctx)?; Scalar::extension_ref(array.ext_dtype().clone(), storage_scalar) } + Canonical::Union(_) => { + todo!("TODO(connor)[Union]: support Union arrays in the scalar_at fuzzer") + } + Canonical::Map(array) => { + let entries = array.entries_at(index)?.execute::(ctx)?; + let keys = entries + .unmasked_field(0) + .clone() + .execute::(ctx)?; + let values = entries + .unmasked_field(1) + .clone() + .execute::(ctx)?; + let pairs = (0..entries.len()) + .map(|entry_index| { + Ok(( + scalar_at_canonical_array(keys.clone(), entry_index, ctx)?, + scalar_at_canonical_array(values.clone(), entry_index, ctx)?, + )) + }) + .collect::>>()?; + Scalar::try_map(array.dtype().clone(), pairs)? + } Canonical::Variant(_) => unreachable!("Variant arrays are not fuzzed"), }) } diff --git a/fuzz/src/array/search_sorted.rs b/fuzz/src/array/search_sorted.rs index 762182b35fb..6eb0d63e93a 100644 --- a/fuzz/src/array/search_sorted.rs +++ b/fuzz/src/array/search_sorted.rs @@ -149,7 +149,11 @@ pub fn search_sorted_canonical_array( .collect::>>()?; scalar_vals.search_sorted(&scalar.cast(array.dtype())?, side) } - d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => { + d @ (DType::Null + | DType::Map(..) + | DType::Union(..) + | DType::Variant(_) + | DType::Extension(_)) => { unreachable!("DType {d} not supported for fuzzing") } } diff --git a/fuzz/src/array/slice.rs b/fuzz/src/array/slice.rs index 2a151002805..e121337e5d3 100644 --- a/fuzz/src/array/slice.rs +++ b/fuzz/src/array/slice.rs @@ -13,8 +13,10 @@ use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::bool::BoolArrayExt; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; -use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::builders::builder_with_capacity_in; use vortex_array::dtype::DType; use vortex_array::match_each_decimal_value_type; use vortex_array::match_each_native_ptype; @@ -124,6 +126,14 @@ pub fn slice_canonical_array( ) .map(|a| a.into_array()) } + DType::Map(..) => { + let mut builder = + builder_with_capacity_in(array.dtype(), stop - start, ctx.allocator()); + for idx in start..stop { + builder.append_scalar(&array.execute_scalar(idx, ctx)?)?; + } + Ok(builder.finish()) + } d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => { unreachable!("DType {d} not supported for fuzzing") } diff --git a/fuzz/src/array/sort.rs b/fuzz/src/array/sort.rs index 3f00fccd06e..b0af8638898 100644 --- a/fuzz/src/array/sort.rs +++ b/fuzz/src/array/sort.rs @@ -102,7 +102,11 @@ pub fn sort_canonical_array(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexR }); take_canonical_array_non_nullable_indices(array, &sort_indices, ctx) } - d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => { + d @ (DType::Null + | DType::Map(..) + | DType::Union(..) + | DType::Variant(_) + | DType::Extension(_)) => { unreachable!("DType {d} not supported for fuzzing") } } diff --git a/fuzz/src/array/take.rs b/fuzz/src/array/take.rs index c67ae6184d8..460cdf1593c 100644 --- a/fuzz/src/array/take.rs +++ b/fuzz/src/array/take.rs @@ -11,7 +11,7 @@ use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::bool::BoolArrayExt; use vortex_array::arrays::struct_::StructArrayExt; -use vortex_array::builders::builder_with_capacity; +use vortex_array::builders::builder_with_capacity_in; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::NativeDecimalType; @@ -115,9 +115,10 @@ pub fn take_canonical_array( .into_array()) } DType::List(..) | DType::FixedSizeList(..) => { - let mut builder = builder_with_capacity( + let mut builder = builder_with_capacity_in( &array.dtype().union_nullability(nullable), indices_slice_non_opt.len(), + ctx.allocator(), ); for idx in indices { if let Some(idx) = idx { @@ -148,6 +149,20 @@ pub fn take_canonical_array( ) .map(|a| a.into_array()) } + DType::Map(..) => { + let result_dtype = array.dtype().union_nullability(nullable); + let mut builder = + builder_with_capacity_in(&result_dtype, indices.len(), ctx.allocator()); + for idx in indices { + if let Some(idx) = idx { + builder + .append_scalar(&array.execute_scalar(*idx, ctx)?.cast(&result_dtype)?)?; + } else { + builder.append_null(); + } + } + Ok(builder.finish()) + } d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => { unreachable!("DType {d} not supported for fuzzing") } diff --git a/fuzz/src/fsst_like.rs b/fuzz/src/fsst_like.rs index 9a8c9ae21ac..92f53cae258 100644 --- a/fuzz/src/fsst_like.rs +++ b/fuzz/src/fsst_like.rs @@ -160,11 +160,8 @@ fn run_like_on_array( len: usize, opts: LikeOptions, ) -> VortexResult { - use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; - let pattern_arr = ConstantArray::new(pattern, len).into_array(); - let result = Like - .try_new_array(len, opts, [array.clone(), pattern_arr])? + let result = Like::try_new(array.clone(), pattern_arr, opts)? .into_array() .execute::(&mut SESSION.create_execution_ctx())?; Ok(result.into_bool()) diff --git a/fuzz/src/gpu/mod.rs b/fuzz/src/gpu/mod.rs index d9ac5c4fa0f..92691949cdf 100644 --- a/fuzz/src/gpu/mod.rs +++ b/fuzz/src/gpu/mod.rs @@ -95,6 +95,7 @@ fn arbitrary_gpu_primitive_dtype(u: &mut Unstructured) -> Result VortexFuzzResult { use vortex_cuda::CanonicalCudaExt; use vortex_cuda::CudaSession; diff --git a/fuzz/src/lib.rs b/fuzz/src/lib.rs index b0abf660045..ca6cd49535e 100644 --- a/fuzz/src/lib.rs +++ b/fuzz/src/lib.rs @@ -7,6 +7,7 @@ mod array; pub mod compress; pub mod error; pub mod fsst_like; +mod row; // File module only available for native builds (requires vortex-file which uses tokio) #[cfg(not(target_arch = "wasm32"))] @@ -31,10 +32,24 @@ pub use fsst_like::run_fsst_like_fuzz; pub use gpu::FuzzCompressGpu; #[cfg(feature = "cuda")] pub use gpu::run_compress_gpu; +pub use row::FuzzRowEncode; +pub use row::run_row_encode; pub const FUZZ_ARRAY_MAX_LEN: usize = 2048; pub const FUZZ_FILE_ARRAY_MAX_LEN: usize = 16_384; +fn enable_latest_core_edition(session: &vortex_session::VortexSession) { + use vortex::editions::CORE_2026_08_3; + use vortex::editions::EditionSessionExt; + use vortex_error::VortexExpect; + use vortex_error::vortex_err; + + session + .enable_edition(CORE_2026_08_3) + .map_err(|error| vortex_err!("{error}")) + .vortex_expect("latest core edition is registered"); +} + // Runtime initialization - platform-specific #[cfg(not(target_arch = "wasm32"))] mod native_runtime { @@ -58,6 +73,7 @@ mod native_runtime { session = session.with::(); vortex_cuda::initialize_cuda(&session.cuda_session()); } + super::enable_latest_core_edition(&session); session }); } @@ -78,8 +94,11 @@ mod wasm_runtime { use vortex_io::session::RuntimeSessionExt; use vortex_session::VortexSession; - pub static SESSION: LazyLock = - LazyLock::new(|| VortexSession::default().with_handle(WasmRuntime::handle())); + pub static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::default().with_handle(WasmRuntime::handle()); + super::enable_latest_core_edition(&session); + session + }); } #[cfg(all(target_arch = "wasm32", target_os = "unknown"))] @@ -94,8 +113,30 @@ mod wasi_runtime { use vortex::VortexSessionDefault; use vortex_session::VortexSession; - pub static SESSION: LazyLock = LazyLock::new(VortexSession::default); + pub static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::default(); + super::enable_latest_core_edition(&session); + session + }); } #[cfg(all(target_arch = "wasm32", not(target_os = "unknown")))] pub use wasi_runtime::SESSION; + +#[cfg(test)] +mod tests { + use vortex::editions::CORE_2026_08_3; + use vortex::editions::EditionSessionExt; + + use super::SESSION; + + #[test] + fn fuzz_session_uses_latest_core_edition() { + assert!( + SESSION + .enabled_editions() + .editions() + .contains(&CORE_2026_08_3) + ); + } +} diff --git a/fuzz/src/row.rs b/fuzz/src/row.rs new file mode 100644 index 00000000000..22508a49a6f --- /dev/null +++ b/fuzz/src/row.rs @@ -0,0 +1,411 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Fuzz harness for the vortex-row byte encoder. +//! +//! Property: comparing two encoded row keys as byte strings yields the same ordering as a +//! logical tuple comparison of the row values under the configured [`RowSortField`]s. This +//! must hold across keys produced by *separate* encode calls over different chunks of the +//! same logical columns — chunks compress to different physical layouts (e.g. decimal value +//! widths), and sort/top-k operators compare keys across chunks. + +use std::cmp::Ordering; + +use arbitrary::Arbitrary; +use arbitrary::Unstructured; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::arbitrary::ArbitraryArray; +use vortex_array::arrays::arbitrary::ArbitraryArrayConfig; +use vortex_array::arrays::arbitrary::ArbitraryWith; +use vortex_array::arrays::bool::BoolArrayExt; +use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; +use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::dtype::BigCast; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalType; +use vortex_array::dtype::PType; +use vortex_array::dtype::half::f16; +use vortex_array::match_each_decimal_value_type; +use vortex_array::match_each_integer_ptype; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_row::RowSortField; +use vortex_row::compute_row_sizes; +use vortex_row::convert_columns; + +use crate::SESSION; +use crate::array::CompressorStrategy; +use crate::array::compress_array; + +/// Keep row counts small: the order property is asserted pairwise (O(rows²)). +const MAX_ROWS: usize = 96; + +#[derive(Debug)] +pub struct FuzzRowEncode { + /// Same-length columns of arbitrary dtypes. + pub columns: Vec, + /// Per-column sort configuration. + pub fields: Vec, + /// Where to split the columns for the cross-chunk property. + pub split_at: usize, + /// Whether to btrblocks-compress each chunk before the cross-chunk encode. + pub compress: bool, +} + +impl<'a> Arbitrary<'a> for FuzzRowEncode { + fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result { + let ncols = u.int_in_range(1..=3)?; + let nrows = u.int_in_range(0..=MAX_ROWS)?; + let mut columns = Vec::with_capacity(ncols); + let mut fields = Vec::with_capacity(ncols); + for _ in 0..ncols { + let array = ArbitraryArray::arbitrary_with_config( + u, + &ArbitraryArrayConfig { + dtype: None, + len: nrows..=nrows, + }, + )? + .0; + columns.push(array); + fields.push(RowSortField::new(u.arbitrary()?, u.arbitrary()?)); + } + Ok(Self { + columns, + fields, + split_at: u.int_in_range(0..=nrows)?, + compress: u.arbitrary()?, + }) + } +} + +/// Run the row-encoding properties for one generated input. +/// +/// Returns `Ok(true)` to keep the input in the corpus, `Ok(false)` to reject inputs that the +/// encoder (correctly) refuses. Property violations panic. +pub fn run_row_encode(input: FuzzRowEncode) -> VortexResult { + let mut ctx = SESSION.create_execution_ctx(); + let FuzzRowEncode { + columns, + fields, + split_at, + compress, + } = input; + + if !columns.iter().all(|c| dtype_row_encodable(c.dtype())) { + assert!( + convert_columns(&columns, &fields, &mut ctx).is_err(), + "row encoder accepted an unsupported dtype: {:?}", + columns.iter().map(|c| c.dtype()).collect::>() + ); + return Ok(false); + } + + let encoded = match convert_columns(&columns, &fields, &mut ctx) { + Ok(encoded) => encoded, + // Values that violate their declared decimal precision are rejected by the encoder; + // such inputs are uninteresting, not property violations. + Err(e) if e.to_string().contains("does not fit") => return Ok(false), + Err(e) => return Err(e), + }; + let rows = collect_row_bytes(&encoded, &mut ctx); + let nrows = columns[0].len(); + + // Property 1: per-row sizes match the encoded key lengths. `RowSize` reports a + // `Struct { fixed, var }` per row; the total size is their sum. + let sizes = compute_row_sizes(&columns, &fields, &mut ctx)?.execute::(&mut ctx)?; + let parts = sizes + .iter_unmasked_fields() + .map(|f| f.clone().execute::(&mut ctx)) + .collect::>>()?; + for (i, row) in rows.iter().enumerate() { + let total: u64 = parts.iter().map(|p| p.as_slice::()[i] as u64).sum(); + assert_eq!( + total, + row.len() as u64, + "row {i}: computed size disagrees with encoded key length" + ); + } + + // Property 2: byte order of whole-array keys equals the logical row order. + let keys = columns + .iter() + .map(|c| row_keys(c, &mut ctx)) + .collect::>>()?; + let logical_cmp = |i: usize, j: usize| -> Ordering { + fields + .iter() + .zip(&keys) + .map(|(f, col)| col[i].compare(&col[j], *f)) + .find(|o| o.is_ne()) + .unwrap_or(Ordering::Equal) + }; + for i in 0..nrows { + for j in 0..nrows { + assert_eq!( + rows[i].cmp(&rows[j]), + logical_cmp(i, j), + "keys for rows {i} and {j} do not compare like their values \ + (fields {fields:?}, dtypes {:?})", + columns.iter().map(|c| c.dtype()).collect::>() + ); + } + } + + // Property 3: keys stay comparable across separately encoded chunks, including chunks + // whose physical layout was changed by compression. + let chunk = |cols: &[ArrayRef], range: std::ops::Range, ctx: &mut ExecutionCtx| { + cols.iter() + .map(|c| { + let sliced = c.slice(range.clone())?; + Ok(if compress { + compress_array(&sliced, CompressorStrategy::Default, ctx) + } else { + sliced + }) + }) + .collect::>>() + }; + let head = chunk(&columns, 0..split_at, &mut ctx)?; + let tail = chunk(&columns, split_at..nrows, &mut ctx)?; + let head_rows = collect_row_bytes(&convert_columns(&head, &fields, &mut ctx)?, &mut ctx); + let tail_rows = collect_row_bytes(&convert_columns(&tail, &fields, &mut ctx)?, &mut ctx); + for (i, head_row) in head_rows.iter().enumerate() { + for (j, tail_row) in tail_rows.iter().enumerate() { + assert_eq!( + head_row.cmp(tail_row), + logical_cmp(i, split_at + j), + "cross-chunk keys for rows {i} and {} do not compare like their values \ + (split at {split_at}, compress {compress}, fields {fields:?}, dtypes {:?})", + split_at + j, + columns.iter().map(|c| c.dtype()).collect::>() + ); + } + } + + Ok(true) +} + +/// Dtypes the row encoder supports; everything else must be rejected with an error. +fn dtype_row_encodable(dtype: &DType) -> bool { + match dtype { + DType::Null | DType::Bool(_) | DType::Primitive(..) | DType::Utf8(_) | DType::Binary(_) => { + true + } + DType::Decimal(dt, _) => !matches!( + DecimalType::smallest_decimal_value_type(dt), + DecimalType::I256 + ), + DType::Struct(fields, _) => fields.fields().all(|f| dtype_row_encodable(&f)), + DType::FixedSizeList(elem, ..) => dtype_row_encodable(elem), + _ => false, + } +} + +fn collect_row_bytes(array: &ListViewArray, ctx: &mut ExecutionCtx) -> Vec> { + (0..array.len()) + .map(|i| { + let slice = array + .list_elements_at(i) + .vortex_expect("row key bytes must be addressable"); + let p = slice + .execute::(ctx) + .vortex_expect("row key bytes must canonicalize"); + p.as_slice::().to_vec() + }) + .collect() +} + +/// A row value lifted into a form whose comparison mirrors the byte encoding's semantics. +#[derive(Debug, Clone, PartialEq)] +enum RowKey { + Null, + Bool(bool), + /// Every integer ptype and decimal unscaled value (Decimal256 is unsupported upstream). + Int(i128), + /// IEEE-754 total-order bit key (sign-flipped bits), matching the encoded byte order. + Float(u64), + Bytes(Vec), + /// Struct fields or fixed-size-list elements, in order. + Composite(Vec), +} + +impl RowKey { + /// Compare as the encoded bytes compare: nulls are positioned by `nulls_first` regardless + /// of direction, and `descending` inverts leaf values only — the sentinel bytes that + /// place nulls are never inverted, at any nesting depth. + fn compare(&self, other: &Self, field: RowSortField) -> Ordering { + use RowKey::*; + match (self, other) { + (Null, Null) => Ordering::Equal, + (Null, _) => { + if field.nulls_first { + Ordering::Less + } else { + Ordering::Greater + } + } + (_, Null) => { + if field.nulls_first { + Ordering::Greater + } else { + Ordering::Less + } + } + (Composite(lhs), Composite(rhs)) => lhs + .iter() + .zip(rhs) + .map(|(l, r)| l.compare(r, field)) + .find(|o| o.is_ne()) + .unwrap_or(Ordering::Equal), + (lhs, rhs) => { + let natural = match (lhs, rhs) { + (Bool(l), Bool(r)) => l.cmp(r), + (Int(l), Int(r)) => l.cmp(r), + (Float(l), Float(r)) => l.cmp(r), + (Bytes(l), Bytes(r)) => l.cmp(r), + _ => unreachable!("mismatched RowKey kinds within one column"), + }; + if field.descending { + natural.reverse() + } else { + natural + } + } + } + } +} + +fn f64_key(v: f64) -> u64 { + let bits = v.to_bits(); + if bits >> 63 == 1 { + !bits + } else { + bits ^ (1 << 63) + } +} + +fn f32_key(v: f32) -> u64 { + let bits = v.to_bits(); + (if bits >> 31 == 1 { + !bits + } else { + bits ^ (1 << 31) + }) as u64 +} + +fn f16_key(v: f16) -> u64 { + let bits = v.to_bits(); + (if bits >> 15 == 1 { + !bits + } else { + bits ^ (1 << 15) + }) as u64 +} + +/// Lift every row of `array` into a [`RowKey`]. +fn row_keys(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { + let len = array.len(); + let with_mask = |keys: Vec, array: &ArrayRef, ctx: &mut ExecutionCtx| { + let mask = array.validity()?.execute_mask(len, ctx)?; + Ok(keys + .into_iter() + .enumerate() + .map(|(i, k)| if mask.value(i) { k } else { RowKey::Null }) + .collect()) + }; + match array.clone().execute::(ctx)? { + Canonical::Null(_) => Ok(vec![RowKey::Null; len]), + Canonical::Bool(a) => { + let bits = a.to_bit_buffer(); + let keys = (0..len).map(|i| RowKey::Bool(bits.value(i))).collect(); + with_mask(keys, &a.into_array(), ctx) + } + Canonical::Primitive(a) => { + let keys = match a.ptype() { + PType::F16 => a + .as_slice::() + .iter() + .map(|v| RowKey::Float(f16_key(*v))) + .collect(), + PType::F32 => a + .as_slice::() + .iter() + .map(|v| RowKey::Float(f32_key(*v))) + .collect(), + PType::F64 => a + .as_slice::() + .iter() + .map(|v| RowKey::Float(f64_key(*v))) + .collect(), + _ => match_each_integer_ptype!(a.ptype(), |P| { + a.as_slice::

() + .iter() + .map(|v| RowKey::Int((*v).into())) + .collect::>() + }), + }; + with_mask(keys, &a.into_array(), ctx) + } + Canonical::Decimal(a) => { + let mask = a.as_ref().validity()?.execute_mask(len, ctx)?; + match_each_decimal_value_type!(a.values_type(), |D| { + let buf = a.buffer::(); + Ok((0..len) + .map(|i| { + if mask.value(i) { + RowKey::Int(::from(buf[i]).vortex_expect( + "valid decimal values fit i128 (Decimal256 dtypes are filtered)", + )) + } else { + RowKey::Null + } + }) + .collect()) + }) + } + Canonical::VarBinView(a) => { + let mask = a.as_ref().validity()?.execute_mask(len, ctx)?; + Ok((0..len) + .map(|i| { + if mask.value(i) { + RowKey::Bytes(a.bytes_at(i).to_vec()) + } else { + RowKey::Null + } + }) + .collect()) + } + Canonical::Struct(a) => { + let children = a + .iter_unmasked_fields() + .map(|child| row_keys(child, ctx)) + .collect::>>()?; + let keys = (0..len) + .map(|i| RowKey::Composite(children.iter().map(|c| c[i].clone()).collect())) + .collect(); + with_mask(keys, &a.into_array(), ctx) + } + Canonical::FixedSizeList(a) => { + let k = a.list_size() as usize; + let elements = row_keys(&a.elements().clone(), ctx)?; + let keys = (0..len) + .map(|i| RowKey::Composite(elements[i * k..(i + 1) * k].to_vec())) + .collect(); + with_mask(keys, &a.into_array(), ctx) + } + c => unreachable!( + "unsupported dtypes are rejected before oracle construction: {:?}", + c.dtype() + ), + } +} diff --git a/java/README.md b/java/README.md index f45f2474f11..e1d682cd7ba 100644 --- a/java/README.md +++ b/java/README.md @@ -11,7 +11,7 @@ We provide two interfaces for working with Vortex from Java: We publish three artifacts out of this repo at CI time to Maven Central Sonatype: -* `vortex-jni` JAR containing the JNI code, plus compiled native libraries for all of the following targets: `aarch64-apple-darwin`, `aarch64-unknown-linux-gnu`, `x86_64-unknown-linux-gnu` +* `vortex-jni` JAR containing the JNI code, plus compiled native libraries for all of the following targets: `aarch64-apple-darwin`, `aarch64-unknown-linux-gnu`, `x86_64-unknown-linux-gnu`, `x86_64-pc-windows-msvc` * `vortex-jni-all` which is the "shadow JAR" containing all of `vortex-jni` as well as all upstream Java dependencies packaged in a single JAR. * `vortex-spark` which is the runtime JAR needed for the Vortex Spark bindings diff --git a/java/gradle/libs.versions.toml b/java/gradle/libs.versions.toml index 1135946d4ab..0b817553b96 100644 --- a/java/gradle/libs.versions.toml +++ b/java/gradle/libs.versions.toml @@ -12,8 +12,8 @@ netty = "4.2.15.Final" nopen = "1.0.1" roaringbitmap = "1.6.14" slf4j = "2.0.18" -spark3 = "3.5.8" -spark4 = "4.1.1" +spark3 = "3.5.9" +spark4 = "4.1.2" s3mock = "5.0.0" testcontainers-jupiter = "1.21.4" diff --git a/java/vortex-jni/.gitignore b/java/vortex-jni/.gitignore index b013b0fca94..0c01ec23f83 100644 --- a/java/vortex-jni/.gitignore +++ b/java/vortex-jni/.gitignore @@ -1,2 +1,3 @@ +*.dll *.dylib *.so diff --git a/java/vortex-jni/build.gradle.kts b/java/vortex-jni/build.gradle.kts index 8f5106a1262..09e2651fca7 100644 --- a/java/vortex-jni/build.gradle.kts +++ b/java/vortex-jni/build.gradle.kts @@ -129,7 +129,11 @@ val libExt = "win" -> ".dll" else -> throw GradleException("Unsupported OS short name: $osShortName") } -val nativeLibrary = rustWorkspaceDir.resolve("target/debug/libvortex_jni$libExt") +// Cargo emits `vortex_jni.dll` on Windows and `libvortex_jni.{so,dylib}` elsewhere. The +// bundled resources always use the `lib` prefix so that NativeLoader can build one name. +val bundledLibraryName = "libvortex_jni$libExt" +val cargoLibraryName = if (osShortName == "win") "vortex_jni$libExt" else bundledLibraryName +val nativeLibrary = rustWorkspaceDir.resolve("target/debug/$cargoLibraryName") val nativeLibraryDir = "src/main/resources/native/$osShortName-$osArch" val buildJniLibrary = @@ -156,7 +160,9 @@ tasks.register("makeTestFiles") { // Only populate the host-arch directory so cross-compiled libs for other // architectures (placed by the publish workflow) are preserved. - from(nativeLibrary) + from(nativeLibrary) { + rename { bundledLibraryName } + } into(nativeLibraryDir) } diff --git a/java/vortex-jni/src/main/java/dev/vortex/api/DataSource.java b/java/vortex-jni/src/main/java/dev/vortex/api/DataSource.java index 9a888f224f9..8658cf5198b 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/api/DataSource.java +++ b/java/vortex-jni/src/main/java/dev/vortex/api/DataSource.java @@ -120,7 +120,7 @@ public static DataSource open(Session session, List readables, i /** Arrow schema of the data source (and of scans produced from it). */ public Schema arrowSchema(BufferAllocator allocator) { try (ArrowSchema schema = ArrowSchema.allocateNew(allocator)) { - NativeDataSource.arrowSchema(pointer, schema.memoryAddress()); + NativeDataSource.arrowSchema(session.nativePointer(), pointer, schema.memoryAddress()); return Data.importSchema(allocator, schema, null); } } diff --git a/java/vortex-jni/src/main/java/dev/vortex/api/Scan.java b/java/vortex-jni/src/main/java/dev/vortex/api/Scan.java index 0e1cf94349a..22a26fed365 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/api/Scan.java +++ b/java/vortex-jni/src/main/java/dev/vortex/api/Scan.java @@ -42,7 +42,7 @@ static Scan fromPointer(Session session, long pointer) { */ public Schema arrowSchema(BufferAllocator allocator) { try (ArrowSchema schema = ArrowSchema.allocateNew(allocator)) { - NativeScan.arrowSchema(pointer, schema.memoryAddress()); + NativeScan.arrowSchema(session.nativePointer(), pointer, schema.memoryAddress()); return Data.importSchema(allocator, schema, null); } } diff --git a/java/vortex-jni/src/main/java/dev/vortex/api/ScanOptions.java b/java/vortex-jni/src/main/java/dev/vortex/api/ScanOptions.java index e9009582e28..25f03e76f7f 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/api/ScanOptions.java +++ b/java/vortex-jni/src/main/java/dev/vortex/api/ScanOptions.java @@ -36,6 +36,8 @@ public interface ScanOptions { /** * Sorted ascending, unique row indices that should be included in (or excluded from) the scan, depending on * {@link #selectionMode()}. + * + *

The native scan validates this contract and rejects indices that are not strictly ascending. */ Optional selectionIndices(); @@ -64,9 +66,6 @@ default void validateSelectionPayload() { if (hasIndices && hasRoaringBitmap) { throw new IllegalArgumentException("row selection must use either indices or roaring bitmap, not both"); } - if (hasIndices) { - validateSelectionIndices(selectionIndices().orElseThrow()); - } switch (selectionMode()) { case INCLUDE_ALL -> { @@ -91,20 +90,6 @@ default void validateSelectionPayload() { } } - private static void validateSelectionIndices(long[] selectionIndices) { - long previous = -1L; - for (int i = 0; i < selectionIndices.length; i++) { - long index = selectionIndices[i]; - if (index < 0) { - throw new IllegalArgumentException("selection indices must be non-negative"); - } - if (i > 0 && index <= previous) { - throw new IllegalArgumentException("selection indices must be sorted ascending and unique"); - } - previous = index; - } - } - static ScanOptions of() { return ImmutableScanOptions.builder().build(); } diff --git a/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java b/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java index a62243aab07..c740a4b70f5 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java +++ b/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java @@ -8,6 +8,7 @@ import dev.vortex.io.NativeWritable; import dev.vortex.jni.NativeWriter; import java.io.IOException; +import java.util.HashMap; import java.util.Map; import java.util.Objects; import java.util.concurrent.atomic.AtomicBoolean; @@ -19,6 +20,9 @@ /** * Writer for Vortex files. * + *

Open one with {@link #builder(Session, String, Schema, BufferAllocator)} to write to a URI, or + * {@link #builder(Session, NativeWritable, Schema, BufferAllocator)} to write into a caller-provided byte sink. + * *

Batches are accepted via the Arrow C Data Interface: callers export an Arrow record batch to an {@code ArrowArray} * / {@code ArrowSchema} pair and pass the addresses to {@link #writeBatch(long, long)}. The writer accepts up to four * in-flight batches on the session's runtime thread before back-pressuring the caller. @@ -44,53 +48,133 @@ private VortexWriter(long pointer) { } /** - * Create a writer that streams records into the file at {@code uri}. The path may be a full URI or a plain local - * filesystem path. The Arrow schema describes the exact layout of every batch written. + * Start configuring a writer that streams records into the file at {@code uri} through a native storage client. The + * path may be a full URI or a plain local filesystem path. + * + * @param arrowSchema describes the exact layout of every batch written */ - public static VortexWriter create( - Session session, String uri, Schema arrowSchema, Map options, BufferAllocator allocator) - throws IOException { - Objects.requireNonNull(session, "session"); - Objects.requireNonNull(uri, "uri"); - Objects.requireNonNull(arrowSchema, "arrowSchema"); - Objects.requireNonNull(allocator, "allocator"); - ArrowSchema ffi = ArrowSchema.allocateNew(allocator); - try { - Data.exportSchema(allocator, arrowSchema, null, ffi); - long ptr = NativeWriter.create(session.nativePointer(), uri, ffi.memoryAddress(), options); - if (ptr <= 0) { - throw new IOException("failed to create writer for uri " + uri + " (ptr=" + ptr + ")"); - } - return new VortexWriter(ptr); - } finally { - ffi.close(); - } + public static Builder builder(Session session, String uri, Schema arrowSchema, BufferAllocator allocator) { + return new Builder(session, Objects.requireNonNull(uri, "uri"), null, arrowSchema, allocator); } /** - * Create a writer that streams the file into a caller-provided byte sink instead of a native storage client. This - * is the integration point for external I/O abstractions (for example Iceberg's {@code PositionOutputStream}). + * Start configuring a writer that streams the file into a caller-provided byte sink instead of a native storage + * client. This is the integration point for external I/O abstractions (for example Iceberg's + * {@code PositionOutputStream}). * *

The native side writes and flushes the sink but never closes it: after {@link #close()} returns, all bytes * have been written and flushed, and the caller must close the sink to finalize the file. + * + * @param arrowSchema describes the exact layout of every batch written */ - public static VortexWriter create( - Session session, NativeWritable writable, Schema arrowSchema, BufferAllocator allocator) - throws IOException { - Objects.requireNonNull(session, "session"); - Objects.requireNonNull(writable, "writable"); - Objects.requireNonNull(arrowSchema, "arrowSchema"); - Objects.requireNonNull(allocator, "allocator"); - ArrowSchema ffi = ArrowSchema.allocateNew(allocator); - try { - Data.exportSchema(allocator, arrowSchema, null, ffi); - long ptr = NativeWriter.createStream(session.nativePointer(), writable, ffi.memoryAddress()); - if (ptr <= 0) { - throw new IOException("failed to create stream writer (ptr=" + ptr + ")"); + public static Builder builder( + Session session, NativeWritable writable, Schema arrowSchema, BufferAllocator allocator) { + return new Builder(session, null, Objects.requireNonNull(writable, "writable"), arrowSchema, allocator); + } + + /** + * Configures and opens a {@link VortexWriter}. Everything required to open one is fixed by + * {@link VortexWriter#builder}; the setters here are all optional. + * + *

Not thread-safe, and not reusable: each {@link #build()} opens one file. + */ + public static final class Builder { + private final Session session; + private final String uri; + private final NativeWritable writable; + private final Schema arrowSchema; + private final BufferAllocator allocator; + private final Map options = new HashMap<>(); + private final Map metadata = new HashMap<>(); + + private Builder( + Session session, String uri, NativeWritable writable, Schema arrowSchema, BufferAllocator allocator) { + this.session = Objects.requireNonNull(session, "session"); + this.uri = uri; + this.writable = writable; + this.arrowSchema = Objects.requireNonNull(arrowSchema, "arrowSchema"); + this.allocator = Objects.requireNonNull(allocator, "allocator"); + } + + /** + * Object-store credentials and options, replacing any set so far. Only the URI destination reaches a storage + * client, so this is rejected on a writer built over a {@link NativeWritable}. + */ + public Builder options(Map newOptions) { + checkOptionsApply(); + Objects.requireNonNull(newOptions, "options"); + options.clear(); + newOptions.forEach(this::putOption); + return this; + } + + /** Add a single object-store option. See {@link #options(Map)}. */ + public Builder putOption(String key, String value) { + checkOptionsApply(); + Objects.requireNonNull(key, "key"); + Objects.requireNonNull(value, "value"); + options.put(key, value); + return this; + } + + /** + * User-defined metadata segments to store in the file footer, replacing any set so far. + * + *

Values are opaque bytes, returned verbatim by {@code NativeFiles.readMetadata}. The native writer caps + * both the number of segments per file and the length of each key; the limits are enforced (and named in the + * error) by {@link #build()}, rather than when the file is finalized. + */ + public Builder metadata(Map newMetadata) { + Objects.requireNonNull(newMetadata, "metadata"); + metadata.clear(); + newMetadata.forEach(this::putMetadata); + return this; + } + + /** Add a single metadata segment. See {@link #metadata(Map)}. */ + public Builder putMetadata(String key, byte[] value) { + Objects.requireNonNull(key, "key"); + Objects.requireNonNull(value, "value"); + metadata.put(key, value); + return this; + } + + /** + * Open the writer. + * + * @throws IOException if the writer cannot be opened, including when the native side rejects the configured + * metadata + */ + public VortexWriter build() throws IOException { + ArrowSchema ffi = ArrowSchema.allocateNew(allocator); + try { + Data.exportSchema(allocator, arrowSchema, null, ffi); + final long pointer; + try { + pointer = uri != null + ? NativeWriter.create(session.nativePointer(), uri, ffi.memoryAddress(), options, metadata) + : NativeWriter.createStream( + session.nativePointer(), writable, ffi.memoryAddress(), metadata); + } catch (RuntimeException e) { + // Native failures arrive as RuntimeException. Wrap them so callers can handle every + // failure to open a writer as an IOException, as writeBatch and finish already do. + throw new IOException("failed to create writer for " + destination(), e); + } + if (pointer <= 0) { + throw new IOException("failed to create writer for " + destination() + " (ptr=" + pointer + ")"); + } + return new VortexWriter(pointer); + } finally { + ffi.close(); } - return new VortexWriter(ptr); - } finally { - ffi.close(); + } + + private String destination() { + return uri != null ? "uri " + uri : "stream"; + } + + private void checkOptionsApply() { + Preconditions.checkState(uri != null, "options apply to uri destinations only"); } } @@ -125,6 +209,23 @@ public synchronized long bytesWritten() { return bytesWritten; } + /** + * Return the number of uncompressed bytes accepted by the writer but not yet written to the sink. + * + *

Together with {@link #bytesWritten()}, this lets callers estimate the in-progress file size: bytes that + * reached the sink are already compressed, while buffered bytes are still uncompressed and will shrink by roughly + * the file's observed compression ratio once flushed. After {@link #finish()}, this is zero. + */ + public synchronized long bufferedBytes() { + if (summary != null) { + return 0; + } + Preconditions.checkState(!closed.get(), "writer closed without a write summary"); + long bufferedBytes = NativeWriter.bufferedBytes(pointer); + Preconditions.checkState(bufferedBytes >= 0, "native writer returned an invalid buffered byte count"); + return bufferedBytes; + } + /** * Flush pending batches, finalize the file, and return its statistics and physical sizes. * diff --git a/java/vortex-jni/src/main/java/dev/vortex/jni/NativeDataSource.java b/java/vortex-jni/src/main/java/dev/vortex/jni/NativeDataSource.java index 76d09baf42e..7ceaf44894b 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/jni/NativeDataSource.java +++ b/java/vortex-jni/src/main/java/dev/vortex/jni/NativeDataSource.java @@ -40,8 +40,13 @@ public static native long openFiles( /** Free a data source pointer. */ public static native void free(long pointer); - /** Export the data source's schema into the Arrow C Data Interface struct at {@code schemaAddress}. */ - public static native void arrowSchema(long pointer, long schemaAddress); + /** + * Export the data source's schema into the Arrow C Data Interface struct at {@code schemaAddress}. Extension dtypes + * are dispatched through the session's registered Arrow export plugins. + * + * @param sessionPointer pointer from {@link NativeSession#newSession()} + */ + public static native void arrowSchema(long sessionPointer, long pointer, long schemaAddress); /** * Populate {@code out} with {@code [rows, cardinality]}. Cardinality is one of {@code 0=unknown}, diff --git a/java/vortex-jni/src/main/java/dev/vortex/jni/NativeFiles.java b/java/vortex-jni/src/main/java/dev/vortex/jni/NativeFiles.java index fae03d50186..8a969276e34 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/jni/NativeFiles.java +++ b/java/vortex-jni/src/main/java/dev/vortex/jni/NativeFiles.java @@ -3,13 +3,16 @@ package dev.vortex.jni; +import com.google.common.base.Preconditions; import dev.vortex.api.Session; +import dev.vortex.io.NativeReadable; import java.util.List; import java.util.Map; +import java.util.Objects; /** - * Static utilities for discovering and deleting Vortex files on an object store. The caller supplies a {@link Session}; - * its runtime handle is forwarded to the underlying object store. + * Static utilities for discovering, deleting, and inspecting the metadata of Vortex files on an object store. The + * caller supplies a {@link Session}; its runtime handle is forwarded to the underlying object store. */ public final class NativeFiles { static { @@ -28,7 +31,42 @@ public static void delete(Session session, String[] uris, Map op delete(session.nativePointer(), uris, options); } + /** + * Read the user-defined metadata segments written into the Vortex file at {@code uri}, as opaque bytes keyed by the + * names the {@code VortexWriter} builder was given. Returns an empty map for a file that carries no metadata. + * + *

This opens the file on its own. Metadata lives in dedicated segments that opening a file does not read by + * default, so a segment that falls outside the file-tail read costs an extra round trip. + */ + public static Map readMetadata(Session session, String uri, Map options) { + return readMetadata(session.nativePointer(), uri, options); + } + + /** + * Read the user-defined metadata segments of a Vortex file through a caller-provided {@link NativeReadable}, + * instead of a native storage client. See {@link #readMetadata(Session, String, Map)}. + * + *

The readable must stay open for the duration of the call; native code never closes it. + * + *

{@link NativeReadable#name()} keys the file in the session's footer cache, exactly as it does when the + * readable is scanned, so reading metadata and then scanning the same file share one footer read. + */ + public static Map readMetadata(Session session, NativeReadable readable) { + Objects.requireNonNull(readable, "readable"); + String name = readable.name(); + Preconditions.checkArgument(name != null && !name.isEmpty(), "readable reported an empty name"); + long length = readable.length(); + Preconditions.checkArgument(length >= 0, "readable for %s reported negative length", name); + return readMetadataFromReadable(session.nativePointer(), readable, name, length); + } + private static native List listFiles(long sessionPointer, String uri, Map options); private static native void delete(long sessionPointer, String[] uris, Map options); + + private static native Map readMetadata( + long sessionPointer, String uri, Map options); + + private static native Map readMetadataFromReadable( + long sessionPointer, Object readable, String name, long length); } diff --git a/java/vortex-jni/src/main/java/dev/vortex/jni/NativeScan.java b/java/vortex-jni/src/main/java/dev/vortex/jni/NativeScan.java index 86b4b2e3359..ff80b30caaa 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/jni/NativeScan.java +++ b/java/vortex-jni/src/main/java/dev/vortex/jni/NativeScan.java @@ -43,8 +43,13 @@ public static native long create( /** Free a scan pointer. */ public static native void free(long pointer); - /** Export the scan's schema into the Arrow C Data Interface struct at {@code schemaAddress}. */ - public static native void arrowSchema(long pointer, long schemaAddress); + /** + * Export the scan's schema into the Arrow C Data Interface struct at {@code schemaAddress}. Extension dtypes are + * dispatched through the session's registered Arrow export plugins. + * + * @param sessionPointer pointer from {@link NativeSession#newSession()} + */ + public static native void arrowSchema(long sessionPointer, long pointer, long schemaAddress); /** Fill {@code out} with {@code [count, cardinality]}. */ public static native void partitionCount(long pointer, long[] out); diff --git a/java/vortex-jni/src/main/java/dev/vortex/jni/NativeWriter.java b/java/vortex-jni/src/main/java/dev/vortex/jni/NativeWriter.java index 4f61a381438..44bbd3b5323 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/jni/NativeWriter.java +++ b/java/vortex-jni/src/main/java/dev/vortex/jni/NativeWriter.java @@ -14,15 +14,26 @@ public final class NativeWriter { private NativeWriter() {} - /** Open a writer at {@code uri} that accepts batches matching the Arrow schema at {@code arrowSchemaAddress}. */ + /** + * Open a writer at {@code uri} that accepts batches matching the Arrow schema at {@code arrowSchemaAddress}. + * + * @param metadata user-defined metadata segments to store in the file footer (may be null) + */ public static native long create( - long sessionPointer, String uri, long arrowSchemaAddress, Map options); + long sessionPointer, + String uri, + long arrowSchemaAddress, + Map options, + Map metadata); /** * Open a writer that streams the file into a caller-provided {@link dev.vortex.io.NativeWritable}. The native side * writes and flushes but never closes the writable; the caller must close it after {@link #close}. + * + * @param metadata user-defined metadata segments to store in the file footer (may be null) */ - public static native long createStream(long sessionPointer, Object writable, long arrowSchemaAddress); + public static native long createStream( + long sessionPointer, Object writable, long arrowSchemaAddress, Map metadata); /** * Write a batch directly from Arrow C Data Interface addresses. @@ -37,6 +48,9 @@ public static native long create( /** Number of bytes successfully written to the underlying sink so far. */ public static native long bytesWritten(long writerPointer); + /** Number of uncompressed bytes buffered by the native writer that have not yet reached the sink. */ + public static native long bufferedBytes(long writerPointer); + /** Flush and close the writer. Must be called exactly once. */ public static native void close(long writerPointer); diff --git a/java/vortex-jni/src/test/java/dev/vortex/api/ExpressionTagParityTest.java b/java/vortex-jni/src/test/java/dev/vortex/api/ExpressionTagParityTest.java new file mode 100644 index 00000000000..00712632434 --- /dev/null +++ b/java/vortex-jni/src/test/java/dev/vortex/api/ExpressionTagParityTest.java @@ -0,0 +1,196 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.api; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import dev.vortex.api.Expression.BinaryOp; +import dev.vortex.api.Expression.DType; +import dev.vortex.api.Expression.DuplicateHandling; +import dev.vortex.api.Expression.TimeUnit; +import dev.vortex.jni.NativeExpression; +import dev.vortex.jni.NativeLoader; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +/** + * Parity tests for the byte tags {@link Expression} hands to the native side. + * + *

Four enums carry a tag that the Rust side switches on: {@link BinaryOp} against {@code parse_op}, + * {@link DuplicateHandling} against {@code parse_duplicate_handling}, {@link TimeUnit} against + * {@code TimeUnit::try_from}, and {@link DType} against the table in {@code literalNull}. Three of the four say in + * their javadoc that the values must match the Rust table, but nothing checked it. Drift compiles on both sides, and + * the two failure modes are not equally loud: a tag past the end of a table reaches the {@code other =>} arm and + * throws, while a tag that collides with a sibling decodes to the wrong operator or the wrong time unit and returns an + * expression that reads valid. + * + *

So each table is pinned twice: the constants against the bytes Rust matches, and every constant against the native + * call that consumes it. Both temporal types are exercised because between them they reject enough units to tell the + * five {@code TimeUnit} tags apart — Timestamp takes everything but {@code DAYS}, Date only {@code DAYS} and + * {@code MILLISECONDS} — and Date names the unit it rejected, so a swapped tag reports the wrong name here instead of + * reading a value at the wrong scale. + */ +public final class ExpressionTagParityTest { + @BeforeAll + public static void loadLibrary() { + NativeLoader.loadJni(); + } + + @Test + public void binaryOpCodesMatchTheParseOpTable() { + assertEquals(0, BinaryOp.EQ.code()); + assertEquals(1, BinaryOp.NOT_EQ.code()); + assertEquals(2, BinaryOp.GT.code()); + assertEquals(3, BinaryOp.GTE.code()); + assertEquals(4, BinaryOp.LT.code()); + assertEquals(5, BinaryOp.LTE.code()); + assertEquals(6, BinaryOp.AND.code()); + assertEquals(7, BinaryOp.OR.code()); + assertEquals(8, BinaryOp.ADD.code()); + assertEquals(9, BinaryOp.SUB.code()); + assertEquals(10, BinaryOp.MUL.code()); + assertEquals(11, BinaryOp.DIV.code()); + assertEquals(12, BinaryOp.values().length); + } + + @Test + public void everyBinaryOpIsAcceptedByTheNativeSide() { + // parse_op bails on a code it does not know, so an operator Java has and Rust lacks fails here. + for (BinaryOp op : BinaryOp.values()) { + assertNotNull( + Expression.binary(op, Expression.column("a"), Expression.literal(1L)), + () -> "native side rejected " + op); + } + } + + @Test + public void aBinaryOpCodePastTheTableIsRejectedByName() { + Expression lhs = Expression.column("a"); + Expression rhs = Expression.literal(1L); + RuntimeException exception = assertThrows( + RuntimeException.class, + () -> NativeExpression.binary( + (byte) BinaryOp.values().length, lhs.nativePointer(), rhs.nativePointer())); + assertTrue( + exception.getMessage().contains("unknown binary operator code: 12"), + () -> "unexpected message: " + exception.getMessage()); + } + + @Test + public void duplicateHandlingTagsMatchTheRustTable() { + assertEquals(0, DuplicateHandling.RIGHT_MOST.tag()); + assertEquals(1, DuplicateHandling.ERROR.tag()); + assertEquals(2, DuplicateHandling.values().length); + } + + @Test + public void everyDuplicateHandlingIsAcceptedByTheNativeSide() { + for (DuplicateHandling handling : DuplicateHandling.values()) { + assertNotNull( + Expression.merge(handling, Expression.column("a"), Expression.column("b")), + () -> "native side rejected " + handling); + } + } + + @Test + public void aDuplicateHandlingTagPastTheTableIsRejectedByName() { + RuntimeException exception = assertThrows( + RuntimeException.class, + () -> NativeExpression.merge(new long[0], (byte) DuplicateHandling.values().length)); + assertTrue( + exception.getMessage().contains("unknown duplicate handling code: 2"), + () -> "unexpected message: " + exception.getMessage()); + } + + @Test + public void timeUnitTagsMatchTheRustTable() { + assertEquals(0, TimeUnit.NANOSECONDS.tag()); + assertEquals(1, TimeUnit.MICROSECONDS.tag()); + assertEquals(2, TimeUnit.MILLISECONDS.tag()); + assertEquals(3, TimeUnit.SECONDS.tag()); + assertEquals(4, TimeUnit.DAYS.tag()); + assertEquals(5, TimeUnit.values().length); + } + + @Test + public void timestampAcceptsEveryUnitExceptDays() { + for (TimeUnit unit : + new TimeUnit[] {TimeUnit.NANOSECONDS, TimeUnit.MICROSECONDS, TimeUnit.MILLISECONDS, TimeUnit.SECONDS}) { + assertNotNull(Expression.literalTimestamp(0L, unit, "UTC"), unit::name); + assertNotNull(Expression.nullLiteralTimestamp(unit, null), unit::name); + } + RuntimeException exception = + assertThrows(RuntimeException.class, () -> Expression.literalTimestamp(0L, TimeUnit.DAYS, "UTC")); + assertTrue( + exception.getMessage().contains("Timestamp does not support Days time unit"), + () -> "unexpected message: " + exception.getMessage()); + } + + @Test + public void dateAcceptsOnlyDaysAndMilliseconds() { + assertNotNull(Expression.literalDate(0L, TimeUnit.DAYS)); + assertNotNull(Expression.literalDate(0L, TimeUnit.MILLISECONDS)); + assertNotNull(Expression.nullLiteralDate(TimeUnit.DAYS)); + assertDateRejects(TimeUnit.NANOSECONDS); + assertDateRejects(TimeUnit.MICROSECONDS); + assertDateRejects(TimeUnit.SECONDS); + } + + @Test + public void aTimeUnitTagPastTheTableIsRejected() { + assertThrows( + RuntimeException.class, + () -> NativeExpression.literalTimestamp(0L, (byte) TimeUnit.values().length, "UTC", false)); + } + + @Test + public void nullLiteralDTypeTagsMatchTheRustTable() { + assertEquals(0, DType.BOOL.tag()); + assertEquals(1, DType.I8.tag()); + assertEquals(2, DType.I16.tag()); + assertEquals(3, DType.I32.tag()); + assertEquals(4, DType.I64.tag()); + assertEquals(5, DType.F32.tag()); + assertEquals(6, DType.F64.tag()); + assertEquals(7, DType.UTF8.tag()); + assertEquals(8, DType.BINARY.tag()); + assertEquals(9, DType.values().length); + } + + @Test + public void everyNullLiteralDTypeIsAcceptedByTheNativeSide() { + for (DType dtype : DType.values()) { + assertNotNull(Expression.nullLiteral(dtype), () -> "native side rejected " + dtype); + } + } + + @Test + public void aDTypeTagPastTheTableIsRejectedByName() { + RuntimeException exception = + assertThrows(RuntimeException.class, () -> NativeExpression.literalNull((byte) DType.values().length)); + assertTrue( + exception.getMessage().contains("unknown null dtype tag: 9"), + () -> "unexpected message: " + exception.getMessage()); + } + + @Test + public void theTagTablesAreReadIndependently() { + // literalNull's javadoc says its tags intentionally do not overlap with parse_time_unit. They do + // overlap numerically — the guarantee is that each tag is only ever read by its own parser, so the + // byte 4 means I64 to literalNull and DAYS to literalDate. + assertEquals(TimeUnit.DAYS.tag(), DType.I64.tag()); + assertNotNull(Expression.nullLiteral(DType.I64)); + assertNotNull(Expression.literalDate(0L, TimeUnit.DAYS)); + } + + private static void assertDateRejects(TimeUnit unit) { + RuntimeException exception = assertThrows(RuntimeException.class, () -> Expression.literalDate(0L, unit)); + assertTrue( + exception.getMessage().contains("Date type does not support time unit"), + () -> "unexpected message for " + unit + ": " + exception.getMessage()); + } +} diff --git a/java/vortex-jni/src/test/java/dev/vortex/api/ScanOptionsTest.java b/java/vortex-jni/src/test/java/dev/vortex/api/ScanOptionsTest.java new file mode 100644 index 00000000000..43be270520d --- /dev/null +++ b/java/vortex-jni/src/test/java/dev/vortex/api/ScanOptionsTest.java @@ -0,0 +1,206 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.api; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import dev.vortex.api.ScanOptions.SelectionMode; +import java.io.ByteArrayInputStream; +import java.io.DataInputStream; +import java.util.OptionalLong; +import org.junit.jupiter.api.Test; +import org.roaringbitmap.longlong.Roaring64NavigableMap; + +/** + * Unit tests for {@link ScanOptions#validateSelectionPayload()} and the factories that build a row selection. + * + *

Validation runs in an Immutables {@code @Value.Check}, so it fires at {@code build()} time on the driver rather + * than inside the native scan. That is the only place a mismatched mode and payload is reported with a message a caller + * can act on: {@code scan.rs} sees the mode as a bare byte alongside two possibly-empty arrays, so an absent payload + * reaches it as an empty selection instead of an error. + */ +public final class ScanOptionsTest { + + @Test + public void defaultsReadEveryRowAndColumn() { + ScanOptions options = ScanOptions.of(); + + assertTrue(options.projection().isEmpty()); + assertTrue(options.filter().isEmpty()); + assertEquals(OptionalLong.empty(), options.rowRangeBegin()); + assertEquals(OptionalLong.empty(), options.rowRangeEnd()); + assertEquals(OptionalLong.empty(), options.limit()); + assertTrue(options.selectionIndices().isEmpty()); + assertTrue(options.selectionRoaringBitmap().isEmpty()); + assertEquals(SelectionMode.INCLUDE_ALL, options.selectionMode()); + assertFalse(options.ordered()); + } + + @Test + public void indexSelectionModesRequireIndices() { + for (SelectionMode mode : new SelectionMode[] {SelectionMode.INCLUDE, SelectionMode.EXCLUDE}) { + IllegalArgumentException exception = assertThrows( + IllegalArgumentException.class, + () -> ScanOptions.builder().selectionMode(mode).build(), + () -> "no exception for " + mode); + assertEquals("selection indices are required for index selection modes", exception.getMessage()); + } + } + + @Test + public void indexSelectionModesRejectARoaringPayload() { + // The payload is present but of the wrong kind, so the indices check is what rejects it. + IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, () -> ScanOptions.builder() + .selectionMode(SelectionMode.INCLUDE) + .selectionRoaringBitmap(new byte[] {1}) + .build()); + assertEquals("selection indices are required for index selection modes", exception.getMessage()); + } + + @Test + public void roaringSelectionModesRequireABitmap() { + for (SelectionMode mode : new SelectionMode[] {SelectionMode.INCLUDE_ROARING, SelectionMode.EXCLUDE_ROARING}) { + IllegalArgumentException exception = assertThrows( + IllegalArgumentException.class, + () -> ScanOptions.builder().selectionMode(mode).build(), + () -> "no exception for " + mode); + assertEquals("selection roaring bitmap is required for roaring selection modes", exception.getMessage()); + } + } + + @Test + public void roaringSelectionModesRejectAnIndexPayload() { + IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, () -> ScanOptions.builder() + .selectionMode(SelectionMode.INCLUDE_ROARING) + .selectionIndices(new long[] {0}) + .build()); + assertEquals("selection roaring bitmap is required for roaring selection modes", exception.getMessage()); + } + + @Test + public void roaringSelectionModesRejectAnEmptyBitmap() { + // `deserialize_roaring_selection` in scan.rs also rejects an empty buffer; catching it here keeps the + // failure on the driver, where the caller still has a stack that names the option. + for (SelectionMode mode : new SelectionMode[] {SelectionMode.INCLUDE_ROARING, SelectionMode.EXCLUDE_ROARING}) { + IllegalArgumentException exception = assertThrows( + IllegalArgumentException.class, + () -> ScanOptions.builder() + .selectionMode(mode) + .selectionRoaringBitmap(new byte[0]) + .build(), + () -> "no exception for " + mode); + assertEquals("selection roaring bitmap must not be empty", exception.getMessage()); + } + } + + @Test + public void anEmptyIndexArrayIsAcceptedForIndexModes() { + // Only the roaring payload has a non-empty requirement; an empty index array selects no rows. + ScanOptions options = ScanOptions.includeRows(); + + assertEquals(SelectionMode.INCLUDE, options.selectionMode()); + assertArrayEquals(new long[0], options.selectionIndices().orElseThrow()); + } + + @Test + public void includeAllRejectsEitherPayload() { + IllegalArgumentException fromIndices = assertThrows( + IllegalArgumentException.class, + () -> ScanOptions.builder().selectionIndices(new long[] {0}).build()); + assertEquals("row selection payload requires a selection mode", fromIndices.getMessage()); + + IllegalArgumentException fromBitmap = assertThrows(IllegalArgumentException.class, () -> ScanOptions.builder() + .selectionRoaringBitmap(new byte[] {1}) + .build()); + assertEquals("row selection payload requires a selection mode", fromBitmap.getMessage()); + } + + @Test + public void bothPayloadsTogetherAreRejectedBeforeTheModeIsConsidered() { + // Checked ahead of the mode switch, so it reports the same message for a roaring mode. + IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, () -> ScanOptions.builder() + .selectionMode(SelectionMode.INCLUDE_ROARING) + .selectionIndices(new long[] {0}) + .selectionRoaringBitmap(new byte[] {1}) + .build()); + assertEquals("row selection must use either indices or roaring bitmap, not both", exception.getMessage()); + } + + @Test + public void indexFactoriesSetTheMatchingMode() { + assertEquals(SelectionMode.INCLUDE, ScanOptions.includeRows(0, 3, 9).selectionMode()); + assertEquals(SelectionMode.EXCLUDE, ScanOptions.excludeRows(0, 9).selectionMode()); + } + + @Test + public void indexFactoriesCopyTheirArgument() { + long[] rowIndices = {0, 3, 9}; + ScanOptions options = ScanOptions.includeRows(rowIndices); + rowIndices[0] = 7; + + assertArrayEquals(new long[] {0, 3, 9}, options.selectionIndices().orElseThrow()); + } + + @Test + public void roaringFactoriesSetTheMatchingModeAndSerializeThePayload() throws Exception { + Roaring64NavigableMap rows = new Roaring64NavigableMap(); + rows.addLong(0); + rows.addLong(3); + rows.addLong(9); + + for (SelectionMode mode : new SelectionMode[] {SelectionMode.INCLUDE_ROARING, SelectionMode.EXCLUDE_ROARING}) { + ScanOptions options = mode == SelectionMode.INCLUDE_ROARING + ? ScanOptions.includeRows(rows) + : ScanOptions.excludeRows(rows); + + assertEquals(mode, options.selectionMode()); + // Written with `serializePortable`, so `deserializePortable` on the same bytes must round trip. + Roaring64NavigableMap parsed = new Roaring64NavigableMap(); + parsed.deserializePortable(new DataInputStream( + new ByteArrayInputStream(options.selectionRoaringBitmap().orElseThrow()))); + assertEquals(3, parsed.getLongCardinality()); + assertTrue(parsed.contains(9L)); + assertFalse(parsed.contains(1L)); + } + } + + @Test + public void anEmptyRoaringBitmapStillSerializesToAHeader() { + // `serializePortable` writes a container count even for an empty map, so the non-empty check above + // guards a malformed buffer rather than an empty selection. + ScanOptions options = ScanOptions.includeRows(new Roaring64NavigableMap()); + + assertTrue(options.selectionRoaringBitmap().orElseThrow().length > 0); + } + + @Test + public void selectionModeCodesMatchTheNativeSwitch() { + // scan.rs matches these bytes directly and bails on anything else. + assertEquals(0, SelectionMode.INCLUDE_ALL.code()); + assertEquals(1, SelectionMode.INCLUDE.code()); + assertEquals(2, SelectionMode.EXCLUDE.code()); + assertEquals(3, SelectionMode.INCLUDE_ROARING.code()); + assertEquals(4, SelectionMode.EXCLUDE_ROARING.code()); + assertEquals(5, SelectionMode.values().length); + } + + @Test + public void rowRangeAndLimitAreCarriedThrough() { + ScanOptions options = ScanOptions.builder() + .rowRangeBegin(10) + .rowRangeEnd(20) + .limit(5) + .ordered(true) + .build(); + + assertEquals(OptionalLong.of(10), options.rowRangeBegin()); + assertEquals(OptionalLong.of(20), options.rowRangeEnd()); + assertEquals(OptionalLong.of(5), options.limit()); + assertTrue(options.ordered()); + } +} diff --git a/java/vortex-jni/src/test/java/dev/vortex/api/SpatialTypesTest.java b/java/vortex-jni/src/test/java/dev/vortex/api/SpatialTypesTest.java new file mode 100644 index 00000000000..6f6977ec70a --- /dev/null +++ b/java/vortex-jni/src/test/java/dev/vortex/api/SpatialTypesTest.java @@ -0,0 +1,164 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.api; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import dev.vortex.arrow.ArrowAllocation; +import dev.vortex.jni.NativeLoader; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import org.apache.arrow.c.ArrowArray; +import org.apache.arrow.c.ArrowSchema; +import org.apache.arrow.c.Data; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ipc.ArrowReader; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +/** + * Round-trips a Vortex spatial extension column ({@code vortex.st.wkb}) through the JNI boundary. Spatial columns cross + * the boundary as Arrow fields tagged with the GeoArrow extension name ({@code geoarrow.wkb}) and JSON metadata + * carrying the CRS. + */ +public final class SpatialTypesTest { + private static final String EXTENSION_NAME_KEY = "ARROW:extension:name"; + private static final String EXTENSION_METADATA_KEY = "ARROW:extension:metadata"; + private static final String GEOARROW_WKB = "geoarrow.wkb"; + private static final String CRS_METADATA = "{\"crs\":\"OGC:CRS84\"}"; + + @TempDir + static Path tempDir; + + static String writePath; + + private static final List WKB_POINTS = new ArrayList<>(); + + @BeforeAll + public static void loadLibrary() { + NativeLoader.loadJni(); + } + + @BeforeAll + static void setup() throws IOException { + writePath = tempDir.resolve("spatial.vortex").toAbsolutePath().toUri().toString(); + + WKB_POINTS.add(wkbPoint(1.0, 2.0)); + WKB_POINTS.add(wkbPoint(-111.7610, 34.8697)); + WKB_POINTS.add(null); + WKB_POINTS.add(wkbPoint(0.0, 0.0)); + + BufferAllocator allocator = ArrowAllocation.rootAllocator(); + Schema schema = new Schema(List.of(wkbField("geom"))); + Session session = Session.create(); + try (VortexWriter writer = VortexWriter.builder(session, writePath, schema, allocator) + .build(); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + VarBinaryVector geomVec = (VarBinaryVector) root.getVector("geom"); + geomVec.allocateNew(WKB_POINTS.size()); + for (int i = 0; i < WKB_POINTS.size(); i++) { + byte[] wkb = WKB_POINTS.get(i); + if (wkb == null) { + geomVec.setNull(i); + } else { + geomVec.setSafe(i, wkb); + } + } + root.setRowCount(WKB_POINTS.size()); + + try (ArrowArray arrowArray = ArrowArray.allocateNew(allocator); + ArrowSchema arrowSchemaFfi = ArrowSchema.allocateNew(allocator)) { + Data.exportVectorSchemaRoot(allocator, root, null, arrowArray, arrowSchemaFfi); + writer.writeBatch(arrowArray.memoryAddress(), arrowSchemaFfi.memoryAddress()); + } + } + } + + @Test + public void testSchemaCarriesGeoArrowExtension() { + BufferAllocator allocator = ArrowAllocation.rootAllocator(); + Session session = Session.create(); + DataSource ds = DataSource.open(session, writePath); + + Schema schema = ds.arrowSchema(allocator); + Field geom = schema.findField("geom"); + assertEquals(GEOARROW_WKB, geom.getMetadata().get(EXTENSION_NAME_KEY)); + assertEquals(CRS_METADATA, geom.getMetadata().get(EXTENSION_METADATA_KEY)); + assertTrue(geom.isNullable()); + } + + @Test + public void testScanSchemaCarriesGeoArrowExtension() { + BufferAllocator allocator = ArrowAllocation.rootAllocator(); + Session session = Session.create(); + DataSource ds = DataSource.open(session, writePath); + + Schema schema = ds.scan(ScanOptions.of()).arrowSchema(allocator); + Field geom = schema.findField("geom"); + assertEquals(GEOARROW_WKB, geom.getMetadata().get(EXTENSION_NAME_KEY)); + assertEquals(CRS_METADATA, geom.getMetadata().get(EXTENSION_METADATA_KEY)); + } + + @Test + public void testWkbValuesRoundTrip() throws Exception { + BufferAllocator allocator = ArrowAllocation.rootAllocator(); + Session session = Session.create(); + DataSource ds = DataSource.open(session, writePath); + + List values = new ArrayList<>(); + Scan scan = ds.scan(ScanOptions.of()); + while (scan.hasNext()) { + Partition partition = scan.next(); + try (ArrowReader arrowReader = partition.scanArrow(allocator)) { + while (arrowReader.loadNextBatch()) { + VectorSchemaRoot root = arrowReader.getVectorSchemaRoot(); + var geomVec = root.getVector("geom"); + for (int i = 0; i < root.getRowCount(); i++) { + values.add(geomVec.isNull(i) ? null : (byte[]) geomVec.getObject(i)); + } + } + } + } + + assertEquals(WKB_POINTS.size(), values.size()); + for (int i = 0; i < WKB_POINTS.size(); i++) { + byte[] expected = WKB_POINTS.get(i); + if (expected == null) { + assertNull(values.get(i)); + } else { + assertArrayEquals(expected, values.get(i)); + } + } + } + + private static Field wkbField(String name) { + Map metadata = Map.of(EXTENSION_NAME_KEY, GEOARROW_WKB, EXTENSION_METADATA_KEY, CRS_METADATA); + return new Field(name, new FieldType(true, ArrowType.Binary.INSTANCE, null, metadata), null); + } + + /** Little-endian WKB encoding of {@code POINT(x y)}. */ + private static byte[] wkbPoint(double x, double y) { + ByteBuffer buffer = ByteBuffer.allocate(21).order(ByteOrder.LITTLE_ENDIAN); + buffer.put((byte) 1); // little-endian marker + buffer.putInt(1); // geometry type: point + buffer.putDouble(x); + buffer.putDouble(y); + return buffer.array(); + } +} diff --git a/java/vortex-jni/src/test/java/dev/vortex/api/TestMinimal.java b/java/vortex-jni/src/test/java/dev/vortex/api/TestMinimal.java index 3a6f9ca0883..c63efe38320 100644 --- a/java/vortex-jni/src/test/java/dev/vortex/api/TestMinimal.java +++ b/java/vortex-jni/src/test/java/dev/vortex/api/TestMinimal.java @@ -6,6 +6,7 @@ import static java.nio.charset.StandardCharsets.UTF_8; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; import dev.vortex.arrow.ArrowAllocation; import dev.vortex.jni.NativeLoader; @@ -14,7 +15,6 @@ import java.nio.file.Files; import java.nio.file.Path; import java.util.ArrayList; -import java.util.HashMap; import java.util.List; import java.util.Objects; import org.apache.arrow.c.ArrowArray; @@ -86,7 +86,8 @@ static void setup() throws IOException { Field.notNullable("Salary", ArrowType.Decimal.createDecimal(9, 2, 128)), Field.nullable("State", ArrowType.Utf8View.INSTANCE))); Session session = Session.create(); - try (VortexWriter writer = VortexWriter.create(session, writePath, schema, new HashMap<>(), allocator); + try (VortexWriter writer = VortexWriter.builder(session, writePath, schema, allocator) + .build(); VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { ViewVarCharVector nameVec = (ViewVarCharVector) root.getVector("Name"); DecimalVector salaryVec = (DecimalVector) root.getVector("Salary"); @@ -226,10 +227,14 @@ public void testRoaringSelectionExcludesRows() throws Exception { } @Test - public void testSelectionIndicesMustBeSortedAndUnique() { - IllegalArgumentException exception = - assertThrows(IllegalArgumentException.class, () -> ScanOptions.includeRows(2, 1)); - assertEquals("selection indices must be sorted ascending and unique", exception.getMessage()); + public void testSelectionIndicesMustBeSortedAndUnique() throws Exception { + Session session = Session.create(); + DataSource ds = DataSource.open(session, writePath); + + RuntimeException exception = assertThrows(RuntimeException.class, () -> ds.scan(ScanOptions.includeRows(2, 1))); + assertTrue( + exception.getMessage().contains("must be strictly increasing"), + "unexpected message: " + exception.getMessage()); } @Test diff --git a/java/vortex-jni/src/test/java/dev/vortex/api/VortexWriteSummaryTest.java b/java/vortex-jni/src/test/java/dev/vortex/api/VortexWriteSummaryTest.java new file mode 100644 index 00000000000..ab1bb029484 --- /dev/null +++ b/java/vortex-jni/src/test/java/dev/vortex/api/VortexWriteSummaryTest.java @@ -0,0 +1,159 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.api; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.math.BigDecimal; +import java.math.BigInteger; +import java.util.List; +import java.util.OptionalLong; +import org.junit.jupiter.api.Test; + +/** + * Unit tests for the value semantics of {@link VortexWriteSummary} and {@link VortexColumnStatistics}. + * + *

Both constructors are public only so {@code writer.rs} can call them through JNI, which means the boundary is + * untyped: counts arrive as bare {@code jlong}s and bounds as {@code Object}. {@code exact_count_jlong} sends + * {@code -1} when Vortex has no exact statistic, and the accessors turn that into an empty {@link OptionalLong} — a + * caller reading the field directly would see a negative count instead. These tests pin that translation and the + * defensive copy that keeps the statistics list immutable. + */ +public final class VortexWriteSummaryTest { + + @Test + public void summaryReportsTheSizeAndRowCountItWasGiven() { + VortexWriteSummary summary = new VortexWriteSummary(4096L, 10L, new VortexColumnStatistics[0]); + + assertEquals(4096L, summary.fileSize()); + assertEquals(10L, summary.rowCount()); + assertTrue(summary.columnStatistics().isEmpty()); + } + + @Test + public void summaryCopiesTheStatisticsArray() { + VortexColumnStatistics first = statistics(0, 1L, 0L); + VortexColumnStatistics[] columns = {first}; + VortexWriteSummary summary = new VortexWriteSummary(1L, 1L, columns); + columns[0] = statistics(1, 2L, 0L); + + assertEquals(1, summary.columnStatistics().size()); + assertSame(first, summary.columnStatistics().get(0)); + } + + @Test + public void statisticsListRejectsModification() { + VortexWriteSummary summary = + new VortexWriteSummary(1L, 1L, new VortexColumnStatistics[] {statistics(0, 1L, 0L)}); + List columns = summary.columnStatistics(); + + assertThrows(UnsupportedOperationException.class, columns::clear); + assertThrows(UnsupportedOperationException.class, () -> columns.add(statistics(1, 1L, 0L))); + } + + @Test + public void statisticsListIsStableAcrossCalls() { + VortexWriteSummary summary = new VortexWriteSummary(1L, 1L, new VortexColumnStatistics[0]); + + assertSame(summary.columnStatistics(), summary.columnStatistics()); + } + + @Test + public void aNegativeCountMeansVortexComputedNone() { + // `exact_count_jlong` returns -1 when the statistic is absent or inexact. + VortexColumnStatistics absent = statistics(0, 5L, 5L, -1L, -1L, null, null); + + assertEquals(OptionalLong.empty(), absent.nullValueCount()); + assertEquals(OptionalLong.empty(), absent.nanValueCount()); + } + + @Test + public void zeroIsAnExactCountAndNotASentinel() { + VortexColumnStatistics none = statistics(0, 5L, 5L, 0L, 0L, null, null); + + assertEquals(OptionalLong.of(0L), none.nullValueCount()); + assertEquals(OptionalLong.of(0L), none.nanValueCount()); + } + + @Test + public void exactCountsAreReportedAsGiven() { + VortexColumnStatistics counted = statistics(2, 128L, 10L, 3L, 1L, null, null); + + assertEquals(2, counted.columnIndex()); + assertEquals(128L, counted.compressedSize()); + assertEquals(10L, counted.valueCount()); + assertEquals(OptionalLong.of(3L), counted.nullValueCount()); + assertEquals(OptionalLong.of(1L), counted.nanValueCount()); + } + + @Test + public void oneCountCanBeExactWhileTheOtherIsAbsent() { + // A non-floating-point column has a null count but no NaN count. + VortexColumnStatistics integral = statistics(0, 64L, 10L, 2L, -1L, null, null); + + assertEquals(OptionalLong.of(2L), integral.nullValueCount()); + assertEquals(OptionalLong.empty(), integral.nanValueCount()); + } + + @Test + public void absentBoundsAreEmpty() { + VortexColumnStatistics unbounded = statistics(0, 1L, 1L, 0L, -1L, null, null); + + assertTrue(unbounded.lowerBound().isEmpty()); + assertTrue(unbounded.upperBound().isEmpty()); + } + + @Test + public void boundsAreReturnedWithoutConversion() { + // `scalar_to_java` picks the Java type per Arrow scalar; the accessor must not narrow or box further. + for (Object[] pair : new Object[][] { + {Integer.valueOf(1), Integer.valueOf(9)}, + {Long.valueOf(1L), Long.valueOf(9L)}, + {BigInteger.ONE, BigInteger.TEN}, + {Float.valueOf(1.5f), Float.valueOf(9.5f)}, + {Double.valueOf(1.5d), Double.valueOf(9.5d)}, + {new BigDecimal("1.50"), new BigDecimal("9.50")}, + {"a", "z"} + }) { + VortexColumnStatistics bounded = statistics(0, 1L, 1L, 0L, -1L, pair[0], pair[1]); + + assertSame(pair[0], bounded.lowerBound().orElseThrow()); + assertSame(pair[1], bounded.upperBound().orElseThrow()); + } + } + + @Test + public void binaryBoundsAreSharedNotCopied() { + // Documented as `byte[]` and handed straight back, so a reader that mutates the array changes what + // every later reader sees. Callers must copy before writing. + VortexColumnStatistics bounded = statistics(0, 1L, 1L, 0L, -1L, new byte[] {1, 2}, new byte[] {3, 4}); + + byte[] lower = (byte[]) bounded.lowerBound().orElseThrow(); + assertArrayEquals(new byte[] {1, 2}, lower); + lower[0] = 9; + + assertSame(lower, bounded.lowerBound().orElseThrow()); + assertArrayEquals(new byte[] {9, 2}, (byte[]) bounded.lowerBound().orElseThrow()); + } + + private static VortexColumnStatistics statistics(int columnIndex, long compressedSize, long valueCount) { + return statistics(columnIndex, compressedSize, valueCount, 0L, -1L, null, null); + } + + private static VortexColumnStatistics statistics( + int columnIndex, + long compressedSize, + long valueCount, + long nullValueCount, + long nanValueCount, + Object lowerBound, + Object upperBound) { + return new VortexColumnStatistics( + columnIndex, compressedSize, valueCount, nullValueCount, nanValueCount, lowerBound, upperBound); + } +} diff --git a/java/vortex-jni/src/test/java/dev/vortex/io/NativeIOBridgeTest.java b/java/vortex-jni/src/test/java/dev/vortex/io/NativeIOBridgeTest.java index 5a6d40e05db..e19bac01e3d 100644 --- a/java/vortex-jni/src/test/java/dev/vortex/io/NativeIOBridgeTest.java +++ b/java/vortex-jni/src/test/java/dev/vortex/io/NativeIOBridgeTest.java @@ -4,6 +4,7 @@ package dev.vortex.io; import static java.nio.charset.StandardCharsets.UTF_8; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -16,6 +17,7 @@ import dev.vortex.api.VortexWriteSummary; import dev.vortex.api.VortexWriter; import dev.vortex.arrow.ArrowAllocation; +import dev.vortex.jni.NativeFiles; import dev.vortex.jni.NativeLoader; import java.io.EOFException; import java.io.IOException; @@ -26,6 +28,8 @@ import java.nio.file.Path; import java.nio.file.StandardOpenOption; import java.util.List; +import java.util.Map; +import java.util.Set; import org.apache.arrow.c.ArrowArray; import org.apache.arrow.c.ArrowSchema; import org.apache.arrow.c.Data; @@ -132,7 +136,8 @@ private void writePeopleFile(Session session, BufferAllocator allocator, Path pa Schema schema = personSchema(); VortexWriteSummary summary; try (StreamWritable writable = new StreamWritable(path)) { - try (VortexWriter writer = VortexWriter.create(session, writable, schema, allocator); + try (VortexWriter writer = VortexWriter.builder(session, writable, schema, allocator) + .build(); VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { VarCharVector nameVec = (VarCharVector) root.getVector("name"); IntVector ageVec = (IntVector) root.getVector("age"); @@ -198,6 +203,49 @@ public void testWritableThenReadableRoundTrip() throws IOException { } } + @Test + public void testMetadataRoundTripThroughBridge() throws IOException { + Path path = tempDir.resolve("bridge_metadata.vortex"); + BufferAllocator allocator = ArrowAllocation.rootAllocator(); + Session session = Session.create(); + Schema schema = personSchema(); + byte[] value = "{\"type\":\"struct\",\"fields\":[]}".getBytes(UTF_8); + + try (StreamWritable writable = new StreamWritable(path)) { + try (VortexWriter writer = VortexWriter.builder(session, writable, schema, allocator) + .putMetadata("iceberg.schema", value) + .build(); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + VarCharVector nameVec = (VarCharVector) root.getVector("name"); + IntVector ageVec = (IntVector) root.getVector("age"); + nameVec.allocateNew(1); + ageVec.allocateNew(1); + nameVec.setSafe(0, "Alice".getBytes(UTF_8)); + ageVec.setSafe(0, 30); + root.setRowCount(1); + try (ArrowArray arrowArray = ArrowArray.allocateNew(allocator); + ArrowSchema arrowSchemaFfi = ArrowSchema.allocateNew(allocator)) { + Data.exportVectorSchemaRoot(allocator, root, null, arrowArray, arrowSchemaFfi); + writer.writeBatch(arrowArray.memoryAddress(), arrowSchemaFfi.memoryAddress()); + } + } + } + + // Both directions go through caller-provided I/O: no native storage client is involved. + try (FileChannelReadable readable = new FileChannelReadable(path)) { + Map metadata = NativeFiles.readMetadata(session, readable); + assertEquals(Set.of("iceberg.schema"), metadata.keySet()); + assertArrayEquals(value, metadata.get("iceberg.schema")); + } + + // The readable's name keys the session footer cache, so this read resolves the metadata + // against the footer cached by the read above rather than reading it again. + try (FileChannelReadable readable = new FileChannelReadable(path)) { + Map metadata = NativeFiles.readMetadata(session, readable); + assertArrayEquals(value, metadata.get("iceberg.schema")); + } + } + @Test public void testMultipleReadables() throws IOException { Path first = tempDir.resolve("bridge_a.vortex"); @@ -256,7 +304,8 @@ public void testLargeRoundTrip() throws IOException { int batches = 20; int rowsPerBatch = 5_000; try (StreamWritable writable = new StreamWritable(path)) { - try (VortexWriter writer = VortexWriter.create(session, writable, schema, allocator); + try (VortexWriter writer = VortexWriter.builder(session, writable, schema, allocator) + .build(); VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { VarCharVector nameVec = (VarCharVector) root.getVector("name"); IntVector ageVec = (IntVector) root.getVector("age"); diff --git a/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java b/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java index 1bb51f8799f..73435033969 100644 --- a/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java +++ b/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java @@ -6,7 +6,9 @@ import static java.nio.charset.StandardCharsets.UTF_8; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import dev.vortex.api.DataSource; @@ -27,18 +29,23 @@ import org.apache.arrow.c.ArrowSchema; import org.apache.arrow.c.Data; import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ViewVarBinaryVector; import org.apache.arrow.vector.ViewVarCharVector; +import org.apache.arrow.vector.complex.MapVector; import org.apache.arrow.vector.complex.StructVector; +import org.apache.arrow.vector.complex.impl.UnionMapWriter; import org.apache.arrow.vector.ipc.ArrowReader; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.types.pojo.Schema; import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; @@ -77,6 +84,40 @@ private static Schema parquetVariantSchema() { return new Schema(List.of(variant)); } + private static Schema mapSchema() { + Field key = Field.notNullable(MapVector.KEY_NAME, new ArrowType.Int(32, true)); + Field value = Field.nullable(MapVector.VALUE_NAME, new ArrowType.Int(64, true)); + Field entries = new Field("entries", FieldType.notNullable(ArrowType.Struct.INSTANCE), List.of(key, value)); + Field map = new Field("map", FieldType.nullable(new ArrowType.Map(false)), List.of(entries)); + return new Schema(List.of(map)); + } + + private static void populateMapRoot(VectorSchemaRoot root) { + MapVector map = (MapVector) root.getVector("map"); + UnionMapWriter writer = map.getWriter(); + writer.allocate(); + + writer.setPosition(0); + writer.startMap(); + writer.startEntry(); + writer.key().integer().writeInt(1); + writer.value().bigInt().writeBigInt(10L); + writer.endEntry(); + writer.startEntry(); + writer.key().integer().writeInt(2); + writer.value().bigInt().writeBigInt(20L); + writer.endEntry(); + writer.endMap(); + + // Position 1 remains null. + writer.setPosition(2); + writer.startMap(); + writer.endMap(); + + map.setValueCount(3); + root.setRowCount(3); + } + private static void populateParquetVariantRoot(VectorSchemaRoot root) { StructVector variant = (StructVector) root.getVector("variant"); VarBinaryVector metadata = variant.getChild("metadata", VarBinaryVector.class); @@ -111,7 +152,9 @@ public void testCreateWriter() throws IOException { Map options = new HashMap<>(); Session session = Session.create(); - try (VortexWriter writer = VortexWriter.create(session, writePath, personSchema(), options, allocator)) { + try (VortexWriter writer = VortexWriter.builder(session, writePath, personSchema(), allocator) + .options(options) + .build()) { assertNotNull(writer); } @@ -127,7 +170,9 @@ public void testCreateWriterPlainLocalPath() throws IOException { Map options = new HashMap<>(); Session session = Session.create(); - try (VortexWriter writer = VortexWriter.create(session, writePath, personSchema(), options, allocator)) { + try (VortexWriter writer = VortexWriter.builder(session, writePath, personSchema(), allocator) + .options(options) + .build()) { assertNotNull(writer); } @@ -143,7 +188,9 @@ public void testCreateWriterCreatesParentDirectories() throws IOException { Map options = new HashMap<>(); Session session = Session.create(); - try (VortexWriter writer = VortexWriter.create(session, writePath, personSchema(), options, allocator)) { + try (VortexWriter writer = VortexWriter.builder(session, writePath, personSchema(), allocator) + .options(options) + .build()) { assertNotNull(writer); } @@ -161,7 +208,8 @@ public void testWriteBatch() throws IOException { Session session = Session.create(); VortexWriteSummary summary; long bytesWhileOpen; - try (VortexWriter writer = VortexWriter.create(session, writePath, schema, new HashMap<>(), allocator); + try (VortexWriter writer = VortexWriter.builder(session, writePath, schema, allocator) + .build(); VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { VarCharVector nameVec = (VarCharVector) root.getVector("name"); IntVector ageVec = (IntVector) root.getVector("age"); @@ -225,6 +273,148 @@ public void testWriteBatch() throws IOException { } @Test + public void testMapRoundTrip() throws IOException { + Path outputPath = tempDir.resolve("test_map.vortex"); + String writePath = outputPath.toAbsolutePath().toUri().toString(); + + BufferAllocator allocator = ArrowAllocation.rootAllocator(); + Schema schema = mapSchema(); + Session session = Session.create(); + + try (VortexWriter writer = VortexWriter.builder(session, writePath, schema, allocator) + .build(); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + populateMapRoot(root); + try (ArrowArray arrowArray = ArrowArray.allocateNew(allocator); + ArrowSchema arrowSchema = ArrowSchema.allocateNew(allocator)) { + Data.exportVectorSchemaRoot(allocator, root, null, arrowArray, arrowSchema); + writer.writeBatch(arrowArray.memoryAddress(), arrowSchema.memoryAddress()); + } + } + + DataSource dataSource = DataSource.open(session, writePath); + assertEquals( + new ArrowType.Map(false), + dataSource.arrowSchema(allocator).findField("map").getType()); + + boolean sawBatch = false; + Scan scan = dataSource.scan(ScanOptions.of()); + while (scan.hasNext()) { + Partition partition = scan.next(); + try (ArrowReader reader = partition.scanArrow(allocator)) { + while (reader.loadNextBatch()) { + sawBatch = true; + MapVector map = (MapVector) reader.getVectorSchemaRoot().getVector("map"); + StructVector entries = (StructVector) map.getDataVector(); + IntVector keys = entries.getChild(MapVector.KEY_NAME, IntVector.class); + BigIntVector values = entries.getChild(MapVector.VALUE_NAME, BigIntVector.class); + + int firstOffset = map.getOffsetBuffer().getInt(0); + assertEquals(2, map.getInnerValueCountAt(0)); + assertEquals(1, keys.get(firstOffset)); + assertEquals(10L, values.get(firstOffset)); + assertEquals(2, keys.get(firstOffset + 1)); + assertEquals(20L, values.get(firstOffset + 1)); + assertTrue(map.isNull(1)); + assertFalse(map.isNull(2)); + assertEquals(0, map.getInnerValueCountAt(2)); + } + } + } + assertTrue(sawBatch); + } + + /** Write a single three-row batch of {@link #personSchema()} data. */ + private static void writePeopleBatch(VortexWriter writer, BufferAllocator allocator) throws IOException { + try (VectorSchemaRoot root = VectorSchemaRoot.create(personSchema(), allocator)) { + VarCharVector nameVec = (VarCharVector) root.getVector("name"); + IntVector ageVec = (IntVector) root.getVector("age"); + nameVec.allocateNew(3); + ageVec.allocateNew(3); + nameVec.setSafe(0, "Alice".getBytes(UTF_8)); + nameVec.setSafe(1, "Bob".getBytes(UTF_8)); + nameVec.setSafe(2, "Carol".getBytes(UTF_8)); + ageVec.setSafe(0, 30); + ageVec.setSafe(1, 25); + ageVec.setSafe(2, 40); + root.setRowCount(3); + + try (ArrowArray arrowArray = ArrowArray.allocateNew(allocator); + ArrowSchema arrowSchemaFfi = ArrowSchema.allocateNew(allocator)) { + Data.exportVectorSchemaRoot(allocator, root, null, arrowArray, arrowSchemaFfi); + writer.writeBatch(arrowArray.memoryAddress(), arrowSchemaFfi.memoryAddress()); + } + } + } + + @Test + public void testFileMetadataRoundTrip() throws IOException { + Path outputPath = tempDir.resolve("test_metadata.vortex"); + String writePath = outputPath.toAbsolutePath().toUri().toString(); + + BufferAllocator allocator = ArrowAllocation.rootAllocator(); + Session session = Session.create(); + byte[] schemaJson = "{\"type\":\"struct\",\"fields\":[]}".getBytes(UTF_8); + byte[] deleteType = "position".getBytes(UTF_8); + Map metadata = Map.of("iceberg.schema", schemaJson, "delete-type", deleteType); + + try (VortexWriter writer = VortexWriter.builder(session, writePath, personSchema(), allocator) + .metadata(metadata) + .build()) { + writePeopleBatch(writer, allocator); + } + + Map read = NativeFiles.readMetadata(session, writePath, new HashMap<>()); + assertEquals(metadata.keySet(), read.keySet()); + assertArrayEquals(schemaJson, read.get("iceberg.schema")); + assertArrayEquals(deleteType, read.get("delete-type")); + + // Metadata is stored out of band: the file still scans as usual. + DataSource ds = DataSource.open(session, writePath); + assertEquals(new DataSource.RowCount.Exact(3L), ds.rowCount()); + } + + @Test + public void testFileWithoutMetadataReadsEmpty() throws IOException { + Path outputPath = tempDir.resolve("test_no_metadata.vortex"); + String writePath = outputPath.toAbsolutePath().toUri().toString(); + + BufferAllocator allocator = ArrowAllocation.rootAllocator(); + Session session = Session.create(); + try (VortexWriter writer = VortexWriter.builder(session, writePath, personSchema(), allocator) + .build()) { + writePeopleBatch(writer, allocator); + } + + assertTrue(NativeFiles.readMetadata(session, writePath, new HashMap<>()).isEmpty()); + } + + @Test + public void testInvalidMetadataKeyRejectedAtCreate() { + Path outputPath = tempDir.resolve("test_bad_metadata.vortex"); + String writePath = outputPath.toAbsolutePath().toUri().toString(); + + BufferAllocator allocator = ArrowAllocation.rootAllocator(); + Session session = Session.create(); + // Keys are capped well below this; the writer must reject the set before any bytes + // are produced, rather than failing the first writeBatch with a send error. + Map metadata = Map.of("k".repeat(1024), new byte[] {1}); + + IOException thrown = assertThrows( + IOException.class, () -> VortexWriter.builder(session, writePath, personSchema(), allocator) + .metadata(metadata) + .build()); + Throwable cause = thrown.getCause(); + assertNotNull(cause, "native failure should be retained as the cause"); + assertTrue( + cause.getMessage().contains("metadata key"), + "error should identify the offending key, got: " + cause.getMessage()); + // Rejected before the sink is touched, so no partial file is left behind. + assertFalse(Files.exists(outputPath)); + } + + @Test + @Disabled public void testParquetVariantRoundTrip() throws IOException { Path outputPath = tempDir.resolve("test_parquet_variant.vortex"); String writePath = outputPath.toAbsolutePath().toUri().toString(); @@ -233,7 +423,8 @@ public void testParquetVariantRoundTrip() throws IOException { Schema schema = parquetVariantSchema(); Session session = Session.create(); - try (VortexWriter writer = VortexWriter.create(session, writePath, schema, new HashMap<>(), allocator); + try (VortexWriter writer = VortexWriter.builder(session, writePath, schema, allocator) + .build(); VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { populateParquetVariantRoot(root); @@ -261,8 +452,9 @@ public void testParquetVariantRoundTrip() throws IOException { assertTrue(reader.loadNextBatch()); VectorSchemaRoot resultRoot = reader.getVectorSchemaRoot(); StructVector variant = (StructVector) resultRoot.getVector("variant"); - VarBinaryVector metadata = variant.getChild("metadata", VarBinaryVector.class); - VarBinaryVector value = variant.getChild("value", VarBinaryVector.class); + // Binary columns cross the boundary as their native view types. + ViewVarBinaryVector metadata = variant.getChild("metadata", ViewVarBinaryVector.class); + ViewVarBinaryVector value = variant.getChild("value", ViewVarBinaryVector.class); assertArrayEquals(VARIANT_METADATA, metadata.get(0)); assertArrayEquals(VARIANT_INT8_42, value.get(0)); diff --git a/java/vortex-spark/README.md b/java/vortex-spark/README.md index c6e8c9345bd..b7ab97daa83 100644 --- a/java/vortex-spark/README.md +++ b/java/vortex-spark/README.md @@ -11,7 +11,7 @@ Two flavors are published to Maven Central: | `dev.vortex:vortex-spark_2.13` | Spark 4.x | 2.13 | | `dev.vortex:vortex-spark_2.12` | Spark 3.5.x | 2.12 | -Use the `all` classifier JAR (e.g. `vortex-spark_2.13-0.78.0-all.jar`). It is self-contained: +Use the `all` classifier JAR (e.g. `vortex-spark_2.13-0.85.0-all.jar`). It is self-contained: it bundles the Vortex JNI bindings, native libraries for Linux (x86_64 and aarch64) and macOS (aarch64), and relocates its Arrow, Guava, and Jackson dependencies to avoid classpath conflicts with Spark. The thin (unclassified) JAR does not work on its own because it @@ -23,7 +23,7 @@ Pass the `all` JAR to `spark-shell`, `spark-submit`, or `pyspark` with `--jars`. either a local path or a URL, so you can point directly at Maven Central: ```shell -spark-shell --jars https://repo1.maven.org/maven2/dev/vortex/vortex-spark_2.13/0.78.0/vortex-spark_2.13-0.78.0-all.jar +spark-shell --jars https://repo1.maven.org/maven2/dev/vortex/vortex-spark_2.13/0.85.0/vortex-spark_2.13-0.85.0-all.jar ``` Or configure it on the session builder, e.g. in PySpark: @@ -31,12 +31,12 @@ Or configure it on the session builder, e.g. in PySpark: ```python spark = ( SparkSession.builder - .config("spark.jars", "/path/to/vortex-spark_2.13-0.78.0-all.jar") + .config("spark.jars", "/path/to/vortex-spark_2.13-0.85.0-all.jar") .getOrCreate() ) ``` -Note that `--packages dev.vortex:vortex-spark_2.13:0.78.0` does not work: `--packages` cannot +Note that `--packages dev.vortex:vortex-spark_2.13:0.85.0` does not work: `--packages` cannot select the `all` classifier and resolves the thin JAR, which fails at runtime with `NoClassDefFoundError: dev/vortex/relocated/...`. @@ -44,7 +44,7 @@ To depend on the connector from a JVM project instead, add the `all` classifier dependency: ```kotlin -implementation("dev.vortex:vortex-spark_2.13:0.78.0:all") +implementation("dev.vortex:vortex-spark_2.13:0.85.0:all") ``` ## Usage diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/VortexCatalog.java b/java/vortex-spark/src/main/java/dev/vortex/spark/VortexCatalog.java index 3d496a4cee0..2bd4fd3273d 100644 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/VortexCatalog.java +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/VortexCatalog.java @@ -12,6 +12,8 @@ import org.apache.spark.sql.connector.expressions.Transform; import org.apache.spark.sql.types.StructType; import org.apache.spark.sql.util.CaseInsensitiveStringMap; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * A path-based Spark catalog for querying Vortex files directly from SQL. @@ -30,6 +32,8 @@ * vortex.`/path/to/data`}), and pushdown all behave identically. The catalog holds no state and supports no DDL. */ public final class VortexCatalog implements TableCatalog { + private static final Logger log = LoggerFactory.getLogger(VortexCatalog.class); + private static final String PATH_KEY = "path"; private String name = "vortex"; @@ -85,8 +89,14 @@ public Table loadTable(Identifier ident) throws NoSuchTableException { schema = provider.inferSchema(options); partitioning = provider.inferPartitioning(options); } catch (RuntimeException e) { - // Missing or unreadable paths surface as "table not found" to SQL users. - throw new NoSuchTableException(ident); + // Missing or unreadable paths surface as "table not found" to SQL users. NoSuchTableException + // has no cause-accepting constructor common to the supported Spark versions, so carry the + // original failure as a suppressed exception and log it: it names the offending path and + // distinguishes an empty directory from a credentials or unsupported-type error. + log.warn("Cannot load {} as a Vortex table, reporting it as not found", path, e); + NoSuchTableException notFound = new NoSuchTableException(ident); + notFound.addSuppressed(e); + throw notFound; } return provider.getTable(schema, partitioning, Map.of(PATH_KEY, path)); } diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/VortexDataSourceV2.java b/java/vortex-spark/src/main/java/dev/vortex/spark/VortexDataSourceV2.java index a5ce4e3b1ce..bd434bcc8d2 100644 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/VortexDataSourceV2.java +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/VortexDataSourceV2.java @@ -62,7 +62,7 @@ public VortexDataSourceV2() { * * @param options the data source options containing file paths * @return the inferred Spark SQL schema - * @throws RuntimeException if required path options are missing + * @throws IllegalArgumentException if no Vortex files can be found under the supplied paths * @throws RuntimeException if there's an error reading the file or converting the schema */ @Override @@ -87,7 +87,11 @@ public StructType inferSchema(CaseInsensitiveStringMap options) { .findFirst(); if (firstFile.isEmpty()) { - throw new RuntimeException(String.format("UNABLE_TO_INFER_SCHEMA format: %s", shortName())); + throw new IllegalArgumentException(String.format( + "Unable to infer schema for %s: no .vortex files found under path %s. " + + "Check that the path is correct and contains at least one Vortex file, " + + "or provide an explicit schema.", + shortName(), pathToInfer)); } else { pathToInfer = firstFile.get(); } diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/VortexSessionCatalog.java b/java/vortex-spark/src/main/java/dev/vortex/spark/VortexSessionCatalog.java index 1d546880cc1..d1162bbf391 100644 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/VortexSessionCatalog.java +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/VortexSessionCatalog.java @@ -47,8 +47,10 @@ public Table loadTable(Identifier ident) throws NoSuchTableException { * Creates the table in the delegate session catalog, then returns it resolved through the Vortex connector when its * provider is {@code vortex}. * - *

The conversion matters for {@code CREATE TABLE ... AS SELECT}: Spark writes the query result into the table - * returned here, which must therefore support V2 writes. + *

Spark does not route table creation through this overload — {@link DelegatingCatalogExtension} sends the + * {@code Column[]} overload straight to the delegate, and that is the one both supported Spark versions call. It is + * kept because the delegate may return {@code null} on the normal path, which {@link #asVortexTableIfVortex} now + * tolerates; {@code CREATE TABLE ... AS SELECT} gets its Vortex table from {@link #loadTable} instead. */ @SuppressWarnings("deprecation") @Override @@ -60,11 +62,15 @@ public Table createTable( /** * Rebuilds a session-catalog table as a Vortex DataSource V2 table when its provider is {@code vortex} and it has a - * location; returns every other table unchanged. The schema and partitioning stored in the catalog are used as-is, - * no file needs to be opened. + * location; returns every other table unchanged, and {@code null} for a {@code null} input. The schema and + * partitioning stored in the catalog are used as-is, no file needs to be opened. */ @SuppressWarnings("deprecation") private static Table asVortexTableIfVortex(Table table) { + if (table == null) { + // V2SessionCatalog.createTable returns null on its normal path. + return null; + } Map properties = table.properties(); VortexDataSourceV2 provider = new VortexDataSourceV2(); if (!provider.shortName().equalsIgnoreCase(properties.get(TableCatalog.PROP_PROVIDER))) { diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/config/HadoopUtils.java b/java/vortex-spark/src/main/java/dev/vortex/spark/config/HadoopUtils.java index d091269df0a..2f3c4d78d90 100644 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/config/HadoopUtils.java +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/config/HadoopUtils.java @@ -35,6 +35,11 @@ public static Map s3PropertiesFromHadoopConf(Configuration hadoo qualified = "https://" + qualified; } properties.setEndpoint(qualified); + // object_store rejects plain-HTTP endpoints (LocalStack, MinIO, S3Mock) + // unless explicitly allowed. + if (qualified.startsWith("http://")) { + properties.setAllowHttp(true); + } break; case FS_S3A_ENDPOINT_REGION: properties.setRegion(entry.getValue()); diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/config/VortexS3Properties.java b/java/vortex-spark/src/main/java/dev/vortex/spark/config/VortexS3Properties.java index f8298891f18..2bd2221e413 100644 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/config/VortexS3Properties.java +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/config/VortexS3Properties.java @@ -14,6 +14,7 @@ public final class VortexS3Properties { private static final String SESSION_TOKEN = "aws_session_token"; private static final String REGION = "aws_region"; private static final String ENDPOINT = "aws_endpoint"; + private static final String ALLOW_HTTP = "aws_allow_http"; private static final String SKIP_SIGNATURE = "aws_skip_signature"; private final Map properties = Maps.newHashMap(); @@ -62,6 +63,10 @@ public void setEndpoint(String endpoint) { properties.put(ENDPOINT, endpoint); } + public void setAllowHttp(boolean allowHttp) { + properties.put(ALLOW_HTTP, Boolean.toString(allowHttp)); + } + public void setSkipSignature(boolean skipSignature) { properties.put(SKIP_SIGNATURE, Boolean.toString(skipSignature)); } diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexPartitionReader.java b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexPartitionReader.java index f9cd2363d59..46616631d00 100644 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexPartitionReader.java +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexPartitionReader.java @@ -95,7 +95,7 @@ public boolean next() { return true; } } catch (IOException e) { - throw new RuntimeException(e); + throw failure("load a batch from", e); } closeCurrentReader(); } @@ -119,7 +119,7 @@ public ColumnarBatch get() { try { root = currentReader.getVectorSchemaRoot(); } catch (IOException e) { - throw new RuntimeException(e); + throw failure("read the loaded batch of", e); } int rowCount = root.getRowCount(); @@ -157,12 +157,17 @@ public void close() { session = null; } + /** Wraps a failure with the paths being read, the one fact a stack trace alone does not give. */ + private RuntimeException failure(String what, IOException cause) { + return new RuntimeException(String.format("Failed to %s %s", what, spark.paths()), cause); + } + private void closeCurrentReader() { if (currentReader != null) { try { currentReader.close(); } catch (IOException e) { - throw new RuntimeException(e); + throw failure("close the reader over", e); } currentReader = null; } diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScanBuilder.java b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScanBuilder.java index 107d22b30f6..62c8085aa0f 100644 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScanBuilder.java +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScanBuilder.java @@ -108,16 +108,17 @@ public VortexScanBuilder addAllColumns(Iterable columns) { /** * Builds a VortexScan with the configured paths and columns. * + *

An empty column list is allowed: aggregates such as {@code count()} need no column data, and the scan then + * reads the minimal schema. + * * @return a new VortexScan instance - * @throws IllegalStateException if no paths or columns have been added + * @throws IllegalStateException if no paths have been added */ @Override public Scan build() { var paths = this.paths.build(); checkState(!paths.isEmpty(), "paths cannot be empty"); - // Allow empty columns for operations like count() that don't need actual column data - // If no columns are specified, we'll read the minimal schema needed return new VortexScan( paths, diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/write/SparkToArrowSchema.java b/java/vortex-spark/src/main/java/dev/vortex/spark/write/SparkToArrowSchema.java index 3ed5a56e52f..879db9fdd31 100644 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/write/SparkToArrowSchema.java +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/write/SparkToArrowSchema.java @@ -3,6 +3,7 @@ package dev.vortex.spark.write; +import dev.vortex.relocated.org.apache.arrow.vector.complex.MapVector; import dev.vortex.relocated.org.apache.arrow.vector.types.DateUnit; import dev.vortex.relocated.org.apache.arrow.vector.types.FloatingPointPrecision; import dev.vortex.relocated.org.apache.arrow.vector.types.TimeUnit; @@ -61,28 +62,29 @@ public static Schema convert(StructType sparkSchema) { * @return the corresponding Arrow field */ private static Field convertField(StructField sparkField) { - ArrowType arrowType = convertType(sparkField.dataType()); - FieldType fieldType = new FieldType(sparkField.nullable(), arrowType, null); + return convertField(sparkField.name(), sparkField.dataType(), sparkField.nullable()); + } - if (sparkField.dataType() instanceof StructType) { - // Handle nested struct - StructType structType = (StructType) sparkField.dataType(); + private static Field convertField(String name, DataType sparkType, boolean nullable) { + if (sparkType instanceof StructType structType) { List children = new ArrayList<>(); for (StructField childField : structType.fields()) { children.add(convertField(childField)); } - return new Field(sparkField.name(), fieldType, children); - } else if (sparkField.dataType() instanceof ArrayType) { - // Handle array type - ArrayType arrayType = (ArrayType) sparkField.dataType(); - Field elementField = new Field( - "element", - new FieldType(arrayType.containsNull(), convertType(arrayType.elementType()), null), - null); - return new Field(sparkField.name(), fieldType, List.of(elementField)); + return new Field(name, new FieldType(nullable, new ArrowType.Struct(), null), children); + } else if (sparkType instanceof ArrayType arrayType) { + Field elementField = convertField("element", arrayType.elementType(), arrayType.containsNull()); + return new Field(name, new FieldType(nullable, new ArrowType.List(), null), List.of(elementField)); + } else if (sparkType instanceof MapType mapType) { + Field keyField = convertField(MapVector.KEY_NAME, mapType.keyType(), false); + Field valueField = convertField(MapVector.VALUE_NAME, mapType.valueType(), mapType.valueContainsNull()); + Field entries = new Field( + MapVector.DATA_VECTOR_NAME, + FieldType.notNullable(new ArrowType.Struct()), + List.of(keyField, valueField)); + return new Field(name, new FieldType(nullable, new ArrowType.Map(false), null), List.of(entries)); } else { - // Primitive type - return new Field(sparkField.name(), fieldType, null); + return new Field(name, new FieldType(nullable, convertType(sparkType), null), null); } } @@ -125,9 +127,6 @@ private static ArrowType convertType(DataType sparkType) { return new ArrowType.List(); } else if (sparkType instanceof StructType) { return new ArrowType.Struct(); - } else if (sparkType instanceof MapType) { - // Map is represented as List> in Arrow - return new ArrowType.List(); } else { throw new UnsupportedOperationException("Unsupported Spark type for Arrow conversion: " + sparkType.getClass().getName()); diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexBatchWrite.java b/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexBatchWrite.java index 03a70c2f676..c9be4159a59 100644 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexBatchWrite.java +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexBatchWrite.java @@ -5,11 +5,7 @@ import dev.vortex.jni.NativeFiles; import dev.vortex.spark.VortexSparkSession; -import java.io.IOException; import java.io.Serializable; -import java.nio.file.Files; -import java.nio.file.Path; -import java.nio.file.Paths; import java.util.Arrays; import java.util.List; import java.util.Map; @@ -137,21 +133,22 @@ public void commit(WriterCommitMessage[] messages) { /** * Aborts the write job due to failures. * - *

This method cleans up any partially written files. + *

Deletes the files the tasks reported, through the same native filesystem layer the writers used, so that URL + * paths ({@code file://}, {@code s3://}) resolve the same way on cleanup as they did on write. * * @param messages commit messages from write tasks (may include failures) */ @Override public void abort(WriterCommitMessage[] messages) { - for (String filePath : extractFilePaths(messages)) { - try { - Path path = Paths.get(filePath); - if (Files.exists(path)) { - Files.delete(path); - } - } catch (IOException e) { - log.error("Failed to clean up file: {}", filePath, e); - } + List filePaths = extractFilePaths(messages); + if (filePaths.isEmpty()) { + return; + } + log.warn("Deleting {} file(s) written before the job failed, under {}", filePaths.size(), outputPath); + try { + NativeFiles.delete(VortexSparkSession.get(options), filePaths.toArray(new String[0]), options); + } catch (RuntimeException e) { + log.error("Failed to clean up {} file(s) under {}", filePaths.size(), outputPath, e); } } diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexDataWriter.java b/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexDataWriter.java index 3a4390cb8b6..48f47284550 100644 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexDataWriter.java +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexDataWriter.java @@ -26,6 +26,7 @@ import dev.vortex.relocated.org.apache.arrow.vector.VarCharVector; import dev.vortex.relocated.org.apache.arrow.vector.VectorSchemaRoot; import dev.vortex.relocated.org.apache.arrow.vector.complex.ListVector; +import dev.vortex.relocated.org.apache.arrow.vector.complex.MapVector; import dev.vortex.relocated.org.apache.arrow.vector.complex.StructVector; import dev.vortex.spark.VortexSparkSession; import java.io.IOException; @@ -36,6 +37,7 @@ import org.apache.spark.sql.catalyst.InternalRow; import org.apache.spark.sql.catalyst.expressions.SpecializedGetters; import org.apache.spark.sql.catalyst.util.ArrayData; +import org.apache.spark.sql.catalyst.util.MapData; import org.apache.spark.sql.connector.write.DataWriter; import org.apache.spark.sql.connector.write.WriterCommitMessage; import org.apache.spark.sql.types.ArrayType; @@ -49,6 +51,7 @@ import org.apache.spark.sql.types.FloatType; import org.apache.spark.sql.types.IntegerType; import org.apache.spark.sql.types.LongType; +import org.apache.spark.sql.types.MapType; import org.apache.spark.sql.types.ShortType; import org.apache.spark.sql.types.StringType; import org.apache.spark.sql.types.StructField; @@ -123,8 +126,9 @@ public final class VortexDataWriter implements DataWriter, AutoClos var arrowSchema = SparkToArrowSchema.convert(schema); this.session = VortexSparkSession.get(options.asCaseSensitiveMap()); - this.vortexWriter = - VortexWriter.create(session, filePath, arrowSchema, options.asCaseSensitiveMap(), allocator); + this.vortexWriter = VortexWriter.builder(session, filePath, arrowSchema, allocator) + .options(options.asCaseSensitiveMap()) + .build(); this.vectorSchemaRoot = VectorSchemaRoot.create(arrowSchema, allocator); logger.debug("Initialized VortexDataWriter for {}", filePath); @@ -188,9 +192,6 @@ private void writeBatch() throws IOException { vectorSchemaRoot.setRowCount(batchRows.size()); // Export via Arrow C Data Interface and write to Vortex - for (FieldVector vector : vectorSchemaRoot.getFieldVectors()) { - bytesWritten += vector.getBufferSize(); - } try (ArrowArray arrowArray = ArrowArray.allocateNew(allocator); ArrowSchema arrowSchema = ArrowSchema.allocateNew(allocator)) { Data.exportVectorSchemaRoot(allocator, vectorSchemaRoot, null, arrowArray, arrowSchema); @@ -249,10 +250,41 @@ private void populateVector( ListVector listVector = ((ListVector) vector); int writtenElements = listVector.getElementEndIndex(listVector.getLastSet()); listVector.startNewValue(rowIndex); + FieldVector elementVector = listVector.getDataVector(); for (int i = 0; i < data.numElements(); i++) { - populateVector(listVector.getDataVector(), arrayType.elementType(), data, i, writtenElements + i); + int elementIndex = writtenElements + i; + if (data.isNullAt(i)) { + // Reading a null slot of a fixed-width element returns the zeroed slot, and the + // typed setters mark it valid, so a null element must be written explicitly. + elementVector.setNull(elementIndex); + } else { + populateVector(elementVector, arrayType.elementType(), data, i, elementIndex); + } } listVector.endValue(rowIndex, data.numElements()); + } else if (dataType instanceof MapType mapType) { + MapData data = row.getMap(fieldIndex); + MapVector mapVector = (MapVector) vector; + int writtenEntries = mapVector.getElementEndIndex(mapVector.getLastSet()); + mapVector.startNewValue(rowIndex); + + StructVector entries = (StructVector) mapVector.getDataVector(); + FieldVector keyVector = entries.getChild(MapVector.KEY_NAME); + FieldVector valueVector = entries.getChild(MapVector.VALUE_NAME); + ArrayData keys = data.keyArray(); + ArrayData values = data.valueArray(); + + for (int i = 0; i < data.numElements(); i++) { + int entryIndex = writtenEntries + i; + entries.setIndexDefined(entryIndex); + populateVector(keyVector, mapType.keyType(), keys, i, entryIndex); + if (values.isNullAt(i)) { + valueVector.setNull(entryIndex); + } else { + populateVector(valueVector, mapType.valueType(), values, i, entryIndex); + } + } + mapVector.endValue(rowIndex, data.numElements()); } else { // For unsupported types, set null throw new IllegalArgumentException("Unsupported data type: " + dataType); @@ -292,10 +324,10 @@ public WriterCommitMessage commit() throws IOException { writeBatch(); } - // Close the Vortex writer to finalize the file + // Finalize the file; the summary carries its physical size if (vortexWriter != null) { try { - vortexWriter.close(); + bytesWritten = vortexWriter.finish().fileSize(); } finally { vortexWriter = null; // Always null out the reference } diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexCatalogTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexCatalogTest.java new file mode 100644 index 00000000000..d08455d376c --- /dev/null +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexCatalogTest.java @@ -0,0 +1,114 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.spark; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.util.Map; +import org.apache.spark.sql.catalyst.analysis.NoSuchTableException; +import org.apache.spark.sql.connector.catalog.Identifier; +import org.apache.spark.sql.connector.catalog.TableChange; +import org.apache.spark.sql.connector.expressions.Transform; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +/** + * Unit tests for {@link VortexCatalog}, the path-based catalog behind {@code SELECT * FROM vortex.`/path`}. + * + *

Characterizes which identifiers the catalog accepts as paths and which it rejects, plus the DDL surface it + * deliberately does not implement. End-to-end querying through a real session is covered by {@code VortexSqlTest}; + * these tests pin the identifier contract without starting Spark. + */ +final class VortexCatalogTest { + + private static final StructType SCHEMA = + new StructType(new StructField[] {new StructField("id", DataTypes.IntegerType, false, Metadata.empty())}); + + private static VortexCatalog catalog() { + VortexCatalog catalog = new VortexCatalog(); + catalog.initialize("vortex", new CaseInsensitiveStringMap(Map.of())); + return catalog; + } + + @Test + @DisplayName("Takes its name from the session config that registered it") + void nameComesFromInitialize() { + VortexCatalog catalog = new VortexCatalog(); + catalog.initialize("my_vortex", new CaseInsensitiveStringMap(Map.of())); + + assertEquals("my_vortex", catalog.name()); + } + + @Test + @DisplayName("Defaults to the name \"vortex\" before initialize is called") + void nameDefaultsToVortex() { + assertEquals("vortex", new VortexCatalog().name()); + } + + @Test + @DisplayName("Lists no tables: the catalog holds no state, tables are addressed by path") + void listTablesIsAlwaysEmpty() { + assertArrayEquals(new Identifier[0], catalog().listTables(new String[0])); + assertArrayEquals(new Identifier[0], catalog().listTables(new String[] {"any", "namespace"})); + } + + @Test + @DisplayName("An identifier without a path separator is not a table") + void identifierWithoutSlashIsRejected() { + assertThrows(NoSuchTableException.class, () -> catalog().loadTable(Identifier.of(new String[0], "not_a_path"))); + } + + @Test + @DisplayName("A namespaced identifier is not a table, even when the name looks like a path") + void namespacedIdentifierIsRejected() { + assertThrows(NoSuchTableException.class, () -> catalog() + .loadTable(Identifier.of(new String[] {"db"}, "/data/a.vortex"))); + } + + @Test + @DisplayName("A path-shaped identifier that cannot be read surfaces as table not found") + void unreadablePathIsReportedAsMissingTable() { + // Schema inference throws for a path with no Vortex files; the catalog translates that into + // NoSuchTableException so SQL users see "table not found" rather than an internal error. + assertThrows(NoSuchTableException.class, () -> catalog() + .loadTable(Identifier.of(new String[0], "/nonexistent/vortex/path"))); + } + + @Test + @DisplayName("Dropping a table never deletes data, it reports that nothing was dropped") + void dropTableReturnsFalse() { + assertFalse(catalog().dropTable(Identifier.of(new String[0], "/data/a.vortex"))); + } + + @Test + @DisplayName("CREATE TABLE is unsupported: write to the path instead") + void createTableIsUnsupported() { + assertThrows(UnsupportedOperationException.class, () -> catalog() + .createTable(Identifier.of(new String[0], "/data/a.vortex"), SCHEMA, new Transform[0], Map.of())); + } + + @Test + @DisplayName("ALTER TABLE is unsupported: there is no metadata to alter") + void alterTableIsUnsupported() { + assertThrows(UnsupportedOperationException.class, () -> catalog() + .alterTable(Identifier.of(new String[0], "/data/a.vortex"), new TableChange[0])); + } + + @Test + @DisplayName("RENAME TABLE is unsupported: there is no metadata to rename") + void renameTableIsUnsupported() { + assertThrows(UnsupportedOperationException.class, () -> catalog() + .renameTable( + Identifier.of(new String[0], "/data/a.vortex"), + Identifier.of(new String[0], "/data/b.vortex"))); + } +} diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceInferSchemaTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceInferSchemaTest.java new file mode 100644 index 00000000000..8c103cea9b9 --- /dev/null +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceInferSchemaTest.java @@ -0,0 +1,67 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.spark; + +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.nio.file.Path; +import org.apache.spark.sql.SparkSession; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; +import org.junit.jupiter.api.io.TempDir; + +/** Tests for schema inference failure reporting in {@link VortexDataSourceV2}. */ +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +public final class VortexDataSourceInferSchemaTest { + + private SparkSession spark; + + @TempDir + Path tempDir; + + @BeforeAll + public void setUp() { + spark = SparkSession.builder() + .appName("VortexInferSchemaTest") + .master("local[1]") + .config("spark.driver.host", "127.0.0.1") + .config("spark.ui.enabled", "false") + .getOrCreate(); + } + + @AfterAll + public void tearDown() { + if (spark != null) { + spark.stop(); + } + } + + @Test + @DisplayName("Reading a directory without Vortex files reports the offending path") + public void inferSchemaFailureNamesThePath() { + String emptyDir = tempDir.toString(); + + Throwable thrown = assertThrows( + Throwable.class, () -> spark.read().format("vortex").load(emptyDir)); + + String allMessages = messagesOf(thrown); + assertTrue( + allMessages.contains("no .vortex files found"), + "error should explain that no Vortex files were found, got: " + allMessages); + assertTrue(allMessages.contains(emptyDir), "error should name the offending path, got: " + allMessages); + } + + /** Concatenates the messages of the whole cause chain, since Spark may wrap data source exceptions. */ + private static String messagesOf(Throwable thrown) { + StringBuilder sb = new StringBuilder(); + for (Throwable t = thrown; t != null; t = t.getCause()) { + sb.append(t.getMessage()).append('\n'); + } + return sb.toString(); + } +} diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceWriteTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceWriteTest.java index 90fb1959ade..14b66c9e10a 100644 --- a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceWriteTest.java +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceWriteTest.java @@ -14,7 +14,9 @@ import java.nio.file.Path; import java.util.Arrays; import java.util.Comparator; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.stream.Collectors; import java.util.stream.Stream; import org.apache.spark.sql.Dataset; @@ -23,6 +25,7 @@ import org.apache.spark.sql.SaveMode; import org.apache.spark.sql.SparkSession; import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.MapType; import org.apache.spark.sql.types.StructField; import org.apache.spark.sql.types.StructType; import org.junit.jupiter.api.AfterAll; @@ -440,6 +443,86 @@ ELSE CAST('2024-02-03 04:05:06.654321' AS TIMESTAMP_NTZ) END assertFalse(findVortexFiles(outputPath).isEmpty(), "TimestampNTZ write should create Vortex files"); } + @Test + @DisplayName("Null array elements round-trip as nulls, not as zeros") + public void testWriteArrayWithNullElements() { + // A fixed-width element is the interesting case: reading a null slot yields the zeroed + // slot, so a missing null check writes 0 with the validity bit set. + Dataset arraysDf = spark.range(0, 1) + .selectExpr( + "cast(id as int) as id", + "array(1, cast(null as int), 3) as ints", + "array(cast(null as double), 2.5) as doubles", + "array('a', cast(null as string), 'c') as strings"); + + Path outputPath = tempDir.resolve("array_nulls_output"); + arraysDf.write() + .format("vortex") + .option("path", outputPath.toUri().toString()) + .mode(SaveMode.Overwrite) + .save(); + + Row read = spark.read() + .format("vortex") + .option("path", outputPath.toUri().toString()) + .load() + .selectExpr("ints", "doubles", "strings") + .collectAsList() + .get(0); + + assertEquals(Arrays.asList(1, null, 3), read.getList(0)); + assertEquals(Arrays.asList(null, 2.5d), read.getList(1)); + assertEquals(Arrays.asList("a", null, "c"), read.getList(2)); + } + + @Test + @DisplayName("Map columns round-trip populated, null, and empty values") + public void testWriteAndReadMapColumns() { + Dataset mapsDf = spark.range(0, 3) + .selectExpr( + "cast(id as int) as id", + """ + CASE + WHEN id = 0 THEN map(1, 'one', 2, cast(null as string)) + WHEN id = 1 THEN cast(null as map) + ELSE map_from_arrays(cast(array() as array), cast(array() as array)) + END AS attrs""", + "named_struct('attrs', map(cast(id as int), cast(id * 10 as bigint))) AS payload"); + + Path outputPath = tempDir.resolve("map_output"); + mapsDf.write() + .format("vortex") + .option("path", outputPath.toUri().toString()) + .mode(SaveMode.Overwrite) + .save(); + + Dataset readDf = spark.read() + .format("vortex") + .option("path", outputPath.toUri().toString()) + .load() + .orderBy("id"); + + MapType mapType = (MapType) readDf.schema().apply("attrs").dataType(); + assertEquals(DataTypes.IntegerType, mapType.keyType()); + assertEquals(DataTypes.StringType, mapType.valueType()); + assertTrue(mapType.valueContainsNull()); + + List rows = readDf.collectAsList(); + Map expected = new HashMap<>(); + expected.put(1, "one"); + expected.put(2, null); + assertEquals(expected, rows.get(0).getJavaMap(1)); + assertTrue(rows.get(1).isNullAt(1)); + assertTrue(rows.get(2).getJavaMap(1).isEmpty()); + + List nestedRows = readDf.selectExpr("id", "payload.attrs[id] AS nested_value") + .orderBy("id") + .collectAsList(); + assertEquals( + List.of(0L, 10L, 20L), + nestedRows.stream().map(row -> row.getLong(1)).toList()); + } + /** Creates a test DataFrame with monotonically increasing integers and their string representations. */ private Dataset createTestDataFrame(int numRows) { // Create DataFrame with monotonically increasing integers diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexTableTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexTableTest.java new file mode 100644 index 00000000000..8c23d6d874b --- /dev/null +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexTableTest.java @@ -0,0 +1,143 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.spark; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.google.common.collect.ImmutableList; +import dev.vortex.spark.read.VortexScanBuilder; +import dev.vortex.spark.write.VortexWriteBuilder; +import java.util.Map; +import java.util.NoSuchElementException; +import org.apache.spark.sql.connector.catalog.TableCapability; +import org.apache.spark.sql.connector.expressions.Expressions; +import org.apache.spark.sql.connector.expressions.Transform; +import org.apache.spark.sql.connector.read.Scan; +import org.apache.spark.sql.connector.write.LogicalWriteInfo; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +/** + * Unit tests for {@link VortexTable}, the Spark V2 table of Vortex files. + * + *

Characterizes the table-level contract Spark relies on: which capabilities are advertised, how the table name is + * rendered from its paths, and the read/write builders it hands back. Notably, writes accept exactly one path, so + * {@link VortexTable#newWriteBuilder} rejects tables built from several paths. + */ +final class VortexTableTest { + + private static final StructType SCHEMA = new StructType(new StructField[] { + new StructField("id", DataTypes.IntegerType, false, Metadata.empty()), + new StructField("name", DataTypes.StringType, true, Metadata.empty()) + }); + + private static VortexTable tableFor(String... paths) { + return new VortexTable(ImmutableList.copyOf(paths), SCHEMA, Map.of(), new Transform[0]); + } + + @Test + @DisplayName("Advertises exactly batch read, batch write and truncate") + void capabilitiesAreBatchReadWriteAndTruncate() { + assertEquals( + java.util.Set.of(TableCapability.BATCH_READ, TableCapability.BATCH_WRITE, TableCapability.TRUNCATE), + tableFor("/data/a.vortex").capabilities()); + } + + @Test + @DisplayName("Table name joins every path with commas under the vortex prefix") + void nameJoinsAllPaths() { + assertEquals("vortex.\"/data/a.vortex\"", tableFor("/data/a.vortex").name()); + assertEquals( + "vortex.\"/data/a.vortex,/data/b.vortex\"", + tableFor("/data/a.vortex", "/data/b.vortex").name()); + } + + @Test + @DisplayName("Schema and partitioning are returned as supplied") + void schemaAndPartitioningRoundTrip() { + Transform[] transforms = new Transform[] {Expressions.identity("year")}; + VortexTable table = new VortexTable(ImmutableList.of("/tbl"), SCHEMA, Map.of(), transforms); + + assertEquals(SCHEMA, table.schema()); + assertArrayEquals(transforms, table.partitioning()); + } + + @Test + @DisplayName("Scan builder projects the full table schema by default") + void scanBuilderStartsFromTableSchema() { + VortexTable table = tableFor("/data/a.vortex"); + + var builder = table.newScanBuilder(new CaseInsensitiveStringMap(Map.of())); + + assertInstanceOf(VortexScanBuilder.class, builder); + Scan scan = builder.build(); + assertEquals(SCHEMA, scan.readSchema()); + } + + @Test + @DisplayName("Scan description carries every table path") + void scanDescriptionCarriesPaths() { + VortexTable table = tableFor("/data/a.vortex", "/data/b.vortex"); + + Scan scan = table.newScanBuilder(new CaseInsensitiveStringMap(Map.of())).build(); + + String description = scan.description(); + assertTrue(description.startsWith("VortexScan{"), description); + assertTrue(description.contains("/data/a.vortex"), description); + assertTrue(description.contains("/data/b.vortex"), description); + assertTrue(description.contains("pushedPredicates=[]"), description); + } + + @Test + @DisplayName("Write builder is created for a single-path table") + void writeBuilderAcceptsSinglePath() { + VortexTable table = tableFor("/data/out"); + + assertInstanceOf(VortexWriteBuilder.class, table.newWriteBuilder(writeInfo())); + } + + @Test + @DisplayName("Writing a table of several paths is rejected: there is no single output path") + void writeBuilderRejectsMultiplePaths() { + VortexTable table = tableFor("/data/a.vortex", "/data/b.vortex"); + + assertThrows(IllegalArgumentException.class, () -> table.newWriteBuilder(writeInfo())); + } + + @Test + @DisplayName("Writing a table with no path is rejected") + void writeBuilderRejectsNoPaths() { + VortexTable table = tableFor(); + + assertThrows(NoSuchElementException.class, () -> table.newWriteBuilder(writeInfo())); + } + + private static LogicalWriteInfo writeInfo() { + return new LogicalWriteInfo() { + @Override + public String queryId() { + return "query-1"; + } + + @Override + public StructType schema() { + return SCHEMA; + } + + @Override + public CaseInsensitiveStringMap options() { + return new CaseInsensitiveStringMap(Map.of()); + } + }; + } +} diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/config/HadoopUtilsTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/config/HadoopUtilsTest.java new file mode 100644 index 00000000000..234db1a5ed6 --- /dev/null +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/config/HadoopUtilsTest.java @@ -0,0 +1,150 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.spark.config; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.Map; +import org.apache.hadoop.conf.Configuration; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +/** + * Unit tests for {@link HadoopUtils}, which forwards S3 and Azure credentials from a Hadoop {@link Configuration} to + * the Vortex object-store property keys. + * + *

Characterizes the {@code fs.s3a.*} to {@code aws_*} key mapping (including https-qualification of bare endpoints), + * the prefix-based Azure account key and SAS token extraction, and the skip-signature fallback for anonymous Azure + * access. + */ +final class HadoopUtilsTest { + /** A Configuration that does not load the default resources, so tests see only the keys they set. */ + private static Configuration emptyConf() { + return new Configuration(/* loadDefaults= */ false); + } + + // --- s3PropertiesFromHadoopConf --- + + @Test + @DisplayName("Maps fs.s3a credentials, session token, and region to the aws_* property keys") + void s3MapsCredentialKeys() { + Configuration conf = emptyConf(); + conf.set(HadoopUtils.FS_S3A_ACCESS_KEY, "AKIAEXAMPLE"); + conf.set(HadoopUtils.FS_S3A_SECRET_KEY, "secret"); + conf.set(HadoopUtils.FS_S3A_SESSION_TOKEN, "token"); + conf.set(HadoopUtils.FS_S3A_ENDPOINT_REGION, "us-west-2"); + + Map properties = HadoopUtils.s3PropertiesFromHadoopConf(conf); + + assertEquals("AKIAEXAMPLE", properties.get("aws_access_key_id")); + assertEquals("secret", properties.get("aws_secret_access_key")); + assertEquals("token", properties.get("aws_session_token")); + assertEquals("us-west-2", properties.get("aws_region")); + } + + @Test + @DisplayName("Qualifies a bare fs.s3a.endpoint with https://") + void s3QualifiesBareEndpoint() { + Configuration conf = emptyConf(); + conf.set(HadoopUtils.FS_S3A_ENDPOINT, "s3.us-west-2.amazonaws.com"); + + Map properties = HadoopUtils.s3PropertiesFromHadoopConf(conf); + + assertEquals("https://s3.us-west-2.amazonaws.com", properties.get("aws_endpoint")); + } + + @Test + @DisplayName("Preserves an fs.s3a.endpoint that already has an http or https scheme") + void s3PreservesSchemedEndpoint() { + Configuration httpConf = emptyConf(); + httpConf.set(HadoopUtils.FS_S3A_ENDPOINT, "http://localhost:9000"); + Map httpProperties = HadoopUtils.s3PropertiesFromHadoopConf(httpConf); + assertEquals("http://localhost:9000", httpProperties.get("aws_endpoint")); + assertEquals("true", httpProperties.get("aws_allow_http")); + + Configuration httpsConf = emptyConf(); + httpsConf.set(HadoopUtils.FS_S3A_ENDPOINT, "https://s3.example.com"); + Map httpsProperties = HadoopUtils.s3PropertiesFromHadoopConf(httpsConf); + assertEquals("https://s3.example.com", httpsProperties.get("aws_endpoint")); + assertNull(httpsProperties.get("aws_allow_http")); + } + + @Test + @DisplayName("Ignores Hadoop keys that are not S3-relevant") + void s3IgnoresUnrelatedKeys() { + Configuration conf = emptyConf(); + conf.set("fs.defaultFS", "hdfs://namenode:8020"); + conf.set("fs.s3a.connection.maximum", "64"); + + assertTrue(HadoopUtils.s3PropertiesFromHadoopConf(conf).isEmpty()); + } + + @Test + @DisplayName("Returns no properties for an empty configuration") + void s3EmptyConfYieldsNoProperties() { + assertTrue(HadoopUtils.s3PropertiesFromHadoopConf(emptyConf()).isEmpty()); + } + + // --- azurePropertiesFromHadoopConf --- + + @Test + @DisplayName("Extracts the storage account key from any fs.azure.account.key-prefixed entry") + void azureExtractsAccountKey() { + Configuration conf = emptyConf(); + conf.set(HadoopUtils.ACCESS_KEY_PREFIX + ".myaccount.dfs.core.windows.net", "account-key"); + + Map properties = HadoopUtils.azurePropertiesFromHadoopConf(conf); + + assertEquals("account-key", properties.get("azure_storage_account_key")); + assertFalse(properties.containsKey("azure_skip_signature")); + } + + @Test + @DisplayName("Extracts the SAS token from any fs.azure.sas.fixed.token-prefixed entry") + void azureExtractsSasToken() { + Configuration conf = emptyConf(); + conf.set(HadoopUtils.FIXED_TOKEN_PREFIX + "myaccount.dfs.core.windows.net", "sas-token"); + + Map properties = HadoopUtils.azurePropertiesFromHadoopConf(conf); + + assertEquals("sas-token", properties.get("azure_storage_sas_key")); + } + + @Test + @DisplayName("Falls back to skipping signatures when no account key is configured") + void azureSkipsSignatureWithoutAccountKey() { + Map properties = HadoopUtils.azurePropertiesFromHadoopConf(emptyConf()); + + assertEquals("true", properties.get("azure_skip_signature")); + assertFalse(properties.containsKey("azure_storage_account_key")); + } + + @Test + @DisplayName("Skips signatures even when only a SAS token is configured") + void azureSasOnlyStillSkipsSignature() { + Configuration conf = emptyConf(); + conf.set(HadoopUtils.FIXED_TOKEN_PREFIX + "myaccount.dfs.core.windows.net", "sas-token"); + + Map properties = HadoopUtils.azurePropertiesFromHadoopConf(conf); + + assertEquals("sas-token", properties.get("azure_storage_sas_key")); + assertEquals("true", properties.get("azure_skip_signature")); + } + + @Test + @DisplayName("Ignores Hadoop keys that are not Azure-relevant") + void azureIgnoresUnrelatedKeys() { + Configuration conf = emptyConf(); + conf.set("fs.azure.io.retry.max.retries", "5"); + conf.set("fs.defaultFS", "abfss://container@account.dfs.core.windows.net/"); + + Map properties = HadoopUtils.azurePropertiesFromHadoopConf(conf); + + assertFalse(properties.containsKey("azure_storage_account_key")); + assertFalse(properties.containsKey("azure_storage_sas_key")); + } +} diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/read/SparkPredicateToVortexExpressionTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/read/SparkPredicateToVortexExpressionTest.java index b7feafd322c..7be565f3553 100644 --- a/java/vortex-spark/src/test/java/dev/vortex/spark/read/SparkPredicateToVortexExpressionTest.java +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/read/SparkPredicateToVortexExpressionTest.java @@ -6,18 +6,36 @@ import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; +import dev.vortex.jni.NativeLoader; +import java.math.BigDecimal; +import java.util.List; import java.util.Map; import org.apache.spark.sql.connector.expressions.Expression; import org.apache.spark.sql.connector.expressions.LiteralValue; import org.apache.spark.sql.connector.expressions.NamedReference; +import org.apache.spark.sql.connector.expressions.filter.AlwaysFalse; +import org.apache.spark.sql.connector.expressions.filter.AlwaysTrue; +import org.apache.spark.sql.connector.expressions.filter.And; +import org.apache.spark.sql.connector.expressions.filter.Not; +import org.apache.spark.sql.connector.expressions.filter.Or; import org.apache.spark.sql.connector.expressions.filter.Predicate; import org.apache.spark.sql.types.DataType; import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Decimal; import org.apache.spark.sql.types.StructType; +import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; -/** Unit tests for {@link SparkPredicateToVortexExpression#isPushable(Predicate, Map)}. */ +/** + * Unit tests for {@link SparkPredicateToVortexExpression#isPushable(Predicate, Map)} and + * {@link SparkPredicateToVortexExpression#convert(Predicate)}. + * + *

{@code isPushable} decides which predicates {@code VortexScanBuilder.pushPredicates} lets Spark drop, while + * {@code convert} builds the filter {@code VortexPartitionReader} actually pushes down. A predicate that passes the + * first but fails the second is silently skipped by the reader, so the scan returns rows the query excluded. The tests + * below pin that {@code isPushable} implies {@code convert().isPresent()} across every accepted shape. + */ final class SparkPredicateToVortexExpressionTest { private static final StructType ADDRESS = DataTypes.createStructType(new org.apache.spark.sql.types.StructField[] { @@ -30,8 +48,48 @@ final class SparkPredicateToVortexExpressionTest { DataTypes.createStructField("address", ADDRESS, true) }); - private static final Map SCHEMA = - Map.of("id", DataTypes.IntegerType, "name", DataTypes.StringType, "profile", PROFILE); + private static final DataType DECIMAL = DataTypes.createDecimalType(10, 2); + + /** One data column per literal type {@code convertLiteral} maps, so every literal meets a same-typed column. */ + private static final Map SCHEMA = Map.ofEntries( + Map.entry("id", DataTypes.IntegerType), + Map.entry("name", DataTypes.StringType), + Map.entry("active", DataTypes.BooleanType), + Map.entry("tiny", DataTypes.ByteType), + Map.entry("small", DataTypes.ShortType), + Map.entry("big", DataTypes.LongType), + Map.entry("ratio", DataTypes.FloatType), + Map.entry("weight", DataTypes.DoubleType), + Map.entry("payload", DataTypes.BinaryType), + Map.entry("birthday", DataTypes.DateType), + Map.entry("createdAt", DataTypes.TimestampType), + Map.entry("createdLocal", DataTypes.TimestampNTZType), + Map.entry("amount", DECIMAL), + Map.entry("profile", PROFILE)); + + private static final List COMPARISON_OPERATORS = List.of("=", "<>", "!=", ">", ">=", "<", "<="); + + /** The columns whose type {@code isPushableLiteral} accepts a {@code null} value for. */ + private static final List NULLABLE_LITERAL_COLUMNS = List.of( + "active", + "tiny", + "small", + "id", + "big", + "ratio", + "weight", + "name", + "payload", + "birthday", + "createdAt", + "createdLocal", + "amount"); + + @BeforeAll + static void loadNativeLibrary() { + // `convert` allocates native expressions; `isPushable` does not. + NativeLoader.loadJni(); + } @Test @DisplayName("Top-level column reference is pushable when present in the schema") @@ -83,6 +141,208 @@ void emptyReferenceIsNotPushable() { assertFalse(SparkPredicateToVortexExpression.isPushable(equality, SCHEMA)); } + @Test + @DisplayName("Every accepted comparison operator converts, with the column on either side") + void everyComparisonOperatorConverts() { + for (String op : COMPARISON_OPERATORS) { + assertPushableAndConvertible(predicate(op, ref("id"), literal(42)), op + " with column on the left"); + // Spark's V2 builder sometimes commutes; `convertComparison` swaps the operator back. + assertPushableAndConvertible(predicate(op, literal(42), ref("id")), op + " with column on the right"); + } + } + + @Test + @DisplayName("Comparison between two columns converts") + void columnToColumnComparisonConverts() { + assertPushableAndConvertible(equality(ref("id"), ref("profile", "address", "zip"))); + } + + @Test + @DisplayName("Comparison between two literals is rejected by both sides") + void literalToLiteralComparisonIsRejected() { + assertNotPushableAndNotConvertible(equality(literal(1), literal(2))); + } + + @Test + @DisplayName("Comparison with the wrong number of children is rejected by both sides") + void comparisonWithWrongArityIsRejected() { + assertNotPushableAndNotConvertible(predicate("=", ref("id"))); + assertNotPushableAndNotConvertible(predicate("=", ref("id"), literal(1), literal(2))); + } + + @Test + @DisplayName("Nested column comparison converts") + void nestedColumnComparisonConverts() { + assertPushableAndConvertible(equality(ref("profile", "email"), literal("a@b.com"))); + } + + @Test + @DisplayName("IS_NULL and IS_NOT_NULL convert for top-level and nested columns") + void nullChecksConvert() { + assertPushableAndConvertible(predicate("IS_NULL", ref("name"))); + assertPushableAndConvertible(predicate("IS_NOT_NULL", ref("name"))); + assertPushableAndConvertible(predicate("IS_NULL", ref("profile", "address", "city"))); + } + + @Test + @DisplayName("IN converts for a single literal and for many literals") + void inConverts() { + assertPushableAndConvertible(predicate("IN", ref("id"), literal(1))); + assertPushableAndConvertible(predicate("IN", ref("id"), literal(1), literal(2), literal(3))); + } + + @Test + @DisplayName("IN with no literals is rejected by both sides") + void inWithoutLiteralsIsRejected() { + assertNotPushableAndNotConvertible(predicate("IN", ref("id"))); + } + + @Test + @DisplayName("String matching predicates convert, including LIKE meta-characters in the needle") + void stringMatchesConvert() { + for (String name : List.of("STARTS_WITH", "ENDS_WITH", "CONTAINS")) { + assertPushableAndConvertible(predicate(name, ref("name"), literal("ali")), name); + // `buildLikePattern` escapes `%`, `_` and `\` so the match stays an exact substring. + assertPushableAndConvertible(predicate(name, ref("name"), literal("100%_a\\b")), name + " with escapes"); + } + } + + @Test + @DisplayName("String matching against a non-string literal is rejected by both sides") + void stringMatchAgainstNonStringLiteralIsRejected() { + assertNotPushableAndNotConvertible(predicate("STARTS_WITH", ref("name"), literal(1))); + } + + @Test + @DisplayName("BOOLEAN_EXPRESSION over a column reference converts") + void bareBooleanColumnConverts() { + assertPushableAndConvertible(predicate("BOOLEAN_EXPRESSION", ref("active"))); + } + + @Test + @DisplayName("An unrecognised predicate name is rejected by both sides") + void unknownPredicateNameIsRejected() { + assertNotPushableAndNotConvertible(predicate("BLOOM_FILTER", ref("id"), literal(1))); + } + + @Test + @DisplayName("AlwaysTrue and AlwaysFalse convert to boolean literals") + void constantPredicatesConvert() { + assertPushableAndConvertible(new AlwaysTrue()); + assertPushableAndConvertible(new AlwaysFalse()); + } + + @Test + @DisplayName("AND, OR and NOT convert when every leaf converts") + void compoundPredicatesConvert() { + Predicate left = equality(ref("id"), literal(1)); + Predicate right = predicate("IS_NOT_NULL", ref("name")); + assertPushableAndConvertible(new And(left, right)); + assertPushableAndConvertible(new Or(left, right)); + assertPushableAndConvertible(new Not(left)); + assertPushableAndConvertible(new Not(new And(left, new Or(right, new AlwaysFalse())))); + } + + @Test + @DisplayName("A compound predicate with one unconvertible leaf is rejected by both sides") + void compoundPredicateWithBadLeafIsRejected() { + Predicate good = equality(ref("id"), literal(1)); + Predicate bad = predicate("BLOOM_FILTER", ref("id"), literal(1)); + assertNotPushableAndNotConvertible(new And(good, bad)); + assertNotPushableAndNotConvertible(new Or(bad, good)); + assertNotPushableAndNotConvertible(new Not(bad)); + } + + @Test + @DisplayName("Every literal type that accepts a null value converts") + void nullLiteralsConvert() { + for (String column : NULLABLE_LITERAL_COLUMNS) { + assertPushableAndConvertible( + equality(ref(column), new LiteralValue<>(null, SCHEMA.get(column))), "null literal for " + column); + } + } + + @Test + @DisplayName("Every non-null literal type the translator maps converts") + void nonNullLiteralsConvert() { + assertPushableAndConvertible(equality(ref("active"), new LiteralValue<>(true, DataTypes.BooleanType))); + assertPushableAndConvertible(equality(ref("tiny"), new LiteralValue<>((byte) 1, DataTypes.ByteType))); + assertPushableAndConvertible(equality(ref("small"), new LiteralValue<>((short) 1, DataTypes.ShortType))); + assertPushableAndConvertible(equality(ref("id"), literal(42))); + assertPushableAndConvertible(equality(ref("big"), new LiteralValue<>(1L, DataTypes.LongType))); + assertPushableAndConvertible(equality(ref("ratio"), new LiteralValue<>(1.5f, DataTypes.FloatType))); + assertPushableAndConvertible(equality(ref("weight"), new LiteralValue<>(1.5d, DataTypes.DoubleType))); + assertPushableAndConvertible(equality(ref("name"), literal("alice"))); + assertPushableAndConvertible( + equality(ref("payload"), new LiteralValue<>(new byte[] {1, 2, 3}, DataTypes.BinaryType))); + // Spark encodes DateType as an epoch-day int and both timestamp types as epoch micros. + assertPushableAndConvertible(equality(ref("birthday"), new LiteralValue<>(19_000, DataTypes.DateType))); + assertPushableAndConvertible( + equality(ref("createdAt"), new LiteralValue<>(1_700_000_000L, DataTypes.TimestampType))); + assertPushableAndConvertible( + equality(ref("createdLocal"), new LiteralValue<>(1_700_000_000L, DataTypes.TimestampNTZType))); + assertPushableAndConvertible(equality(ref("amount"), decimalLiteral("12.34"))); + } + + @Test + @DisplayName("A literal type with no Vortex representation is rejected by both sides") + void unrepresentableLiteralIsRejected() { + assertNotPushableAndNotConvertible(equality(ref("id"), new LiteralValue<>(null, DataTypes.NullType))); + } + + @Test + @DisplayName("A decimal literal that does not fit the declared scale is rejected by both sides") + void decimalThatDoesNotFitTheScaleIsRejected() { + // `unscaledValueOf` calls `setScale(2)` without a rounding mode, so 12.345 throws and + // `isPushableLiteral` falls through to `literalOf`, which is empty. + assertNotPushableAndNotConvertible(equality(ref("amount"), decimalLiteral("12.345"))); + } + + @Test + @DisplayName("An empty named reference is rejected by convert as well as by isPushable") + void emptyReferenceIsAlsoNotConvertible() { + // `isFieldRefExpr` on the convert path only checks `instanceof NamedReference`, so the + // zero-part guard lives in `columnOf`. Without it a pushable-looking predicate would reach + // the reader and be silently dropped. + assertNotPushableAndNotConvertible(equality(ref(), literal(1))); + assertNotPushableAndNotConvertible(predicate("IS_NULL", ref())); + assertNotPushableAndNotConvertible(predicate("IN", ref(), literal(1))); + assertNotPushableAndNotConvertible(predicate("STARTS_WITH", ref(), literal("a"))); + assertNotPushableAndNotConvertible(predicate("BOOLEAN_EXPRESSION", ref())); + } + + /** + * Asserts the invariant documented on {@link SparkPredicateToVortexExpression#isPushable(Predicate, Map)}: a + * predicate Spark is allowed to drop must produce a Vortex expression. + */ + private static void assertPushableAndConvertible(Predicate predicate) { + assertPushableAndConvertible(predicate, predicate.name()); + } + + private static void assertPushableAndConvertible(Predicate predicate, String what) { + assertTrue(SparkPredicateToVortexExpression.isPushable(predicate, SCHEMA), () -> "not pushable: " + what); + assertTrue( + SparkPredicateToVortexExpression.convert(predicate).isPresent(), + () -> "pushable but not convertible: " + what); + } + + private static void assertNotPushableAndNotConvertible(Predicate predicate) { + assertNotPushableAndNotConvertible(predicate, predicate.name()); + } + + private static void assertNotPushableAndNotConvertible(Predicate predicate, String what) { + assertFalse(SparkPredicateToVortexExpression.isPushable(predicate, SCHEMA), () -> "pushable: " + what); + assertFalse(SparkPredicateToVortexExpression.convert(predicate).isPresent(), () -> "convertible: " + what); + } + + private static Predicate predicate(String name, Expression... children) { + return new Predicate(name, children); + } + + private static LiteralValue decimalLiteral(String value) { + return new LiteralValue<>(Decimal.apply(new BigDecimal(value)), DECIMAL); + } + private static Predicate equality(Expression left, Expression right) { return new Predicate("=", new Expression[] {left, right}); } diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexBatchExecTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexBatchExecTest.java new file mode 100644 index 00000000000..967940567b2 --- /dev/null +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexBatchExecTest.java @@ -0,0 +1,95 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.spark.read; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; + +import dev.vortex.spark.VortexFilePartition; +import java.util.List; +import java.util.Map; +import org.apache.spark.sql.connector.catalog.Column; +import org.apache.spark.sql.connector.expressions.filter.Predicate; +import org.apache.spark.sql.connector.read.InputPartition; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +/** + * Unit tests for {@link VortexBatchExec} input-partition planning. + * + *

Characterizes how explicit {@code .vortex} paths are planned: one {@link VortexFilePartition} per file, each + * carrying the requested read schema and the Hive-style partition values parsed from its own path. All paths used here + * end in {@code .vortex} so planning never lists directories through native I/O. + */ +final class VortexBatchExecTest { + + private static final List COLUMNS = List.of( + Column.create("id", org.apache.spark.sql.types.DataTypes.IntegerType), + Column.create("name", org.apache.spark.sql.types.DataTypes.StringType)); + + private static VortexBatchExec execFor(List paths) { + return new VortexBatchExec(paths, COLUMNS, Map.of(), new Predicate[0]); + } + + @Test + @DisplayName("Plans one input partition per .vortex file") + void plansOnePartitionPerFile() { + VortexBatchExec exec = execFor(List.of("/data/a.vortex", "/data/b.vortex", "/data/c.vortex")); + + InputPartition[] partitions = exec.planInputPartitions(); + + assertEquals(3, partitions.length); + for (InputPartition partition : partitions) { + assertInstanceOf(VortexFilePartition.class, partition); + assertEquals(1, ((VortexFilePartition) partition).paths().size()); + } + } + + @Test + @DisplayName("Each partition carries exactly its own file path, in input order") + void partitionsKeepInputOrder() { + VortexBatchExec exec = execFor(List.of("/data/first.vortex", "/data/second.vortex")); + + InputPartition[] partitions = exec.planInputPartitions(); + + assertEquals(List.of("/data/first.vortex"), ((VortexFilePartition) partitions[0]).paths()); + assertEquals(List.of("/data/second.vortex"), ((VortexFilePartition) partitions[1]).paths()); + } + + @Test + @DisplayName("Partitions carry the requested read schema") + void partitionsCarryReadSchema() { + VortexBatchExec exec = execFor(List.of("/data/a.vortex")); + + VortexFilePartition partition = (VortexFilePartition) exec.planInputPartitions()[0]; + + assertEquals(2, partition.readSchema().size()); + assertEquals("id", partition.readSchema().fields()[0].name()); + assertEquals("name", partition.readSchema().fields()[1].name()); + } + + @Test + @DisplayName("Hive-style partition values are parsed from each file's own path") + void partitionValuesParsedPerFile() { + VortexBatchExec exec = execFor( + List.of("/tbl/year=2024/month=01/a.vortex", "/tbl/year=2025/month=02/b.vortex", "/tbl/plain.vortex")); + + InputPartition[] partitions = exec.planInputPartitions(); + + assertEquals(Map.of("year", "2024", "month", "01"), ((VortexFilePartition) partitions[0]).partitionValues()); + assertEquals(Map.of("year", "2025", "month", "02"), ((VortexFilePartition) partitions[1]).partitionValues()); + assertEquals(Map.of(), ((VortexFilePartition) partitions[2]).partitionValues()); + } + + @Test + @DisplayName("Format options are propagated to every partition") + void formatOptionsPropagated() { + Map options = Map.of("vortex.workerThreads", "8"); + VortexBatchExec exec = new VortexBatchExec(List.of("/data/a.vortex"), COLUMNS, options, new Predicate[0]); + + VortexFilePartition partition = (VortexFilePartition) exec.planInputPartitions()[0]; + + assertEquals("8", partition.formatOptions().get("vortex.workerThreads")); + } +} diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexScanBuilderTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexScanBuilderTest.java new file mode 100644 index 00000000000..9cc841b631b --- /dev/null +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexScanBuilderTest.java @@ -0,0 +1,176 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.spark.read; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.util.Map; +import org.apache.spark.sql.connector.catalog.Column; +import org.apache.spark.sql.connector.expressions.Expression; +import org.apache.spark.sql.connector.expressions.Expressions; +import org.apache.spark.sql.connector.expressions.LiteralValue; +import org.apache.spark.sql.connector.expressions.NamedReference; +import org.apache.spark.sql.connector.expressions.Transform; +import org.apache.spark.sql.connector.expressions.filter.Predicate; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.unsafe.types.UTF8String; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +/** + * Unit tests for {@link VortexScanBuilder}, focusing on how {@link VortexScanBuilder#pushPredicates(Predicate[])} + * splits predicates between Vortex pushdown and Spark post-scan evaluation. + * + *

Characterizes the partition-column exclusion rule (partition columns live in directory paths, not in the Vortex + * files, so predicates referencing them must stay with Spark), the handling of unsupported operators, and the defensive + * copy returned by {@link VortexScanBuilder#pushedPredicates()}. + */ +final class VortexScanBuilderTest { + + private static VortexScanBuilder builderWithColumns(Transform[] partitionTransforms, Column... columns) { + VortexScanBuilder builder = new VortexScanBuilder(Map.of(), partitionTransforms); + builder.addPath("/tmp/data.vortex"); + for (Column column : columns) { + builder.addColumn(column); + } + return builder; + } + + private static Predicate equality(String column, int value) { + return new Predicate("=", new Expression[] {ref(column), new LiteralValue<>(value, DataTypes.IntegerType)}); + } + + private static NamedReference ref(String name) { + return Expressions.column(name); + } + + // --- pushPredicates: data columns --- + + @Test + @DisplayName("Predicate on a data column is pushed and removed from post-scan predicates") + void dataColumnPredicateIsPushed() { + VortexScanBuilder builder = builderWithColumns( + new Transform[0], + Column.create("id", DataTypes.IntegerType), + Column.create("name", DataTypes.StringType)); + + Predicate predicate = equality("id", 42); + Predicate[] postScan = builder.pushPredicates(new Predicate[] {predicate}); + + assertEquals(0, postScan.length); + assertArrayEquals(new Predicate[] {predicate}, builder.pushedPredicates()); + } + + @Test + @DisplayName("Predicate on a column missing from the read schema is left to Spark") + void unknownColumnPredicateIsNotPushed() { + VortexScanBuilder builder = builderWithColumns(new Transform[0], Column.create("id", DataTypes.IntegerType)); + + Predicate predicate = equality("missing", 1); + Predicate[] postScan = builder.pushPredicates(new Predicate[] {predicate}); + + assertArrayEquals(new Predicate[] {predicate}, postScan); + assertEquals(0, builder.pushedPredicates().length); + } + + @Test + @DisplayName("Predicate with an unsupported operator is left to Spark") + void unsupportedOperatorIsNotPushed() { + VortexScanBuilder builder = builderWithColumns(new Transform[0], Column.create("name", DataTypes.StringType)); + + // LIKE with arbitrary user pattern is not a translatable V2 predicate name. + Predicate predicate = new Predicate( + "LIKE", + new Expression[] {ref("name"), new LiteralValue<>(UTF8String.fromString("%a%"), DataTypes.StringType)}); + Predicate[] postScan = builder.pushPredicates(new Predicate[] {predicate}); + + assertArrayEquals(new Predicate[] {predicate}, postScan); + assertEquals(0, builder.pushedPredicates().length); + } + + // --- pushPredicates: partition columns --- + + @Test + @DisplayName("Predicate on a partition column is left to Spark") + void partitionColumnPredicateIsNotPushed() { + Transform[] transforms = new Transform[] {Expressions.identity("year")}; + VortexScanBuilder builder = builderWithColumns( + transforms, Column.create("id", DataTypes.IntegerType), Column.create("year", DataTypes.IntegerType)); + + Predicate predicate = equality("year", 2024); + Predicate[] postScan = builder.pushPredicates(new Predicate[] {predicate}); + + assertArrayEquals(new Predicate[] {predicate}, postScan); + assertEquals(0, builder.pushedPredicates().length); + } + + @Test + @DisplayName("Mixed predicates split into pushed (data) and post-scan (partition)") + void mixedPredicatesAreSplit() { + Transform[] transforms = new Transform[] {Expressions.identity("year")}; + VortexScanBuilder builder = builderWithColumns( + transforms, Column.create("id", DataTypes.IntegerType), Column.create("year", DataTypes.IntegerType)); + + Predicate onData = equality("id", 7); + Predicate onPartition = equality("year", 2024); + Predicate[] postScan = builder.pushPredicates(new Predicate[] {onData, onPartition}); + + assertArrayEquals(new Predicate[] {onPartition}, postScan); + assertArrayEquals(new Predicate[] {onData}, builder.pushedPredicates()); + } + + @Test + @DisplayName("AND spanning a data column and a partition column is left to Spark as a whole") + void conjunctionSpanningPartitionColumnIsNotPushed() { + Transform[] transforms = new Transform[] {Expressions.identity("year")}; + VortexScanBuilder builder = builderWithColumns( + transforms, Column.create("id", DataTypes.IntegerType), Column.create("year", DataTypes.IntegerType)); + + Predicate and = + new org.apache.spark.sql.connector.expressions.filter.And(equality("id", 7), equality("year", 2024)); + Predicate[] postScan = builder.pushPredicates(new Predicate[] {and}); + + assertArrayEquals(new Predicate[] {and}, postScan); + assertEquals(0, builder.pushedPredicates().length); + } + + // --- pushedPredicates --- + + @Test + @DisplayName("pushedPredicates is empty before any pushPredicates call") + void pushedPredicatesEmptyByDefault() { + VortexScanBuilder builder = builderWithColumns(new Transform[0], Column.create("id", DataTypes.IntegerType)); + assertEquals(0, builder.pushedPredicates().length); + } + + @Test + @DisplayName("pushedPredicates returns a defensive copy") + void pushedPredicatesReturnsCopy() { + VortexScanBuilder builder = builderWithColumns(new Transform[0], Column.create("id", DataTypes.IntegerType)); + builder.pushPredicates(new Predicate[] {equality("id", 1)}); + + Predicate[] first = builder.pushedPredicates(); + Predicate[] second = builder.pushedPredicates(); + assertNotSame(first, second); + + // Mutating one returned array must not leak into subsequent calls. + first[0] = null; + assertEquals(1, builder.pushedPredicates().length); + assertNotNull(builder.pushedPredicates()[0]); + } + + // --- build --- + + @Test + @DisplayName("build without paths fails fast") + void buildWithoutPathsThrows() { + VortexScanBuilder builder = new VortexScanBuilder(Map.of()); + builder.addColumn(Column.create("id", DataTypes.IntegerType)); + assertThrows(IllegalStateException.class, builder::build); + } +} diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/write/PartitionedVortexDataWriterTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/write/PartitionedVortexDataWriterTest.java new file mode 100644 index 00000000000..7d3318d61b9 --- /dev/null +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/write/PartitionedVortexDataWriterTest.java @@ -0,0 +1,169 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.spark.write; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import dev.vortex.spark.write.PartitionedVortexDataWriter.ResolvedTransform; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.ObjectInputStream; +import java.io.ObjectOutputStream; +import java.util.List; +import org.apache.spark.sql.connector.expressions.Expressions; +import org.apache.spark.sql.connector.expressions.Transform; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +/** + * Unit tests for {@link PartitionedVortexDataWriter#resolveTransforms}, which turns Spark partition transforms into the + * directory keys and column indices a writer uses. + * + *

Resolution happens eagerly on the driver because Spark's {@code Transform} objects are Scala case classes that are + * not Java-serializable, so only the resolved form may cross to the executors; the serialization round trip below is + * what that eager step exists for. + */ +final class PartitionedVortexDataWriterTest { + + private static final StructType SCHEMA = new StructType(new StructField[] { + new StructField("id", DataTypes.IntegerType, false, Metadata.empty()), + new StructField("name", DataTypes.StringType, true, Metadata.empty()), + new StructField("event_date", DataTypes.DateType, true, Metadata.empty()), + new StructField("event_ts", DataTypes.TimestampType, true, Metadata.empty()) + }); + + private static ResolvedTransform resolveOne(Transform transform) { + ResolvedTransform[] resolved = + PartitionedVortexDataWriter.resolveTransforms(new Transform[] {transform}, SCHEMA); + assertEquals(1, resolved.length); + return resolved[0]; + } + + @Test + @DisplayName("No transforms resolve to no partitioning") + void emptyInput() { + assertEquals(0, PartitionedVortexDataWriter.resolveTransforms(new Transform[0], SCHEMA).length); + } + + @Test + @DisplayName("An identity transform keeps the column name as its directory key") + void identity() { + ResolvedTransform resolved = resolveOne(Expressions.identity("name")); + + assertEquals("name", resolved.directoryKey()); + assertEquals("identity", resolved.transformName()); + assertEquals(1, resolved.columnIndices().get(0)); + assertEquals(DataTypes.StringType, resolved.columnTypes().get(0)); + assertEquals(-1, resolved.bucketCount(), "only bucket transforms carry a bucket count"); + } + + @Test + @DisplayName("Temporal transforms suffix the directory key with the unit they truncate to") + void temporalDirectoryKeys() { + assertEquals( + "event_date_year", resolveOne(Expressions.years("event_date")).directoryKey()); + assertEquals( + "event_date_month", resolveOne(Expressions.months("event_date")).directoryKey()); + assertEquals( + "event_date_day", resolveOne(Expressions.days("event_date")).directoryKey()); + assertEquals("event_ts_hour", resolveOne(Expressions.hours("event_ts")).directoryKey()); + } + + @Test + @DisplayName("Temporal transforms reject a non-temporal column, naming the transform and the type") + void temporalTransformsRejectNonTemporalColumns() { + IllegalArgumentException thrown = + assertThrows(IllegalArgumentException.class, () -> resolveOne(Expressions.years("id"))); + + assertTrue(thrown.getMessage().contains("years"), thrown.getMessage()); + assertTrue(thrown.getMessage().contains("IntegerType"), thrown.getMessage()); + } + + @Test + @DisplayName("The hours transform additionally rejects a date column, which has no hour to truncate to") + void hoursRejectsDateColumn() { + IllegalArgumentException thrown = + assertThrows(IllegalArgumentException.class, () -> resolveOne(Expressions.hours("event_date"))); + + assertTrue(thrown.getMessage().contains("hours"), thrown.getMessage()); + } + + @Test + @DisplayName("A bucket transform carries its bucket count and joins its columns with underscores") + void bucket() { + ResolvedTransform single = resolveOne(Expressions.bucket(8, "id")); + assertEquals("id_bucket", single.directoryKey()); + assertEquals(8, single.bucketCount()); + + ResolvedTransform multi = resolveOne(Expressions.bucket(4, "id", "name")); + assertEquals("id_name_bucket", multi.directoryKey()); + assertEquals(4, multi.bucketCount()); + assertEquals(0, multi.columnIndices().get(0)); + assertEquals(1, multi.columnIndices().get(1)); + assertEquals(List.of(DataTypes.IntegerType, DataTypes.StringType), multi.columnTypes()); + } + + @Test + @DisplayName("A bucket transform without a numBuckets argument is rejected") + void bucketWithoutCount() { + Transform noCount = Expressions.apply("bucket", Expressions.column("id")); + + IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class, () -> resolveOne(noCount)); + + assertTrue(thrown.getMessage().contains("numBuckets"), thrown.getMessage()); + } + + @Test + @DisplayName("An unsupported transform name is rejected") + void unsupportedTransform() { + IllegalArgumentException thrown = assertThrows( + IllegalArgumentException.class, + () -> resolveOne(Expressions.apply("truncate", Expressions.column("name")))); + + assertTrue(thrown.getMessage().contains("truncate"), thrown.getMessage()); + } + + @Test + @DisplayName("A transform that references no column is rejected") + void transformWithoutReferences() { + IllegalArgumentException thrown = + assertThrows(IllegalArgumentException.class, () -> resolveOne(Expressions.apply("identity"))); + + assertTrue(thrown.getMessage().contains("no column references"), thrown.getMessage()); + } + + @Test + @DisplayName("A transform on a column outside the schema is rejected") + void transformOnUnknownColumn() { + assertThrows(IllegalArgumentException.class, () -> resolveOne(Expressions.identity("missing"))); + } + + @Test + @DisplayName("Resolved transforms survive Java serialization, which is why they are resolved eagerly") + void resolvedTransformsAreSerializable() throws IOException, ClassNotFoundException { + ResolvedTransform original = resolveOne(Expressions.bucket(4, "id", "name")); + + ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + try (ObjectOutputStream out = new ObjectOutputStream(bytes)) { + out.writeObject(original); + } + ResolvedTransform restored; + try (ObjectInputStream in = new ObjectInputStream(new ByteArrayInputStream(bytes.toByteArray()))) { + restored = (ResolvedTransform) in.readObject(); + } + + assertEquals(original.directoryKey(), restored.directoryKey()); + assertEquals(original.transformName(), restored.transformName()); + assertEquals(original.bucketCount(), restored.bucketCount()); + assertEquals(original.columnTypes(), restored.columnTypes()); + assertEquals(original.columnIndices().asList(), restored.columnIndices().asList()); + } +} diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/write/SparkToArrowSchemaTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/write/SparkToArrowSchemaTest.java index cd9b3ca184c..d65f0403800 100644 --- a/java/vortex-spark/src/test/java/dev/vortex/spark/write/SparkToArrowSchemaTest.java +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/write/SparkToArrowSchemaTest.java @@ -173,4 +173,45 @@ void unsupportedTypeIsRejected() { assertThrows(UnsupportedOperationException.class, () -> SparkToArrowSchema.convert(schema)); assertTrue(e.getMessage().contains("CalendarIntervalType")); } + + @Test + @DisplayName("MapType converts to a Map field with non-null keys and value nullability") + void mapConvertsWithEntries() { + Field map = single("m", DataTypes.createMapType(DataTypes.StringType, DataTypes.IntegerType, true), true); + + assertEquals(new ArrowType.Map(false), map.getType()); + assertEquals(1, map.getChildren().size()); + + Field entries = map.getChildren().get(0); + assertEquals("entries", entries.getName()); + assertEquals(new ArrowType.Struct(), entries.getType()); + assertFalse(entries.isNullable()); + + Field key = entries.getChildren().get(0); + assertEquals("key", key.getName()); + assertEquals(new ArrowType.Utf8(), key.getType()); + assertFalse(key.isNullable()); + + Field value = entries.getChildren().get(1); + assertEquals("value", value.getName()); + assertEquals(new ArrowType.Int(32, true), value.getType()); + assertTrue(value.isNullable()); + } + + @Test + @DisplayName("Nested MapType children retain their Arrow structure") + void nestedMapConvertsRecursively() { + Field map = single( + "m", + DataTypes.createMapType( + DataTypes.StringType, + DataTypes.createArrayType(DataTypes.createMapType(DataTypes.IntegerType, DataTypes.LongType))), + true); + + Field value = map.getChildren().get(0).getChildren().get(1); + Field element = value.getChildren().get(0); + assertEquals(new ArrowType.List(), value.getType()); + assertEquals(new ArrowType.Map(false), element.getType()); + assertEquals(2, element.getChildren().get(0).getChildren().size()); + } } diff --git a/lang/cpp/CMakeLists.txt b/lang/cpp/CMakeLists.txt index 8949c1977cc..7884228cf70 100644 --- a/lang/cpp/CMakeLists.txt +++ b/lang/cpp/CMakeLists.txt @@ -1,147 +1,115 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright the Vortex contributors -cmake_minimum_required(VERSION 3.10) -project(VortexCXX - VERSION 0.0.1 - LANGUAGES CXX) +# Entry point for the static Vortex C++ wrapper. It consumes the Rust FFI +# archive from vortex-ffi and exports Vortex::cpp_static. It supports standalone +# configuration and add_subdirectory() from the repository root or a parent. -include(FetchContent) +# CMake 3.25 is required for block() and FetchContent SYSTEM. +cmake_minimum_required(VERSION 3.25) -set(CMAKE_EXPORT_COMPILE_COMMANDS ON) -set(CMAKE_CXX_STANDARD 20) -set(CMAKE_CXX_STANDARD_REQUIRED ON) +project(VortexCXX DESCRIPTION "Vortex C++ wrapper" LANGUAGES C CXX) -find_program(SCCACHE_PROGRAM sccache) -if (SCCACHE_PROGRAM) - set(CMAKE_C_COMPILER_LAUNCHER "${SCCACHE_PROGRAM}") - set(CMAKE_CXX_COMPILER_LAUNCHER "${SCCACHE_PROGRAM}") - message(STATUS "Sccache found: ${SCCACHE_PROGRAM}") -else () - message(STATUS "Sccache not found") -endif () - -option(BUILD_TESTS "Build tests" OFF) - -set(SANITIZER "" CACHE STRING "Build with sanitizers") -set(TARGET_TRIPLE "" CACHE STRING "Rust target triple for FFI library") -set(RUST_BUILD_PROFILE "" CACHE STRING "Cargo profile name for Rust FFI library") - -if (NOT CMAKE_BUILD_TYPE) - set(CMAKE_BUILD_TYPE Debug) +# Standalone entry points own their defaults; otherwise inherit Vortex's marker. +if(PROJECT_IS_TOP_LEVEL) + set(_vortex_top_level ON) +elseif(NOT DEFINED _vortex_top_level) + set(_vortex_top_level OFF) endif() -if (NOT SANITIZER STREQUAL "") - message(NOTICE "Sanitizer: ${SANITIZER}") - if (NOT CMAKE_BUILD_TYPE STREQUAL "Debug") - message(FATAL_ERROR "Only debug build is supported for sanitizer builds") - endif() +# Vortex does not use C++ modules, so it does not require compiler-specific +# dependency scanners such as clang-scan-deps. +set(CMAKE_CXX_SCAN_FOR_MODULES OFF) - if (SANITIZER STREQUAL "asan") - set(SANITIZER_FLAGS "-fsanitize=address,undefined,leak") - elseif (SANITIZER STREQUAL "tsan") - set(SANITIZER_FLAGS "-fsanitize=thread") - else() - message(FATAL_ERROR "Unknown sanitizer ${SANITIZER}") - endif() - set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${SANITIZER_FLAGS}") - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${SANITIZER_FLAGS}") - set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} ${SANITIZER_FLAGS}") +# A standalone configure writes compile_commands.json for editor tooling. +if(PROJECT_IS_TOP_LEVEL) + set(CMAKE_EXPORT_COMPILE_COMMANDS ON) endif() -if (RUST_BUILD_PROFILE STREQUAL "") - if (CMAKE_BUILD_TYPE STREQUAL "RelWithDebInfo") - set(RUST_BUILD_PROFILE "release_debug") - else() - string(TOLOWER "${CMAKE_BUILD_TYPE}" RUST_BUILD_PROFILE) - endif() -endif() - -set(VORTEX_FFI_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../../vortex-ffi") -set(VORTEX_FFI_LIB_DIR "${VORTEX_FFI_DIR}/../target/${TARGET_TRIPLE}/${RUST_BUILD_PROFILE}") -set(VORTEX_FFI_HEADERS "${VORTEX_FFI_DIR}/cinclude") - -if(APPLE) - set(VORTEX_FFI_STATIC "${VORTEX_FFI_LIB_DIR}/libvortex_ffi.a") - set(VORTEX_FFI_SHARED "${VORTEX_FFI_LIB_DIR}/libvortex_ffi.dylib") -elseif(WIN32) - set(VORTEX_FFI_STATIC "${VORTEX_FFI_LIB_DIR}/libvortex_ffi.lib") - set(VORTEX_FFI_SHARED "${VORTEX_FFI_LIB_DIR}/libvortex_ffi.dll") -else() - set(VORTEX_FFI_STATIC "${VORTEX_FFI_LIB_DIR}/libvortex_ffi.a") - set(VORTEX_FFI_SHARED "${VORTEX_FFI_LIB_DIR}/libvortex_ffi.so") +if(_vortex_top_level AND NOT CMAKE_BUILD_TYPE) + set(CMAKE_BUILD_TYPE Debug) endif() -if(NOT EXISTS "${VORTEX_FFI_STATIC}") - message(FATAL_ERROR - "vortex-ffi static library not found at ${VORTEX_FFI_STATIC}. " - "Run: cargo build --profile -p vortex-ffi") -endif() +option(VORTEX_BUILD_TESTS "Build Vortex tests" OFF) +option(VORTEX_BUILD_EXAMPLES "Build Vortex examples" OFF) +option(VORTEX_WARNINGS_AS_ERRORS + "Treat warnings in Vortex-owned sources as errors" + "${_vortex_top_level}") -add_library(vortex_ffi STATIC IMPORTED) -set_target_properties(vortex_ffi PROPERTIES - IMPORTED_LOCATION "${VORTEX_FFI_STATIC}" - INTERFACE_INCLUDE_DIRECTORIES "${VORTEX_FFI_HEADERS}") - -if(EXISTS "${VORTEX_FFI_SHARED}") - add_library(vortex_ffi_shared SHARED IMPORTED) - set_target_properties(vortex_ffi_shared PROPERTIES - IMPORTED_LOCATION "${VORTEX_FFI_SHARED}" - INTERFACE_INCLUDE_DIRECTORIES "${VORTEX_FFI_HEADERS}" - INTERFACE_LINK_OPTIONS "LINKER:-rpath,${VORTEX_FFI_LIB_DIR}") +# The Rust FFI archive comes from vortex-ffi. A standalone configure of this +# directory adds it here; the repository root or a parent may already have. +if(NOT TARGET Vortex::ffi_static) + add_subdirectory(../../vortex-ffi ffi) endif() -file(GLOB VORTEX_CXX_SOURCES CONFIGURE_DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/src/*.cpp") -add_library(vortex_cxx STATIC ${VORTEX_CXX_SOURCES}) -target_include_directories(vortex_cxx PUBLIC - ${CMAKE_CURRENT_SOURCE_DIR} - ${CMAKE_CURRENT_SOURCE_DIR}/include - ${CMAKE_CURRENT_SOURCE_DIR}/../vortex-ffi/cinclude -) -target_link_libraries(vortex_cxx PUBLIC vortex_ffi) -add_library(vortex::cxx ALIAS vortex_cxx) +# The C++ wrapper over the C FFI. Public headers live under include/. +file(GLOB _vortex_cpp_sources CONFIGURE_DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/src/*.cpp") + +# Consumers inherit the FFI archive, its headers, and its system libraries +# through this target. Hidden visibility keeps wrapper symbols out of a parent +# shared library's exports. +add_library(vortex_cxx STATIC ${_vortex_cpp_sources}) +set_target_properties(vortex_cxx PROPERTIES + POSITION_INDEPENDENT_CODE ON + CXX_VISIBILITY_PRESET hidden + VISIBILITY_INLINES_HIDDEN ON) target_compile_features(vortex_cxx PUBLIC cxx_std_20) - -set(APPLE_LINK_FLAGS "-framework CoreFoundation -framework Security") - -if(APPLE) - target_link_libraries(vortex_cxx PRIVATE ${APPLE_LINK_FLAGS}) +target_include_directories(vortex_cxx PUBLIC "${CMAKE_CURRENT_SOURCE_DIR}/include") +target_link_libraries(vortex_cxx PUBLIC Vortex::ffi_static) + +# Warning policy is private to Vortex-owned sources so embedding this target +# never injects warning flags into consumers. +target_compile_options(vortex_cxx PRIVATE -Wall -Wextra -Wpedantic) +if(VORTEX_WARNINGS_AS_ERRORS) + target_compile_options(vortex_cxx PRIVATE -Werror) endif() -if(TARGET vortex_ffi_shared) - add_library(vortex_cxx_shared SHARED ${VORTEX_CXX_SOURCES}) - set_target_properties(vortex_cxx_shared PROPERTIES POSITION_INDEPENDENT_CODE ON) - target_include_directories(vortex_cxx_shared PUBLIC - ${CMAKE_CURRENT_SOURCE_DIR} - ${CMAKE_CURRENT_SOURCE_DIR}/include - ${CMAKE_CURRENT_SOURCE_DIR}/../vortex-ffi/cinclude - ) - target_link_libraries(vortex_cxx_shared PUBLIC vortex_ffi_shared) - add_library(vortex::cxx_shared ALIAS vortex_cxx_shared) - target_compile_features(vortex_cxx_shared PUBLIC cxx_std_20) - - if(APPLE) - target_link_libraries(vortex_cxx_shared PRIVATE ${APPLE_LINK_FLAGS}) - endif() -endif() +# The only supported consumer-facing name for the C++ API. +add_library(Vortex::cpp_static ALIAS vortex_cxx) + +# Test and example dependencies are fetched at configure time and never installed. +if(VORTEX_BUILD_TESTS OR VORTEX_BUILD_EXAMPLES) + include(FetchContent) -if (BUILD_TESTS) + # Keep dependency warnings out of Vortex's warning policy. + set(NANOARROW_DEBUG_EXTRA_WARNINGS OFF) FetchContent_Declare( - Nanoarrow - GIT_REPOSITORY https://github.com/apache/arrow-nanoarrow - GIT_TAG apache-arrow-nanoarrow-0.8.0 - ) + Nanoarrow + SYSTEM + GIT_REPOSITORY https://github.com/apache/arrow-nanoarrow + GIT_TAG apache-arrow-nanoarrow-0.8.0) FetchContent_MakeAvailable(Nanoarrow) + # Suppress Nanoarrow warnings. + target_compile_options(nanoarrow_static PRIVATE -w) + target_compile_options(nanoarrow_shared PRIVATE -w) + + # arrow-nanoarrow creates a self-referential symlink: + # python/subprojects/arrow-nanoarrow -> ../.. + if(IS_SYMLINK "${nanoarrow_SOURCE_DIR}/python/subprojects/arrow-nanoarrow") + file(REMOVE "${nanoarrow_SOURCE_DIR}/python/subprojects/arrow-nanoarrow") + endif() +endif() +if(VORTEX_BUILD_TESTS) FetchContent_Declare( Catch + SYSTEM GIT_REPOSITORY https://github.com/catchorg/Catch2.git - GIT_TAG v3.8.1 - ) + GIT_TAG v3.8.1) FetchContent_MakeAvailable(Catch) - include(Catch) + # By path: Catch2 exports its module directory only to the directory that + # populated it, which may be vortex-ffi/test. + include("${catch_SOURCE_DIR}/extras/Catch.cmake") + + # Catch2's headers and library must agree on standard-dependent features. + target_compile_features(Catch2 PRIVATE cxx_std_20) target_compile_definitions(Catch2 PRIVATE CATCH_CONFIG_NO_POSIX_SIGNALS) - include(CTest) + # Embedding parents must also call enable_testing() for ctest to find these. + enable_testing() add_subdirectory(tests) endif() + +if(VORTEX_BUILD_EXAMPLES) + add_subdirectory(examples) +endif() diff --git a/lang/cpp/README.md b/lang/cpp/README.md index fe9d2527e63..5159bb9207b 100644 --- a/lang/cpp/README.md +++ b/lang/cpp/README.md @@ -1,45 +1,178 @@ # Vortex C++ bindings -For a usage guide, see docs/api/cpp/index.rst. +A C++20 API for reading and writing Vortex files. See the +[C++ quickstart](../../docs/getting-started/cpp.rst) for API examples. -## Requirements +## Build -- CMake 3.10+ -- C++20 compiler (C++23 compiler for tests). -- Rust toolchain for building `vortex-ffi`. +Requires CMake 3.25+, C/C++ compilers, and Cargo/rustc 1.95+ on `PATH`. -## Build +From a Vortex checkout: + +```sh +cmake -S lang/cpp -B build/cpp -DCMAKE_BUILD_TYPE=Release +cmake --build build/cpp --parallel +``` + +CMake runs Cargo for you. Configuration and builds download uncached dependencies. + +**Native builds only:** GNU/Linux x86_64 and aarch64, plus macOS arm64 for standalone development. +Cross-compilation, universal binaries, Windows, musl, and shared Vortex targets are unsupported. + +## Embed in a CMake project + +Vendor or fetch a pinned, complete checkout: + +```cmake +add_subdirectory(path/to/vortex vortex) +target_link_libraries(my_cpp_target PRIVATE Vortex::cpp_static) +target_link_libraries(my_c_target PRIVATE Vortex::ffi_static) +``` + +Vortex leaves parent build settings unchanged. Its archives are position-independent. + +## Build options + +Pass `-D

") summary_text = ( - f"{engine} / {file_format} ({group_performance}, {significant_improvements}↑ {significant_regressions}↓)" + f"{engine} / {file_format} / {unit} " + f"({group_performance}, {significant_improvements}↑ {significant_regressions}↓)" ) print(f"{summary_text}") print("") @@ -938,7 +1314,8 @@ def main() -> None: display_df.to_markdown( index=False, tablefmt="github", - floatfmt=".2f", + disable_numparse=True, + colalign=tuple("left" for _ in display_df.columns), ) ) print("") diff --git a/scripts/compare-file-sizes.py b/scripts/compare-file-sizes.py index a4a40538e47..6808269980c 100644 --- a/scripts/compare-file-sizes.py +++ b/scripts/compare-file-sizes.py @@ -39,7 +39,7 @@ def format_pct_change(pct: float) -> str: return f"{sign}{pct:.1f}%" -def main(): +def main() -> None: parser = argparse.ArgumentParser(description="Compare file sizes between base and HEAD") parser.add_argument("base_file", help="Base JSONL file") parser.add_argument("head_file", help="HEAD JSONL file") diff --git a/scripts/compress-split.py b/scripts/compress-split.py new file mode 100755 index 00000000000..e216ff52bcf --- /dev/null +++ b/scripts/compress-split.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +""" +Run compress-bench once per dataset, drop OS cache between datasets, +merte outputs. +""" + +import argparse +import glob +import re +import subprocess +from pathlib import Path + +SCRIPT_DIR = Path(__file__).resolve().parent +BINARY = "target/release_debug/compress-bench" +PARTS_DIR = Path("parts") + +# TODO(myrrc): we should also drop CUDA allocator caches + + +def drop_os_caches() -> None: + try: + subprocess.run(["sync"], check=True) + subprocess.run( + ["sudo", "-n", "sh", "-c", "echo 3 > /proc/sys/vm/drop_caches"], + check=True, + capture_output=True, + ) + except (OSError, subprocess.CalledProcessError): + pass + + +def list_datasets(gpu_decompress: bool) -> list[str]: + cmd = [BINARY, "--print-datasets"] + if gpu_decompress: + cmd.append("--gpu-decompress") + result = subprocess.run(cmd, check=True, capture_output=True, text=True) + return [line.strip() for line in result.stdout.splitlines() if line.strip()] + + +def run_datasets(formats: str, emit_ingest_records: bool, gpu_decompress: bool) -> list[str]: + PARTS_DIR.mkdir(parents=True, exist_ok=True) + failures: list[str] = [] + for i, dataset in enumerate(list_datasets(gpu_decompress)): + drop_os_caches() + + args = [ + "bash", + str(SCRIPT_DIR / "bench-taskset.sh"), + BINARY, + "--datasets", + f"^{re.escape(dataset)}$", + ] + if gpu_decompress: + # GPU mode fixes its own formats + args.append("--gpu-decompress") + else: + args += ["--formats", formats] + args += ["-d", "gh-json", "-o", str(PARTS_DIR / f"{i}.gh.json")] + if emit_ingest_records: + args += ["--ingest-jsonl", str(PARTS_DIR / f"{i}.ingest.jsonl")] + print("+", " ".join(args), flush=True) + + result = subprocess.run(args, check=not gpu_decompress) + if result.returncode != 0: + failures.append(dataset) + return failures + + +def merge(pattern: str, out_path: str) -> None: + lines: list[str] = [] + for path in sorted(glob.glob(pattern)): + with open(path, encoding="utf-8") as handle: + for line in handle: + line = line.strip() + if line: + lines.append(line) + Path(out_path).write_text("".join(line + "\n" for line in lines), encoding="utf-8") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--formats", + default="arrow-ipc,parquet,vortex", + help="comma-separated formats to forward to compress-bench", + ) + parser.add_argument( + "--emit-ingest-records", + action="store_true", + help="merge --ingest-jsonl records into results.ingest.jsonl", + ) + parser.add_argument( + "--gpu-decompress", + action="store_true", + help="run the GPU decompression suite (forwards --gpu-decompress, ignores --formats)", + ) + args = parser.parse_args() + + failures = run_datasets(args.formats, args.emit_ingest_records, args.gpu_decompress) + merge(f"{PARTS_DIR}/*.gh.json", "results.json") + if args.emit_ingest_records: + merge(f"{PARTS_DIR}/*.ingest.jsonl", "results.ingest.jsonl") + + if failures: + raise SystemExit("GPU decompression failed for: " + ", ".join(failures)) + + +if __name__ == "__main__": + main() diff --git a/scripts/cudf-parquet-read.py b/scripts/cudf-parquet-read.py new file mode 100644 index 00000000000..abe14532c2a --- /dev/null +++ b/scripts/cudf-parquet-read.py @@ -0,0 +1,121 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +"""Times a full GPU Parquet read with cuDF. + +`cudf.read_parquet` performs the entire read on the device: page header decode, +codec decompression, dictionary/RLE/plain decoding and column assembly. That makes it +the like-for-like opponent for the Vortex GPU backend, which also decodes all the way +to canonical arrays on device. + +Timing excludes interpreter start, `import cudf`, CUDA context creation and any JIT +warm-up, all of which are paid once per process and are not part of a read. A warm-up +read runs first for exactly that reason. + +Emits one JSON object on stdout so the benchmark can parse it. +""" + +import argparse +import json +import sys +import time +from datetime import date +from typing import TypeVar + +_T = TypeVar("_T") + + +def synchronize() -> None: + """Block until queued device work finishes. + + `cudf.read_parquet` returns a materialized DataFrame, but synchronizing explicitly + keeps the measurement honest if that ever stops being true. + """ + try: + import cupy + + cupy.cuda.runtime.deviceSynchronize() + except ImportError: + pass + + +def normalize(frame: _T) -> _T: + """Collapses representation differences that are not value differences. + + A Parquet DATE column comes back from pyarrow as a column of `datetime.date` + objects but from cuDF as `datetime64[s]`. Those hold the same instants, yet + `check_dtype=False` does not bridge them because one side is `object`, so the + comparison reports every row as different. Coercing both sides to datetime64 + compares the dates themselves. + """ + import pandas as pd + + for name in frame.columns: + column = frame[name] + if column.dtype == object and len(column) and isinstance(column.iloc[0], date): + frame[name] = pd.to_datetime(column) + return frame + + +def verify(path: str, frame) -> None: + """Fails unless the GPU read matches a CPU Parquet read of the same file.""" + import pandas as pd + from pandas.testing import assert_frame_equal + + expected = normalize(pd.read_parquet(path)) + actual = normalize(frame.to_pandas()) + + # cuDF and pyarrow can land on different-but-equivalent dtypes (nullable vs numpy + # backed, for instance), so compare values and leave dtype policy out of it. + assert_frame_equal(actual, expected, check_dtype=False) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("path", help="Parquet file to read") + parser.add_argument("--iterations", type=int, default=1, help="timed reads to perform") + parser.add_argument( + "--verify", + action="store_true", + help="cross-check the GPU read against a CPU Parquet read", + ) + args = parser.parse_args() + + import cudf + + # Warm-up: pays CUDA context creation and any first-call JIT so they stay out of + # the timed reads below. + warmup = cudf.read_parquet(args.path) + synchronize() + + if args.verify: + verify(args.path, warmup) + + rows, columns = warmup.shape + del warmup + + runs_ns = [] + for _ in range(max(args.iterations, 1)): + start = time.perf_counter_ns() + frame = cudf.read_parquet(args.path) + synchronize() + runs_ns.append(time.perf_counter_ns() - start) + del frame + + json.dump( + { + "min_ns": min(runs_ns), + "runs_ns": runs_ns, + "rows": int(rows), + "columns": int(columns), + "verified": bool(args.verify), + }, + sys.stdout, + ) + sys.stdout.write("\n") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/post-ingest.py b/scripts/post-ingest.py index e349048f04b..00b71ae2ff6 100755 --- a/scripts/post-ingest.py +++ b/scripts/post-ingest.py @@ -7,45 +7,18 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright the Vortex contributors -"""Ingest a `--gh-json-v3` JSONL file into the benchmarks store. - -Reads bare v3 records from a JSONL file produced by `vortex-bench --gh-json-v3` -and fills the `commit` fields by shelling out to `git show`. Two mutually -exclusive ingest modes select the destination substrate: - -- `--server ` (v3): wraps the records in an envelope and POSTs to - `/api/ingest` with a bearer token. Standard library only -- urllib, - json, subprocess. -- `--postgres ` (v4): computes the server-internal `measurement_id` - locally (via `_measurement_id.py`) and upserts directly into the RDS - Postgres tables with `INSERT ... ON CONFLICT (measurement_id) DO UPDATE`, - over a verify-full TLS connection, authenticating with an RDS IAM auth token - when the DSN carries no password. Requires `psycopg`, `boto3`, and `xxhash` - from the project environment; these are imported lazily inside the - `--postgres` code path so the v3 `--server` path stays standard-library-only - under a bare `python3` (CI invokes `python3 scripts/post-ingest.py`, not - `uv run`, and the v3 path is in production until the Phase 5 cutover). The - PEP 723 metadata block below intentionally keeps `dependencies = []` for that - reason; run the `--postgres` mode from the repo's uv environment. - -The default envelope size (60 MiB, just under the server's 64 MiB body limit) -is sized so a single JSONL run normally posts in one envelope -- preserving the -"per-file all-or-nothing" contract the server documents. If the JSONL is large -enough that splitting kicks in, the script emits a warning and proceeds with the -chunked semantics (per-chunk commit, mid-chunk failure leaves earlier chunks -ingested; subsequent retries re-upsert via the server's ON CONFLICT idempotency -on `measurement_id`). The `--postgres` mode applies a whole JSONL file in one -transaction (all-or-nothing), with no chunking. - -Wire-contract pointers (kept in sync as a coordinated change per -the `vortex-data/benchmarks-website` repo's `AGENTS.md`): - -- the `vortex-data/benchmarks-website` repo's `server/src/records.rs` - envelope + per-record wire - shapes that the server deserializes. -- `vortex-bench/src/v3.rs` - bare-record producer that writes the JSONL this - script wraps. -- the `vortex-data/benchmarks-website` repo's `server/src/schema.rs` - `SCHEMA_VERSION` source of - truth that the `SCHEMA_VERSION` constant below MUST equal at every bump. +"""Ingest benchmark JSONL records into the v4 Postgres store. + +`vortex-bench --ingest-jsonl` writes the current versioned JSONL record shape. This script +fills commit metadata from `git show`, computes the internal `measurement_id`, +then atomically upserts the whole file into the RDS fact tables. + +The wire contract is shared with the benchmarks-website repository: + +- `vortex-bench/src/v3.rs` owns the producer record shape. +- `benchmarks-website/CONTRACT.md` documents the transport and schema contract. +- `benchmarks-website/web/lib/schema-version.ts` is the `SCHEMA_VERSION` source + of truth that the constant below must match. """ from __future__ import annotations @@ -56,62 +29,43 @@ import os import subprocess import sys -import urllib.error import urllib.request +from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor -from datetime import UTC, datetime from pathlib import Path +from types import ModuleType +from typing import TypeVar -# MUST equal the `vortex-data/benchmarks-website` repo's `server/src/schema.rs::SCHEMA_VERSION`. -# Bumping this is a coordinated change across schema.rs, records.rs, v3.rs, -# and this script. See the `vortex-data/benchmarks-website` repo's `AGENTS.md` ("Wire shapes are a -# coordinated change") for the full list of coupled sites. -SCHEMA_VERSION = 1 -# Default sized to fit comfortably under the server's 64 MiB ingest body -# limit while leaving headroom for HTTP and JSON framing overhead. The -# point is to keep a normal JSONL run in one envelope so the documented -# "per-file all-or-nothing" contract holds. -DEFAULT_MAX_ENVELOPE_BYTES = 60 * 1024 * 1024 +# MUST equal `benchmarks-website/web/lib/schema-version.ts::SCHEMA_VERSION`. +# Bumping this is a coordinated change across the website contract, v3.rs, and +# this script. +SCHEMA_VERSION = 3 def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description=( - "Ingest a v3 JSONL records file: --server POSTs a v3 envelope to " - "/api/ingest; --postgres upserts directly into the RDS Postgres tables." - ), - ) + parser = argparse.ArgumentParser(description="Ingest benchmark JSONL records into RDS Postgres.") parser.add_argument( "jsonl_path", type=Path, - help="Path to the JSONL file written by vortex-bench --gh-json-v3.", + help="Path to the JSONL file written by vortex-bench --ingest-jsonl.", ) - mode = parser.add_mutually_exclusive_group(required=True) - mode.add_argument( - "--server", - help=( - "v3 mode: server base URL, e.g. http://localhost:8080. Wraps the " - "records in an envelope and POSTs to /api/ingest with a " - "bearer token. Mutually exclusive with --postgres." - ), - ) - mode.add_argument( + parser.add_argument( "--postgres", metavar="DSN", + required=True, help=( - "v4 mode: libpq DSN for the RDS Postgres ingest target, e.g. " + "Libpq DSN for the RDS Postgres ingest target, e.g. " "'postgresql://bench_ingest@host:5432/benchmarks?sslmode=verify-full" "&sslrootcert=/path/rds-ca.pem'. Computes measurement_id locally and " "upserts via INSERT ... ON CONFLICT DO UPDATE. If the DSN carries no " - "password, an RDS IAM auth token is minted for the DSN's user. " - "Mutually exclusive with --server." + "password, an RDS IAM auth token is minted for the DSN's user." ), ) parser.add_argument( "--region", default=None, help=( - "AWS region for RDS IAM token minting (--postgres mode). Precedence: " + "AWS region for RDS IAM token minting. Precedence: " "this explicit --region, then the boto3 session region, then the region " "parsed from the RDS hostname." ), @@ -121,18 +75,6 @@ def parse_args() -> argparse.Namespace: required=True, help="40-hex commit SHA. Usually ${{ github.sha }} in CI.", ) - parser.add_argument( - "--benchmark-id", - default=None, - help="Run identifier echoed back in run_meta.benchmark_id. Required for " - "--server mode; unused by --postgres mode (the v4 tables have no " - "benchmark_id column).", - ) - parser.add_argument( - "--token-env", - default="INGEST_BEARER_TOKEN", - help="Env var holding the bearer token (default: INGEST_BEARER_TOKEN).", - ) parser.add_argument( "--repo-url", default="https://github.com/vortex-data/vortex", @@ -148,13 +90,7 @@ def parse_args() -> argparse.Namespace: "--timeout", type=float, default=30.0, - help="HTTP timeout in seconds (default: 30).", - ) - parser.add_argument( - "--max-envelope-bytes", - type=int, - default=DEFAULT_MAX_ENVELOPE_BYTES, - help=(f"Maximum encoded JSON bytes per POST before splitting records (default: {DEFAULT_MAX_ENVELOPE_BYTES})."), + help="Site cache refresh timeout in seconds (default: 30).", ) return parser.parse_args() @@ -216,120 +152,13 @@ def build_commit(sha: str, repo_url: str, git_dir: Path | None) -> dict: } -def build_envelope(run_meta: dict, commit: dict, records: list[dict]) -> dict: - return { - "run_meta": run_meta, - "commit": commit, - "records": records, - } - - -def encode_envelope(envelope: dict) -> bytes: - return json.dumps(envelope, separators=(",", ":")).encode("utf-8") - - -def chunk_envelopes( - run_meta: dict, - commit: dict, - records: list[dict], - max_envelope_bytes: int, -) -> list[tuple[dict, bytes]]: - if max_envelope_bytes <= 0: - raise SystemExit("--max-envelope-bytes must be positive") - if not records: - envelope = build_envelope(run_meta, commit, []) - return [(envelope, encode_envelope(envelope))] - - # Cost model: re-encoding the growing envelope per record was O(N^2) in - # the size of the JSONL on the CI hot path. Instead track each record's - # encoded size once (`json.dumps(record)`) and reason about the - # cumulative chunk size via that plus per-record separator overhead - # plus a one-time envelope shell. Cross-check at chunk flush time by - # encoding the actual chunk and `assert len(body) <= cap` so a - # misestimation surfaces here, not at the server. - shell = build_envelope(run_meta, commit, []) - shell_bytes = len(encode_envelope(shell)) - # `,` between records inside the JSON array. - record_sep_bytes = 1 - - def encoded_size(records_chunk: list[dict]) -> int: - env = build_envelope(run_meta, commit, records_chunk) - return len(encode_envelope(env)) - - encoded_records: list[bytes] = [json.dumps(r, separators=(",", ":")).encode("utf-8") for r in records] - - chunks: list[tuple[dict, bytes]] = [] - batch: list[dict] = [] - batch_payload_bytes = 0 - - def flush_batch() -> None: - env = build_envelope(run_meta, commit, batch) - body = encode_envelope(env) - assert len(body) <= max_envelope_bytes, ( - f"chunk_envelopes invariant violated: {len(body)} > {max_envelope_bytes}" - ) - chunks.append((env, body)) - - for i, record in enumerate(records): - record_bytes = len(encoded_records[i]) - # First record in a fresh batch has no separator before it. - delta = record_bytes if not batch else record_bytes + record_sep_bytes - projected = shell_bytes + batch_payload_bytes + delta - - if not batch and projected > max_envelope_bytes: - # First record in a fresh batch ALONE exceeds the cap. Refuse - # rather than ship an over-cap chunk that the server will 413. - raise SystemExit( - f"single record exceeds --max-envelope-bytes " - f"(record {i} encodes to {record_bytes} bytes; " - f"envelope shell adds {shell_bytes}; cap is {max_envelope_bytes}). " - f"Raise --max-envelope-bytes, or trim the record before posting." - ) +# `psycopg`, `boto3`, and `_measurement_id` (which pulls in `xxhash`) are +# imported lazily below so `--help` and syntax checks need no optional packages. - if batch and projected > max_envelope_bytes: - flush_batch() - batch = [record] - batch_payload_bytes = record_bytes - else: - batch.append(record) - batch_payload_bytes += delta - - if batch: - flush_batch() - return chunks - - -def post(server: str, body: bytes, token: str, timeout: float) -> tuple[int, bytes]: - url = f"{server.rstrip('/')}/api/ingest" - request = urllib.request.Request( - url, - data=body, - method="POST", - headers={ - "Content-Type": "application/json", - "Authorization": f"Bearer {token}", - }, - ) - try: - with urllib.request.urlopen(request, timeout=timeout) as response: - return response.status, response.read() - except urllib.error.HTTPError as exc: - return exc.code, exc.read() - - -# -------------------------------------------------------------------------- -# Postgres dual-write mode (--postgres) -# -------------------------------------------------------------------------- -# The v4 ingest path. `psycopg`, `boto3`, and `_measurement_id` (which pulls in -# `xxhash`) are imported lazily inside the functions below so the v3 `--server` -# path above stays standard-library-only under a bare `python3`. - -# Per-`kind` record field sets, mirroring the `#[serde(deny_unknown_fields)]` -# structs in the `vortex-data/benchmarks-website` repo's `server/src/records.rs`. `required` plus -# `optional` is the exact set of fields a record of that `kind` may carry -# (besides `kind` itself); an unknown field is rejected loudly, preserving the -# v3 server's deny_unknown_fields behavior so producer/schema drift surfaces -# instead of silently dropping data. `optional` fields default to NULL. +# Per-`kind` record field sets. `required` plus `optional` is the exact set of +# fields a record of that `kind` may carry (besides `kind` itself); unknown +# fields fail loudly so producer/schema drift cannot silently drop data. +# `optional` fields default to NULL. _RECORD_FIELDS: dict[str, tuple[frozenset[str], frozenset[str]]] = { "query_measurement": ( frozenset( @@ -361,11 +190,11 @@ def post(server: str, body: bytes, token: str, timeout: float) -> tuple[int, byt frozenset({"dataset_variant", "env_triple"}), ), "compression_size": ( - frozenset({"commit_sha", "dataset", "format", "value_bytes"}), + frozenset({"commit_sha", "dataset", "format", "value_bytes", "uncompressed_bytes"}), frozenset({"dataset_variant"}), ), "random_access_time": ( - frozenset({"commit_sha", "dataset", "format", "value_ns", "all_runtimes_ns"}), + frozenset({"commit_sha", "dataset", "format", "open_mode", "value_ns", "all_runtimes_ns"}), frozenset({"env_triple"}), ), "vector_search_run": ( @@ -388,10 +217,12 @@ def post(server: str, body: bytes, token: str, timeout: float) -> tuple[int, byt ), } -_MEASUREMENT_ID_MODULE = None +_T = TypeVar("_T") +_MEASUREMENT_ID_MODULE: ModuleType | None = None -def _measurement_id_module(): + +def _measurement_id_module() -> ModuleType: """Lazily load `scripts/_measurement_id.py` by path (cached). Loaded by file path rather than `import _measurement_id` so it resolves @@ -415,8 +246,7 @@ def _measurement_id_module(): def _validate_record_fields(record: object, index: int) -> str: """Validate a record's `kind` and field set; return the `kind`. - Mirrors the server's deny_unknown_fields + required-field deserialization: - a non-object record, an unknown `kind`, an unknown field, or a missing + A non-object record, an unknown `kind`, an unknown field, or a missing required field is a loud error keyed by the record's index. """ if not isinstance(record, dict): @@ -431,7 +261,7 @@ def _validate_record_fields(record: object, index: int) -> str: present = set(record) - {"kind"} unknown = present - required - optional if unknown: - raise SystemExit(f"record {index} ({kind}): unknown field(s) {sorted(unknown)} (deny_unknown_fields)") + raise SystemExit(f"record {index} ({kind}): unknown field(s) {sorted(unknown)}") missing = required - present if missing: raise SystemExit(f"record {index} ({kind}): missing required field(s) {sorted(missing)}") @@ -457,7 +287,7 @@ def _require_finite(value: object, field: str, kind: str, index: int) -> None: ) -# i32/i64 bounds enforced by the Postgres INTEGER/BIGINT columns and the v3 serde. +# i32/i64 bounds enforced by the Postgres INTEGER/BIGINT columns. _INT32_MIN, _INT32_MAX = -(2**31), 2**31 - 1 _INT64_MIN, _INT64_MAX = -(2**63), 2**63 - 1 @@ -467,11 +297,10 @@ def _require_int(value: object, field: str, kind: str, index: int, *, bits: int) The integer columns bind straight to INTEGER/BIGINT; psycopg adapts a Python float to float8 and Postgres assignment-casts (rounds) it, so a JSON float - would silently persist a rounded value where the v3 server's serde `i32`/`i64` - rejects it. An out-of-range integer would otherwise fail late as an uncaught + would silently persist a rounded value. An out-of-range integer would otherwise fail late as an uncaught `struct.error` (i32 hash dims via `_write_i32`) or a raw Postgres 22003 overflow; validate the type AND width here so a malformed scalar fails loud - (record-indexed) instead, matching the v3 serde boundary. + with its record index. """ lo, hi = (_INT32_MIN, _INT32_MAX) if bits == 32 else (_INT64_MIN, _INT64_MAX) if isinstance(value, bool) or not isinstance(value, int): @@ -484,11 +313,11 @@ def _require_int_list(value: object, field: str, kind: str, index: int) -> None: """Raise loudly unless `value` is a JSON array of plain (non-bool) i64 integers. Guards the `all_runtimes_ns` -> `bigint[]` bind: the explicit `::bigint[]` - cast is permissive in ways the v3 server's `Vec` serde is not. psycopg + cast accepts values that are not valid JSON integer arrays. psycopg sends the string `"{}"` as text and the cast parses it into an empty array, a `[1, null]` list adapts to `{1,NULL}`, and an out-of-i64 element hits a raw - Postgres 22003 -- each diverges from or fails differently than the v3 path. - Validate the element type AND i64 range so a malformed value fails loud. + Postgres 22003. Validate the element type AND i64 range so a malformed value + fails loud. """ if not isinstance(value, list) or any(isinstance(x, bool) or not isinstance(x, int) for x in value): raise SystemExit(f"record {index} ({kind}): {field} must be a JSON array of integers, got {value!r}") @@ -497,7 +326,7 @@ def _require_int_list(value: object, field: str, kind: str, index: int) -> None: def _require_str(value: object, field: str, kind: str, index: int) -> None: - """Raise loudly unless `value` is a string. Mirrors the v3 serde `String` fields.""" + """Raise loudly unless `value` is a string.""" if not isinstance(value, str): raise SystemExit(f"record {index} ({kind}): {field} must be a string, got {value!r}") @@ -511,8 +340,7 @@ def _require_opt_str(value: object, field: str, kind: str, index: int) -> None: def _memory_quartet_consistent(r: dict) -> bool: """The four `query_measurements` memory columns are all set or all absent. - Mirrors `ingest.rs::memory_quartet_consistent`: a partial quartet is a - validation error, not a half-populated row. + A partial quartet is a validation error, not a half-populated row. """ present = [ r.get("peak_physical") is not None, @@ -523,9 +351,8 @@ def _memory_quartet_consistent(r: dict) -> bool: return not any(present) or all(present) -# Per-kind field -> type token, mirroring the v3 server's typed serde boundary in -# `records.rs`. The direct-to-Postgres writer can no longer lean on server-side -# deserialization, so every field is type/range-validated here. Tokens: +# Per-kind field -> type token. Every field is type/range-validated before the +# direct Postgres write. Tokens: # "str" required String "opt_str" Option # "i32" required i32 "i64" required i64 # "opt_i64" Option "i64_list" Vec (all_runtimes_ns) @@ -566,11 +393,13 @@ def _memory_quartet_consistent(r: dict) -> bool: ("dataset_variant", "opt_str"), ("format", "str"), ("value_bytes", "i64"), + ("uncompressed_bytes", "i64"), ), "random_access_time": ( ("commit_sha", "str"), ("dataset", "str"), ("format", "str"), + ("open_mode", "str"), ("value_ns", "i64"), ("all_runtimes_ns", "i64_list"), ("env_triple", "opt_str"), @@ -593,13 +422,11 @@ def _memory_quartet_consistent(r: dict) -> bool: def _validate_record_values(record: dict, kind: str, index: int) -> None: - """Validate every field's type/range against the v3 server's serde boundary. + """Validate every field's type/range before the Postgres write. - Runs in `ingest_postgres`'s loop (where the record index is known) so every - failure is reported as `record {index} ({kind}): ...`, matching the v3 - server's indexed per-record errors. Drives field type/range checks off - `_FIELD_TYPES`, then applies the semantic checks the type alone does not cover - (the storage enum + memory quartet for query_measurements). + Runs in `ingest_postgres`'s loop, where the record index is known. It drives + type/range checks from `_FIELD_TYPES`, then applies semantic checks the type + alone does not cover (the storage enum + memory quartet for query_measurements). """ for field, typ in _FIELD_TYPES[kind]: if typ == "str": @@ -627,6 +454,10 @@ def _validate_record_values(record: dict, kind: str, index: int) -> None: raise SystemExit( f"record {index} (query_measurement): memory fields must be populated together (all four or none)" ) + elif kind == "random_access_time" and record["open_mode"] not in ("cached", "reopen"): + raise SystemExit( + f"record {index} (random_access_time): open_mode must be 'cached' or 'reopen', got {record['open_mode']!r}" + ) def _upsert_returning_was_update(conn, sql: str, params: tuple) -> bool: @@ -641,7 +472,7 @@ def _upsert_returning_was_update(conn, sql: str, params: tuple) -> bool: the flag from the single atomic statement removes that window. A duplicate `measurement_id` within ONE transaction is also handled correctly: the second upsert locks the tuple the first created (stamping its xmax), so it is - classified as an update, matching the v3 server. `sql` must end with + classified as an update. `sql` must end with `RETURNING (xmax = 0) AS inserted`. """ row = conn.execute(sql, params).fetchone() @@ -649,7 +480,7 @@ def _upsert_returning_was_update(conn, sql: str, params: tuple) -> bool: def _insert_query_measurement(conn, mid_mod, r: dict) -> bool: - """Upsert a `query_measurements` row. Mirrors `ingest.rs::insert_query_measurement`. + """Upsert a `query_measurements` row. Record values are validated by `_validate_record_values` before dispatch. """ @@ -712,7 +543,7 @@ def _insert_query_measurement(conn, mid_mod, r: dict) -> bool: def _insert_compression_time(conn, mid_mod, r: dict) -> bool: - """Upsert a `compression_times` row. Mirrors `ingest.rs::insert_compression_time`.""" + """Upsert a `compression_times` row.""" mid = mid_mod.measurement_id_compression_time( commit_sha=r["commit_sha"], dataset=r["dataset"], @@ -749,7 +580,7 @@ def _insert_compression_time(conn, mid_mod, r: dict) -> bool: def _insert_compression_size(conn, mid_mod, r: dict) -> bool: - """Upsert a `compression_sizes` row. Mirrors `ingest.rs::insert_compression_size`.""" + """Upsert a `compression_sizes` row.""" mid = mid_mod.measurement_id_compression_size( commit_sha=r["commit_sha"], dataset=r["dataset"], @@ -761,11 +592,12 @@ def _insert_compression_size(conn, mid_mod, r: dict) -> bool: """ INSERT INTO compression_sizes ( measurement_id, commit_sha, dataset, dataset_variant, - format, value_bytes - ) VALUES (%s, %s, %s, %s, %s, %s) + format, value_bytes, uncompressed_bytes + ) VALUES (%s, %s, %s, %s, %s, %s, %s) ON CONFLICT (measurement_id) DO UPDATE SET - commit_sha = excluded.commit_sha, - value_bytes = excluded.value_bytes + commit_sha = excluded.commit_sha, + value_bytes = excluded.value_bytes, + uncompressed_bytes = excluded.uncompressed_bytes RETURNING (xmax = 0) AS inserted """, ( @@ -775,26 +607,29 @@ def _insert_compression_size(conn, mid_mod, r: dict) -> bool: r.get("dataset_variant"), r["format"], r["value_bytes"], + r["uncompressed_bytes"], ), ) def _insert_random_access(conn, mid_mod, r: dict) -> bool: - """Upsert a `random_access_times` row. Mirrors `ingest.rs::insert_random_access`.""" + """Upsert a `random_access_times` row.""" mid = mid_mod.measurement_id_random_access( commit_sha=r["commit_sha"], dataset=r["dataset"], format=r["format"], + open_mode=r["open_mode"], ) return _upsert_returning_was_update( conn, """ INSERT INTO random_access_times ( - measurement_id, commit_sha, dataset, format, + measurement_id, commit_sha, dataset, format, open_mode, value_ns, all_runtimes_ns, env_triple - ) VALUES (%s, %s, %s, %s, %s, %s::bigint[], %s) + ) VALUES (%s, %s, %s, %s, %s, %s, %s::bigint[], %s) ON CONFLICT (measurement_id) DO UPDATE SET commit_sha = excluded.commit_sha, + open_mode = excluded.open_mode, value_ns = excluded.value_ns, all_runtimes_ns = excluded.all_runtimes_ns, env_triple = excluded.env_triple @@ -805,6 +640,7 @@ def _insert_random_access(conn, mid_mod, r: dict) -> bool: r["commit_sha"], r["dataset"], r["format"], + r["open_mode"], r["value_ns"], r["all_runtimes_ns"], r.get("env_triple"), @@ -813,7 +649,7 @@ def _insert_random_access(conn, mid_mod, r: dict) -> bool: def _insert_vector_search(conn, mid_mod, r: dict) -> bool: - """Upsert a `vector_search_runs` row. Mirrors `ingest.rs::insert_vector_search`. + """Upsert a `vector_search_runs` row. `threshold` is validated finite by `_validate_record_values` before dispatch. """ @@ -874,7 +710,7 @@ def _insert_vector_search(conn, mid_mod, r: dict) -> bool: def _upsert_commit(conn, commit: dict) -> None: - """Upsert the `commits` dim row. Mirrors `ingest.rs::upsert_commit`.""" + """Upsert the `commits` dimension row.""" conn.execute( """ INSERT INTO commits ( @@ -905,20 +741,17 @@ def _upsert_commit(conn, commit: dict) -> None: ) -# Mirrors the v3 server's `WRITE_CONFLICT_ATTEMPTS` -# (the `vortex-data/benchmarks-website` repo's `server/src/ingest.rs`). The DuckDB ingest wrapped -# `apply_envelope_once` in `retry_write_conflicts`; the DuckDB -> Postgres substrate change -# must preserve that behavior, since the dual-write CI runs ~14 concurrent writers. +# CI runs concurrent writers whose upserts can touch commits and dimensions in +# conflicting orders. Retrying transaction-level deadlocks and serialization +# failures keeps each JSONL file all-or-nothing. _WRITE_CONFLICT_ATTEMPTS = 128 -def _retry_write_conflicts(op): - """Retry `op` on a Postgres write conflict, mirroring the v3 server's `retry_write_conflicts`. +def _retry_write_conflicts(op: Callable[[], _T]) -> _T: + """Retry `op` on a Postgres write conflict. - The v3 DuckDB ingest retried on write conflicts up to `WRITE_CONFLICT_ATTEMPTS` times; the - Postgres writer must preserve that, because the dual-write CI runs many concurrent writers - whose row-level `ON CONFLICT DO UPDATE` upserts touching the same `commits` / dim rows in - conflicting orders can deadlock. The retryable analogs on Postgres are deadlock + Row-level `ON CONFLICT DO UPDATE` upserts touching the same commits or + dimensions in conflicting orders can deadlock. The retryable Postgres errors are deadlock (`SQLSTATE 40P01`) and serialization failure (`40001`); both abort one transaction cleanly, so re-running the whole transaction is safe. A non-retryable error (e.g. a validation `SystemExit`) propagates immediately. Returns `op`'s value on the first success. @@ -931,31 +764,23 @@ def _retry_write_conflicts(op): except (pg_errors.DeadlockDetected, pg_errors.SerializationFailure): # The failing `op`'s `with conn.transaction()` block already rolled back, so the # connection is idle and the whole transaction can be retried. Re-raise on the - # final attempt (mirrors v3 returning the error after the last try; the bare - # retry loop mirrors v3's `std::thread::yield_now()` between attempts). + # final attempt. if attempt >= _WRITE_CONFLICT_ATTEMPTS: raise raise AssertionError("unreachable: _retry_write_conflicts exited without return or raise") def ingest_postgres(conn, commit: dict, records: list[dict]) -> tuple[int, int]: - """Upsert a commit and its records into Postgres, retrying on write conflicts. - - Wraps `_ingest_postgres_once` in `_retry_write_conflicts` (mirroring the v3 server's - `apply_envelope = retry_write_conflicts(apply_envelope_once)`). Returns `(inserted, updated)` - aggregated across all fact tables. - """ + """Upsert a commit and its records into Postgres, retrying on write conflicts.""" mid_mod = _measurement_id_module() return _retry_write_conflicts(lambda: _ingest_postgres_once(conn, commit, records, mid_mod)) def _ingest_postgres_once(conn, commit: dict, records: list[dict], mid_mod) -> tuple[int, int]: - """Upsert a commit and its records into Postgres in one transaction (a single attempt). + """Upsert a commit and its records in one transaction (a single attempt). - Mirrors the v3 server's `apply_envelope_once`: upsert `commits` first, then each fact - record, classifying each as inserted or updated. Any validation failure rolls the whole - transaction back (all-or-nothing). Returns `(inserted, updated)` aggregated across all - fact tables. + Upsert `commits` first, then each fact record while classifying it as inserted + or updated. Any validation failure rolls the whole transaction back. """ inserted = 0 updated = 0 @@ -966,7 +791,7 @@ def _ingest_postgres_once(conn, commit: dict, records: list[dict], mid_mod) -> t if record["commit_sha"] != commit["sha"]: raise SystemExit( f"record {idx} ({kind}): commit_sha {record['commit_sha']!r} does not " - f"match envelope commit.sha {commit['sha']!r}" + f"match the requested commit SHA {commit['sha']!r}" ) _validate_record_values(record, kind, idx) if _APPLY_RECORD[kind](conn, mid_mod, record): @@ -1011,7 +836,7 @@ def _rds_iam_token(*, host: str, port: int, user: str, region: str | None) -> st _INGEST_ROLE = "bench_ingest" -def connect_postgres(dsn: str, region: str | None): +def connect_postgres(dsn: str, region: str | None) -> object: """Open a psycopg connection to the RDS Postgres ingest target. Enforces the ingest contract: verify-full TLS, and authentication only as the @@ -1160,11 +985,8 @@ def _main_postgres(args: argparse.Namespace) -> int: separators=(",", ":"), ) ) - # Best-effort site-cache refresh after a successful write. No-op unless both - # env vars are set (so the script stays inert until the ops wiring lands), - # and it can never fail the ingest. The ops prerequisite (setting the two env - # vars in Vercel and as GitHub secrets/vars) is documented in the "Ops - # prerequisite" section of .big-plans/ct__bench-v4-uiux-r3-design.md. + # Best-effort site-cache refresh after a successful write. It runs only + # when both environment variables are configured and can never fail ingest. base_url = os.environ.get("BENCH_SITE_BASE_URL") revalidate_token = os.environ.get("BENCH_REVALIDATE_TOKEN") if base_url and revalidate_token: @@ -1172,83 +994,8 @@ def _main_postgres(args: argparse.Namespace) -> int: return 0 -def _main_server(args: argparse.Namespace) -> int: - if args.benchmark_id is None: - print("error: --benchmark-id is required in --server mode", file=sys.stderr) - return 2 - token = os.environ.get(args.token_env) - if not token: - print( - f"error: env var {args.token_env} is not set", - file=sys.stderr, - ) - return 2 - - records = read_records(args.jsonl_path) - commit = build_commit(args.commit_sha, args.repo_url, args.git_dir) - run_meta = { - "benchmark_id": args.benchmark_id, - "schema_version": SCHEMA_VERSION, - "started_at": datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ"), - } - - chunks = chunk_envelopes(run_meta, commit, records, args.max_envelope_bytes) - if len(chunks) > 1: - print( - f"warning: JSONL exceeds --max-envelope-bytes ({args.max_envelope_bytes} B); " - f"splitting into {len(chunks)} envelopes. Each chunk commits independently on " - "the server, so a mid-stream failure leaves earlier chunks ingested. " - "Re-run on failure to re-upsert via measurement_id ON CONFLICT.", - file=sys.stderr, - ) - inserted = 0 - updated = 0 - raw_bodies: list[str] = [] - for idx, (envelope, body) in enumerate(chunks, start=1): - if len(chunks) > 1: - print( - f"POST chunk {idx}/{len(chunks)} records={len(envelope['records'])} bytes={len(body)}", - file=sys.stderr, - ) - status, response = post(args.server, body, token, args.timeout) - body_text = response.decode("utf-8", errors="replace") - - if status >= 400: - print( - f"error: POST {args.server}/api/ingest chunk {idx}/{len(chunks)} -> {status}\n{body_text}", - file=sys.stderr, - ) - return 1 - - try: - parsed = json.loads(body_text) - inserted += int(parsed.get("inserted", 0)) - updated += int(parsed.get("updated", 0)) - except (TypeError, ValueError, json.JSONDecodeError): - raw_bodies.append(body_text) - - if raw_bodies: - print("\n".join(raw_bodies)) - else: - print( - json.dumps( - { - "chunks": len(chunks), - "records": len(records), - "inserted": inserted, - "updated": updated, - }, - separators=(",", ":"), - ) - ) - return 0 - - def main() -> int: - args = parse_args() - if args.postgres is not None: - return _main_postgres(args) - return _main_server(args) + return _main_postgres(parse_args()) if __name__ == "__main__": diff --git a/scripts/random-access-split.py b/scripts/random-access-split.py index 92b44aee7e0..3a5a41740b7 100755 --- a/scripts/random-access-split.py +++ b/scripts/random-access-split.py @@ -18,18 +18,32 @@ PARTS_DIR = Path("parts") DATASETS = ["taxi", "feature-vectors", "nested-lists", "nested-structs"] -FORMATS = ["parquet", "lance", "vortex"] +FORMATS = ["arrow-ipc", "parquet", "lance", "vortex"] PATTERNS = ["correlated", "uniform"] OPEN_MODES = ["cached", "reopen"] -def run_combinations(emit_v3: bool) -> None: +def drop_os_caches() -> None: + try: + subprocess.run(["sync"], check=True) + subprocess.run( + ["sudo", "-n", "sh", "-c", "echo 3 > /proc/sys/vm/drop_caches"], + check=True, + capture_output=True, + ) + except (OSError, subprocess.CalledProcessError): + pass + + +def run_combinations(emit_ingest_records: bool) -> None: PARTS_DIR.mkdir(parents=True, exist_ok=True) i = 0 for dataset in DATASETS: for fmt in FORMATS: for pattern in PATTERNS: for open_mode in OPEN_MODES: + drop_os_caches() + args = [ "bash", str(SCRIPT_DIR / "bench-taskset.sh"), @@ -47,8 +61,8 @@ def run_combinations(emit_v3: bool) -> None: "-o", str(PARTS_DIR / f"{i}.gh.json"), ] - if emit_v3: - args += ["--gh-json-v3", str(PARTS_DIR / f"{i}.v3.jsonl")] + if emit_ingest_records: + args += ["--ingest-jsonl", str(PARTS_DIR / f"{i}.ingest.jsonl")] print("+", " ".join(args), flush=True) subprocess.run(args, check=True) i += 1 @@ -79,22 +93,31 @@ def merge(pattern: str, key: Callable[[dict], object], out_path: str) -> None: Path(out_path).write_text("".join(line + "\n" for line in lines), encoding="utf-8") +def ingest_identity(record: dict) -> tuple[object, object, object, object]: + return ( + record["kind"], + record["dataset"], + record["format"], + record["open_mode"], + ) + + def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( - "--v3", + "--emit-ingest-records", action="store_true", - help="merge --gh-json-v3 records into results.v3.jsonl", + help="merge --ingest-jsonl records into results.ingest.jsonl", ) args = parser.parse_args() - run_combinations(args.v3) + run_combinations(args.emit_ingest_records) merge(f"{PARTS_DIR}/*.gh.json", lambda record: record["name"], "results.json") - if args.v3: + if args.emit_ingest_records: merge( - f"{PARTS_DIR}/*.v3.jsonl", - lambda record: (record["kind"], record["dataset"], record["format"]), - "results.v3.jsonl", + f"{PARTS_DIR}/*.ingest.jsonl", + ingest_identity, + "results.ingest.jsonl", ) diff --git a/scripts/s3-download.py b/scripts/s3-download.py index ed2343efd8f..ed8cc297acb 100755 --- a/scripts/s3-download.py +++ b/scripts/s3-download.py @@ -10,7 +10,7 @@ import time -def main(): +def main() -> None: parser = argparse.ArgumentParser(description="Download a file from S3 with retry") parser.add_argument("s3_url", help="S3 URL to download (e.g. s3://bucket/key)") parser.add_argument("output", help="Local output file path") diff --git a/scripts/s3-upload.py b/scripts/s3-upload.py index 6215e0531d4..0fa4deeef05 100755 --- a/scripts/s3-upload.py +++ b/scripts/s3-upload.py @@ -65,7 +65,7 @@ def put_object( return result.returncode == 0 -def main(): +def main() -> None: parser = argparse.ArgumentParser(description="Upload a file to S3 with retry and optional optimistic locking") parser.add_argument("--bucket", required=True, help="S3 bucket name") parser.add_argument("--key", required=True, help="S3 object key") diff --git a/scripts/setup-benchmark.sh b/scripts/setup-benchmark.sh index b9049071478..c02441ef9f7 100755 --- a/scripts/setup-benchmark.sh +++ b/scripts/setup-benchmark.sh @@ -10,41 +10,60 @@ if [ "$EUID" -ne 0 ]; then exit 0 fi +# Write a value to a sysfs file, skipping hosts where the file is absent or read-only +# (e.g. virtualized guests without frequency control). +sysfs_write() { + [[ -n "$2" && -w "$1" ]] && echo "$2" > "$1" || true +} + +# Expand a cpulist like "0-2,8" into one CPU index per line. +expand_cpulist() { + [[ -n "$1" ]] || return 0 + local parts part + IFS=, read -ra parts <<< "$1" + for part in "${parts[@]}"; do + seq "${part%-*}" "${part#*-}" + done +} + # Disable turbo boost so benchmark runs stay at a more stable clock rate. -[[ -w /sys/devices/system/cpu/intel_pstate/no_turbo ]] && echo 1 > /sys/devices/system/cpu/intel_pstate/no_turbo || true -[[ -w /sys/devices/system/cpu/cpufreq/boost ]] && echo 0 > /sys/devices/system/cpu/cpufreq/boost || true +sysfs_write /sys/devices/system/cpu/intel_pstate/no_turbo 1 +sysfs_write /sys/devices/system/cpu/cpufreq/boost 0 + +# Pin the frequency governor to performance and raise the floor to the policy max so +# clocks hold at a constant base clock (turbo is capped above) instead of ramping +# under powersave/schedutil. +for policy in /sys/devices/system/cpu/cpufreq/policy*; do + sysfs_write "$policy/scaling_governor" performance + sysfs_write "$policy/scaling_min_freq" "$(cat "$policy/scaling_max_freq" 2>/dev/null)" +done # Really discourage swapping to disk -sysctl vm.swappiness=0 +sysctl -w vm.swappiness=0 || true swapoff -a || true # Might be worse if a single application uses the OS # https://www.intel.com/content/www/us/en/developer/articles/technical/measuring-impact-of-numa-migrations-on-performance.html -sysctl -w kernel.numa_balancing=0 +sysctl -w kernel.numa_balancing=0 || true # Disable ASLR - https://docs.kernel.org/admin-guide/sysctl/kernel.html#randomize-va-space -sysctl kernel.randomize_va_space=0 +sysctl -w kernel.randomize_va_space=0 || true # This is a desktop optimization, making sure its disabled -sysctl -w kernel.sched_autogroup_enabled=0 +sysctl -w kernel.sched_autogroup_enabled=0 || true + +# Keep armed timers on the CPU that armed them instead of migrating them onto busy +# (i.e. benchmark) CPUs. +sysctl -w kernel.timer_migration=0 || true # Reduce kernel logging to minimum dmesg -n 1 -# Disable some unused services and features -systemctl stop apparmor ModemManager -systemctl disable apparmor ModemManager - -# mask prevents them from being started by other services -systemctl mask ModemManager - -# For apparmor specifically, also teardown loaded profiles. Runners that do not -# have apparmor fully available fail this, which is not a reason to abandon the -# rest of the tuning. -aa-teardown || true - -# Reduce background activity (Ubuntu-specific) +# Stop background services that could wake up mid-benchmark (Ubuntu-specific list). +# Masking prevents them from being restarted by other services. for unit in \ + apparmor \ + ModemManager \ irqbalance \ apt-daily.service \ apt-daily-upgrade.service \ @@ -54,22 +73,79 @@ for unit in \ motd-news.timer \ apport do - # Not every unit exists on every image; a missing one is not a failure. systemctl disable --now "$unit" 2>/dev/null || true + systemctl mask "$unit" 2>/dev/null || true done -systemctl mask irqbalance 2>/dev/null || true -CPU_COUNT="$(nproc)" -HOUSEKEEPING_CPUS="0-1" -BENCH_CPUS="2-$((CPU_COUNT - 1))" +# For apparmor, also teardown already-loaded profiles +aa-teardown || true + +# Split CPUs between the benchmark and everything else ("housekeeping"). On machines +# with multiple NUMA nodes the benchmark owns all of node 0 (bench-taskset.sh binds +# its memory to node 0) and housekeeping moves to the other nodes entirely. On +# single-node machines housekeeping gets cores 0-1 plus their SMT siblings, so +# background work never shares a physical core with the benchmark. +node_dirs=(/sys/devices/system/node/node[0-9]*) +if [[ ${#node_dirs[@]} -gt 1 && -r "${node_dirs[1]:-}/cpulist" ]]; then + BENCH_CPUS="$(cat /sys/devices/system/node/node0/cpulist)" + HOUSEKEEPING_CPUS="$(cat /sys/devices/system/node/node[1-9]*/cpulist | paste -sd, -)" +else + HOUSEKEEPING_CPUS="$(expand_cpulist "$(cat /sys/devices/system/cpu/cpu{0,1}/topology/thread_siblings_list 2>/dev/null | paste -sd, -)" | sort -un | paste -sd, -)" + bench=() + for cpu in $(expand_cpulist "$(cat /sys/devices/system/cpu/online)"); do + [[ ",$HOUSEKEEPING_CPUS," == *",$cpu,"* ]] || bench+=("$cpu") + done + BENCH_CPUS="$(IFS=,; echo "${bench[*]}")" +fi + +# Don't pin NIC queue IRQs. Pinning them serializes packet processing, and for +# S3 tests like fineweb-s3 network is the bottlenenck. +declare -A NET_IRQS=() +bench_cpu_arr=() +while IFS= read -r cpu; do bench_cpu_arr+=("$cpu"); done < <(expand_cpulist "$BENCH_CPUS") +i=0 +for msi in /sys/class/net/*/device/msi_irqs/*; do + [[ -e "$msi" && ${#bench_cpu_arr[@]} -gt 0 ]] || continue + irq="${msi##*/}" + NET_IRQS["$irq"]=1 + echo "${bench_cpu_arr[i % ${#bench_cpu_arr[@]}]}" > "/proc/irq/$irq/smp_affinity_list" 2>/dev/null || true + i=$((i + 1)) +done -# Pin all IRQs to housekeeping CPUs +# Pin all non-NIC queue IRQs to housekeeping CPUs. Some IRQs are kernel-managed +# and reject writes with EPERM even as root. for f in /proc/irq/[0-9]*/smp_affinity_list; do - if [[ -w "$f" ]]; then - # Some IRQs are kernel-managed and reject writes with EPERM even as root. - echo "$HOUSEKEEPING_CPUS" > "$f" 2>/dev/null || true - fi + irq="${f#/proc/irq/}" + irq="${irq%%/*}" + [[ -n "${NET_IRQS[$irq]:-}" ]] && continue + echo "$HOUSEKEEPING_CPUS" > "$f" 2>/dev/null || true +done + +# Steer unbound kernel workqueues (writeback etc.) to the housekeeping CPUs. The +# sysfs file takes a hex cpumask rather than a cpulist. +if command -v python3 >/dev/null 2>&1; then + wq_mask="$(expand_cpulist "$HOUSEKEEPING_CPUS" \ + | python3 -c 'import sys; print(f"{sum(1 << int(l) for l in sys.stdin):x}")' || true)" + sysfs_write /sys/devices/virtual/workqueue/cpumask "$wq_mask" +fi + +# Keep everything that is not explicitly re-pinned on the housekeeping CPUs. Two +# layers: a systemd manager default so services (re)started from now on land there, +# and a one-shot re-pin of everything already running (daemons, the CI runner agent, +# movable kernel threads). Plain affinity is a default rather than a ceiling (unlike +# cpusets), so bench-taskset.sh can still move the benchmark onto BENCH_CPUS. +mkdir -p /etc/systemd/system.conf.d +cat > /etc/systemd/system.conf.d/99-vortex-benchmark.conf </dev/null # the per-PID loop is far too noisy for trace output +for pid_dir in /proc/[0-9]*; do + taskset -a -pc "$HOUSEKEEPING_CPUS" "${pid_dir##*/}" >/dev/null 2>&1 || true done +set -x # Persist CPU affinity ranges for non-root benchmark steps in CI. cat > /tmp/vortex-benchmark.env < None: + try: + subprocess.run(["sync"], check=True) + subprocess.run( + ["sudo", "-n", "sh", "-c", "echo 3 > /proc/sys/vm/drop_caches"], + check=True, + capture_output=True, + ) + except (OSError, subprocess.CalledProcessError): + pass + + +def run_combinations() -> None: + PARTS_DIR.mkdir(parents=True, exist_ok=True) + i = 0 + for column in COLUMNS: + for encoder in ENCODERS: + drop_os_caches() + + args = [ + "bash", + str(SCRIPT_DIR / "bench-taskset.sh"), + BINARY, + "--columns", + f"^{re.escape(column)}$", + "--encoders", + encoder, + "-d", + "gh-json", + "-o", + str(PARTS_DIR / f"{i}.gh.json"), + ] + print("+", " ".join(args), flush=True) + subprocess.run(args, check=True) + i += 1 + + +def merge(pattern: str, out_path: str) -> None: + lines: list[str] = [] + for path in sorted(glob.glob(pattern)): + with open(path, encoding="utf-8") as handle: + for line in handle: + line = line.strip() + if line: + lines.append(line) + Path(out_path).write_text("".join(line + "\n" for line in lines), encoding="utf-8") + + +def main() -> None: + run_combinations() + merge(f"{PARTS_DIR}/*.gh.json", "results.json") + + +if __name__ == "__main__": + main() diff --git a/scripts/tests/test_benchmark_reporting.py b/scripts/tests/test_benchmark_reporting.py index 7925161e91e..46c7b4b1a1d 100644 --- a/scripts/tests/test_benchmark_reporting.py +++ b/scripts/tests/test_benchmark_reporting.py @@ -3,18 +3,21 @@ import importlib.util import json +import math import subprocess import sys from pathlib import Path +from types import ModuleType import pandas as pd REPO_ROOT = Path(__file__).resolve().parents[2] COMPARE_SCRIPT = REPO_ROOT / "scripts" / "compare-benchmark-jsons.py" CAPTURE_SCRIPT = REPO_ROOT / "scripts" / "capture-file-sizes.py" +RANDOM_ACCESS_SPLIT_SCRIPT = REPO_ROOT / "scripts" / "random-access-split.py" -def load_compare_module(): +def load_compare_module() -> ModuleType: spec = importlib.util.spec_from_file_location("compare_benchmark_jsons", COMPARE_SCRIPT) assert spec is not None module = importlib.util.module_from_spec(spec) @@ -23,13 +26,46 @@ def load_compare_module(): return module +def load_random_access_split_module() -> ModuleType: + spec = importlib.util.spec_from_file_location("random_access_split", RANDOM_ACCESS_SPLIT_SCRIPT) + assert spec is not None + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def test_random_access_merge_preserves_each_open_mode(tmp_path: Path) -> None: + split = load_random_access_split_module() + parts = tmp_path / "parts" + parts.mkdir() + cached = { + "kind": "random_access_time", + "dataset": "taxi", + "format": "parquet", + "open_mode": "cached", + } + reopen = cached | {"open_mode": "reopen"} + (parts / "0.ingest.jsonl").write_text(f"{json.dumps(cached)}\n", encoding="utf-8") + (parts / "1.ingest.jsonl").write_text( + f"{json.dumps(cached)}\n{json.dumps(reopen)}\n", + encoding="utf-8", + ) + output = tmp_path / "results.ingest.jsonl" + + split.merge(str(parts / "*.ingest.jsonl"), split.ingest_identity, str(output)) + + records = [json.loads(line) for line in output.read_text(encoding="utf-8").splitlines()] + assert records == [cached, reopen] + + def timing_row(name: str, base: int, pr: int) -> dict[str, object]: return { "name": name, - "value_base": base, - "value_pr": pr, - "all_runtimes_base": [base, base, base], - "all_runtimes_pr": [pr, pr, pr], + "hot_value_base": base, + "hot_value_pr": pr, + "hot_runtimes_base": [base, base, base], + "hot_runtimes_pr": [pr, pr, pr], } @@ -39,13 +75,23 @@ def stored_timing_row( value: int, storage: str | None = None, dataset: dict[str, object] | None = None, + engine: str = "datafusion", + file_format: str = "parquet", + cold: int | None = None, ) -> dict[str, object]: + """Build a runner result row whose first run is the cold one. + + `value` stays the median across all runs, which is what the benchmark binary + writes, so the reporter has to derive the hot median from `all_runtimes`. + """ + row: dict[str, object] = { "name": name, "unit": "ns", "value": value, - "all_runtimes": [value, value, value], + "all_runtimes": [value if cold is None else cold, value, value], "commit_id": commit, + "target": {"engine": engine, "format": file_format}, } if storage is not None: row["storage"] = storage @@ -54,6 +100,77 @@ def stored_timing_row( return row +def stored_custom_row( + commit: str, + name: str, + unit: str, + value: float, + engine: str = "vortex", + file_format: str = "vortex", +) -> dict[str, object]: + return { + "name": name, + "unit": unit, + "value": value, + "commit_id": commit, + "target": {"engine": engine, "format": file_format}, + } + + +def git_repo_with_one_commit(path: Path) -> Path: + """Create a throwaway repository the reporter can walk. + + The reporter resolves its baseline from `git rev-list HEAD`, so the tests + need a repository of their own rather than whatever checkout they happen to + run inside — a contributor's shallow clone must not change the result. + """ + + path.mkdir(parents=True) + identity = ["-c", "user.name=Vortex Test", "-c", "user.email=test@example.com"] + subprocess.run(["git", "init", "--quiet"], cwd=path, check=True) + (path / "README.md").write_text("benchmark reporting fixture\n", encoding="utf-8") + subprocess.run(["git", "add", "README.md"], cwd=path, check=True) + subprocess.run(["git", *identity, "commit", "--quiet", "-m", "fixture"], cwd=path, check=True) + return path + + +def render_report( + tmp_path: Path, + base_rows: list[dict[str, object]], + pr_rows: list[dict[str, object]], + benchmark_name: str, +) -> str: + repo = git_repo_with_one_commit(tmp_path / "repo") + head_commit = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=repo, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + base_rows = [row | {"commit_id": head_commit} for row in base_rows] + base_path = tmp_path / "base.jsonl" + pr_path = tmp_path / "pr.jsonl" + base_path.write_text("".join(f"{json.dumps(row)}\n" for row in base_rows), encoding="utf-8") + pr_path.write_text("".join(f"{json.dumps(row)}\n" for row in pr_rows), encoding="utf-8") + + result = subprocess.run( + [sys.executable, str(COMPARE_SCRIPT), str(base_path), str(pr_path), benchmark_name], + cwd=repo, + check=False, + capture_output=True, + text=True, + ) + + assert result.returncode == 0, result.stderr + return result.stdout + + +def markdown_row(report: str, name: str) -> list[str]: + line = next(line for line in report.splitlines() if line.startswith(f"| {name} ")) + return [cell.strip() for cell in line.strip("|").split("|")] + + def test_select_latest_baseline_rows_uses_latest_matching_benchmark_commit() -> None: compare = load_compare_module() history = pd.DataFrame( @@ -89,7 +206,7 @@ def test_select_latest_baseline_rows_uses_latest_matching_benchmark_commit() -> ] ) - selected = compare.select_latest_baseline_rows(history, pr) + selected = compare.select_latest_baseline_rows(history, pr, {"base-old", "base-current"}) assert set(selected["commit_id"]) == {"base-current"} assert len(selected) == 2 @@ -115,6 +232,22 @@ def test_read_latest_baseline_rows_streams_latest_matching_benchmark_commit(tmp_ {"scale_factor": "1.0"}, ), file_size_record_for("base-current", 120, "tpch", "1.0", "vortex-file-compressed", "part-0.vortex"), + stored_timing_row( + "base-wrong-unit", + "tpch_q01/datafusion:parquet", + 111, + "nvme", + {"scale_factor": "1.0"}, + ) + | {"unit": "ms"}, + stored_timing_row( + "base-wrong-target", + "tpch_q01/datafusion:parquet", + 112, + "nvme", + {"scale_factor": "1.0"}, + file_format="vortex", + ), stored_timing_row("base-other", "clickbench_q01/datafusion:parquet", 200, "nvme"), ] history_path.write_text( @@ -133,12 +266,82 @@ def test_read_latest_baseline_rows_streams_latest_matching_benchmark_commit(tmp_ ] ) - selected = compare.read_latest_baseline_rows(history_path, pr) + selected = compare.read_latest_baseline_rows(history_path, pr, {"base-old", "base-current"}) assert set(selected["commit_id"]) == {"base-current"} assert len(selected) == 2 +def test_read_latest_baseline_rows_skips_commits_outside_git_tree(tmp_path: Path) -> None: + compare = load_compare_module() + history_path = tmp_path / "history.jsonl" + history_rows = [ + stored_timing_row("base-reachable", "tpch_q01/datafusion:parquet", 100), + file_size_record_for( + "base-reachable", + 100, + "tpch", + "1.0", + "vortex-file-compressed", + "part-0.vortex", + ), + stored_timing_row("base-off-tree", "tpch_q01/datafusion:parquet", 110), + file_size_record_for( + "base-off-tree", + 120, + "tpch", + "1.0", + "vortex-file-compressed", + "part-0.vortex", + ), + ] + history_path.write_text( + "".join(f"{json.dumps(row)}\n" for row in history_rows), + encoding="utf-8", + ) + pr = pd.DataFrame([stored_timing_row("pr", "tpch_q01/datafusion:parquet", 105)]) + + selected = compare.read_latest_baseline_rows(history_path, pr, {"base-reachable"}) + + assert set(selected["commit_id"]) == {"base-reachable"} + assert len(selected) == 2 + + +def test_read_latest_baseline_rows_returns_empty_without_reachable_match(tmp_path: Path) -> None: + compare = load_compare_module() + history_path = tmp_path / "history.jsonl" + history_path.write_text( + f"{json.dumps(stored_timing_row('base-off-tree', 'tpch_q01/datafusion:parquet', 100))}\n", + encoding="utf-8", + ) + pr = pd.DataFrame([stored_timing_row("pr", "tpch_q01/datafusion:parquet", 105)]) + + selected = compare.read_latest_baseline_rows(history_path, pr, set()) + + assert selected.empty + + +def test_read_latest_baseline_rows_uses_last_result_from_rerun(tmp_path: Path) -> None: + compare = load_compare_module() + history_path = tmp_path / "history.jsonl" + history_path.write_text( + "".join( + f"{json.dumps(row)}\n" + for row in [ + stored_timing_row("base", "tpch_q01/datafusion:parquet", 100), + stored_timing_row("base", "tpch_q01/datafusion:parquet", 110), + ] + ), + encoding="utf-8", + ) + pr = pd.DataFrame([stored_timing_row("pr", "tpch_q01/datafusion:parquet", 105)]) + + selected = compare.read_latest_baseline_rows(history_path, pr, {"base"}) + + assert len(selected) == 1 + assert selected.iloc[0]["value"] == 110 + + def test_within_engine_analysis_uses_each_engines_own_parquet_control() -> None: compare = load_compare_module() rows = [ @@ -157,6 +360,257 @@ def test_within_engine_analysis_uses_each_engines_own_parquet_control() -> None: assert compare.build_verdict(analyses["duckdb"])["impact"] == "+20.0%" +def test_random_access_attribution_excludes_lance_rows() -> None: + compare = load_compare_module() + rows = [ + timing_row("random-access/taxi/correlated/parquet-tokio-local-disk", 100, 110), + timing_row("random-access/taxi/correlated/vortex-tokio-local-disk", 100, 99), + timing_row("random-access/taxi/correlated/lance-tokio-local-disk", 100, 200), + timing_row("random-access/taxi/uniform/parquet-tokio-local-disk", 100, 110), + timing_row("random-access/taxi/uniform/vortex-tokio-local-disk", 100, 99), + timing_row("random-access/taxi/uniform/lance-tokio-local-disk", 100, 200), + ] + df = pd.DataFrame(rows) + df[["engine", "file_format", "query"]] = df["name"].apply(compare.extract_target_fields) + + analyses = compare.build_within_engine_statistical_analyses(df, threshold_pct=5) + + assert df.loc[df["file_format"] == "lance", "query"].isna().all() + assert set(analyses) == {"random-access"} + assert set(analyses["random-access"]["detail_df"]["file_format"]) == { + "parquet", + "vortex-file-compressed", + } + assert compare.build_verdict(analyses["random-access"])["status"] == "Likely improvement" + + +def test_random_access_report_keeps_lance_details_without_attribution(tmp_path: Path) -> None: + names_and_values = [ + ("random-access/taxi/correlated/parquet-tokio-local-disk", 110), + ("random-access/taxi/correlated/vortex-tokio-local-disk", 99), + ("random-access/taxi/correlated/lance-tokio-local-disk", 200), + ("random-access/taxi/uniform/parquet-tokio-local-disk", 110), + ("random-access/taxi/uniform/vortex-tokio-local-disk", 99), + ("random-access/taxi/uniform/lance-tokio-local-disk", 200), + ] + base_rows = [stored_timing_row("base-sha", name, 100) for name, _pr_value in names_and_values] + pr_rows = [stored_timing_row("pr-sha", name, pr_value) for name, pr_value in names_and_values] + + report = render_report(tmp_path, base_rows, pr_rows, "Random Access") + + assert "**Attributed Vortex impact**: -10.0%" in report + assert "random-access / lance / ns " in report + + +def test_comparison_report_groups_by_target_and_unit(tmp_path: Path) -> None: + base_rows = [ + stored_custom_row("base-sha", "timing/fixture", "ms", 12.5), + stored_custom_row("base-sha", "size/fixture", "%", 45.25), + stored_custom_row("base-sha", "ratio/fixture", "ratio", 0.75), + stored_custom_row("base-sha", "parquet timing/fixture", "ms", 20.0, file_format="parquet"), + ] + pr_rows = [ + stored_custom_row("pr-sha", "timing/fixture", "ms", 11.75), + stored_custom_row("pr-sha", "new timing/fixture", "ms", 3.125), + stored_custom_row("pr-sha", "size/fixture", "%", 44.5), + stored_custom_row("pr-sha", "ratio/fixture", "ratio", 0.8), + stored_custom_row("pr-sha", "parquet timing/fixture", "ms", 21.0, file_format="parquet"), + ] + + report = render_report(tmp_path, base_rows, pr_rows, "Mixed metrics") + + assert "unknown / unknown" not in report + assert "How to read Verdict and Engines" not in report + # These measurements report a single number, so there is no cold run to split out. + assert "cold" not in report + assert "hot" not in report + assert "vortex / vortex-file-compressed / ms " in report + assert "vortex / vortex-file-compressed / % " in report + assert "vortex / vortex-file-compressed / ratio " in report + assert "vortex / parquet / ms " in report + assert markdown_row(report, "timing/fixture") == ["timing/fixture", "11.75 / 12.5 / -6.0%"] + assert markdown_row(report, "size/fixture") == ["size/fixture", "44.5 / 45.25 / -1.7%"] + assert markdown_row(report, "ratio/fixture") == ["ratio/fixture", "0.8 / 0.75 / +6.7%"] + assert markdown_row(report, "new timing/fixture") == [ + "new timing/fixture", + "3.125 / — / no baseline", + ] + + +def test_comparison_report_handles_missing_benchmark_baseline(tmp_path: Path) -> None: + base_rows = [stored_custom_row("base-sha", "other timing/fixture", "ms", 12.5)] + pr_rows = [stored_custom_row("pr-sha", "new timing/fixture", "ms", 3.125)] + + report = render_report(tmp_path, base_rows, pr_rows, "New benchmark") + + assert "No baseline is available for this benchmark yet" in report + assert "vs base `none`" in report + assert markdown_row(report, "new timing/fixture") == [ + "new timing/fixture", + "3.125 / — / no baseline", + ] + + +def test_comparison_report_handles_mixed_query_types_in_baseline(tmp_path: Path) -> None: + base_rows = [ + stored_custom_row("base-sha", "random-access/fixture", "ms", 12.5), + stored_timing_row("base-sha", "tpch_q01/datafusion:parquet", 100), + ] + pr_rows = [stored_custom_row("pr-sha", "random-access/fixture", "ms", 13.0)] + + report = render_report(tmp_path, base_rows, pr_rows, "Random Access") + + assert markdown_row(report, "random-access/fixture") == [ + "random-access/fixture", + "13 / 12.5 / +4.0%", + ] + + +def test_comparison_report_retains_sql_analysis(tmp_path: Path) -> None: + targets = [ + ("parquet", "parquet", 100, 105), + ("vortex-file-compressed", "vortex", 80, 70), + ] + base_rows = [ + stored_timing_row("base-sha", f"tpch_q01/datafusion:{name}", base, file_format=target) + for name, target, base, _pr in targets + ] + pr_rows = [ + stored_timing_row("pr-sha", f"tpch_q01/datafusion:{name}", pr, file_format=target) + for name, target, _base, pr in targets + ] + + report = render_report(tmp_path, base_rows, pr_rows, "TPC-H") + + assert "**Verdict**:" in report + assert "**Vortex (geomean)**: hot " in report + assert "**Parquet (geomean)**: hot " in report + assert "How to read Verdict and Engines" in report + assert "datafusion / vortex-file-compressed / ns " in report + assert "datafusion / parquet / ns " in report + assert "unknown / unknown" not in report + + +def test_large_hot_medians_render_without_scientific_notation() -> None: + compare = load_compare_module() + + # A median over two runs can land on a half nanosecond; at this magnitude "%g" + # would switch to scientific notation and stop matching the baseline column. + assert compare.format_measurement_value(1478968300.5) == "1478968300" + assert compare.format_measurement_value(8228437.5) == "8228438" + assert compare.format_measurement_value(12.5) == "12.5" + assert compare.format_measurement_value(100) == "100" + + +def test_cold_and_hot_runtimes_split_the_first_run_from_the_rest() -> None: + compare = load_compare_module() + + assert compare.cold_runtime([7, 9, 11]) == 7 + assert compare.hot_runtimes([7, 9, 11]) == [9, 11] + assert compare.hot_runtime([7, 9, 11], 9) == 10 + + +def test_hot_runtime_falls_back_to_the_reported_value_without_later_runs() -> None: + compare = load_compare_module() + + assert math.isnan(compare.cold_runtime(None)) + assert compare.hot_runtimes(None) == [] + assert compare.hot_runtime(None, 42) == 42 + assert compare.hot_runtime([7], 7) == 7 + assert math.isnan(compare.hot_runtime(None, None)) + + +def test_only_changes_past_the_suite_threshold_are_marked(tmp_path: Path) -> None: + # TPC-H runs at a 10% threshold, so 9% is noise here and 11% is not. + base_rows = [ + stored_timing_row("base-sha", "tpch_q01/datafusion:vortex-file-compressed", 100, file_format="vortex"), + stored_timing_row("base-sha", "tpch_q02/datafusion:vortex-file-compressed", 100, file_format="vortex"), + stored_timing_row("base-sha", "tpch_q03/datafusion:vortex-file-compressed", 100, file_format="vortex"), + ] + pr_rows = [ + stored_timing_row("pr-sha", "tpch_q01/datafusion:vortex-file-compressed", 109, file_format="vortex"), + stored_timing_row("pr-sha", "tpch_q02/datafusion:vortex-file-compressed", 111, file_format="vortex"), + stored_timing_row("pr-sha", "tpch_q03/datafusion:vortex-file-compressed", 89, file_format="vortex"), + ] + + report = render_report(tmp_path, base_rows, pr_rows, "TPC-H") + + assert markdown_row(report, "tpch_q01/datafusion:vortex-file-compressed")[1] == "109 / 100 / +9.0%" + assert markdown_row(report, "tpch_q02/datafusion:vortex-file-compressed")[1] == "111 / 100 / +11.0% 🔴" + assert markdown_row(report, "tpch_q03/datafusion:vortex-file-compressed")[1] == "89 / 100 / -11.0% 🟢" + + +def test_report_splits_cold_and_hot_runs(tmp_path: Path) -> None: + base_rows = [ + stored_timing_row("base-sha", "tpch_q01/datafusion:parquet", 100, cold=1000), + stored_timing_row( + "base-sha", + "tpch_q01/datafusion:vortex-file-compressed", + 100, + file_format="vortex", + cold=1000, + ), + ] + pr_rows = [ + stored_timing_row("pr-sha", "tpch_q01/datafusion:parquet", 100, cold=1000), + stored_timing_row( + "pr-sha", + "tpch_q01/datafusion:vortex-file-compressed", + 50, + file_format="vortex", + cold=2000, + ), + ] + + report = render_report(tmp_path, base_rows, pr_rows, "TPC-H") + + # The marks sit on the numbers that moved, not on the benchmark name. + assert markdown_row(report, "tpch_q01/datafusion:vortex-file-compressed") == [ + "tpch_q01/datafusion:vortex-file-compressed", + "50 / 100 / -50.0% 🟢", + "2000 / 1000 / +100.0% 🔴", + "0.03 / 0.10 / -75.0%", + ] + assert "**Attributed Vortex impact**: -50.0%" in report + assert "**Vortex (geomean)**: hot 0.500x ✅ · cold 2.000x ❌" in report + assert "**Parquet (geomean)**: hot 1.000x ➖ · cold 1.000x ➖" in report + assert "**Commits**: PR `pr-sha` vs base " in report + + +def test_verdict_ignores_a_cold_start_only_regression(tmp_path: Path) -> None: + base_rows = [ + stored_timing_row("base-sha", "tpch_q01/datafusion:parquet", 100, cold=100), + stored_timing_row( + "base-sha", + "tpch_q01/datafusion:vortex-file-compressed", + 100, + file_format="vortex", + cold=100, + ), + ] + pr_rows = [ + stored_timing_row("pr-sha", "tpch_q01/datafusion:parquet", 100, cold=100), + stored_timing_row( + "pr-sha", + "tpch_q01/datafusion:vortex-file-compressed", + 100, + file_format="vortex", + cold=400, + ), + ] + + report = render_report(tmp_path, base_rows, pr_rows, "TPC-H") + + assert "**Attributed Vortex impact**: +0.0%" in report + assert "**Vortex (geomean)**: hot 1.000x ➖ · cold 4.000x ❌" in report + assert markdown_row(report, "tpch_q01/datafusion:vortex-file-compressed") == [ + "tpch_q01/datafusion:vortex-file-compressed", + "100 / 100 / +0.0%", + "400 / 100 / +300.0% 🔴", + "0.25 / 1.00 / -75.0%", + ] + + def file_size_record(commit: str, size: int) -> dict[str, object]: return file_size_record_for(commit, size, "tpch", "10", "vortex-file-compressed", "part-0.vortex") @@ -195,6 +649,31 @@ def test_file_size_report_reads_shared_benchmark_rows() -> None: assert "| part-0.vortex | 10 | vortex-file-compressed | 100 B | 125 B | +25 B | +25.0% |" in report +def test_file_size_report_ignores_file_identities_with_a_zero_byte_side() -> None: + compare = load_compare_module() + + report = compare.format_file_size_report( + pd.DataFrame( + [ + file_size_record_for("base-sha", 100, "tpch", "10", "vortex-file-compressed", "head-empty"), + file_size_record_for("base-sha", 0, "tpch", "10", "vortex-file-compressed", "base-empty"), + file_size_record_for("base-sha", 100, "tpch", "10", "vortex-file-compressed", "changed"), + ] + ), + pd.DataFrame( + [ + file_size_record_for("pr-sha", 0, "tpch", "10", "vortex-file-compressed", "head-empty"), + file_size_record_for("pr-sha", 125, "tpch", "10", "vortex-file-compressed", "base-empty"), + file_size_record_for("pr-sha", 125, "tpch", "10", "vortex-file-compressed", "changed"), + ] + ), + ) + + assert "File Size Changes (1 files changed, +25.0% overall, 1↑ 0↓)" in report + assert "head-empty" not in report + assert "base-empty" not in report + + def test_file_size_report_ignores_baseline_rows_outside_pr_scope() -> None: compare = load_compare_module() @@ -218,6 +697,29 @@ def test_file_size_report_ignores_baseline_rows_outside_pr_scope() -> None: assert "| part-0.vortex | 1.0 |" not in report +def test_file_size_report_omits_formats_the_pr_run_skipped() -> None: + compare = load_compare_module() + + report = compare.format_file_size_report( + pd.DataFrame( + [ + file_size_record_for("base-sha", 100, "tpch", "10", "vortex-file-compressed", "part-0.vortex"), + file_size_record_for("base-sha", 80, "tpch", "10", "vortex-compact", "part-0.vortex"), + file_size_record_for("base-sha", 5, "tpch", "10", "vortex-compact", "duckdb.db"), + ] + ), + pd.DataFrame( + [ + file_size_record_for("pr-sha", 125, "tpch", "10", "vortex-file-compressed", "part-0.vortex"), + ] + ), + ) + + assert "File Size Changes (1 files changed, +25.0% overall, 1↑ 0↓)" in report + assert "vortex-compact" not in report + assert "-100.0%" not in report + + def test_capture_file_sizes_emits_shared_benchmark_rows(tmp_path: Path) -> None: data_dir = tmp_path / "data" format_dir = data_dir / "tpch" / "10" / "vortex-file-compressed" @@ -258,3 +760,63 @@ def test_capture_file_sizes_emits_shared_benchmark_rows(tmp_path: Path) -> None: }, } ] + + +def test_shallow_checkout_fails_instead_of_reporting_no_baseline(tmp_path: Path) -> None: + """A shallow clone must fail loudly, not render a baseline-free report. + + `git rev-list HEAD` reaches only the tip of a shallow clone, so every + recorded commit looks unreachable. Reporting that as "no baseline yet" + hides a real comparison behind a report that merely looks new. + """ + + origin = git_repo_with_one_commit(tmp_path / "origin") + shallow = tmp_path / "shallow" + subprocess.run( + ["git", "clone", "--quiet", "--depth", "1", "--no-local", f"file://{origin}", str(shallow)], + check=True, + capture_output=True, + ) + base_path = tmp_path / "base.jsonl" + pr_path = tmp_path / "pr.jsonl" + base_path.write_text( + f"{json.dumps(stored_timing_row('recorded', 'tpch_q01/datafusion:parquet', 100))}\n", + encoding="utf-8", + ) + pr_path.write_text( + f"{json.dumps(stored_timing_row('pr', 'tpch_q01/datafusion:parquet', 105))}\n", + encoding="utf-8", + ) + + result = subprocess.run( + [sys.executable, str(COMPARE_SCRIPT), str(base_path), str(pr_path), "TPC-H"], + cwd=shallow, + check=False, + capture_output=True, + text=True, + ) + + assert result.returncode != 0 + assert "shallow checkout" in result.stderr + assert "fetch-depth: 0" in result.stderr + assert "No baseline is available" not in result.stdout + + +def test_summary_reports_hot_alone_when_no_row_recorded_its_runs(tmp_path: Path) -> None: + # Random Access, String Encoding and Compression report one value per row, so + # there is no cold geomean to sit beside the hot one. + base_rows = [ + stored_custom_row("base-sha", "random-access/vortex-tokio-local-disk", "ns", 100.0), + stored_custom_row("base-sha", "random-access/parquet-tokio-local-disk", "ns", 200.0, file_format="parquet"), + ] + pr_rows = [ + stored_custom_row("pr-sha", "random-access/vortex-tokio-local-disk", "ns", 50.0), + stored_custom_row("pr-sha", "random-access/parquet-tokio-local-disk", "ns", 200.0, file_format="parquet"), + ] + + report = render_report(tmp_path, base_rows, pr_rows, "Random Access") + + summary = report.splitlines()[2] + assert "**Vortex (geomean)**: hot 0.500x ✅" in summary + assert "**Parquet (geomean)**: hot 1.000x ➖" in summary + assert "cold" not in summary diff --git a/scripts/tests/test_measurement_id.py b/scripts/tests/test_measurement_id.py index c8baf0cfb8d..14f74141d75 100644 --- a/scripts/tests/test_measurement_id.py +++ b/scripts/tests/test_measurement_id.py @@ -23,13 +23,14 @@ import importlib.util import json from pathlib import Path +from types import ModuleType REPO_ROOT = Path(__file__).resolve().parents[2] MODULE_PATH = REPO_ROOT / "scripts" / "_measurement_id.py" GOLDEN_PATH = REPO_ROOT / "scripts" / "measurement_id_golden.json" -def load_measurement_id_module(): +def load_measurement_id_module() -> ModuleType: spec = importlib.util.spec_from_file_location("_measurement_id", MODULE_PATH) assert spec is not None module = importlib.util.module_from_spec(spec) @@ -38,7 +39,7 @@ def load_measurement_id_module(): return module -def test_every_golden_vector_reproduces(): +def test_every_golden_vector_reproduces() -> None: module = load_measurement_id_module() golden = json.loads(GOLDEN_PATH.read_text()) assert golden["seed"] == 0 @@ -52,7 +53,7 @@ def test_every_golden_vector_reproduces(): ) -def test_every_fact_table_is_covered(): +def test_every_fact_table_is_covered() -> None: """Every port function must have at least one golden vector, so a new fact table cannot be wired into `MEASUREMENT_ID_BY_TABLE` without also pinning its hash.""" module = load_measurement_id_module() @@ -60,3 +61,19 @@ def test_every_fact_table_is_covered(): covered_tables = {vector["table"] for vector in golden["vectors"]} assert covered_tables == set(module.MEASUREMENT_ID_BY_TABLE) + + +def test_random_access_open_mode_preserves_cached_ids_and_separates_reopen() -> None: + module = load_measurement_id_module() + dimensions = { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "taxi", + "format": "parquet", + } + + historical = module.measurement_id_random_access(**dimensions) + cached = module.measurement_id_random_access(**dimensions, open_mode="cached") + reopen = module.measurement_id_random_access(**dimensions, open_mode="reopen") + + assert cached == historical + assert reopen != cached diff --git a/scripts/vortex_gdb.py b/scripts/vortex_gdb.py index 791bbdd868e..05f78bb079f 100644 --- a/scripts/vortex_gdb.py +++ b/scripts/vortex_gdb.py @@ -135,7 +135,7 @@ def _str_to_python(s: gdb.Value) -> str | None: class _PinnedThread: """Make a pretty-printer's call safe with multiple threads""" - def __enter__(self): + def __enter__(self) -> _PinnedThread: try: self._frame = gdb.selected_frame() except gdb.error: @@ -148,7 +148,7 @@ def __enter__(self): self._prev_lock = None return self - def __exit__(self, exc_type, exc, tb): + def __exit__(self, exc_type, exc, tb) -> None: if self._prev_lock is not None: try: gdb.execute(f"set scheduler-locking {self._prev_lock}", to_string=True) @@ -247,7 +247,7 @@ def _format_array_ref(val: gdb.Value) -> str: if inner is None: return f"ArrayRef(?) @ {arc_inner_base:#x}" - def _try(fn, default): + def _try(fn, default) -> object: try: return fn() except (gdb.error, RuntimeError): @@ -276,7 +276,7 @@ def _format_typed_array(val: gdb.Value) -> str: class _Printer: - def __init__(self, val: gdb.Value, fmt: Callable[[gdb.Value], str]): + def __init__(self, val: gdb.Value, fmt: Callable[[gdb.Value], str]) -> None: self._val, self._fmt = val, fmt def to_string(self) -> str: diff --git a/uv.lock b/uv.lock index 645780c261a..ad6507d0851 100644 --- a/uv.lock +++ b/uv.lock @@ -26,6 +26,146 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8d/3f/95338030883d8c8b91223b4e21744b04d11b161a3ef117295d8241f50ab4/accessible_pygments-0.0.5-py3-none-any.whl", hash = "sha256:88ae3211e68a1d0b011504b2ffc1691feafce124b845bd072ab6f9f66f34d4b7", size = 1395903, upload-time = "2024-05-10T11:23:08.421Z" }, ] +[[package]] +name = "aiohappyeyeballs" +version = "2.7.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ce/f4/eec0465c2f67b2664688d0240b3212d5196fd89e741df67ddb81f8d35658/aiohappyeyeballs-2.7.1.tar.gz", hash = "sha256:065665c041c42a5938ed220bdcd7230f22527fbec085e1853d2402c8a3615d9d", size = 24757, upload-time = "2026-07-01T17:11:55.501Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/71/43/1947f06babed6b3f1d7f38b0c767f52df66bfb2bc10b468c4a7de9eceff2/aiohappyeyeballs-2.7.1-py3-none-any.whl", hash = "sha256:9243213661e29250eb41368e5daa826fc017156c3b8a11440826b2e3ed376472", size = 15038, upload-time = "2026-07-01T17:11:54.055Z" }, +] + +[[package]] +name = "aiohttp" +version = "3.14.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiohappyeyeballs" }, + { name = "aiosignal" }, + { name = "attrs" }, + { name = "frozenlist" }, + { name = "multidict" }, + { name = "propcache" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "yarl" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/58/d9/22ce5786ac0c1653ae8b6c23bded02c1686d11f0dbb45b31ce128e0df985/aiohttp-3.14.3.tar.gz", hash = "sha256:9491196535a88924a60afd5b5f434b5b203b6cc616250878dbdb223a8f7844bc", size = 7971213, upload-time = "2026-07-23T01:57:27.037Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f8/5c/b3e4ff8ad43a8afef9602c5e90285936da1beaea8b029016b793891f03c3/aiohttp-3.14.3-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:e568e14940c09955aa51f4e645b6daa18a581c5dcfcd73744dcc86a856e3ced3", size = 764250, upload-time = "2026-07-23T01:52:48.525Z" }, + { url = "https://files.pythonhosted.org/packages/0e/da/f1b384465e51449d844056b75070461da03a9a23e6c1747003695bf4172a/aiohttp-3.14.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:54cfcdee2770dac994417cbb0ee1f3eb0e7cb6b30c79bf44f2c02ff79ec5124a", size = 516281, upload-time = "2026-07-23T01:52:51.047Z" }, + { url = "https://files.pythonhosted.org/packages/b9/3f/01264f820ee2e3712a827892b1cd6ff80f3300c1fcbffbb45714a915d47a/aiohttp-3.14.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:21c016079415ed3fd676963e9793700a566d85dbbd6bfc564b9b2d209147dcc8", size = 514742, upload-time = "2026-07-23T01:52:53.779Z" }, + { url = "https://files.pythonhosted.org/packages/9e/8d/a71c6f2db52ac1ed142b133f7feddaa6b70539c3f4de24d7e226c95b794c/aiohttp-3.14.3-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d6088ec9894113802bddb3c09e974929aed2c7b3a8c456219b8aab4481f1a239", size = 1780613, upload-time = "2026-07-23T01:52:56.948Z" }, + { url = "https://files.pythonhosted.org/packages/a5/11/3dd9b3fb3a170f6ec9011b5291d876a6fab4086714c9e158600edf01b4fd/aiohttp-3.14.3-cp311-cp311-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:16ea7e24c309fb7c0bbd505d149abe4fe4dccfb8db911db7dbec0921bc889a6f", size = 1737688, upload-time = "2026-07-23T01:52:59.294Z" }, + { url = "https://files.pythonhosted.org/packages/6d/3e/834c26918be7d88068822b40e0db30fca50b5f4fe79104aa16a93f1d74e6/aiohttp-3.14.3-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:56f355e79f71aef2a85c80305cc915f894b170dba76de5fe84f6351939b83c06", size = 1845742, upload-time = "2026-07-23T01:53:01.641Z" }, + { url = "https://files.pythonhosted.org/packages/cc/c9/49ab8572df7d66bc13d11e31f781292badb04180dd87ba98733066c6aed7/aiohttp-3.14.3-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:18c441d0a8fca6de8d1f546849b9f0ab20d435993e2c5b59562b2fae6be2f929", size = 1928412, upload-time = "2026-07-23T01:53:04.018Z" }, + { url = "https://files.pythonhosted.org/packages/a5/b9/2b8f0c0ce09c87a1daf80fd483431b56b1435d3f62789bc86f572e1245de/aiohttp-3.14.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:53e7b4ce82b54a8bcc71b3b67a5cbd177ca1d7f592cbc92cd38b7349f73482db", size = 1786220, upload-time = "2026-07-23T01:53:06.481Z" }, + { url = "https://files.pythonhosted.org/packages/85/00/9c45f81de11710460edfa1dc81317b6e882703b160926c879a9d20da9fcc/aiohttp-3.14.3-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f55119f7bf25f49ed210f6096090715da24f2943c62102448915fde3c62877ce", size = 1637231, upload-time = "2026-07-23T01:53:10.258Z" }, + { url = "https://files.pythonhosted.org/packages/19/ce/967d628e910756f3539c6107cb7844a1b69440dcb3029a5ee7871b09ab63/aiohttp-3.14.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:9aa6e61fdf20105c4144e755bd586008ff450791d67b1c8146fdc15959c4d51c", size = 1753161, upload-time = "2026-07-23T01:53:13.817Z" }, + { url = "https://files.pythonhosted.org/packages/11/b2/0c3d4114f0aee4f580f5b3b4eb71b24d7a23b834ea506a4dfebe76513f35/aiohttp-3.14.3-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:ccd4893707b3e2a13e39c90d43cf80edf2e4d0457935bcc103bf2346214c3f15", size = 1756356, upload-time = "2026-07-23T01:53:16.211Z" }, + { url = "https://files.pythonhosted.org/packages/63/5d/99e7d91c82f1399d1ae2a854e080bd1493fbc31e5e959dbc4ec33dac3bec/aiohttp-3.14.3-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:b2466434105a4e03113c36ec775cc2ebe6676b62eae326fa670bb607ef788c1c", size = 1819846, upload-time = "2026-07-23T01:53:18.289Z" }, + { url = "https://files.pythonhosted.org/packages/ad/05/d5e1cb6480eeffd3f901d40a2c5e2d1e7effdc797837da3b490272699f13/aiohttp-3.14.3-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:ba59d59aba08ac02fc03b0c8983ccd5ee39a199d0552ce9e6d2b4845b34d59ae", size = 1628531, upload-time = "2026-07-23T01:53:23.86Z" }, + { url = "https://files.pythonhosted.org/packages/c9/90/b934682bcaefae18a9e04f3dff5b68522ba810906358ae5029b68110ea3b/aiohttp-3.14.3-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:ed099d105449c4f9e84f24af203cd131349d4761d8813fa7e02c32e7128cd910", size = 1832712, upload-time = "2026-07-23T01:53:27.551Z" }, + { url = "https://files.pythonhosted.org/packages/21/df/6061679faaf81fac746e7307c7adb71e858071a5d34c27583afefc64f543/aiohttp-3.14.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:152516815ef926786a0b6ae2b8f1fd2e0c71582dee0b435636865316fd4891b7", size = 1775014, upload-time = "2026-07-23T01:53:30.223Z" }, + { url = "https://files.pythonhosted.org/packages/8a/1d/f854878bbc69b88faefe924b619a34a6f59ec05fd387c77690667eaa75eb/aiohttp-3.14.3-cp311-cp311-win32.whl", hash = "sha256:a4af35c443e0b1a1bd6a8af3f3485d7fda15c142751a00f3ff8090f0b93346fa", size = 456006, upload-time = "2026-07-23T01:53:34.97Z" }, + { url = "https://files.pythonhosted.org/packages/73/0c/2af9d1674baccd1dbd47282a93d660a22e57ef6167c856deb24b4214fbab/aiohttp-3.14.3-cp311-cp311-win_amd64.whl", hash = "sha256:e1e74298bab6ee0d6e749ed4fd1901c7e604bdda32c03d787a2cc71c46d0433d", size = 481069, upload-time = "2026-07-23T01:53:39.673Z" }, + { url = "https://files.pythonhosted.org/packages/8e/76/88401ff3fc95e85c5fc38d588f36f55e61ecb64343b2bc8d69326f453cc0/aiohttp-3.14.3-cp311-cp311-win_arm64.whl", hash = "sha256:03cd2bde3d7f085b64e549c985f4bb928cad7e8ecf5323bfca320db548d81b39", size = 453021, upload-time = "2026-07-23T01:53:43.749Z" }, + { url = "https://files.pythonhosted.org/packages/18/d4/eb96299230e20acf2efae207cb8d69051f1f68e357e5ea5e479bf6fb097a/aiohttp-3.14.3-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:39aded8c7f3b935b54aab1d8d73c70ec0ee2d3ec3b943e0e86611bc150ba47f5", size = 754690, upload-time = "2026-07-23T01:53:47.332Z" }, + { url = "https://files.pythonhosted.org/packages/88/11/e7a70a209eb9a067c0d3212b518a0134e3484f5178c7533878b6b514d469/aiohttp-3.14.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:5bcb6ff3fdab1258a192679ff1a05d44f59626430aa05cd1a9d2447423599228", size = 509484, upload-time = "2026-07-23T01:53:51.159Z" }, + { url = "https://files.pythonhosted.org/packages/30/07/4bbc222cc8dbe31d4c3e8a5baad2286e4d42026ac0c570027b89afce6344/aiohttp-3.14.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:617105e2c3018ee38d0c8ce5ee3c84f621a6d8b9f723202aacaff28449ca91ee", size = 511949, upload-time = "2026-07-23T01:53:55.083Z" }, + { url = "https://files.pythonhosted.org/packages/54/b9/42e74c46b7b7c794b995bbc1f573fb48950c38b19d8600c62a6804ee2d67/aiohttp-3.14.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f631fe87a6f30df5fbe6d79640b25e4cffb38c31c7fb6f10871517b84b0f8c1a", size = 1765282, upload-time = "2026-07-23T01:53:59.662Z" }, + { url = "https://files.pythonhosted.org/packages/6b/ed/62bc4d74363ad346d518e0720363a949f63e2e23439a79eb5813d4d29bb3/aiohttp-3.14.3-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:a94dbaae5ae27bd849c93570669bff91e0510f33a80805738e3de72a7be0447b", size = 1741511, upload-time = "2026-07-23T01:54:04.063Z" }, + { url = "https://files.pythonhosted.org/packages/d0/9f/181e8a8bc79e47d13c7fc4540bd7a3b729d9505609c61f392a8dd2fbfe55/aiohttp-3.14.3-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:8f2f1c4c032c7cedd7d8da6f54c97b70266c6570c3108d3fdffee7188bb70529", size = 1810680, upload-time = "2026-07-23T01:54:09.882Z" }, + { url = "https://files.pythonhosted.org/packages/5c/9a/dec94d6ad694552fe3424e3f1928d7a606a5d9d9433a04e7ecdd9d38ae7f/aiohttp-3.14.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:ea05e1f97ceea523942d9b2a7d7c0359d781d683d6b043f5943a602b14da4787", size = 1905646, upload-time = "2026-07-23T01:54:13.475Z" }, + { url = "https://files.pythonhosted.org/packages/52/b7/7cd31f29d6055bd711ae6e669367fba6f5ae9de463910a793e30556a8db7/aiohttp-3.14.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:543906c127fb1d929b95076db19b83fa2d46751006ff1e23b093aa5ac4d8db42", size = 1792122, upload-time = "2026-07-23T01:54:15.752Z" }, + { url = "https://files.pythonhosted.org/packages/66/73/10b1ef93afa61f4963c746257b70ced619cf31a4798671de5fdb2608501d/aiohttp-3.14.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:0a5ff2dfbb9ce645fa5b8ef3e02c6c0b9cc3f6030ff863d0c51fffc50cb5541b", size = 1591127, upload-time = "2026-07-23T01:54:19.489Z" }, + { url = "https://files.pythonhosted.org/packages/49/ed/3b203fa6de1b338c14acdc06bf6ca9b043b7944f005966958c2ced932cde/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:041badb8f84396357c4d3ad26de6afd7a32b112f43d3c63045c0c8278cfd2043", size = 1725210, upload-time = "2026-07-23T01:54:24.129Z" }, + { url = "https://files.pythonhosted.org/packages/28/b7/1c2aab8c706436dcc28598452488ac9cd7c409da815237c28c27d58993e6/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:530125ee1163c4219af35dc3aa1206e541e7b31b6efc1a3f93b70a136f65d427", size = 1764848, upload-time = "2026-07-23T01:54:27.973Z" }, + { url = "https://files.pythonhosted.org/packages/54/50/94c28f08b131c4bf10984ea2c7a536c9920608bb2d6e7f95642c30cc87b7/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:c8653fd547c93a61aadc612007790f5555cdd18946fa48cf45e26d8ea4ea473d", size = 1777102, upload-time = "2026-07-23T01:54:31.775Z" }, + { url = "https://files.pythonhosted.org/packages/13/d4/e7d09ba7d345fb2d74440fd2fa033c5e079fac05552927705986f41a364f/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:89176250f686cb9853c0fb7ead90e639e915b84a6f43eedc2a4e7ec21f1037f0", size = 1580205, upload-time = "2026-07-23T01:54:34.518Z" }, + { url = "https://files.pythonhosted.org/packages/a3/84/072a91d68e1e1eb587985b54baab94221277f877e8ef274fc213a0ceae28/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:3a26434dafe408229ff3403458ca58de24fb51936504decac49ce6755f77e59d", size = 1797219, upload-time = "2026-07-23T01:54:36.995Z" }, + { url = "https://files.pythonhosted.org/packages/e0/eb/aad34e897e668424d6e995da5dff8a4a09af93363d3392488772957a63aa/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:d1558173930a5a8d3069cee5c92fc91c87c4dbcb099debbb3622053717145a19", size = 1768629, upload-time = "2026-07-23T01:54:40.103Z" }, + { url = "https://files.pythonhosted.org/packages/b6/2b/6bb88ddba0fecd9122aa3ebcad25996cf6c083a4a7040dbb3a4f97972af6/aiohttp-3.14.3-cp312-cp312-win32.whl", hash = "sha256:16100ad3ab8d649fdfbee87602d9d2dcdca9df0b9eda8a1b5fdc0d41f96da559", size = 451481, upload-time = "2026-07-23T01:54:42.547Z" }, + { url = "https://files.pythonhosted.org/packages/76/9b/f2f8f108da17ecef2cc3efc424e8b7ad3782b1a8360f7b8eae8ced84f6ea/aiohttp-3.14.3-cp312-cp312-win_amd64.whl", hash = "sha256:33a2d7c28d33797a2e99923dffa63f83d908a19b6bf26cfe80fa790aa5e1a75a", size = 476845, upload-time = "2026-07-23T01:54:44.853Z" }, + { url = "https://files.pythonhosted.org/packages/3e/44/28dac80a8941b604f4da10ce21097614ca1bf905ce93dca28d8d7de9c1e7/aiohttp-3.14.3-cp312-cp312-win_arm64.whl", hash = "sha256:362a3fd481769cac1a824514bcd86fda51c65e8fe6e051099e008fddde6db17c", size = 448050, upload-time = "2026-07-23T01:54:47.087Z" }, + { url = "https://files.pythonhosted.org/packages/57/be/5afd201cc0ab139029aadb75392efe85a293403d9dd3a3226161c21ce00c/aiohttp-3.14.3-cp313-cp313-android_21_arm64_v8a.whl", hash = "sha256:2e9878ae68e4a5f1c0abe4dd497dbc3d51946f5837b56759e2a02e78fa90ef86", size = 506269, upload-time = "2026-07-23T01:54:49.075Z" }, + { url = "https://files.pythonhosted.org/packages/22/09/dec8189d62b45ade009f6792a2264b942a90cb88aeaf181239933cd72c3c/aiohttp-3.14.3-cp313-cp313-android_21_x86_64.whl", hash = "sha256:f3d2669fe7dec7fc359ecdb5984b29b50d85d5d00f8c1cb61de4f4a24ee42627", size = 515166, upload-time = "2026-07-23T01:54:51.894Z" }, + { url = "https://files.pythonhosted.org/packages/28/24/2854869d29ed8a8b19d74f9ec6629515f7e04d02dd329d9d179201e58e47/aiohttp-3.14.3-cp313-cp313-ios_13_0_arm64_iphoneos.whl", hash = "sha256:cc7cb243a68167172f48c1fd43cee91ec4b1d40cefd190edd43369d1a6bc9c82", size = 486263, upload-time = "2026-07-23T01:54:54.223Z" }, + { url = "https://files.pythonhosted.org/packages/d4/dd/57187c8be2a35aea65eaee3bd2c3dcbbcf0204f5106c89637e3610380cd1/aiohttp-3.14.3-cp313-cp313-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:78253b573e6ffab5028924fc98bc281aae05445969982a10864bc360dea2016c", size = 492299, upload-time = "2026-07-23T01:54:56.236Z" }, + { url = "https://files.pythonhosted.org/packages/b9/11/06ae6ed8f0d414edf4068861e233d8fe23ee699bfd4b3ceb8663db948a62/aiohttp-3.14.3-cp313-cp313-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:7041d52c3a7fa20c9e8c182b534704abb19502c8bdcbde7ab23bfda6f642394f", size = 502235, upload-time = "2026-07-23T01:54:58.377Z" }, + { url = "https://files.pythonhosted.org/packages/7e/a3/559639c34a345d2cf7c52dff6838119f2eaf29eb508227b5b83f573af813/aiohttp-3.14.3-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:ac74facc01463f138b0da5580329cfcc82818dea5656e83ddcd11268fc12ff80", size = 750883, upload-time = "2026-07-23T01:55:00.65Z" }, + { url = "https://files.pythonhosted.org/packages/91/cd/41e131f13afd1e7b0172a9d9eda085ef90eb8439f41f0d279db81ed3ae60/aiohttp-3.14.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:d6218d92e450824e9b4881f44e8c09f1853b490f9a64130801024a4793b1b3b0", size = 508473, upload-time = "2026-07-23T01:55:02.945Z" }, + { url = "https://files.pythonhosted.org/packages/bc/6b/e7f13410d391c6e55b4c007a8de024355389d7d459e3d64c42b2d33617e5/aiohttp-3.14.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:11fb37ef075669eee52ab1928fbf6e1741fada40409fa309ebde9607a962aebf", size = 509190, upload-time = "2026-07-23T01:55:05.173Z" }, + { url = "https://files.pythonhosted.org/packages/97/21/6464573e53d69672cc1eada3e5c5cb2d2efa82701e8305a0f2047a576967/aiohttp-3.14.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:55bdcc472aafe2de4a253045cc128007a64f1e0264fb675791e132ea5edaa3bd", size = 1761478, upload-time = "2026-07-23T01:55:07.383Z" }, + { url = "https://files.pythonhosted.org/packages/1a/81/d217043a4c17fbce360905e3b2bdd20139ebc9a2de836d035d179c4da006/aiohttp-3.14.3-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:c39846c3aad97a8530c89d7a3869a8f8e9e3762c6ac0504481e5c80948f7e807", size = 1735092, upload-time = "2026-07-23T01:55:09.803Z" }, + { url = "https://files.pythonhosted.org/packages/a1/66/e13a02d0eeb1a9a502402a977abb4e4abff9fe4051c26f80558c57a7c975/aiohttp-3.14.3-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:5895ef58c4620afe02fa16044f023dc4dafec08158f9d08874a46a7dbc0341b8", size = 1800546, upload-time = "2026-07-23T01:55:12.012Z" }, + { url = "https://files.pythonhosted.org/packages/26/5e/57d42fca1d18cb5acc1cad945d017fabc5d6ae71d8a08ad66be8dc3ee544/aiohttp-3.14.3-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:fa9467a8113aa69d3d7c55a70ef0b7c636010a40993f3df9d9d0d73b3eb7ef24", size = 1895250, upload-time = "2026-07-23T01:55:14.357Z" }, + { url = "https://files.pythonhosted.org/packages/ca/1c/7da8d08e74d56f00070822f9638ff3f1c563f8ad87d1efa996c87bfc8644/aiohttp-3.14.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d7d2deec16eeedf55f2c7cf75b521ea3856a5177e123844f8fd0f114ce252cb5", size = 1789289, upload-time = "2026-07-23T01:55:16.668Z" }, + { url = "https://files.pythonhosted.org/packages/cd/0f/cf16bcf56896981c1a0319f5d5db9337994b5165730c48a8fa07e9b34be6/aiohttp-3.14.3-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:dd54d0e8717de95939766febac482ac0474d8ac3b048115f9f2b1d23a16e7db4", size = 1586706, upload-time = "2026-07-23T01:55:18.913Z" }, + { url = "https://files.pythonhosted.org/packages/fe/6f/76eac12a7f2480e1e304f842efdb07db33256b0d9165b866b6ef0806c202/aiohttp-3.14.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:df82f3787c940c94986b34222d59c9e38843fba85139f36e85255a82ad5355a9", size = 1724652, upload-time = "2026-07-23T01:55:21.296Z" }, + { url = "https://files.pythonhosted.org/packages/39/b6/19c8c592baeeb94b75f966547d40c02ac7590902306ec5863d5c027cf506/aiohttp-3.14.3-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:42a67efc36300d052fb4508a53e8b6901b9284b599ae63945c377569c5fcc1e1", size = 1756239, upload-time = "2026-07-23T01:55:23.705Z" }, + { url = "https://files.pythonhosted.org/packages/dc/c9/4e9383150296f97f873b680c4de8fb2cd88608fb9f48c79edcb111611abc/aiohttp-3.14.3-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:7a75aa63cbf9b21cfaf60dc2657e19df2c2867d91707d653fee171ffeedd1371", size = 1769161, upload-time = "2026-07-23T01:55:26.082Z" }, + { url = "https://files.pythonhosted.org/packages/aa/1e/147bdc6cc5de5f3ab011be8bf5d6e786633249f22c20bae06f85e45f5387/aiohttp-3.14.3-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:e92eb8acc45eb6a9f4935071a77edf5b85cc6f8dfad5cd99e97653c26593cdde", size = 1578759, upload-time = "2026-07-23T01:55:28.846Z" }, + { url = "https://files.pythonhosted.org/packages/fd/31/78388a9d6040ece2e11df62ea229a822cf5e52d238374b220ae9975b2623/aiohttp-3.14.3-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:b014a6ed7cf912e787149fdc529166d3ceabac23f26efeea3158c9aba2354e7e", size = 1792025, upload-time = "2026-07-23T01:55:31.457Z" }, + { url = "https://files.pythonhosted.org/packages/03/51/a3d29fdf2c25d796746af8ad6fe56a45d6256c38b0a8a2ed752e1160b3a2/aiohttp-3.14.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:3d4f72af88ac2474bb5bca640030320e3d38a0163a1d7533500e87be458eef71", size = 1768477, upload-time = "2026-07-23T01:55:33.87Z" }, + { url = "https://files.pythonhosted.org/packages/29/a6/442e18b5afeade534d877a2dc3c3e392aff8d49787890b0cf84790410267/aiohttp-3.14.3-cp313-cp313-win32.whl", hash = "sha256:5f08ec777f35ee70720233b8b9811d3bb5d728137f30ac91b7457709c3261ac0", size = 451069, upload-time = "2026-07-23T01:55:36.121Z" }, + { url = "https://files.pythonhosted.org/packages/9d/69/3d876ac02659f271cf7f6769f14a8e3de5b6e888ed8b5a7e998086a4cec8/aiohttp-3.14.3-cp313-cp313-win_amd64.whl", hash = "sha256:dff9461ec275f22135650d5ba4b4931a11f3958df7dfbb8db630000d4dee0883", size = 476518, upload-time = "2026-07-23T01:55:38.303Z" }, + { url = "https://files.pythonhosted.org/packages/b2/0e/50d6e6471cd31edce8b282bdec59375a3a69124d8a989a0b1313355cae52/aiohttp-3.14.3-cp313-cp313-win_arm64.whl", hash = "sha256:ddcac3c6b382e81f1dd0499199d4136b877beb4cb5ef770bbbfba56c4b8f55d2", size = 447676, upload-time = "2026-07-23T01:55:40.451Z" }, + { url = "https://files.pythonhosted.org/packages/c8/20/887fdcf832326571b370ffc347b3e70abe101096f3720126aac161b1d872/aiohttp-3.14.3-cp314-cp314-android_24_arm64_v8a.whl", hash = "sha256:49f7325beb0f85ef4aef5f48f490269575f83e6e2acad00a1d80b807eb027062", size = 509067, upload-time = "2026-07-23T01:55:42.618Z" }, + { url = "https://files.pythonhosted.org/packages/ad/a3/92cec936f78cc4bf0fa5554ebe593b73459d94e3c62303e1902a4cccb6f7/aiohttp-3.14.3-cp314-cp314-android_24_x86_64.whl", hash = "sha256:e3be98a7c30b8c25d573dafba7171d66dfb05ee6a9070fc46535464ff97700a6", size = 514774, upload-time = "2026-07-23T01:55:44.937Z" }, + { url = "https://files.pythonhosted.org/packages/29/ba/2a0c38df3fc557620b6a5acd98364af050053b6285b4dc7ee74100c63c18/aiohttp-3.14.3-cp314-cp314-ios_13_0_arm64_iphoneos.whl", hash = "sha256:614c61d478b83953e261d02bb2df750f17227cd33ef8002945bf5aebbde21919", size = 488134, upload-time = "2026-07-23T01:55:47.135Z" }, + { url = "https://files.pythonhosted.org/packages/48/d6/d51b7d4bf309af3693940d8ffd2b9ed0b682434ef85959b7c9c137f60cf8/aiohttp-3.14.3-cp314-cp314-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:1caa7b0d05f3e3a36f87788c59e970a7ee1cefcfcbb924a9f138c4a6551c9cb7", size = 494201, upload-time = "2026-07-23T01:55:49.451Z" }, + { url = "https://files.pythonhosted.org/packages/3f/5a/8f624384e5f1efabb5229b94157eb966b021e97bdb188c62860c2ae243c2/aiohttp-3.14.3-cp314-cp314-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:dfa68deb2a443bdaa3ea5297b0699c1464f08aef3812b486d1348eee61b07dc0", size = 502766, upload-time = "2026-07-23T01:55:51.656Z" }, + { url = "https://files.pythonhosted.org/packages/a6/26/4ff0164370deec18fb19254ee4ab10b7a73304ac0c860b13f5f84663759b/aiohttp-3.14.3-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:e72ee89e28d907a18f46959b4eb0bb06701cc7f8cf4366e00029e2ccfaaf5924", size = 756557, upload-time = "2026-07-23T01:55:53.964Z" }, + { url = "https://files.pythonhosted.org/packages/97/a3/7056b86dc0d9ec709ea9777eae3b0161428f943372f8b98c01c11593b682/aiohttp-3.14.3-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:ad4c8b7488d745d2ca4838ebd8ae5ba9b56341d30b1da43640e4ce87f9f49646", size = 510168, upload-time = "2026-07-23T01:55:56.22Z" }, + { url = "https://files.pythonhosted.org/packages/85/ed/0357a015892fd68058bf2d39d3fd1958e459b997a7db30aaa6aaa434ae96/aiohttp-3.14.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:db332af25642007330fca8be5c4d194caf2bea7a7fc84415aff3497af5dfee6b", size = 512957, upload-time = "2026-07-23T01:55:58.437Z" }, + { url = "https://files.pythonhosted.org/packages/47/d1/8aba53f15ccb2238405f5e9d30e2a8ca44f93878c26e7165ade00d374b1c/aiohttp-3.14.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:25bd2708db6bdf6a6630dd37bdcdfcb47c4434d22ac69c64665b802910140b30", size = 1750149, upload-time = "2026-07-23T01:56:00.856Z" }, + { url = "https://files.pythonhosted.org/packages/49/bd/40c3fee327529284375c6701cbb0fa4600cc2e8432af1378f897e2ef7d3a/aiohttp-3.14.3-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:cef89a58e628c4efcac3275c2d68083f82426dcdc89c1492a6f654f9f7ea6ab9", size = 1707685, upload-time = "2026-07-23T01:56:03.371Z" }, + { url = "https://files.pythonhosted.org/packages/2a/a3/ca0cc6724cca8114b05694abd916060758c79894c3aa5b012cdadc1bc28e/aiohttp-3.14.3-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c23ec8ee9d5ab2f5421f9c7fffce208435607af27fd46d4a44e031954352838f", size = 1803911, upload-time = "2026-07-23T01:56:05.817Z" }, + { url = "https://files.pythonhosted.org/packages/95/b5/85b099c299c3ffd38ad9b3e43694c8a346934e4a30c88c4fd5a841234f77/aiohttp-3.14.3-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:e2667f0bbe7eb6c74eae5e9691441ad186e5845ca3cff63230fc09c4e7514f5d", size = 1876929, upload-time = "2026-07-23T01:56:08.413Z" }, + { url = "https://files.pythonhosted.org/packages/d5/b7/1da684a04175473fa4cddbf9a2f572e79514c3fd27a74597f43057d4f3da/aiohttp-3.14.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:18cb43369747b2ae007bd2655fb8e63a099c2ff1d207962943636dac989b3147", size = 1761112, upload-time = "2026-07-23T01:56:10.918Z" }, + { url = "https://files.pythonhosted.org/packages/d1/16/bc4b55e3e5cb175fd69c53c90d60d2f47797cb343da5106e23863dc4dba4/aiohttp-3.14.3-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:d77640cc618c1d99fc4f8589c0f24a730adfa54eb1e57ef7bf0c8dfb78da898c", size = 1583500, upload-time = "2026-07-23T01:56:13.613Z" }, + { url = "https://files.pythonhosted.org/packages/2a/e8/13a9d957a1ee40837f46aa30f0f4c657e673ad86a2e6362a9f9be20d26d9/aiohttp-3.14.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:53e5179d8abb5710f8e83ba207c41c8d1261fcffd4616500e15ca2b7a33be10a", size = 1713940, upload-time = "2026-07-23T01:56:15.969Z" }, + { url = "https://files.pythonhosted.org/packages/38/05/d33c680c1bcf1c7e130f9cbfc1fc02fe8bb0c4af2a94a53dd5fb56131e5c/aiohttp-3.14.3-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:cd817772b2fcf2b8c0905795318485f9ec16eae60b29feb7f4c77085311637f0", size = 1724413, upload-time = "2026-07-23T01:56:18.591Z" }, + { url = "https://files.pythonhosted.org/packages/85/1d/af798d306f7a74b6a632dbcabcf62a4c91391b7582d2a8c6d7712e2cc54e/aiohttp-3.14.3-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:4e3ac92d90e92773b2362d506068e9a948192bd553e743c5b2429e28527c8661", size = 1770748, upload-time = "2026-07-23T01:56:21.074Z" }, + { url = "https://files.pythonhosted.org/packages/a8/92/ad720d472556a995049206867765e9410969684f86ee09423ff9969044c1/aiohttp-3.14.3-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:3f42e9b78301f11c8f861746175d8b9c1ccef713fcad9eab396e2f6db8ed4a22", size = 1577564, upload-time = "2026-07-23T01:56:23.475Z" }, + { url = "https://files.pythonhosted.org/packages/60/ad/0ed7586cbef7a884e23a752fa2bb987a122e6a5dd50dab109258d0a95193/aiohttp-3.14.3-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:9d9edccfe496b476db5f398d97b865e9a6752bcf8aec4eef8390ce20fb64bb41", size = 1782080, upload-time = "2026-07-23T01:56:25.994Z" }, + { url = "https://files.pythonhosted.org/packages/97/ea/dbaed0d73e8a69aad653b045dab451c67c2454bb731a37b45a86593e9422/aiohttp-3.14.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:1c5ec8fb1bcc31a8466f74aaf26c345d5c386fa4bd08a3f0eb9c7a4a3fe8b5bf", size = 1745813, upload-time = "2026-07-23T01:56:28.604Z" }, + { url = "https://files.pythonhosted.org/packages/81/1b/6893d4bc57e434fc93a6c9217c637d967a0b651d989f6e3265179375754a/aiohttp-3.14.3-cp314-cp314-win32.whl", hash = "sha256:38901a84da3ce22249f6e860bf8f90d141bcab7da090cc398f8bb58c0e44b7da", size = 455872, upload-time = "2026-07-23T01:56:31.031Z" }, + { url = "https://files.pythonhosted.org/packages/f5/8b/c7baa1ba1eda4db6989baefe5de6d99834921b84ebd7918624febcb9f290/aiohttp-3.14.3-cp314-cp314-win_amd64.whl", hash = "sha256:8b3b60de05f3dcb6f6a00f818bb2ec781cee4de0645f59ccaf99b1d1823b6100", size = 481030, upload-time = "2026-07-23T01:56:33.365Z" }, + { url = "https://files.pythonhosted.org/packages/22/8c/c29d067df825a2df88ca432db848aa2fe8199598359cc06c12b09320cac9/aiohttp-3.14.3-cp314-cp314-win_arm64.whl", hash = "sha256:1576145bdceeb92382d899751e12743a3a5b8e460a841e3e50543859e54864dc", size = 453669, upload-time = "2026-07-23T01:56:35.731Z" }, + { url = "https://files.pythonhosted.org/packages/6a/a4/9c033beb355d39b6147980597ec9645e4729243f686ee4dc73945de72030/aiohttp-3.14.3-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:8800c996b01c2772a783e3e46f3e1abd5823029adca0df54231960de9bfefa5b", size = 791403, upload-time = "2026-07-23T01:56:37.972Z" }, + { url = "https://files.pythonhosted.org/packages/80/ca/87c32a0a7704583cfc49660bd817889bae5b830bf53b5dcb4e92145ac2da/aiohttp-3.14.3-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:ebe8e504f058fe91223351cecd2d9d6946c9d241bb0250d898ffbdf584cc72b0", size = 526413, upload-time = "2026-07-23T01:56:40.523Z" }, + { url = "https://files.pythonhosted.org/packages/9e/d8/8ec0e471248c500acdce2be3f46db8fb62b5eb60efef072529cc85ee1d26/aiohttp-3.14.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:30402d03a7c0ff52bce290b57e564e9079fd9d0cb545c8aba73f86a103162d2e", size = 532135, upload-time = "2026-07-23T01:56:42.876Z" }, + { url = "https://files.pythonhosted.org/packages/fe/45/f8919fd936e8b79fcd9bda7b6d8e62613462a713f4f17987fd7c34399142/aiohttp-3.14.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9fc7b5bfec6573f3ae844f457fdde5adeb713f8b8e4a81ad64fc207b49383716", size = 1922742, upload-time = "2026-07-23T01:56:45.528Z" }, + { url = "https://files.pythonhosted.org/packages/f6/ec/9ca76b28a27525b0cc53e20842e0228b022f301ce1f436b7d814b4aaf2df/aiohttp-3.14.3-cp314-cp314t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:8a5fd34f7f7410d1730d5c2ba873cacb2eed3fede366feb268a70ba22581ed8f", size = 1787371, upload-time = "2026-07-23T01:56:48.045Z" }, + { url = "https://files.pythonhosted.org/packages/b1/04/6acdbf17315f7b55f1937e3387acb89a3cddeb4995689553d064af8e92ab/aiohttp-3.14.3-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:270d3dace9ca2f10f0da5d8ebe519b7a310fc6112ed916e32df5866df0888553", size = 1912623, upload-time = "2026-07-23T01:56:50.605Z" }, + { url = "https://files.pythonhosted.org/packages/86/e6/438b0c79ca6f45eb9fd9817dd4c01a91919a38c0de5ee9e05e2b4dc0ece7/aiohttp-3.14.3-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:3ae5b3a59436d089b5395d910121a390feed4d00578eb95a0fd1a329fe963100", size = 2005515, upload-time = "2026-07-23T01:56:53.153Z" }, + { url = "https://files.pythonhosted.org/packages/bb/6b/62cbd6577758699525f5c712d1ddef57d9875fbab0ae8d5f5a202fd598f8/aiohttp-3.14.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2498f0fe69ead802f9675beca44a7c21c62fdaa4ec5145ea1c3ad6edbee29f85", size = 1879906, upload-time = "2026-07-23T01:56:55.818Z" }, + { url = "https://files.pythonhosted.org/packages/00/95/18bcbf830a21dc3aae24d8f6b6feaf3db1d2090242d00a7868db2ffb0b67/aiohttp-3.14.3-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:a0dc483c00da8b673abbb367eb6f8d8f4bcec30eb58529ea13cb42e7fd2dfa33", size = 1675849, upload-time = "2026-07-23T01:56:58.861Z" }, + { url = "https://files.pythonhosted.org/packages/a9/19/47f4968659c5e23606c3790c80fc624e691c153d036148449ee84d31b287/aiohttp-3.14.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:c7d3a97c678d34fc5b59da671ee9cd630096ddc643e7b5a30d54a2a6f3574d3f", size = 1843496, upload-time = "2026-07-23T01:57:01.591Z" }, + { url = "https://files.pythonhosted.org/packages/64/af/38c33c4dd82fddcb4e56c4653b6f1072a8edbc6b7fa15809f14932c41e2d/aiohttp-3.14.3-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:f8fb78a83c9e5f741ca3a68cfb455c1f5bb83b4e7249a3848b3cd78d0a8563b0", size = 1827746, upload-time = "2026-07-23T01:57:05.131Z" }, + { url = "https://files.pythonhosted.org/packages/a1/9d/0537cda4885ac8f5b7053d164dd06312f4c483a4edcb8ee5b8aaf2a989bf/aiohttp-3.14.3-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:74ab5b6a9fb13e873e5a90946588baecaf488745e1db1a4a5c433f971f035098", size = 1853810, upload-time = "2026-07-23T01:57:08.043Z" }, + { url = "https://files.pythonhosted.org/packages/19/fe/26f9c5e6458385aa86497836b0dea6fb2f027827d63f37c7856cce9286ee/aiohttp-3.14.3-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:bd52f811e65f6fb634b1047159657c98f52b407f8efec907bcfc09da9a4c0a25", size = 1668895, upload-time = "2026-07-23T01:57:10.837Z" }, + { url = "https://files.pythonhosted.org/packages/ec/4c/618b1db9b9ba079b8875d2cdf78e7c4a3bf72903bd5850fee7dd9544600a/aiohttp-3.14.3-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:f0f177d1b195b9e06376cfd7d308d8a1b920909a609d03ac82a8c73bbb16d3b9", size = 1883833, upload-time = "2026-07-23T01:57:13.672Z" }, + { url = "https://files.pythonhosted.org/packages/94/c6/bd959bd1e4771f9fd944e9e436224c48c77b018b73b519b5aad346335bcc/aiohttp-3.14.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:498c6c623134f8e09a3c4e60bcd607a0b4590dd7dbf08dd40851b27cbb520ccb", size = 1844251, upload-time = "2026-07-23T01:57:16.593Z" }, + { url = "https://files.pythonhosted.org/packages/5e/19/08d41839658bdd44a0ed2480f3891705ecb487ce28c0dde62c9040c997e0/aiohttp-3.14.3-cp314-cp314t-win32.whl", hash = "sha256:b304db572b4368edd8dda8a2274f73156fe15558fca4a917cb8a09fc47af5963", size = 474180, upload-time = "2026-07-23T01:57:19.306Z" }, + { url = "https://files.pythonhosted.org/packages/99/5d/3cd6ef0a2b2851f7ab913b5b079334781bd50ff56a323e4454063377a080/aiohttp-3.14.3-cp314-cp314t-win_amd64.whl", hash = "sha256:b20032766aedf6261c7a566585a40867d092ac03a0d81592d5370ef9b054f99b", size = 500528, upload-time = "2026-07-23T01:57:21.762Z" }, + { url = "https://files.pythonhosted.org/packages/a4/37/cfd1ed540a4d318da025590d96b728e63713c09e9377950fc655dadeb856/aiohttp-3.14.3-cp314-cp314t-win_arm64.whl", hash = "sha256:2e1161602f45a54de2ce0905243a95f58cb42dcd378402f3697f5e0b21e9d2e7", size = 469280, upload-time = "2026-07-23T01:57:24.241Z" }, +] + +[[package]] +name = "aiosignal" +version = "1.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "frozenlist" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/61/62/06741b579156360248d1ec624842ad0edf697050bbaf7c3e46394e106ad1/aiosignal-1.4.0.tar.gz", hash = "sha256:f47eecd9468083c2029cc99945502cb7708b082c232f9aca65da147157b251c7", size = 25007, upload-time = "2025-07-03T22:54:43.528Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/76/641ae371508676492379f16e2fa48f4e2c11741bd63c48be4b12a6b09cba/aiosignal-1.4.0-py3-none-any.whl", hash = "sha256:053243f8b92b990551949e63930a839ff0cf0b0ebbe0597b0f3fb19e1a0fe82e", size = 7490, upload-time = "2025-07-03T22:54:42.156Z" }, +] + [[package]] name = "alabaster" version = "1.0.0" @@ -37,16 +177,15 @@ wheels = [ [[package]] name = "anyio" -version = "4.11.0" +version = "4.14.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "idna" }, - { name = "sniffio" }, { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/c6/78/7d432127c41b50bccba979505f272c16cbcadcc33645d5fa3a738110ae75/anyio-4.11.0.tar.gz", hash = "sha256:82a8d0b81e318cc5ce71a5f1f8b5c4e63619620b63141ef8c995fa0db95a57c4", size = 219094, upload-time = "2025-09-23T09:19:12.58Z" } +sdist = { url = "https://files.pythonhosted.org/packages/61/cc/a381afa6efea9f496eff839d4a6a1aed3bfafc7b3ab4b0d1b243a12573dd/anyio-4.14.2.tar.gz", hash = "sha256:cfa139f3ed1a23ee8f88a145ddb5ac7605b8bbfd8592baacd7ce3d8bb4313c7f", size = 260176, upload-time = "2026-07-12T20:29:07.082Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/15/b3/9b1a8074496371342ec1e796a96f99c82c945a339cd81a8e73de28b4cf9e/anyio-4.11.0-py3-none-any.whl", hash = "sha256:0287e96f4d26d4149305414d4e3bc32f0dcd0862365a4bddea19d7a1ec38c4fc", size = 109097, upload-time = "2025-09-23T09:19:10.601Z" }, + { url = "https://files.pythonhosted.org/packages/da/35/f2287558c17e29fafc8ef3daf819bb9834061cfa43bff8014f7df7f63bdc/anyio-4.14.2-py3-none-any.whl", hash = "sha256:9f505dda5ac9f0c8309b5e8bd445a8c2bf7246f3ce950121e45ea15bc41d1494", size = 125813, upload-time = "2026-07-12T20:29:05.763Z" }, ] [[package]] @@ -76,31 +215,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b7/b8/3fe70c75fe32afc4bb507f75563d39bc5642255d1d94f1f23604725780bf/babel-2.17.0-py3-none-any.whl", hash = "sha256:4d0b53093fdfb4b21c92b5213dba5a1b23885afa8383709427046b21c366e5f2", size = 10182537, upload-time = "2025-02-01T15:17:37.39Z" }, ] -[[package]] -name = "basedpyright" -version = "1.31.7" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "nodejs-wheel-binaries" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/c6/ba/ed69e8df732a09c8ca469f592c8e08707fe29149735b834c276d94d4a3da/basedpyright-1.31.7.tar.gz", hash = "sha256:394f334c742a19bcc5905b2455c9f5858182866b7679a6f057a70b44b049bceb", size = 22710948, upload-time = "2025-10-11T05:12:48.3Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f8/90/ce01ad2d0afdc1b82b8b5aaba27e60d2e138e39d887e71c35c55d8f1bfcd/basedpyright-1.31.7-py3-none-any.whl", hash = "sha256:7c54beb7828c9ed0028630aaa6904f395c27e5a9f5a313aa9e91fc1d11170831", size = 11817571, upload-time = "2025-10-11T05:12:45.432Z" }, -] - -[[package]] -name = "beautifulsoup4" -version = "4.14.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "soupsieve" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/77/e9/df2358efd7659577435e2177bfa69cba6c33216681af51a707193dec162a/beautifulsoup4-4.14.2.tar.gz", hash = "sha256:2a98ab9f944a11acee9cc848508ec28d9228abfd522ef0fad6a02a72e0ded69e", size = 625822, upload-time = "2025-09-29T10:05:42.613Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/94/fe/3aed5d0be4d404d12d36ab97e2f1791424d9ca39c2f754a6285d59a3b01d/beautifulsoup4-4.14.2-py3-none-any.whl", hash = "sha256:5ef6fa3a8cbece8488d66985560f97ed091e22bbc4e9c2338508a9d5de6d4515", size = 106392, upload-time = "2025-09-29T10:05:43.771Z" }, -] - [[package]] name = "bench-orchestrator" version = "0.1.0" @@ -126,19 +240,6 @@ requires-dist = [ [package.metadata.requires-dev] dev = [{ name = "pandas-stubs", specifier = ">=2.3.3.251201" }] -[[package]] -name = "breathe" -version = "4.36.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "sphinx", version = "9.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, - { name = "sphinx", version = "9.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/01/56/99bf7d0799d95ad485d95596dc01c2a5b3dda58ebf50a94f6f73b33bacdf/breathe-4.36.0.tar.gz", hash = "sha256:14860b73118ac140b7a3f55446890c777d1b67149cb024279fe3710dad7f535c", size = 154842, upload-time = "2025-02-22T18:36:03.36Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/2c/bc/d67ef1e11ed6e6343c135bf605aa9d5734ff0cc77eb42a2a41f182abc9d9/breathe-4.36.0-py3-none-any.whl", hash = "sha256:af85436f1f09e842bd1fd95617281211c635f8768d245ff830c59b979888d1d5", size = 97231, upload-time = "2025-02-22T18:36:01.087Z" }, -] - [[package]] name = "certifi" version = "2025.10.5" @@ -203,14 +304,14 @@ wheels = [ [[package]] name = "click" -version = "8.2.1" +version = "8.4.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "colorama", marker = "sys_platform == 'win32'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/60/6c/8ca2efa64cf75a977a0d7fac081354553ebe483345c734fb6b6515d96bbc/click-8.2.1.tar.gz", hash = "sha256:27c491cc05d968d271d5a1db13e3b5a184636d9d930f148c50b038f0d0646202", size = 286342, upload-time = "2025-05-20T23:19:49.832Z" } +sdist = { url = "https://files.pythonhosted.org/packages/76/d4/81420972a676e8ffea40450d8c8c92943e7218a78fe9b64359836cc9876b/click-8.4.2.tar.gz", hash = "sha256:9a6cea6e60b17ebe0a44c5cc636d94f09bd66142c1cd7d8b4cd731c4917a15f6", size = 338000, upload-time = "2026-06-24T17:45:15.148Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/85/32/10bb5764d90a8eee674e9dc6f4db6a0ab47c8c4d0d83c27f7c39ac415a4d/click-8.2.1-py3-none-any.whl", hash = "sha256:61a3265b914e850b85317d0b3109c7f8cd35a670f963866005d6ef1d5175a12b", size = 102215, upload-time = "2025-05-20T23:19:47.796Z" }, + { url = "https://files.pythonhosted.org/packages/fb/e2/79c688af8b210d232694e31e59da9f6ec747bae31c3f5946e4e9b98860d5/click-8.4.2-py3-none-any.whl", hash = "sha256:e6f9f66136c816745b9d65817da91d61d957fb16e02e4dcd0552553c5a197b76", size = 119243, upload-time = "2026-06-24T17:45:13.73Z" }, ] [[package]] @@ -222,6 +323,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, ] +[[package]] +name = "datasets" +version = "5.0.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "dill" }, + { name = "filelock" }, + { name = "fsspec", extra = ["http"] }, + { name = "httpx" }, + { name = "huggingface-hub" }, + { name = "multiprocess" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "pandas" }, + { name = "pyarrow" }, + { name = "pyyaml" }, + { name = "requests" }, + { name = "tqdm" }, + { name = "xxhash" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0a/5b/836516269d4f618efe621661cfb6f9acc57e6f95265db3efaee48a5ffe04/datasets-5.0.1.tar.gz", hash = "sha256:ce22bb851efd7494f08aad33b940803784434f6e77763d00679a0dc45fcf686a", size = 641498, upload-time = "2026-07-28T11:09:12.016Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/44/0b/98fc6eb83333508ca5f44c52b3e287ea8137a0ad582714e2cbc67a02154b/datasets-5.0.1-py3-none-any.whl", hash = "sha256:9fbf73688f8c18f7529b4fe592abd04015f81d1e58001e4bac73ffb2b39d7cc4", size = 559079, upload-time = "2026-07-28T11:09:10.266Z" }, +] + [[package]] name = "decorator" version = "5.2.1" @@ -231,47 +357,38 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4e/8c/f3147f5c4b73e7550fe5f9352eaa956ae838d5c51eb58e7a25b9f3e2643b/decorator-5.2.1-py3-none-any.whl", hash = "sha256:d316bb415a2d9e2d2b3abcc4084c6502fc09240e292cd76a76afc106a1c8e04a", size = 9190, upload-time = "2025-02-24T04:41:32.565Z" }, ] +[[package]] +name = "dill" +version = "0.4.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/81/e1/56027a71e31b02ddc53c7d65b01e68edf64dea2932122fe7746a516f75d5/dill-0.4.1.tar.gz", hash = "sha256:423092df4182177d4d8ba8290c8a5b640c66ab35ec7da59ccfa00f6fa3eea5fa", size = 187315, upload-time = "2026-01-19T02:36:56.85Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1e/77/dc8c558f7593132cf8fefec57c4f60c83b16941c574ac5f619abb3ae7933/dill-0.4.1-py3-none-any.whl", hash = "sha256:1e1ce33e978ae97fcfcff5638477032b801c46c7c65cf717f95fbc2248f79a9d", size = 120019, upload-time = "2026-01-19T02:36:55.663Z" }, +] + [[package]] name = "docs" version = "0.1.0" source = { virtual = "docs" } dependencies = [ - { name = "breathe" }, + { name = "accessible-pygments" }, { name = "hawkmoth" }, { name = "libclang" }, { name = "myst-parser" }, - { name = "pydata-sphinx-theme" }, - { name = "setuptools" }, + { name = "rcssmin" }, { name = "sphinx", version = "9.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, { name = "sphinx", version = "9.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, - { name = "sphinx-autobuild" }, - { name = "sphinx-copybutton" }, - { name = "sphinx-design" }, - { name = "sphinx-inline-tabs" }, - { name = "sphinxcontrib-bibtex" }, - { name = "sphinxcontrib-mermaid" }, - { name = "sphinxext-opengraph" }, - { name = "starlette" }, { name = "vortex-data" }, ] [package.metadata] requires-dist = [ - { name = "breathe", specifier = ">=4.35.0" }, + { name = "accessible-pygments", specifier = ">=0.0.5" }, { name = "hawkmoth", specifier = ">=0.20.0" }, { name = "libclang", specifier = ">=18.1.1" }, { name = "myst-parser", specifier = ">=4.0.0" }, - { name = "pydata-sphinx-theme", specifier = ">=0.16.0" }, - { name = "setuptools", specifier = ">=75.8.0" }, + { name = "rcssmin", specifier = ">=1.1.2" }, { name = "sphinx", specifier = ">=9.0.0" }, - { name = "sphinx-autobuild", specifier = ">=2024.10.3" }, - { name = "sphinx-copybutton", specifier = ">=0.5.2" }, - { name = "sphinx-design", specifier = ">=0.6.0" }, - { name = "sphinx-inline-tabs", specifier = ">=2023.4.21" }, - { name = "sphinxcontrib-bibtex", specifier = ">=2.6.3" }, - { name = "sphinxcontrib-mermaid", specifier = ">=1.0.0" }, - { name = "sphinxext-opengraph", specifier = ">=0.9.1" }, - { name = "starlette", specifier = ">=0.49.1" }, { name = "vortex-data", editable = "vortex-python" }, ] @@ -347,6 +464,125 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b5/36/7fb70f04bf00bc646cd5bb45aa9eddb15e19437a28b8fb2b4a5249fac770/filelock-3.20.3-py3-none-any.whl", hash = "sha256:4b0dda527ee31078689fc205ec4f1c1bf7d56cf88b6dc9426c4f230e46c2dce1", size = 16701, upload-time = "2026-01-09T17:55:04.334Z" }, ] +[[package]] +name = "frozenlist" +version = "1.8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2d/f5/c831fac6cc817d26fd54c7eaccd04ef7e0288806943f7cc5bbf69f3ac1f0/frozenlist-1.8.0.tar.gz", hash = "sha256:3ede829ed8d842f6cd48fc7081d7a41001a56f1f38603f9d49bf3020d59a31ad", size = 45875, upload-time = "2025-10-06T05:38:17.865Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bc/03/077f869d540370db12165c0aa51640a873fb661d8b315d1d4d67b284d7ac/frozenlist-1.8.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:09474e9831bc2b2199fad6da3c14c7b0fbdd377cce9d3d77131be28906cb7d84", size = 86912, upload-time = "2025-10-06T05:35:45.98Z" }, + { url = "https://files.pythonhosted.org/packages/df/b5/7610b6bd13e4ae77b96ba85abea1c8cb249683217ef09ac9e0ae93f25a91/frozenlist-1.8.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:17c883ab0ab67200b5f964d2b9ed6b00971917d5d8a92df149dc2c9779208ee9", size = 50046, upload-time = "2025-10-06T05:35:47.009Z" }, + { url = "https://files.pythonhosted.org/packages/6e/ef/0e8f1fe32f8a53dd26bdd1f9347efe0778b0fddf62789ea683f4cc7d787d/frozenlist-1.8.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:fa47e444b8ba08fffd1c18e8cdb9a75db1b6a27f17507522834ad13ed5922b93", size = 50119, upload-time = "2025-10-06T05:35:48.38Z" }, + { url = "https://files.pythonhosted.org/packages/11/b1/71a477adc7c36e5fb628245dfbdea2166feae310757dea848d02bd0689fd/frozenlist-1.8.0-cp311-cp311-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:2552f44204b744fba866e573be4c1f9048d6a324dfe14475103fd51613eb1d1f", size = 231067, upload-time = "2025-10-06T05:35:49.97Z" }, + { url = "https://files.pythonhosted.org/packages/45/7e/afe40eca3a2dc19b9904c0f5d7edfe82b5304cb831391edec0ac04af94c2/frozenlist-1.8.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:957e7c38f250991e48a9a73e6423db1bb9dd14e722a10f6b8bb8e16a0f55f695", size = 233160, upload-time = "2025-10-06T05:35:51.729Z" }, + { url = "https://files.pythonhosted.org/packages/a6/aa/7416eac95603ce428679d273255ffc7c998d4132cfae200103f164b108aa/frozenlist-1.8.0-cp311-cp311-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:8585e3bb2cdea02fc88ffa245069c36555557ad3609e83be0ec71f54fd4abb52", size = 228544, upload-time = "2025-10-06T05:35:53.246Z" }, + { url = "https://files.pythonhosted.org/packages/8b/3d/2a2d1f683d55ac7e3875e4263d28410063e738384d3adc294f5ff3d7105e/frozenlist-1.8.0-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:edee74874ce20a373d62dc28b0b18b93f645633c2943fd90ee9d898550770581", size = 243797, upload-time = "2025-10-06T05:35:54.497Z" }, + { url = "https://files.pythonhosted.org/packages/78/1e/2d5565b589e580c296d3bb54da08d206e797d941a83a6fdea42af23be79c/frozenlist-1.8.0-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:c9a63152fe95756b85f31186bddf42e4c02c6321207fd6601a1c89ebac4fe567", size = 247923, upload-time = "2025-10-06T05:35:55.861Z" }, + { url = "https://files.pythonhosted.org/packages/aa/c3/65872fcf1d326a7f101ad4d86285c403c87be7d832b7470b77f6d2ed5ddc/frozenlist-1.8.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:b6db2185db9be0a04fecf2f241c70b63b1a242e2805be291855078f2b404dd6b", size = 230886, upload-time = "2025-10-06T05:35:57.399Z" }, + { url = "https://files.pythonhosted.org/packages/a0/76/ac9ced601d62f6956f03cc794f9e04c81719509f85255abf96e2510f4265/frozenlist-1.8.0-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:f4be2e3d8bc8aabd566f8d5b8ba7ecc09249d74ba3c9ed52e54dc23a293f0b92", size = 245731, upload-time = "2025-10-06T05:35:58.563Z" }, + { url = "https://files.pythonhosted.org/packages/b9/49/ecccb5f2598daf0b4a1415497eba4c33c1e8ce07495eb07d2860c731b8d5/frozenlist-1.8.0-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:c8d1634419f39ea6f5c427ea2f90ca85126b54b50837f31497f3bf38266e853d", size = 241544, upload-time = "2025-10-06T05:35:59.719Z" }, + { url = "https://files.pythonhosted.org/packages/53/4b/ddf24113323c0bbcc54cb38c8b8916f1da7165e07b8e24a717b4a12cbf10/frozenlist-1.8.0-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:1a7fa382a4a223773ed64242dbe1c9c326ec09457e6b8428efb4118c685c3dfd", size = 241806, upload-time = "2025-10-06T05:36:00.959Z" }, + { url = "https://files.pythonhosted.org/packages/a7/fb/9b9a084d73c67175484ba2789a59f8eebebd0827d186a8102005ce41e1ba/frozenlist-1.8.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:11847b53d722050808926e785df837353bd4d75f1d494377e59b23594d834967", size = 229382, upload-time = "2025-10-06T05:36:02.22Z" }, + { url = "https://files.pythonhosted.org/packages/95/a3/c8fb25aac55bf5e12dae5c5aa6a98f85d436c1dc658f21c3ac73f9fa95e5/frozenlist-1.8.0-cp311-cp311-win32.whl", hash = "sha256:27c6e8077956cf73eadd514be8fb04d77fc946a7fe9f7fe167648b0b9085cc25", size = 39647, upload-time = "2025-10-06T05:36:03.409Z" }, + { url = "https://files.pythonhosted.org/packages/0a/f5/603d0d6a02cfd4c8f2a095a54672b3cf967ad688a60fb9faf04fc4887f65/frozenlist-1.8.0-cp311-cp311-win_amd64.whl", hash = "sha256:ac913f8403b36a2c8610bbfd25b8013488533e71e62b4b4adce9c86c8cea905b", size = 44064, upload-time = "2025-10-06T05:36:04.368Z" }, + { url = "https://files.pythonhosted.org/packages/5d/16/c2c9ab44e181f043a86f9a8f84d5124b62dbcb3a02c0977ec72b9ac1d3e0/frozenlist-1.8.0-cp311-cp311-win_arm64.whl", hash = "sha256:d4d3214a0f8394edfa3e303136d0575eece0745ff2b47bd2cb2e66dd92d4351a", size = 39937, upload-time = "2025-10-06T05:36:05.669Z" }, + { url = "https://files.pythonhosted.org/packages/69/29/948b9aa87e75820a38650af445d2ef2b6b8a6fab1a23b6bb9e4ef0be2d59/frozenlist-1.8.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:78f7b9e5d6f2fdb88cdde9440dc147259b62b9d3b019924def9f6478be254ac1", size = 87782, upload-time = "2025-10-06T05:36:06.649Z" }, + { url = "https://files.pythonhosted.org/packages/64/80/4f6e318ee2a7c0750ed724fa33a4bdf1eacdc5a39a7a24e818a773cd91af/frozenlist-1.8.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:229bf37d2e4acdaf808fd3f06e854a4a7a3661e871b10dc1f8f1896a3b05f18b", size = 50594, upload-time = "2025-10-06T05:36:07.69Z" }, + { url = "https://files.pythonhosted.org/packages/2b/94/5c8a2b50a496b11dd519f4a24cb5496cf125681dd99e94c604ccdea9419a/frozenlist-1.8.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f833670942247a14eafbb675458b4e61c82e002a148f49e68257b79296e865c4", size = 50448, upload-time = "2025-10-06T05:36:08.78Z" }, + { url = "https://files.pythonhosted.org/packages/6a/bd/d91c5e39f490a49df14320f4e8c80161cfcce09f1e2cde1edd16a551abb3/frozenlist-1.8.0-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:494a5952b1c597ba44e0e78113a7266e656b9794eec897b19ead706bd7074383", size = 242411, upload-time = "2025-10-06T05:36:09.801Z" }, + { url = "https://files.pythonhosted.org/packages/8f/83/f61505a05109ef3293dfb1ff594d13d64a2324ac3482be2cedc2be818256/frozenlist-1.8.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:96f423a119f4777a4a056b66ce11527366a8bb92f54e541ade21f2374433f6d4", size = 243014, upload-time = "2025-10-06T05:36:11.394Z" }, + { url = "https://files.pythonhosted.org/packages/d8/cb/cb6c7b0f7d4023ddda30cf56b8b17494eb3a79e3fda666bf735f63118b35/frozenlist-1.8.0-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:3462dd9475af2025c31cc61be6652dfa25cbfb56cbbf52f4ccfe029f38decaf8", size = 234909, upload-time = "2025-10-06T05:36:12.598Z" }, + { url = "https://files.pythonhosted.org/packages/31/c5/cd7a1f3b8b34af009fb17d4123c5a778b44ae2804e3ad6b86204255f9ec5/frozenlist-1.8.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c4c800524c9cd9bac5166cd6f55285957fcfc907db323e193f2afcd4d9abd69b", size = 250049, upload-time = "2025-10-06T05:36:14.065Z" }, + { url = "https://files.pythonhosted.org/packages/c0/01/2f95d3b416c584a1e7f0e1d6d31998c4a795f7544069ee2e0962a4b60740/frozenlist-1.8.0-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:d6a5df73acd3399d893dafc71663ad22534b5aa4f94e8a2fabfe856c3c1b6a52", size = 256485, upload-time = "2025-10-06T05:36:15.39Z" }, + { url = "https://files.pythonhosted.org/packages/ce/03/024bf7720b3abaebcff6d0793d73c154237b85bdf67b7ed55e5e9596dc9a/frozenlist-1.8.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:405e8fe955c2280ce66428b3ca55e12b3c4e9c336fb2103a4937e891c69a4a29", size = 237619, upload-time = "2025-10-06T05:36:16.558Z" }, + { url = "https://files.pythonhosted.org/packages/69/fa/f8abdfe7d76b731f5d8bd217827cf6764d4f1d9763407e42717b4bed50a0/frozenlist-1.8.0-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:908bd3f6439f2fef9e85031b59fd4f1297af54415fb60e4254a95f75b3cab3f3", size = 250320, upload-time = "2025-10-06T05:36:17.821Z" }, + { url = "https://files.pythonhosted.org/packages/f5/3c/b051329f718b463b22613e269ad72138cc256c540f78a6de89452803a47d/frozenlist-1.8.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:294e487f9ec720bd8ffcebc99d575f7eff3568a08a253d1ee1a0378754b74143", size = 246820, upload-time = "2025-10-06T05:36:19.046Z" }, + { url = "https://files.pythonhosted.org/packages/0f/ae/58282e8f98e444b3f4dd42448ff36fa38bef29e40d40f330b22e7108f565/frozenlist-1.8.0-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:74c51543498289c0c43656701be6b077f4b265868fa7f8a8859c197006efb608", size = 250518, upload-time = "2025-10-06T05:36:20.763Z" }, + { url = "https://files.pythonhosted.org/packages/8f/96/007e5944694d66123183845a106547a15944fbbb7154788cbf7272789536/frozenlist-1.8.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:776f352e8329135506a1d6bf16ac3f87bc25b28e765949282dcc627af36123aa", size = 239096, upload-time = "2025-10-06T05:36:22.129Z" }, + { url = "https://files.pythonhosted.org/packages/66/bb/852b9d6db2fa40be96f29c0d1205c306288f0684df8fd26ca1951d461a56/frozenlist-1.8.0-cp312-cp312-win32.whl", hash = "sha256:433403ae80709741ce34038da08511d4a77062aa924baf411ef73d1146e74faf", size = 39985, upload-time = "2025-10-06T05:36:23.661Z" }, + { url = "https://files.pythonhosted.org/packages/b8/af/38e51a553dd66eb064cdf193841f16f077585d4d28394c2fa6235cb41765/frozenlist-1.8.0-cp312-cp312-win_amd64.whl", hash = "sha256:34187385b08f866104f0c0617404c8eb08165ab1272e884abc89c112e9c00746", size = 44591, upload-time = "2025-10-06T05:36:24.958Z" }, + { url = "https://files.pythonhosted.org/packages/a7/06/1dc65480ab147339fecc70797e9c2f69d9cea9cf38934ce08df070fdb9cb/frozenlist-1.8.0-cp312-cp312-win_arm64.whl", hash = "sha256:fe3c58d2f5db5fbd18c2987cba06d51b0529f52bc3a6cdc33d3f4eab725104bd", size = 40102, upload-time = "2025-10-06T05:36:26.333Z" }, + { url = "https://files.pythonhosted.org/packages/2d/40/0832c31a37d60f60ed79e9dfb5a92e1e2af4f40a16a29abcc7992af9edff/frozenlist-1.8.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:8d92f1a84bb12d9e56f818b3a746f3efba93c1b63c8387a73dde655e1e42282a", size = 85717, upload-time = "2025-10-06T05:36:27.341Z" }, + { url = "https://files.pythonhosted.org/packages/30/ba/b0b3de23f40bc55a7057bd38434e25c34fa48e17f20ee273bbde5e0650f3/frozenlist-1.8.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:96153e77a591c8adc2ee805756c61f59fef4cf4073a9275ee86fe8cba41241f7", size = 49651, upload-time = "2025-10-06T05:36:28.855Z" }, + { url = "https://files.pythonhosted.org/packages/0c/ab/6e5080ee374f875296c4243c381bbdef97a9ac39c6e3ce1d5f7d42cb78d6/frozenlist-1.8.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:f21f00a91358803399890ab167098c131ec2ddd5f8f5fd5fe9c9f2c6fcd91e40", size = 49417, upload-time = "2025-10-06T05:36:29.877Z" }, + { url = "https://files.pythonhosted.org/packages/d5/4e/e4691508f9477ce67da2015d8c00acd751e6287739123113a9fca6f1604e/frozenlist-1.8.0-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:fb30f9626572a76dfe4293c7194a09fb1fe93ba94c7d4f720dfae3b646b45027", size = 234391, upload-time = "2025-10-06T05:36:31.301Z" }, + { url = "https://files.pythonhosted.org/packages/40/76/c202df58e3acdf12969a7895fd6f3bc016c642e6726aa63bd3025e0fc71c/frozenlist-1.8.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:eaa352d7047a31d87dafcacbabe89df0aa506abb5b1b85a2fb91bc3faa02d822", size = 233048, upload-time = "2025-10-06T05:36:32.531Z" }, + { url = "https://files.pythonhosted.org/packages/f9/c0/8746afb90f17b73ca5979c7a3958116e105ff796e718575175319b5bb4ce/frozenlist-1.8.0-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:03ae967b4e297f58f8c774c7eabcce57fe3c2434817d4385c50661845a058121", size = 226549, upload-time = "2025-10-06T05:36:33.706Z" }, + { url = "https://files.pythonhosted.org/packages/7e/eb/4c7eefc718ff72f9b6c4893291abaae5fbc0c82226a32dcd8ef4f7a5dbef/frozenlist-1.8.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:f6292f1de555ffcc675941d65fffffb0a5bcd992905015f85d0592201793e0e5", size = 239833, upload-time = "2025-10-06T05:36:34.947Z" }, + { url = "https://files.pythonhosted.org/packages/c2/4e/e5c02187cf704224f8b21bee886f3d713ca379535f16893233b9d672ea71/frozenlist-1.8.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:29548f9b5b5e3460ce7378144c3010363d8035cea44bc0bf02d57f5a685e084e", size = 245363, upload-time = "2025-10-06T05:36:36.534Z" }, + { url = "https://files.pythonhosted.org/packages/1f/96/cb85ec608464472e82ad37a17f844889c36100eed57bea094518bf270692/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:ec3cc8c5d4084591b4237c0a272cc4f50a5b03396a47d9caaf76f5d7b38a4f11", size = 229314, upload-time = "2025-10-06T05:36:38.582Z" }, + { url = "https://files.pythonhosted.org/packages/5d/6f/4ae69c550e4cee66b57887daeebe006fe985917c01d0fff9caab9883f6d0/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:517279f58009d0b1f2e7c1b130b377a349405da3f7621ed6bfae50b10adf20c1", size = 243365, upload-time = "2025-10-06T05:36:40.152Z" }, + { url = "https://files.pythonhosted.org/packages/7a/58/afd56de246cf11780a40a2c28dc7cbabbf06337cc8ddb1c780a2d97e88d8/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:db1e72ede2d0d7ccb213f218df6a078a9c09a7de257c2fe8fcef16d5925230b1", size = 237763, upload-time = "2025-10-06T05:36:41.355Z" }, + { url = "https://files.pythonhosted.org/packages/cb/36/cdfaf6ed42e2644740d4a10452d8e97fa1c062e2a8006e4b09f1b5fd7d63/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:b4dec9482a65c54a5044486847b8a66bf10c9cb4926d42927ec4e8fd5db7fed8", size = 240110, upload-time = "2025-10-06T05:36:42.716Z" }, + { url = "https://files.pythonhosted.org/packages/03/a8/9ea226fbefad669f11b52e864c55f0bd57d3c8d7eb07e9f2e9a0b39502e1/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:21900c48ae04d13d416f0e1e0c4d81f7931f73a9dfa0b7a8746fb2fe7dd970ed", size = 233717, upload-time = "2025-10-06T05:36:44.251Z" }, + { url = "https://files.pythonhosted.org/packages/1e/0b/1b5531611e83ba7d13ccc9988967ea1b51186af64c42b7a7af465dcc9568/frozenlist-1.8.0-cp313-cp313-win32.whl", hash = "sha256:8b7b94a067d1c504ee0b16def57ad5738701e4ba10cec90529f13fa03c833496", size = 39628, upload-time = "2025-10-06T05:36:45.423Z" }, + { url = "https://files.pythonhosted.org/packages/d8/cf/174c91dbc9cc49bc7b7aab74d8b734e974d1faa8f191c74af9b7e80848e6/frozenlist-1.8.0-cp313-cp313-win_amd64.whl", hash = "sha256:878be833caa6a3821caf85eb39c5ba92d28e85df26d57afb06b35b2efd937231", size = 43882, upload-time = "2025-10-06T05:36:46.796Z" }, + { url = "https://files.pythonhosted.org/packages/c1/17/502cd212cbfa96eb1388614fe39a3fc9ab87dbbe042b66f97acb57474834/frozenlist-1.8.0-cp313-cp313-win_arm64.whl", hash = "sha256:44389d135b3ff43ba8cc89ff7f51f5a0bb6b63d829c8300f79a2fe4fe61bcc62", size = 39676, upload-time = "2025-10-06T05:36:47.8Z" }, + { url = "https://files.pythonhosted.org/packages/d2/5c/3bbfaa920dfab09e76946a5d2833a7cbdf7b9b4a91c714666ac4855b88b4/frozenlist-1.8.0-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:e25ac20a2ef37e91c1b39938b591457666a0fa835c7783c3a8f33ea42870db94", size = 89235, upload-time = "2025-10-06T05:36:48.78Z" }, + { url = "https://files.pythonhosted.org/packages/d2/d6/f03961ef72166cec1687e84e8925838442b615bd0b8854b54923ce5b7b8a/frozenlist-1.8.0-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:07cdca25a91a4386d2e76ad992916a85038a9b97561bf7a3fd12d5d9ce31870c", size = 50742, upload-time = "2025-10-06T05:36:49.837Z" }, + { url = "https://files.pythonhosted.org/packages/1e/bb/a6d12b7ba4c3337667d0e421f7181c82dda448ce4e7ad7ecd249a16fa806/frozenlist-1.8.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:4e0c11f2cc6717e0a741f84a527c52616140741cd812a50422f83dc31749fb52", size = 51725, upload-time = "2025-10-06T05:36:50.851Z" }, + { url = "https://files.pythonhosted.org/packages/bc/71/d1fed0ffe2c2ccd70b43714c6cab0f4188f09f8a67a7914a6b46ee30f274/frozenlist-1.8.0-cp313-cp313t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:b3210649ee28062ea6099cfda39e147fa1bc039583c8ee4481cb7811e2448c51", size = 284533, upload-time = "2025-10-06T05:36:51.898Z" }, + { url = "https://files.pythonhosted.org/packages/c9/1f/fb1685a7b009d89f9bf78a42d94461bc06581f6e718c39344754a5d9bada/frozenlist-1.8.0-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:581ef5194c48035a7de2aefc72ac6539823bb71508189e5de01d60c9dcd5fa65", size = 292506, upload-time = "2025-10-06T05:36:53.101Z" }, + { url = "https://files.pythonhosted.org/packages/e6/3b/b991fe1612703f7e0d05c0cf734c1b77aaf7c7d321df4572e8d36e7048c8/frozenlist-1.8.0-cp313-cp313t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:3ef2d026f16a2b1866e1d86fc4e1291e1ed8a387b2c333809419a2f8b3a77b82", size = 274161, upload-time = "2025-10-06T05:36:54.309Z" }, + { url = "https://files.pythonhosted.org/packages/ca/ec/c5c618767bcdf66e88945ec0157d7f6c4a1322f1473392319b7a2501ded7/frozenlist-1.8.0-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:5500ef82073f599ac84d888e3a8c1f77ac831183244bfd7f11eaa0289fb30714", size = 294676, upload-time = "2025-10-06T05:36:55.566Z" }, + { url = "https://files.pythonhosted.org/packages/7c/ce/3934758637d8f8a88d11f0585d6495ef54b2044ed6ec84492a91fa3b27aa/frozenlist-1.8.0-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:50066c3997d0091c411a66e710f4e11752251e6d2d73d70d8d5d4c76442a199d", size = 300638, upload-time = "2025-10-06T05:36:56.758Z" }, + { url = "https://files.pythonhosted.org/packages/fc/4f/a7e4d0d467298f42de4b41cbc7ddaf19d3cfeabaf9ff97c20c6c7ee409f9/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:5c1c8e78426e59b3f8005e9b19f6ff46e5845895adbde20ece9218319eca6506", size = 283067, upload-time = "2025-10-06T05:36:57.965Z" }, + { url = "https://files.pythonhosted.org/packages/dc/48/c7b163063d55a83772b268e6d1affb960771b0e203b632cfe09522d67ea5/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_armv7l.whl", hash = "sha256:eefdba20de0d938cec6a89bd4d70f346a03108a19b9df4248d3cf0d88f1b0f51", size = 292101, upload-time = "2025-10-06T05:36:59.237Z" }, + { url = "https://files.pythonhosted.org/packages/9f/d0/2366d3c4ecdc2fd391e0afa6e11500bfba0ea772764d631bbf82f0136c9d/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:cf253e0e1c3ceb4aaff6df637ce033ff6535fb8c70a764a8f46aafd3d6ab798e", size = 289901, upload-time = "2025-10-06T05:37:00.811Z" }, + { url = "https://files.pythonhosted.org/packages/b8/94/daff920e82c1b70e3618a2ac39fbc01ae3e2ff6124e80739ce5d71c9b920/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:032efa2674356903cd0261c4317a561a6850f3ac864a63fc1583147fb05a79b0", size = 289395, upload-time = "2025-10-06T05:37:02.115Z" }, + { url = "https://files.pythonhosted.org/packages/e3/20/bba307ab4235a09fdcd3cc5508dbabd17c4634a1af4b96e0f69bfe551ebd/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:6da155091429aeba16851ecb10a9104a108bcd32f6c1642867eadaee401c1c41", size = 283659, upload-time = "2025-10-06T05:37:03.711Z" }, + { url = "https://files.pythonhosted.org/packages/fd/00/04ca1c3a7a124b6de4f8a9a17cc2fcad138b4608e7a3fc5877804b8715d7/frozenlist-1.8.0-cp313-cp313t-win32.whl", hash = "sha256:0f96534f8bfebc1a394209427d0f8a63d343c9779cda6fc25e8e121b5fd8555b", size = 43492, upload-time = "2025-10-06T05:37:04.915Z" }, + { url = "https://files.pythonhosted.org/packages/59/5e/c69f733a86a94ab10f68e496dc6b7e8bc078ebb415281d5698313e3af3a1/frozenlist-1.8.0-cp313-cp313t-win_amd64.whl", hash = "sha256:5d63a068f978fc69421fb0e6eb91a9603187527c86b7cd3f534a5b77a592b888", size = 48034, upload-time = "2025-10-06T05:37:06.343Z" }, + { url = "https://files.pythonhosted.org/packages/16/6c/be9d79775d8abe79b05fa6d23da99ad6e7763a1d080fbae7290b286093fd/frozenlist-1.8.0-cp313-cp313t-win_arm64.whl", hash = "sha256:bf0a7e10b077bf5fb9380ad3ae8ce20ef919a6ad93b4552896419ac7e1d8e042", size = 41749, upload-time = "2025-10-06T05:37:07.431Z" }, + { url = "https://files.pythonhosted.org/packages/f1/c8/85da824b7e7b9b6e7f7705b2ecaf9591ba6f79c1177f324c2735e41d36a2/frozenlist-1.8.0-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:cee686f1f4cadeb2136007ddedd0aaf928ab95216e7691c63e50a8ec066336d0", size = 86127, upload-time = "2025-10-06T05:37:08.438Z" }, + { url = "https://files.pythonhosted.org/packages/8e/e8/a1185e236ec66c20afd72399522f142c3724c785789255202d27ae992818/frozenlist-1.8.0-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:119fb2a1bd47307e899c2fac7f28e85b9a543864df47aa7ec9d3c1b4545f096f", size = 49698, upload-time = "2025-10-06T05:37:09.48Z" }, + { url = "https://files.pythonhosted.org/packages/a1/93/72b1736d68f03fda5fdf0f2180fb6caaae3894f1b854d006ac61ecc727ee/frozenlist-1.8.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:4970ece02dbc8c3a92fcc5228e36a3e933a01a999f7094ff7c23fbd2beeaa67c", size = 49749, upload-time = "2025-10-06T05:37:10.569Z" }, + { url = "https://files.pythonhosted.org/packages/a7/b2/fabede9fafd976b991e9f1b9c8c873ed86f202889b864756f240ce6dd855/frozenlist-1.8.0-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:cba69cb73723c3f329622e34bdbf5ce1f80c21c290ff04256cff1cd3c2036ed2", size = 231298, upload-time = "2025-10-06T05:37:11.993Z" }, + { url = "https://files.pythonhosted.org/packages/3a/3b/d9b1e0b0eed36e70477ffb8360c49c85c8ca8ef9700a4e6711f39a6e8b45/frozenlist-1.8.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:778a11b15673f6f1df23d9586f83c4846c471a8af693a22e066508b77d201ec8", size = 232015, upload-time = "2025-10-06T05:37:13.194Z" }, + { url = "https://files.pythonhosted.org/packages/dc/94/be719d2766c1138148564a3960fc2c06eb688da592bdc25adcf856101be7/frozenlist-1.8.0-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:0325024fe97f94c41c08872db482cf8ac4800d80e79222c6b0b7b162d5b13686", size = 225038, upload-time = "2025-10-06T05:37:14.577Z" }, + { url = "https://files.pythonhosted.org/packages/e4/09/6712b6c5465f083f52f50cf74167b92d4ea2f50e46a9eea0523d658454ae/frozenlist-1.8.0-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:97260ff46b207a82a7567b581ab4190bd4dfa09f4db8a8b49d1a958f6aa4940e", size = 240130, upload-time = "2025-10-06T05:37:15.781Z" }, + { url = "https://files.pythonhosted.org/packages/f8/d4/cd065cdcf21550b54f3ce6a22e143ac9e4836ca42a0de1022da8498eac89/frozenlist-1.8.0-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:54b2077180eb7f83dd52c40b2750d0a9f175e06a42e3213ce047219de902717a", size = 242845, upload-time = "2025-10-06T05:37:17.037Z" }, + { url = "https://files.pythonhosted.org/packages/62/c3/f57a5c8c70cd1ead3d5d5f776f89d33110b1addae0ab010ad774d9a44fb9/frozenlist-1.8.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:2f05983daecab868a31e1da44462873306d3cbfd76d1f0b5b69c473d21dbb128", size = 229131, upload-time = "2025-10-06T05:37:18.221Z" }, + { url = "https://files.pythonhosted.org/packages/6c/52/232476fe9cb64f0742f3fde2b7d26c1dac18b6d62071c74d4ded55e0ef94/frozenlist-1.8.0-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:33f48f51a446114bc5d251fb2954ab0164d5be02ad3382abcbfe07e2531d650f", size = 240542, upload-time = "2025-10-06T05:37:19.771Z" }, + { url = "https://files.pythonhosted.org/packages/5f/85/07bf3f5d0fb5414aee5f47d33c6f5c77bfe49aac680bfece33d4fdf6a246/frozenlist-1.8.0-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:154e55ec0655291b5dd1b8731c637ecdb50975a2ae70c606d100750a540082f7", size = 237308, upload-time = "2025-10-06T05:37:20.969Z" }, + { url = "https://files.pythonhosted.org/packages/11/99/ae3a33d5befd41ac0ca2cc7fd3aa707c9c324de2e89db0e0f45db9a64c26/frozenlist-1.8.0-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:4314debad13beb564b708b4a496020e5306c7333fa9a3ab90374169a20ffab30", size = 238210, upload-time = "2025-10-06T05:37:22.252Z" }, + { url = "https://files.pythonhosted.org/packages/b2/60/b1d2da22f4970e7a155f0adde9b1435712ece01b3cd45ba63702aea33938/frozenlist-1.8.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:073f8bf8becba60aa931eb3bc420b217bb7d5b8f4750e6f8b3be7f3da85d38b7", size = 231972, upload-time = "2025-10-06T05:37:23.5Z" }, + { url = "https://files.pythonhosted.org/packages/3f/ab/945b2f32de889993b9c9133216c068b7fcf257d8595a0ac420ac8677cab0/frozenlist-1.8.0-cp314-cp314-win32.whl", hash = "sha256:bac9c42ba2ac65ddc115d930c78d24ab8d4f465fd3fc473cdedfccadb9429806", size = 40536, upload-time = "2025-10-06T05:37:25.581Z" }, + { url = "https://files.pythonhosted.org/packages/59/ad/9caa9b9c836d9ad6f067157a531ac48b7d36499f5036d4141ce78c230b1b/frozenlist-1.8.0-cp314-cp314-win_amd64.whl", hash = "sha256:3e0761f4d1a44f1d1a47996511752cf3dcec5bbdd9cc2b4fe595caf97754b7a0", size = 44330, upload-time = "2025-10-06T05:37:26.928Z" }, + { url = "https://files.pythonhosted.org/packages/82/13/e6950121764f2676f43534c555249f57030150260aee9dcf7d64efda11dd/frozenlist-1.8.0-cp314-cp314-win_arm64.whl", hash = "sha256:d1eaff1d00c7751b7c6662e9c5ba6eb2c17a2306ba5e2a37f24ddf3cc953402b", size = 40627, upload-time = "2025-10-06T05:37:28.075Z" }, + { url = "https://files.pythonhosted.org/packages/c0/c7/43200656ecc4e02d3f8bc248df68256cd9572b3f0017f0a0c4e93440ae23/frozenlist-1.8.0-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:d3bb933317c52d7ea5004a1c442eef86f426886fba134ef8cf4226ea6ee1821d", size = 89238, upload-time = "2025-10-06T05:37:29.373Z" }, + { url = "https://files.pythonhosted.org/packages/d1/29/55c5f0689b9c0fb765055629f472c0de484dcaf0acee2f7707266ae3583c/frozenlist-1.8.0-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:8009897cdef112072f93a0efdce29cd819e717fd2f649ee3016efd3cd885a7ed", size = 50738, upload-time = "2025-10-06T05:37:30.792Z" }, + { url = "https://files.pythonhosted.org/packages/ba/7d/b7282a445956506fa11da8c2db7d276adcbf2b17d8bb8407a47685263f90/frozenlist-1.8.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:2c5dcbbc55383e5883246d11fd179782a9d07a986c40f49abe89ddf865913930", size = 51739, upload-time = "2025-10-06T05:37:32.127Z" }, + { url = "https://files.pythonhosted.org/packages/62/1c/3d8622e60d0b767a5510d1d3cf21065b9db874696a51ea6d7a43180a259c/frozenlist-1.8.0-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:39ecbc32f1390387d2aa4f5a995e465e9e2f79ba3adcac92d68e3e0afae6657c", size = 284186, upload-time = "2025-10-06T05:37:33.21Z" }, + { url = "https://files.pythonhosted.org/packages/2d/14/aa36d5f85a89679a85a1d44cd7a6657e0b1c75f61e7cad987b203d2daca8/frozenlist-1.8.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:92db2bf818d5cc8d9c1f1fc56b897662e24ea5adb36ad1f1d82875bd64e03c24", size = 292196, upload-time = "2025-10-06T05:37:36.107Z" }, + { url = "https://files.pythonhosted.org/packages/05/23/6bde59eb55abd407d34f77d39a5126fb7b4f109a3f611d3929f14b700c66/frozenlist-1.8.0-cp314-cp314t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:2dc43a022e555de94c3b68a4ef0b11c4f747d12c024a520c7101709a2144fb37", size = 273830, upload-time = "2025-10-06T05:37:37.663Z" }, + { url = "https://files.pythonhosted.org/packages/d2/3f/22cff331bfad7a8afa616289000ba793347fcd7bc275f3b28ecea2a27909/frozenlist-1.8.0-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:cb89a7f2de3602cfed448095bab3f178399646ab7c61454315089787df07733a", size = 294289, upload-time = "2025-10-06T05:37:39.261Z" }, + { url = "https://files.pythonhosted.org/packages/a4/89/5b057c799de4838b6c69aa82b79705f2027615e01be996d2486a69ca99c4/frozenlist-1.8.0-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:33139dc858c580ea50e7e60a1b0ea003efa1fd42e6ec7fdbad78fff65fad2fd2", size = 300318, upload-time = "2025-10-06T05:37:43.213Z" }, + { url = "https://files.pythonhosted.org/packages/30/de/2c22ab3eb2a8af6d69dc799e48455813bab3690c760de58e1bf43b36da3e/frozenlist-1.8.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:168c0969a329b416119507ba30b9ea13688fafffac1b7822802537569a1cb0ef", size = 282814, upload-time = "2025-10-06T05:37:45.337Z" }, + { url = "https://files.pythonhosted.org/packages/59/f7/970141a6a8dbd7f556d94977858cfb36fa9b66e0892c6dd780d2219d8cd8/frozenlist-1.8.0-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:28bd570e8e189d7f7b001966435f9dac6718324b5be2990ac496cf1ea9ddb7fe", size = 291762, upload-time = "2025-10-06T05:37:46.657Z" }, + { url = "https://files.pythonhosted.org/packages/c1/15/ca1adae83a719f82df9116d66f5bb28bb95557b3951903d39135620ef157/frozenlist-1.8.0-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:b2a095d45c5d46e5e79ba1e5b9cb787f541a8dee0433836cea4b96a2c439dcd8", size = 289470, upload-time = "2025-10-06T05:37:47.946Z" }, + { url = "https://files.pythonhosted.org/packages/ac/83/dca6dc53bf657d371fbc88ddeb21b79891e747189c5de990b9dfff2ccba1/frozenlist-1.8.0-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:eab8145831a0d56ec9c4139b6c3e594c7a83c2c8be25d5bcf2d86136a532287a", size = 289042, upload-time = "2025-10-06T05:37:49.499Z" }, + { url = "https://files.pythonhosted.org/packages/96/52/abddd34ca99be142f354398700536c5bd315880ed0a213812bc491cff5e4/frozenlist-1.8.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:974b28cf63cc99dfb2188d8d222bc6843656188164848c4f679e63dae4b0708e", size = 283148, upload-time = "2025-10-06T05:37:50.745Z" }, + { url = "https://files.pythonhosted.org/packages/af/d3/76bd4ed4317e7119c2b7f57c3f6934aba26d277acc6309f873341640e21f/frozenlist-1.8.0-cp314-cp314t-win32.whl", hash = "sha256:342c97bf697ac5480c0a7ec73cd700ecfa5a8a40ac923bd035484616efecc2df", size = 44676, upload-time = "2025-10-06T05:37:52.222Z" }, + { url = "https://files.pythonhosted.org/packages/89/76/c615883b7b521ead2944bb3480398cbb07e12b7b4e4d073d3752eb721558/frozenlist-1.8.0-cp314-cp314t-win_amd64.whl", hash = "sha256:06be8f67f39c8b1dc671f5d83aaefd3358ae5cdcf8314552c57e7ed3e6475bdd", size = 49451, upload-time = "2025-10-06T05:37:53.425Z" }, + { url = "https://files.pythonhosted.org/packages/e0/a3/5982da14e113d07b325230f95060e2169f5311b1017ea8af2a29b374c289/frozenlist-1.8.0-cp314-cp314t-win_arm64.whl", hash = "sha256:102e6314ca4da683dca92e3b1355490fed5f313b768500084fbe6371fddfdb79", size = 42507, upload-time = "2025-10-06T05:37:54.513Z" }, + { url = "https://files.pythonhosted.org/packages/9a/9a/e35b4a917281c0b8419d4207f4334c8e8c5dbf4f3f5f9ada73958d937dcc/frozenlist-1.8.0-py3-none-any.whl", hash = "sha256:0c18a16eab41e82c295618a77502e17b195883241c563b00f0aa5106fc4eaa0d", size = 13409, upload-time = "2025-10-06T05:38:16.721Z" }, +] + +[[package]] +name = "fsspec" +version = "2026.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d5/8d/1c51c094345df128ca4a990d633fe1a0ff28726c9e6b3c41ba65087bba1d/fsspec-2026.4.0.tar.gz", hash = "sha256:301d8ac70ae90ef3ad05dcf94d6c3754a097f9b5fe4667d2787aa359ec7df7e4", size = 312760, upload-time = "2026-04-29T20:42:38.635Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d5/0c/043d5e551459da400957a1395e0febbf771446ff34291afcbe3d8be2a279/fsspec-2026.4.0-py3-none-any.whl", hash = "sha256:11ef7bb35dab8a394fde6e608221d5cf3e8499401c249bebaeaad760a1a8dec2", size = 203402, upload-time = "2026-04-29T20:42:36.842Z" }, +] + +[package.optional-dependencies] +http = [ + { name = "aiohttp" }, +] + [[package]] name = "h11" version = "0.16.0" @@ -369,6 +605,78 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/12/8d/2abf4487fdb55cf0aef7b8458d17cc2ed2db4446489d52ad4454c970166e/hawkmoth-0.21.0-py3-none-any.whl", hash = "sha256:49d289f1d61c6ab8b8ce1221e49915056b985c8665bd5cdb6cccbfb5a6038ef1", size = 31952, upload-time = "2025-05-10T12:35:20.72Z" }, ] +[[package]] +name = "hf-xet" +version = "1.5.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/63/39/67be8d71f900d9a55761b6022821d6679fb56c64f1b6063d5af2c2606727/hf_xet-1.5.2.tar.gz", hash = "sha256:73044bd31bae33c984af832d19c752a0dffb67518fee9ddbd91d616e1101cf47", size = 903674, upload-time = "2026-07-16T17:29:56.833Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/29/be/525eabac5d1736b679c39e342ecd4292534012546a2d18f0043c8e3b6021/hf_xet-1.5.2-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:4a5ecb9cda8512ba2aa8ee5d37c87a1422992165892d653098c7b90247481c3b", size = 4064284, upload-time = "2026-07-16T17:29:29.907Z" }, + { url = "https://files.pythonhosted.org/packages/c5/3f/699749dd78442480eda4e4fca494284b0e3542e4063cc37654d5fdc929e6/hf_xet-1.5.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:8764488197c1d7b1378c8438c18d2eea902e150dbca0b0f0d2d32603fb9b5576", size = 3828537, upload-time = "2026-07-16T17:29:31.549Z" }, + { url = "https://files.pythonhosted.org/packages/22/d7/2658ac0a5b9f4664ca27ce31bd015044fe9dea50ed455fb5197aba819c11/hf_xet-1.5.2-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8d7446f72abbf7e01ca5ff131786bc2e74a56393462c17a6bf1e303fbab81db4", size = 4417133, upload-time = "2026-07-16T17:29:33.391Z" }, + { url = "https://files.pythonhosted.org/packages/d9/58/8343f3cb63c8fa058d576136df3871550f7d5214a8f048a7ea2eab6ac906/hf_xet-1.5.2-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:580e59e29bf37aece1f2b68537de1e3fb04f43a23d910dcf6f128280b5bfbba4", size = 4212613, upload-time = "2026-07-16T17:29:34.989Z" }, + { url = "https://files.pythonhosted.org/packages/0c/33/a968f4e4535037b36941ec00714625fb60e026302407e7e26ca9f3e65f4e/hf_xet-1.5.2-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:bee28c619622d36968056532fd49cf2b35ca75099b1d616c31a618a893491380", size = 4412710, upload-time = "2026-07-16T17:29:36.646Z" }, + { url = "https://files.pythonhosted.org/packages/7b/d9/9e33981173dbaf194ba0015202b02d467b624d44d4eba89e1bf06c0d2995/hf_xet-1.5.2-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:e396ab0faf6298199ad7a95305c3ca8498cb825978a6485be6d00587ee4ec577", size = 4628455, upload-time = "2026-07-16T17:29:38.352Z" }, + { url = "https://files.pythonhosted.org/packages/e9/4b/cc682832de4264a03880a2d1b5ec3e1fab3bf307f508817250baafdb9996/hf_xet-1.5.2-cp314-cp314t-win_amd64.whl", hash = "sha256:fd3add255549e8ef58fa35b2e42dc016961c050600444e7d77d030ba6b57120e", size = 3979044, upload-time = "2026-07-16T17:29:40.329Z" }, + { url = "https://files.pythonhosted.org/packages/ea/09/b2cdf2a0fb39a08af3222b96092a36bd3b40c54123eef07de4422e870971/hf_xet-1.5.2-cp314-cp314t-win_arm64.whl", hash = "sha256:d6f9c58549407b84b9a5383afd68db0acc42345326a3159990b36a5ca8a20e4e", size = 3808037, upload-time = "2026-07-16T17:29:42.357Z" }, + { url = "https://files.pythonhosted.org/packages/de/ba/2b70603c7552db82baeb2623e2336898304a17328845151be4fe1f48d420/hf_xet-1.5.2-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:f922b8f5fb84f1dd3d7ab7a1316354a1bca9b1c73ecfc19c76e51a2a49d29799", size = 4033760, upload-time = "2026-07-16T17:29:43.884Z" }, + { url = "https://files.pythonhosted.org/packages/60/ac/b097a86a1e4a6098f3a79382643ab09d5733d87ccc864877ad1e12b49b70/hf_xet-1.5.2-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:045f84440c55cdeb659cf1a1dd48c77bcd0d2e93632e2fea8f2c3bdee79f38ed", size = 3841438, upload-time = "2026-07-16T17:29:45.539Z" }, + { url = "https://files.pythonhosted.org/packages/d3/35/db860aa3a0780660324a506ad4b3d322ddc6ecbba4b9340aed0942cbf21c/hf_xet-1.5.2-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:db78c39c83d6279daddc98e2238f373ab8980685556d42472b4ec51abcf03e8c", size = 4428006, upload-time = "2026-07-16T17:29:46.996Z" }, + { url = "https://files.pythonhosted.org/packages/af/6b/832dd980af4b0c3ae0660e309285f2ffcdff2faa38129390dbb47aa4a3f9/hf_xet-1.5.2-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:7db73c810500c54c6760be8c39d4b2e476974de85424c50063efc22fdda13025", size = 4221099, upload-time = "2026-07-16T17:29:48.525Z" }, + { url = "https://files.pythonhosted.org/packages/9e/05/ae50f0d34e3254e6c3e208beb2519f6b8673016fc4b3643badaf6450d186/hf_xet-1.5.2-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:6395cfe3c9cbead4f16b31808b0e67eac428b66c656f856e99636adaddea878f", size = 4420766, upload-time = "2026-07-16T17:29:50.092Z" }, + { url = "https://files.pythonhosted.org/packages/07/a9/c050bc2743a2bcd68928bfee157b08681667a164a24ec95fbfcfcd717e08/hf_xet-1.5.2-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:cde8cd167126bb6109b2ceb19b844433a4988643e8f3e01dd9dd0e4a34535097", size = 4636716, upload-time = "2026-07-16T17:29:51.62Z" }, + { url = "https://files.pythonhosted.org/packages/e9/f8/68b01c5c2edb56ac9a67b3d076ffddcb90867abaee923923eb34e7a14e76/hf_xet-1.5.2-cp38-abi3-win_amd64.whl", hash = "sha256:ecf63d1cb69a9a7319910f8f83fcf9b46e7a32dfcf4b8f8eeddb55f647306e65", size = 3988373, upload-time = "2026-07-16T17:29:53.395Z" }, + { url = "https://files.pythonhosted.org/packages/39/c6/988383e9dc17294d536fcbcd6fd16eed882e411ad16c954984a53e47b09c/hf_xet-1.5.2-cp38-abi3-win_arm64.whl", hash = "sha256:1da28519496eb7c8094c11e4d25509b4a468457a0302d58136099db2fd9a671d", size = 3816957, upload-time = "2026-07-16T17:29:54.991Z" }, +] + +[[package]] +name = "httpcore" +version = "1.0.9" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "h11" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/06/94/82699a10bca87a5556c9c59b5963f2d039dbd239f25bc2a63907a05a14cb/httpcore-1.0.9.tar.gz", hash = "sha256:6e34463af53fd2ab5d807f399a9b45ea31c3dfa2276f15a2c3f00afff6e176e8", size = 85484, upload-time = "2025-04-24T22:06:22.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl", hash = "sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55", size = 78784, upload-time = "2025-04-24T22:06:20.566Z" }, +] + +[[package]] +name = "httpx" +version = "0.28.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "certifi" }, + { name = "httpcore" }, + { name = "idna" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b1/df/48c586a5fe32a0f01324ee087459e112ebb7224f646c0b5023f5e79e9956/httpx-0.28.1.tar.gz", hash = "sha256:75e98c5f16b0f35b567856f597f06ff2270a374470a5c2392242528e3e3e42fc", size = 141406, upload-time = "2024-12-06T15:37:23.222Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517, upload-time = "2024-12-06T15:37:21.509Z" }, +] + +[[package]] +name = "huggingface-hub" +version = "1.25.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "filelock" }, + { name = "fsspec" }, + { name = "hf-xet", marker = "platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'" }, + { name = "httpx" }, + { name = "packaging" }, + { name = "pyyaml" }, + { name = "tqdm" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/4b/50/db3771a6e4fad4bd28fb055d4363b51cb0ae98c1aa504b79d41fdcab5483/huggingface_hub-1.25.1.tar.gz", hash = "sha256:21129595ca7a753be479b319913e22cc8808361ac118bd76cc413db831b28a99", size = 928426, upload-time = "2026-07-27T09:24:10.117Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f7/3f/21e816831c6d16f88a6c784974413fa0421ce8a5d04380c2666ed5b503e5/huggingface_hub-1.25.1-py3-none-any.whl", hash = "sha256:004d4e70350517e24c68a7dbb7dc5e40b2b6aefef8f94bf7a85f6f9835102ea5", size = 774909, upload-time = "2026-07-27T09:24:08.079Z" }, +] + [[package]] name = "idna" version = "3.11" @@ -481,15 +789,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/41/45/1a4ed80516f02155c51f51e8cedb3c1902296743db0bbc66608a0db2814f/jsonschema_specifications-2025.9.1-py3-none-any.whl", hash = "sha256:98802fee3a11ee76ecaca44429fda8a41bff98b00a0f2838151b113f210cc6fe", size = 18437, upload-time = "2025-09-08T01:34:57.871Z" }, ] -[[package]] -name = "latexcodec" -version = "3.0.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/27/dd/4270b2c5e2ee49316c3859e62293bd2ea8e382339d63ab7bbe9f39c0ec3b/latexcodec-3.0.1.tar.gz", hash = "sha256:e78a6911cd72f9dec35031c6ec23584de6842bfbc4610a9678868d14cdfb0357", size = 31222, upload-time = "2025-06-17T18:47:34.051Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b5/40/23569737873cc9637fd488606347e9dd92b9fa37ba4fcda1f98ee5219a97/latexcodec-3.0.1-py3-none-any.whl", hash = "sha256:a9eb8200bff693f0437a69581f7579eb6bca25c4193515c09900ce76451e452e", size = 18532, upload-time = "2025-06-17T18:47:30.726Z" }, -] - [[package]] name = "libclang" version = "18.1.1" @@ -607,23 +906,23 @@ wheels = [ [[package]] name = "maturin" -version = "1.13.1" +version = "1.15.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/39/16/b284a7bc4af3dd87717c784278c1b8cb18606ad1f6f7a671c47bfd9c3df0/maturin-1.13.1.tar.gz", hash = "sha256:9a87ff3b8e4d1c6eac33ebfe8e261e8236516d98d45c0323550621819b5a1a2f", size = 340369, upload-time = "2026-04-09T15:14:07.026Z" } +sdist = { url = "https://files.pythonhosted.org/packages/b9/c8/22e5e21b2679c9bce6415ca578034ca2cc9316be0642ae21e051a2d5198c/maturin-1.15.0.tar.gz", hash = "sha256:94b26cc8e8aba61a5f2099715fe640e18c5f678e9a500408b38761263954228a", size = 385504, upload-time = "2026-08-24T12:11:22.665Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/43/4d/a23fc95be881aa8c7a6ea353410417872e4d7065df03d7f3db8f0dbed4a7/maturin-1.13.1-py3-none-linux_armv6l.whl", hash = "sha256:416e4e01cb88b798e606ee43929df897e42c1647b722ef68283816cca99a8742", size = 10102444, upload-time = "2026-04-09T15:13:48.393Z" }, - { url = "https://files.pythonhosted.org/packages/a6/1e/65c385d65bae95cf04895d52f39dbed8b1453ae55da2903d252ade40a774/maturin-1.13.1-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:72888e87819ce546d0d2df900e4b385e4ef299077d92ee37b48923a5602dae94", size = 19576043, upload-time = "2026-04-09T15:14:08.685Z" }, - { url = "https://files.pythonhosted.org/packages/8f/13/f6bc868d0bfecd9314870b97f530a167e31f7878ac4945c78245c6eef69c/maturin-1.13.1-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:98b5fcf1a186c217830a8295ecc2989c6b1cf50945417adfc15252107b9475b7", size = 10117339, upload-time = "2026-04-09T15:13:40.559Z" }, - { url = "https://files.pythonhosted.org/packages/51/58/279e081305c11c1c1c4fccacf77df8959646c5d4de7a57ec7e787653e270/maturin-1.13.1-py3-none-manylinux_2_12_i686.manylinux2010_i686.musllinux_1_1_i686.whl", hash = "sha256:3da18cccf2f683c0977bff9146a0908d6ffce836d600665736ac01679f588cb9", size = 10139689, upload-time = "2026-04-09T15:13:38.291Z" }, - { url = "https://files.pythonhosted.org/packages/00/94/69391af5396c6aab723932240803f49e5f3de3dd7c57d32f02d237a0ce32/maturin-1.13.1-py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.musllinux_1_1_x86_64.whl", hash = "sha256:6b1e5916a253243e8f5f9e847b62bbc98420eec48c9ce2e2e8724c6da89d359b", size = 10551141, upload-time = "2026-04-09T15:13:42.887Z" }, - { url = "https://files.pythonhosted.org/packages/9e/bf/4edac2667b49e3733438062ae416413b8fc8d42e1bd499ba15e1fb02fc55/maturin-1.13.1-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.musllinux_1_1_aarch64.whl", hash = "sha256:dc91031e0619c1e28730279ef9ee5f106c9b9ec806b013f888676b242f892eb7", size = 9983094, upload-time = "2026-04-09T15:13:56.868Z" }, - { url = "https://files.pythonhosted.org/packages/79/94/a6d651cfe8fc6bf2e892c90e3cdbb25c06d81c9115140d03ea1a68a97575/maturin-1.13.1-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.musllinux_1_1_armv7l.whl", hash = "sha256:001741c6cff56aa8ea59a0d78ae990c0550d0e3e82b00b683eedb4158a8ef7e6", size = 9949980, upload-time = "2026-04-09T15:13:59.185Z" }, - { url = "https://files.pythonhosted.org/packages/b5/d1/82c067464f848e38af9910bce55eb54302b1c1284a279d515dbfcf5994f5/maturin-1.13.1-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.musllinux_1_1_ppc64le.whl", hash = "sha256:01c845825c917c07c1d0b2c9032c59c16a7d383d1e649a46481d3e5693c2750f", size = 13186276, upload-time = "2026-04-09T15:13:45.725Z" }, - { url = "https://files.pythonhosted.org/packages/7c/f4/25367baf1025580f047f9b37598bb3fadc416e24536afd4f28e190335c73/maturin-1.13.1-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:f69093ed4a0e6464e52a7fc26d714f859ce15630ec8070743398c6bf41f38a9e", size = 10891837, upload-time = "2026-04-09T15:13:35.68Z" }, - { url = "https://files.pythonhosted.org/packages/af/be/caafad8ce74974b7deafdf144d12f758993dfea4c66c9905b138f51a7792/maturin-1.13.1-py3-none-manylinux_2_31_riscv64.musllinux_1_1_riscv64.whl", hash = "sha256:c1490584f3c70af45466ee99065b49e6657ebdccac6b10571bb44681309c9396", size = 10351032, upload-time = "2026-04-09T15:14:01.632Z" }, - { url = "https://files.pythonhosted.org/packages/66/0e/970a721d27cfa410e8bfa0a1e32e6ef52cb8169692110a5fdabe1af3f570/maturin-1.13.1-py3-none-win32.whl", hash = "sha256:c6a720b252c99de072922dbe4432ab19662b6f80045b0355fec23bdfccb450da", size = 8855465, upload-time = "2026-04-09T15:13:51.122Z" }, - { url = "https://files.pythonhosted.org/packages/88/70/7c1e0d65fa147d5479055a171541c82b8cdfc1c825d85a82240470f14176/maturin-1.13.1-py3-none-win_amd64.whl", hash = "sha256:a2017d2281203d0c6570240e7d746564d766d756105823b7de68bda6ae722711", size = 10230471, upload-time = "2026-04-09T15:13:53.89Z" }, - { url = "https://files.pythonhosted.org/packages/c5/2a/afe0193b673a79ffd2e01ad999511b7e9e6b49af02bb3759d82a78c3043d/maturin-1.13.1-py3-none-win_arm64.whl", hash = "sha256:2839024dcd65776abb4759e5bca29941971e095574162a4d335191da4be9ff24", size = 8905575, upload-time = "2026-04-09T15:14:03.891Z" }, + { url = "https://files.pythonhosted.org/packages/14/69/5c01b461044eb1f45ddcce006706eb88110c793cdb11c7ae0b5e08492e94/maturin-1.15.0-py3-none-linux_armv6l.whl", hash = "sha256:6bf6dc62e22d4dcfd5a51244ff0d58975fa4979c48209fe84159617648956d82", size = 10206220, upload-time = "2026-08-24T12:10:53.327Z" }, + { url = "https://files.pythonhosted.org/packages/eb/1f/2b431554e11687cdb1077e0cdadcc118c53f611086b3af00c8545a67c6a5/maturin-1.15.0-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:cd35772633f489841132bc8e71d6fc7f842df30b9c05cd5cdf1ee1ddcb744cc7", size = 19416513, upload-time = "2026-08-24T12:10:56.126Z" }, + { url = "https://files.pythonhosted.org/packages/51/36/e23a21cb34a648b711036b9b2fe1d4f3f4ee24f8db54215d73f1a9a3a3ec/maturin-1.15.0-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:c40b4eae7bf5ef1f4b1af8d623fe4105016f93578fb15b764e741d08ec3b92dd", size = 10014962, upload-time = "2026-08-24T12:10:58.486Z" }, + { url = "https://files.pythonhosted.org/packages/8c/80/33b15cb2d8f30f12c807955e8f2fd775027692904e30ec0784744ce8cd83/maturin-1.15.0-py3-none-manylinux_2_12_i686.manylinux2010_i686.musllinux_1_1_i686.whl", hash = "sha256:7eb066372f541f8eb4909c79c5d9bd0b9e8125980bdf1ec9e8aba23c6c8d6c55", size = 10196223, upload-time = "2026-08-24T12:11:00.696Z" }, + { url = "https://files.pythonhosted.org/packages/fe/91/b495e19e2f5c503b540452b2039115e7b2363867e8c5ad4179eb752fa92c/maturin-1.15.0-py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.musllinux_1_1_x86_64.whl", hash = "sha256:653020a63525bb224e5ab0adf02e17a2e08bc86dbea7fc1399c9a56d7529b99e", size = 10541186, upload-time = "2026-08-24T12:11:02.857Z" }, + { url = "https://files.pythonhosted.org/packages/5d/2b/2abff58037188d852b124871b1f0d720e1c2bfb3d4f1b03d87c52cd66488/maturin-1.15.0-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.musllinux_1_1_aarch64.whl", hash = "sha256:0ebf9767892725083138e671c34482c660317a2f3d6a29fc0e0f34e9d8c99136", size = 10083468, upload-time = "2026-08-24T12:11:05.012Z" }, + { url = "https://files.pythonhosted.org/packages/3f/07/b7e9f8be99a6627849e81ac7b6694876bce8f50a92995fe17e3cf2610f0a/maturin-1.15.0-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.musllinux_1_1_armv7l.whl", hash = "sha256:7ab7eebffd7b8debca2265985de4eaeb332141276d24b9560b5ad484d4b3add1", size = 10047786, upload-time = "2026-08-24T12:11:07.042Z" }, + { url = "https://files.pythonhosted.org/packages/ca/ab/167e3cb7accee11b507dbe53e0e87aeccb376d44ae66284c96ee4df3a9fd/maturin-1.15.0-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.musllinux_1_1_ppc64le.whl", hash = "sha256:126e12e618b4db42f68c779a56d41f82a390145ba36ac3f621d057eb34f5ad9d", size = 13315332, upload-time = "2026-08-24T12:11:09.433Z" }, + { url = "https://files.pythonhosted.org/packages/14/4d/801379f646cbc6b00998e5289b0630a886be3a4ee4c75b6bc9b87478a7f1/maturin-1.15.0-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:4f9d33e6c3f9615c8caceecbbbd440f8eb25a3ddeb687077682cd5eca2e9ae15", size = 10807183, upload-time = "2026-08-24T12:11:11.73Z" }, + { url = "https://files.pythonhosted.org/packages/89/27/2e612e1cbd1580e9e94d4722c227b5180dca27b32b955a34b79918aa1292/maturin-1.15.0-py3-none-manylinux_2_31_riscv64.musllinux_1_1_riscv64.whl", hash = "sha256:bf29beddd0c6708f112db51d5275fc28b28b9e9c9c5faae387eaef662918b176", size = 10413274, upload-time = "2026-08-24T12:11:14.04Z" }, + { url = "https://files.pythonhosted.org/packages/70/d8/202a7b4d75a51f20f84ec9ce3b7345b12b822207072164e2b1c6ef665125/maturin-1.15.0-py3-none-win32.whl", hash = "sha256:da649988be98e87e009e51b1bf0d301b6a301bc0cecbdd60d40d8ba60748d1ca", size = 8928744, upload-time = "2026-08-24T12:11:16.269Z" }, + { url = "https://files.pythonhosted.org/packages/40/dc/4e90da594986ba78dd3bc8a5921ecdcdb11085b22b02a412caab3b225601/maturin-1.15.0-py3-none-win_amd64.whl", hash = "sha256:552c2be4afd43fe8d5c9f3ec8d4c4756d973b8dcbe94c14084390301f50243e1", size = 10335085, upload-time = "2026-08-24T12:11:18.326Z" }, + { url = "https://files.pythonhosted.org/packages/8b/10/15d4314edf130955edf2dc237aa393a8a7c10f2b9b57b89fa2f61f915659/maturin-1.15.0-py3-none-win_arm64.whl", hash = "sha256:c7dc0c66c78d3debdd9c5aa807e861fbcbf07f3505d34b125df74c03986b0f48", size = 9713795, upload-time = "2026-08-24T12:11:20.83Z" }, ] [[package]] @@ -700,6 +999,143 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/81/f2/08ace4142eb281c12701fc3b93a10795e4d4dc7f753911d836675050f886/msgpack-1.1.2-cp314-cp314t-win_arm64.whl", hash = "sha256:d99ef64f349d5ec3293688e91486c5fdb925ed03807f64d98d205d2713c60b46", size = 70868, upload-time = "2025-10-08T09:15:44.959Z" }, ] +[[package]] +name = "multidict" +version = "6.7.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1a/c2/c2d94cbe6ac1753f3fc980da97b3d930efe1da3af3c9f5125354436c073d/multidict-6.7.1.tar.gz", hash = "sha256:ec6652a1bee61c53a3e5776b6049172c53b6aaba34f18c9ad04f82712bac623d", size = 102010, upload-time = "2026-01-26T02:46:45.979Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ce/f1/a90635c4f88fb913fbf4ce660b83b7445b7a02615bda034b2f8eb38fd597/multidict-6.7.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:7ff981b266af91d7b4b3793ca3382e53229088d193a85dfad6f5f4c27fc73e5d", size = 76626, upload-time = "2026-01-26T02:43:26.485Z" }, + { url = "https://files.pythonhosted.org/packages/a6/9b/267e64eaf6fc637a15b35f5de31a566634a2740f97d8d094a69d34f524a4/multidict-6.7.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:844c5bca0b5444adb44a623fb0a1310c2f4cd41f402126bb269cd44c9b3f3e1e", size = 44706, upload-time = "2026-01-26T02:43:27.607Z" }, + { url = "https://files.pythonhosted.org/packages/dd/a4/d45caf2b97b035c57267791ecfaafbd59c68212004b3842830954bb4b02e/multidict-6.7.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:f2a0a924d4c2e9afcd7ec64f9de35fcd96915149b2216e1cb2c10a56df483855", size = 44356, upload-time = "2026-01-26T02:43:28.661Z" }, + { url = "https://files.pythonhosted.org/packages/fd/d2/0a36c8473f0cbaeadd5db6c8b72d15bbceeec275807772bfcd059bef487d/multidict-6.7.1-cp311-cp311-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:8be1802715a8e892c784c0197c2ace276ea52702a0ede98b6310c8f255a5afb3", size = 244355, upload-time = "2026-01-26T02:43:31.165Z" }, + { url = "https://files.pythonhosted.org/packages/5d/16/8c65be997fd7dd311b7d39c7b6e71a0cb449bad093761481eccbbe4b42a2/multidict-6.7.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2e2d2ed645ea29f31c4c7ea1552fcfd7cb7ba656e1eafd4134a6620c9f5fdd9e", size = 246433, upload-time = "2026-01-26T02:43:32.581Z" }, + { url = "https://files.pythonhosted.org/packages/01/fb/4dbd7e848d2799c6a026ec88ad39cf2b8416aa167fcc903baa55ecaa045c/multidict-6.7.1-cp311-cp311-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:95922cee9a778659e91db6497596435777bd25ed116701a4c034f8e46544955a", size = 225376, upload-time = "2026-01-26T02:43:34.417Z" }, + { url = "https://files.pythonhosted.org/packages/b6/8a/4a3a6341eac3830f6053062f8fbc9a9e54407c80755b3f05bc427295c2d0/multidict-6.7.1-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:6b83cabdc375ffaaa15edd97eb7c0c672ad788e2687004990074d7d6c9b140c8", size = 257365, upload-time = "2026-01-26T02:43:35.741Z" }, + { url = "https://files.pythonhosted.org/packages/f7/a2/dd575a69c1aa206e12d27d0770cdf9b92434b48a9ef0cd0d1afdecaa93c4/multidict-6.7.1-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:38fb49540705369bab8484db0689d86c0a33a0a9f2c1b197f506b71b4b6c19b0", size = 254747, upload-time = "2026-01-26T02:43:36.976Z" }, + { url = "https://files.pythonhosted.org/packages/5a/56/21b27c560c13822ed93133f08aa6372c53a8e067f11fbed37b4adcdac922/multidict-6.7.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:439cbebd499f92e9aa6793016a8acaa161dfa749ae86d20960189f5398a19144", size = 246293, upload-time = "2026-01-26T02:43:38.258Z" }, + { url = "https://files.pythonhosted.org/packages/5a/a4/23466059dc3854763423d0ad6c0f3683a379d97673b1b89ec33826e46728/multidict-6.7.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:6d3bc717b6fe763b8be3f2bee2701d3c8eb1b2a8ae9f60910f1b2860c82b6c49", size = 242962, upload-time = "2026-01-26T02:43:40.034Z" }, + { url = "https://files.pythonhosted.org/packages/1f/67/51dd754a3524d685958001e8fa20a0f5f90a6a856e0a9dcabff69be3dbb7/multidict-6.7.1-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:619e5a1ac57986dbfec9f0b301d865dddf763696435e2962f6d9cf2fdff2bb71", size = 237360, upload-time = "2026-01-26T02:43:41.752Z" }, + { url = "https://files.pythonhosted.org/packages/64/3f/036dfc8c174934d4b55d86ff4f978e558b0e585cef70cfc1ad01adc6bf18/multidict-6.7.1-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:0b38ebffd9be37c1170d33bc0f36f4f262e0a09bc1aac1c34c7aa51a7293f0b3", size = 245940, upload-time = "2026-01-26T02:43:43.042Z" }, + { url = "https://files.pythonhosted.org/packages/3d/20/6214d3c105928ebc353a1c644a6ef1408bc5794fcb4f170bb524a3c16311/multidict-6.7.1-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:10ae39c9cfe6adedcdb764f5e8411d4a92b055e35573a2eaa88d3323289ef93c", size = 253502, upload-time = "2026-01-26T02:43:44.371Z" }, + { url = "https://files.pythonhosted.org/packages/b1/e2/c653bc4ae1be70a0f836b82172d643fcf1dade042ba2676ab08ec08bff0f/multidict-6.7.1-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:25167cc263257660290fba06b9318d2026e3c910be240a146e1f66dd114af2b0", size = 247065, upload-time = "2026-01-26T02:43:45.745Z" }, + { url = "https://files.pythonhosted.org/packages/c8/11/a854b4154cd3bd8b1fd375e8a8ca9d73be37610c361543d56f764109509b/multidict-6.7.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:128441d052254f42989ef98b7b6a6ecb1e6f708aa962c7984235316db59f50fa", size = 241870, upload-time = "2026-01-26T02:43:47.054Z" }, + { url = "https://files.pythonhosted.org/packages/13/bf/9676c0392309b5fdae322333d22a829715b570edb9baa8016a517b55b558/multidict-6.7.1-cp311-cp311-win32.whl", hash = "sha256:d62b7f64ffde3b99d06b707a280db04fb3855b55f5a06df387236051d0668f4a", size = 41302, upload-time = "2026-01-26T02:43:48.753Z" }, + { url = "https://files.pythonhosted.org/packages/c9/68/f16a3a8ba6f7b6dc92a1f19669c0810bd2c43fc5a02da13b1cbf8e253845/multidict-6.7.1-cp311-cp311-win_amd64.whl", hash = "sha256:bdbf9f3b332abd0cdb306e7c2113818ab1e922dc84b8f8fd06ec89ed2a19ab8b", size = 45981, upload-time = "2026-01-26T02:43:49.921Z" }, + { url = "https://files.pythonhosted.org/packages/ac/ad/9dd5305253fa00cd3c7555dbef69d5bf4133debc53b87ab8d6a44d411665/multidict-6.7.1-cp311-cp311-win_arm64.whl", hash = "sha256:b8c990b037d2fff2f4e33d3f21b9b531c5745b33a49a7d6dbe7a177266af44f6", size = 43159, upload-time = "2026-01-26T02:43:51.635Z" }, + { url = "https://files.pythonhosted.org/packages/8d/9c/f20e0e2cf80e4b2e4b1c365bf5fe104ee633c751a724246262db8f1a0b13/multidict-6.7.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:a90f75c956e32891a4eda3639ce6dd86e87105271f43d43442a3aedf3cddf172", size = 76893, upload-time = "2026-01-26T02:43:52.754Z" }, + { url = "https://files.pythonhosted.org/packages/fe/cf/18ef143a81610136d3da8193da9d80bfe1cb548a1e2d1c775f26b23d024a/multidict-6.7.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:3fccb473e87eaa1382689053e4a4618e7ba7b9b9b8d6adf2027ee474597128cd", size = 45456, upload-time = "2026-01-26T02:43:53.893Z" }, + { url = "https://files.pythonhosted.org/packages/a9/65/1caac9d4cd32e8433908683446eebc953e82d22b03d10d41a5f0fefe991b/multidict-6.7.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b0fa96985700739c4c7853a43c0b3e169360d6855780021bfc6d0f1ce7c123e7", size = 43872, upload-time = "2026-01-26T02:43:55.041Z" }, + { url = "https://files.pythonhosted.org/packages/cf/3b/d6bd75dc4f3ff7c73766e04e705b00ed6dbbaccf670d9e05a12b006f5a21/multidict-6.7.1-cp312-cp312-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:cb2a55f408c3043e42b40cc8eecd575afa27b7e0b956dfb190de0f8499a57a53", size = 251018, upload-time = "2026-01-26T02:43:56.198Z" }, + { url = "https://files.pythonhosted.org/packages/fd/80/c959c5933adedb9ac15152e4067c702a808ea183a8b64cf8f31af8ad3155/multidict-6.7.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:eb0ce7b2a32d09892b3dd6cc44877a0d02a33241fafca5f25c8b6b62374f8b75", size = 258883, upload-time = "2026-01-26T02:43:57.499Z" }, + { url = "https://files.pythonhosted.org/packages/86/85/7ed40adafea3d4f1c8b916e3b5cc3a8e07dfcdcb9cd72800f4ed3ca1b387/multidict-6.7.1-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:c3a32d23520ee37bf327d1e1a656fec76a2edd5c038bf43eddfa0572ec49c60b", size = 242413, upload-time = "2026-01-26T02:43:58.755Z" }, + { url = "https://files.pythonhosted.org/packages/d2/57/b8565ff533e48595503c785f8361ff9a4fde4d67de25c207cd0ba3befd03/multidict-6.7.1-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:9c90fed18bffc0189ba814749fdcc102b536e83a9f738a9003e569acd540a733", size = 268404, upload-time = "2026-01-26T02:44:00.216Z" }, + { url = "https://files.pythonhosted.org/packages/e0/50/9810c5c29350f7258180dfdcb2e52783a0632862eb334c4896ac717cebcb/multidict-6.7.1-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:da62917e6076f512daccfbbde27f46fed1c98fee202f0559adec8ee0de67f71a", size = 269456, upload-time = "2026-01-26T02:44:02.202Z" }, + { url = "https://files.pythonhosted.org/packages/f3/8d/5e5be3ced1d12966fefb5c4ea3b2a5b480afcea36406559442c6e31d4a48/multidict-6.7.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bfde23ef6ed9db7eaee6c37dcec08524cb43903c60b285b172b6c094711b3961", size = 256322, upload-time = "2026-01-26T02:44:03.56Z" }, + { url = "https://files.pythonhosted.org/packages/31/6e/d8a26d81ac166a5592782d208dd90dfdc0a7a218adaa52b45a672b46c122/multidict-6.7.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3758692429e4e32f1ba0df23219cd0b4fc0a52f476726fff9337d1a57676a582", size = 253955, upload-time = "2026-01-26T02:44:04.845Z" }, + { url = "https://files.pythonhosted.org/packages/59/4c/7c672c8aad41534ba619bcd4ade7a0dc87ed6b8b5c06149b85d3dd03f0cd/multidict-6.7.1-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:398c1478926eca669f2fd6a5856b6de9c0acf23a2cb59a14c0ba5844fa38077e", size = 251254, upload-time = "2026-01-26T02:44:06.133Z" }, + { url = "https://files.pythonhosted.org/packages/7b/bd/84c24de512cbafbdbc39439f74e967f19570ce7924e3007174a29c348916/multidict-6.7.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:c102791b1c4f3ab36ce4101154549105a53dc828f016356b3e3bcae2e3a039d3", size = 252059, upload-time = "2026-01-26T02:44:07.518Z" }, + { url = "https://files.pythonhosted.org/packages/fa/ba/f5449385510825b73d01c2d4087bf6d2fccc20a2d42ac34df93191d3dd03/multidict-6.7.1-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:a088b62bd733e2ad12c50dad01b7d0166c30287c166e137433d3b410add807a6", size = 263588, upload-time = "2026-01-26T02:44:09.382Z" }, + { url = "https://files.pythonhosted.org/packages/d7/11/afc7c677f68f75c84a69fe37184f0f82fce13ce4b92f49f3db280b7e92b3/multidict-6.7.1-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:3d51ff4785d58d3f6c91bdbffcb5e1f7ddfda557727043aa20d20ec4f65e324a", size = 259642, upload-time = "2026-01-26T02:44:10.73Z" }, + { url = "https://files.pythonhosted.org/packages/2b/17/ebb9644da78c4ab36403739e0e6e0e30ebb135b9caf3440825001a0bddcb/multidict-6.7.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fc5907494fccf3e7d3f94f95c91d6336b092b5fc83811720fae5e2765890dfba", size = 251377, upload-time = "2026-01-26T02:44:12.042Z" }, + { url = "https://files.pythonhosted.org/packages/ca/a4/840f5b97339e27846c46307f2530a2805d9d537d8b8bd416af031cad7fa0/multidict-6.7.1-cp312-cp312-win32.whl", hash = "sha256:28ca5ce2fd9716631133d0e9a9b9a745ad7f60bac2bccafb56aa380fc0b6c511", size = 41887, upload-time = "2026-01-26T02:44:14.245Z" }, + { url = "https://files.pythonhosted.org/packages/80/31/0b2517913687895f5904325c2069d6a3b78f66cc641a86a2baf75a05dcbb/multidict-6.7.1-cp312-cp312-win_amd64.whl", hash = "sha256:fcee94dfbd638784645b066074b338bc9cc155d4b4bffa4adce1615c5a426c19", size = 46053, upload-time = "2026-01-26T02:44:15.371Z" }, + { url = "https://files.pythonhosted.org/packages/0c/5b/aba28e4ee4006ae4c7df8d327d31025d760ffa992ea23812a601d226e682/multidict-6.7.1-cp312-cp312-win_arm64.whl", hash = "sha256:ba0a9fb644d0c1a2194cf7ffb043bd852cea63a57f66fbd33959f7dae18517bf", size = 43307, upload-time = "2026-01-26T02:44:16.852Z" }, + { url = "https://files.pythonhosted.org/packages/f2/22/929c141d6c0dba87d3e1d38fbdf1ba8baba86b7776469f2bc2d3227a1e67/multidict-6.7.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:2b41f5fed0ed563624f1c17630cb9941cf2309d4df00e494b551b5f3e3d67a23", size = 76174, upload-time = "2026-01-26T02:44:18.509Z" }, + { url = "https://files.pythonhosted.org/packages/c7/75/bc704ae15fee974f8fccd871305e254754167dce5f9e42d88a2def741a1d/multidict-6.7.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:84e61e3af5463c19b67ced91f6c634effb89ef8bfc5ca0267f954451ed4bb6a2", size = 45116, upload-time = "2026-01-26T02:44:19.745Z" }, + { url = "https://files.pythonhosted.org/packages/79/76/55cd7186f498ed080a18440c9013011eb548f77ae1b297206d030eb1180a/multidict-6.7.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:935434b9853c7c112eee7ac891bc4cb86455aa631269ae35442cb316790c1445", size = 43524, upload-time = "2026-01-26T02:44:21.571Z" }, + { url = "https://files.pythonhosted.org/packages/e9/3c/414842ef8d5a1628d68edee29ba0e5bcf235dbfb3ccd3ea303a7fe8c72ff/multidict-6.7.1-cp313-cp313-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:432feb25a1cb67fe82a9680b4d65fb542e4635cb3166cd9c01560651ad60f177", size = 249368, upload-time = "2026-01-26T02:44:22.803Z" }, + { url = "https://files.pythonhosted.org/packages/f6/32/befed7f74c458b4a525e60519fe8d87eef72bb1e99924fa2b0f9d97a221e/multidict-6.7.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e82d14e3c948952a1a85503817e038cba5905a3352de76b9a465075d072fba23", size = 256952, upload-time = "2026-01-26T02:44:24.306Z" }, + { url = "https://files.pythonhosted.org/packages/03/d6/c878a44ba877f366630c860fdf74bfb203c33778f12b6ac274936853c451/multidict-6.7.1-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:4cfb48c6ea66c83bcaaf7e4dfa7ec1b6bbcf751b7db85a328902796dfde4c060", size = 240317, upload-time = "2026-01-26T02:44:25.772Z" }, + { url = "https://files.pythonhosted.org/packages/68/49/57421b4d7ad2e9e60e25922b08ceb37e077b90444bde6ead629095327a6f/multidict-6.7.1-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:1d540e51b7e8e170174555edecddbd5538105443754539193e3e1061864d444d", size = 267132, upload-time = "2026-01-26T02:44:27.648Z" }, + { url = "https://files.pythonhosted.org/packages/b7/fe/ec0edd52ddbcea2a2e89e174f0206444a61440b40f39704e64dc807a70bd/multidict-6.7.1-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:273d23f4b40f3dce4d6c8a821c741a86dec62cded82e1175ba3d99be128147ed", size = 268140, upload-time = "2026-01-26T02:44:29.588Z" }, + { url = "https://files.pythonhosted.org/packages/b0/73/6e1b01cbeb458807aa0831742232dbdd1fa92bfa33f52a3f176b4ff3dc11/multidict-6.7.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9d624335fd4fa1c08a53f8b4be7676ebde19cd092b3895c421045ca87895b429", size = 254277, upload-time = "2026-01-26T02:44:30.902Z" }, + { url = "https://files.pythonhosted.org/packages/6a/b2/5fb8c124d7561a4974c342bc8c778b471ebbeb3cc17df696f034a7e9afe7/multidict-6.7.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:12fad252f8b267cc75b66e8fc51b3079604e8d43a75428ffe193cd9e2195dfd6", size = 252291, upload-time = "2026-01-26T02:44:32.31Z" }, + { url = "https://files.pythonhosted.org/packages/5a/96/51d4e4e06bcce92577fcd488e22600bd38e4fd59c20cb49434d054903bd2/multidict-6.7.1-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:03ede2a6ffbe8ef936b92cb4529f27f42be7f56afcdab5ab739cd5f27fb1cbf9", size = 250156, upload-time = "2026-01-26T02:44:33.734Z" }, + { url = "https://files.pythonhosted.org/packages/db/6b/420e173eec5fba721a50e2a9f89eda89d9c98fded1124f8d5c675f7a0c0f/multidict-6.7.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:90efbcf47dbe33dcf643a1e400d67d59abeac5db07dc3f27d6bdeae497a2198c", size = 249742, upload-time = "2026-01-26T02:44:35.222Z" }, + { url = "https://files.pythonhosted.org/packages/44/a3/ec5b5bd98f306bc2aa297b8c6f11a46714a56b1e6ef5ebda50a4f5d7c5fb/multidict-6.7.1-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:5c4b9bfc148f5a91be9244d6264c53035c8a0dcd2f51f1c3c6e30e30ebaa1c84", size = 262221, upload-time = "2026-01-26T02:44:36.604Z" }, + { url = "https://files.pythonhosted.org/packages/cd/f7/e8c0d0da0cd1e28d10e624604e1a36bcc3353aaebdfdc3a43c72bc683a12/multidict-6.7.1-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:401c5a650f3add2472d1d288c26deebc540f99e2fb83e9525007a74cd2116f1d", size = 258664, upload-time = "2026-01-26T02:44:38.008Z" }, + { url = "https://files.pythonhosted.org/packages/52/da/151a44e8016dd33feed44f730bd856a66257c1ee7aed4f44b649fb7edeb3/multidict-6.7.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:97891f3b1b3ffbded884e2916cacf3c6fc87b66bb0dde46f7357404750559f33", size = 249490, upload-time = "2026-01-26T02:44:39.386Z" }, + { url = "https://files.pythonhosted.org/packages/87/af/a3b86bf9630b732897f6fc3f4c4714b90aa4361983ccbdcd6c0339b21b0c/multidict-6.7.1-cp313-cp313-win32.whl", hash = "sha256:e1c5988359516095535c4301af38d8a8838534158f649c05dd1050222321bcb3", size = 41695, upload-time = "2026-01-26T02:44:41.318Z" }, + { url = "https://files.pythonhosted.org/packages/b2/35/e994121b0e90e46134673422dd564623f93304614f5d11886b1b3e06f503/multidict-6.7.1-cp313-cp313-win_amd64.whl", hash = "sha256:960c83bf01a95b12b08fd54324a4eb1d5b52c88932b5cba5d6e712bb3ed12eb5", size = 45884, upload-time = "2026-01-26T02:44:42.488Z" }, + { url = "https://files.pythonhosted.org/packages/ca/61/42d3e5dbf661242a69c97ea363f2d7b46c567da8eadef8890022be6e2ab0/multidict-6.7.1-cp313-cp313-win_arm64.whl", hash = "sha256:563fe25c678aaba333d5399408f5ec3c383ca5b663e7f774dd179a520b8144df", size = 43122, upload-time = "2026-01-26T02:44:43.664Z" }, + { url = "https://files.pythonhosted.org/packages/6d/b3/e6b21c6c4f314bb956016b0b3ef2162590a529b84cb831c257519e7fde44/multidict-6.7.1-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:c76c4bec1538375dad9d452d246ca5368ad6e1c9039dadcf007ae59c70619ea1", size = 83175, upload-time = "2026-01-26T02:44:44.894Z" }, + { url = "https://files.pythonhosted.org/packages/fb/76/23ecd2abfe0957b234f6c960f4ade497f55f2c16aeb684d4ecdbf1c95791/multidict-6.7.1-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:57b46b24b5d5ebcc978da4ec23a819a9402b4228b8a90d9c656422b4bdd8a963", size = 48460, upload-time = "2026-01-26T02:44:46.106Z" }, + { url = "https://files.pythonhosted.org/packages/c4/57/a0ed92b23f3a042c36bc4227b72b97eca803f5f1801c1ab77c8a212d455e/multidict-6.7.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:e954b24433c768ce78ab7929e84ccf3422e46deb45a4dc9f93438f8217fa2d34", size = 46930, upload-time = "2026-01-26T02:44:47.278Z" }, + { url = "https://files.pythonhosted.org/packages/b5/66/02ec7ace29162e447f6382c495dc95826bf931d3818799bbef11e8f7df1a/multidict-6.7.1-cp313-cp313t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:3bd231490fa7217cc832528e1cd8752a96f0125ddd2b5749390f7c3ec8721b65", size = 242582, upload-time = "2026-01-26T02:44:48.604Z" }, + { url = "https://files.pythonhosted.org/packages/58/18/64f5a795e7677670e872673aca234162514696274597b3708b2c0d276cce/multidict-6.7.1-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:253282d70d67885a15c8a7716f3a73edf2d635793ceda8173b9ecc21f2fb8292", size = 250031, upload-time = "2026-01-26T02:44:50.544Z" }, + { url = "https://files.pythonhosted.org/packages/c8/ed/e192291dbbe51a8290c5686f482084d31bcd9d09af24f63358c3d42fd284/multidict-6.7.1-cp313-cp313t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:0b4c48648d7649c9335cf1927a8b87fa692de3dcb15faa676c6a6f1f1aabda43", size = 228596, upload-time = "2026-01-26T02:44:51.951Z" }, + { url = "https://files.pythonhosted.org/packages/1e/7e/3562a15a60cf747397e7f2180b0a11dc0c38d9175a650e75fa1b4d325e15/multidict-6.7.1-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:98bc624954ec4d2c7cb074b8eefc2b5d0ce7d482e410df446414355d158fe4ca", size = 257492, upload-time = "2026-01-26T02:44:53.902Z" }, + { url = "https://files.pythonhosted.org/packages/24/02/7d0f9eae92b5249bb50ac1595b295f10e263dd0078ebb55115c31e0eaccd/multidict-6.7.1-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:1b99af4d9eec0b49927b4402bcbb58dea89d3e0db8806a4086117019939ad3dd", size = 255899, upload-time = "2026-01-26T02:44:55.316Z" }, + { url = "https://files.pythonhosted.org/packages/00/e3/9b60ed9e23e64c73a5cde95269ef1330678e9c6e34dd4eb6b431b85b5a10/multidict-6.7.1-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6aac4f16b472d5b7dc6f66a0d49dd57b0e0902090be16594dc9ebfd3d17c47e7", size = 247970, upload-time = "2026-01-26T02:44:56.783Z" }, + { url = "https://files.pythonhosted.org/packages/3e/06/538e58a63ed5cfb0bd4517e346b91da32fde409d839720f664e9a4ae4f9d/multidict-6.7.1-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:21f830fe223215dffd51f538e78c172ed7c7f60c9b96a2bf05c4848ad49921c3", size = 245060, upload-time = "2026-01-26T02:44:58.195Z" }, + { url = "https://files.pythonhosted.org/packages/b2/2f/d743a3045a97c895d401e9bd29aaa09b94f5cbdf1bd561609e5a6c431c70/multidict-6.7.1-cp313-cp313t-musllinux_1_2_armv7l.whl", hash = "sha256:f5dd81c45b05518b9aa4da4aa74e1c93d715efa234fd3e8a179df611cc85e5f4", size = 235888, upload-time = "2026-01-26T02:44:59.57Z" }, + { url = "https://files.pythonhosted.org/packages/38/83/5a325cac191ab28b63c52f14f1131f3b0a55ba3b9aa65a6d0bf2a9b921a0/multidict-6.7.1-cp313-cp313t-musllinux_1_2_i686.whl", hash = "sha256:eb304767bca2bb92fb9c5bd33cedc95baee5bb5f6c88e63706533a1c06ad08c8", size = 243554, upload-time = "2026-01-26T02:45:01.054Z" }, + { url = "https://files.pythonhosted.org/packages/20/1f/9d2327086bd15da2725ef6aae624208e2ef828ed99892b17f60c344e57ed/multidict-6.7.1-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:c9035dde0f916702850ef66460bc4239d89d08df4d02023a5926e7446724212c", size = 252341, upload-time = "2026-01-26T02:45:02.484Z" }, + { url = "https://files.pythonhosted.org/packages/e8/2c/2a1aa0280cf579d0f6eed8ee5211c4f1730bd7e06c636ba2ee6aafda302e/multidict-6.7.1-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:af959b9beeb66c822380f222f0e0a1889331597e81f1ded7f374f3ecb0fd6c52", size = 246391, upload-time = "2026-01-26T02:45:03.862Z" }, + { url = "https://files.pythonhosted.org/packages/e5/03/7ca022ffc36c5a3f6e03b179a5ceb829be9da5783e6fe395f347c0794680/multidict-6.7.1-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:41f2952231456154ee479651491e94118229844dd7226541788be783be2b5108", size = 243422, upload-time = "2026-01-26T02:45:05.296Z" }, + { url = "https://files.pythonhosted.org/packages/dc/1d/b31650eab6c5778aceed46ba735bd97f7c7d2f54b319fa916c0f96e7805b/multidict-6.7.1-cp313-cp313t-win32.whl", hash = "sha256:df9f19c28adcb40b6aae30bbaa1478c389efd50c28d541d76760199fc1037c32", size = 47770, upload-time = "2026-01-26T02:45:06.754Z" }, + { url = "https://files.pythonhosted.org/packages/ac/5b/2d2d1d522e51285bd61b1e20df8f47ae1a9d80839db0b24ea783b3832832/multidict-6.7.1-cp313-cp313t-win_amd64.whl", hash = "sha256:d54ecf9f301853f2c5e802da559604b3e95bb7a3b01a9c295c6ee591b9882de8", size = 53109, upload-time = "2026-01-26T02:45:08.044Z" }, + { url = "https://files.pythonhosted.org/packages/3d/a3/cc409ba012c83ca024a308516703cf339bdc4b696195644a7215a5164a24/multidict-6.7.1-cp313-cp313t-win_arm64.whl", hash = "sha256:5a37ca18e360377cfda1d62f5f382ff41f2b8c4ccb329ed974cc2e1643440118", size = 45573, upload-time = "2026-01-26T02:45:09.349Z" }, + { url = "https://files.pythonhosted.org/packages/91/cc/db74228a8be41884a567e88a62fd589a913708fcf180d029898c17a9a371/multidict-6.7.1-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:8f333ec9c5eb1b7105e3b84b53141e66ca05a19a605368c55450b6ba208cb9ee", size = 75190, upload-time = "2026-01-26T02:45:10.651Z" }, + { url = "https://files.pythonhosted.org/packages/d5/22/492f2246bb5b534abd44804292e81eeaf835388901f0c574bac4eeec73c5/multidict-6.7.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:a407f13c188f804c759fc6a9f88286a565c242a76b27626594c133b82883b5c2", size = 44486, upload-time = "2026-01-26T02:45:11.938Z" }, + { url = "https://files.pythonhosted.org/packages/f1/4f/733c48f270565d78b4544f2baddc2fb2a245e5a8640254b12c36ac7ac68e/multidict-6.7.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:0e161ddf326db5577c3a4cc2d8648f81456e8a20d40415541587a71620d7a7d1", size = 43219, upload-time = "2026-01-26T02:45:14.346Z" }, + { url = "https://files.pythonhosted.org/packages/24/bb/2c0c2287963f4259c85e8bcbba9182ced8d7fca65c780c38e99e61629d11/multidict-6.7.1-cp314-cp314-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:1e3a8bb24342a8201d178c3b4984c26ba81a577c80d4d525727427460a50c22d", size = 245132, upload-time = "2026-01-26T02:45:15.712Z" }, + { url = "https://files.pythonhosted.org/packages/a7/f9/44d4b3064c65079d2467888794dea218d1601898ac50222ab8a9a8094460/multidict-6.7.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:97231140a50f5d447d3164f994b86a0bed7cd016e2682f8650d6a9158e14fd31", size = 252420, upload-time = "2026-01-26T02:45:17.293Z" }, + { url = "https://files.pythonhosted.org/packages/8b/13/78f7275e73fa17b24c9a51b0bd9d73ba64bb32d0ed51b02a746eb876abe7/multidict-6.7.1-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:6b10359683bd8806a200fd2909e7c8ca3a7b24ec1d8132e483d58e791d881048", size = 233510, upload-time = "2026-01-26T02:45:19.356Z" }, + { url = "https://files.pythonhosted.org/packages/4b/25/8167187f62ae3cbd52da7893f58cb036b47ea3fb67138787c76800158982/multidict-6.7.1-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:283ddac99f7ac25a4acadbf004cb5ae34480bbeb063520f70ce397b281859362", size = 264094, upload-time = "2026-01-26T02:45:20.834Z" }, + { url = "https://files.pythonhosted.org/packages/a1/e7/69a3a83b7b030cf283fb06ce074a05a02322359783424d7edf0f15fe5022/multidict-6.7.1-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:538cec1e18c067d0e6103aa9a74f9e832904c957adc260e61cd9d8cf0c3b3d37", size = 260786, upload-time = "2026-01-26T02:45:22.818Z" }, + { url = "https://files.pythonhosted.org/packages/fe/3b/8ec5074bcfc450fe84273713b4b0a0dd47c0249358f5d82eb8104ffe2520/multidict-6.7.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7eee46ccb30ff48a1e35bb818cc90846c6be2b68240e42a78599166722cea709", size = 248483, upload-time = "2026-01-26T02:45:24.368Z" }, + { url = "https://files.pythonhosted.org/packages/48/5a/d5a99e3acbca0e29c5d9cba8f92ceb15dce78bab963b308ae692981e3a5d/multidict-6.7.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:fa263a02f4f2dd2d11a7b1bb4362aa7cb1049f84a9235d31adf63f30143469a0", size = 248403, upload-time = "2026-01-26T02:45:25.982Z" }, + { url = "https://files.pythonhosted.org/packages/35/48/e58cd31f6c7d5102f2a4bf89f96b9cf7e00b6c6f3d04ecc44417c00a5a3c/multidict-6.7.1-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:2e1425e2f99ec5bd36c15a01b690a1a2456209c5deed58f95469ffb46039ccbb", size = 240315, upload-time = "2026-01-26T02:45:27.487Z" }, + { url = "https://files.pythonhosted.org/packages/94/33/1cd210229559cb90b6786c30676bb0c58249ff42f942765f88793b41fdce/multidict-6.7.1-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:497394b3239fc6f0e13a78a3e1b61296e72bf1c5f94b4c4eb80b265c37a131cd", size = 245528, upload-time = "2026-01-26T02:45:28.991Z" }, + { url = "https://files.pythonhosted.org/packages/64/f2/6e1107d226278c876c783056b7db43d800bb64c6131cec9c8dfb6903698e/multidict-6.7.1-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:233b398c29d3f1b9676b4b6f75c518a06fcb2ea0b925119fb2c1bc35c05e1601", size = 258784, upload-time = "2026-01-26T02:45:30.503Z" }, + { url = "https://files.pythonhosted.org/packages/4d/c1/11f664f14d525e4a1b5327a82d4de61a1db604ab34c6603bb3c2cc63ad34/multidict-6.7.1-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:93b1818e4a6e0930454f0f2af7dfce69307ca03cdcfb3739bf4d91241967b6c1", size = 251980, upload-time = "2026-01-26T02:45:32.603Z" }, + { url = "https://files.pythonhosted.org/packages/e1/9f/75a9ac888121d0c5bbd4ecf4eead45668b1766f6baabfb3b7f66a410e231/multidict-6.7.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:f33dc2a3abe9249ea5d8360f969ec7f4142e7ac45ee7014d8f8d5acddf178b7b", size = 243602, upload-time = "2026-01-26T02:45:34.043Z" }, + { url = "https://files.pythonhosted.org/packages/9a/e7/50bf7b004cc8525d80dbbbedfdc7aed3e4c323810890be4413e589074032/multidict-6.7.1-cp314-cp314-win32.whl", hash = "sha256:3ab8b9d8b75aef9df299595d5388b14530839f6422333357af1339443cff777d", size = 40930, upload-time = "2026-01-26T02:45:36.278Z" }, + { url = "https://files.pythonhosted.org/packages/e0/bf/52f25716bbe93745595800f36fb17b73711f14da59ed0bb2eba141bc9f0f/multidict-6.7.1-cp314-cp314-win_amd64.whl", hash = "sha256:5e01429a929600e7dab7b166062d9bb54a5eed752384c7384c968c2afab8f50f", size = 45074, upload-time = "2026-01-26T02:45:37.546Z" }, + { url = "https://files.pythonhosted.org/packages/97/ab/22803b03285fa3a525f48217963da3a65ae40f6a1b6f6cf2768879e208f9/multidict-6.7.1-cp314-cp314-win_arm64.whl", hash = "sha256:4885cb0e817aef5d00a2e8451d4665c1808378dc27c2705f1bf4ef8505c0d2e5", size = 42471, upload-time = "2026-01-26T02:45:38.889Z" }, + { url = "https://files.pythonhosted.org/packages/e0/6d/f9293baa6146ba9507e360ea0292b6422b016907c393e2f63fc40ab7b7b5/multidict-6.7.1-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:0458c978acd8e6ea53c81eefaddbbee9c6c5e591f41b3f5e8e194780fe026581", size = 82401, upload-time = "2026-01-26T02:45:40.254Z" }, + { url = "https://files.pythonhosted.org/packages/7a/68/53b5494738d83558d87c3c71a486504d8373421c3e0dbb6d0db48ad42ee0/multidict-6.7.1-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:c0abd12629b0af3cf590982c0b413b1e7395cd4ec026f30986818ab95bfaa94a", size = 48143, upload-time = "2026-01-26T02:45:41.635Z" }, + { url = "https://files.pythonhosted.org/packages/37/e8/5284c53310dcdc99ce5d66563f6e5773531a9b9fe9ec7a615e9bc306b05f/multidict-6.7.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:14525a5f61d7d0c94b368a42cff4c9a4e7ba2d52e2672a7b23d84dc86fb02b0c", size = 46507, upload-time = "2026-01-26T02:45:42.99Z" }, + { url = "https://files.pythonhosted.org/packages/e4/fc/6800d0e5b3875568b4083ecf5f310dcf91d86d52573160834fb4bfcf5e4f/multidict-6.7.1-cp314-cp314t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:17307b22c217b4cf05033dabefe68255a534d637c6c9b0cc8382718f87be4262", size = 239358, upload-time = "2026-01-26T02:45:44.376Z" }, + { url = "https://files.pythonhosted.org/packages/41/75/4ad0973179361cdf3a113905e6e088173198349131be2b390f9fa4da5fc6/multidict-6.7.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7a7e590ff876a3eaf1c02a4dfe0724b6e69a9e9de6d8f556816f29c496046e59", size = 246884, upload-time = "2026-01-26T02:45:47.167Z" }, + { url = "https://files.pythonhosted.org/packages/c3/9c/095bb28b5da139bd41fb9a5d5caff412584f377914bd8787c2aa98717130/multidict-6.7.1-cp314-cp314t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:5fa6a95dfee63893d80a34758cd0e0c118a30b8dcb46372bf75106c591b77889", size = 225878, upload-time = "2026-01-26T02:45:48.698Z" }, + { url = "https://files.pythonhosted.org/packages/07/d0/c0a72000243756e8f5a277b6b514fa005f2c73d481b7d9e47cd4568aa2e4/multidict-6.7.1-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a0543217a6a017692aa6ae5cc39adb75e587af0f3a82288b1492eb73dd6cc2a4", size = 253542, upload-time = "2026-01-26T02:45:50.164Z" }, + { url = "https://files.pythonhosted.org/packages/c0/6b/f69da15289e384ecf2a68837ec8b5ad8c33e973aa18b266f50fe55f24b8c/multidict-6.7.1-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:f99fe611c312b3c1c0ace793f92464d8cd263cc3b26b5721950d977b006b6c4d", size = 252403, upload-time = "2026-01-26T02:45:51.779Z" }, + { url = "https://files.pythonhosted.org/packages/a2/76/b9669547afa5a1a25cd93eaca91c0da1c095b06b6d2d8ec25b713588d3a1/multidict-6.7.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9004d8386d133b7e6135679424c91b0b854d2d164af6ea3f289f8f2761064609", size = 244889, upload-time = "2026-01-26T02:45:53.27Z" }, + { url = "https://files.pythonhosted.org/packages/7e/a9/a50d2669e506dad33cfc45b5d574a205587b7b8a5f426f2fbb2e90882588/multidict-6.7.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e628ef0e6859ffd8273c69412a2465c4be4a9517d07261b33334b5ec6f3c7489", size = 241982, upload-time = "2026-01-26T02:45:54.919Z" }, + { url = "https://files.pythonhosted.org/packages/c5/bb/1609558ad8b456b4827d3c5a5b775c93b87878fd3117ed3db3423dfbce1b/multidict-6.7.1-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:841189848ba629c3552035a6a7f5bf3b02eb304e9fea7492ca220a8eda6b0e5c", size = 232415, upload-time = "2026-01-26T02:45:56.981Z" }, + { url = "https://files.pythonhosted.org/packages/d8/59/6f61039d2aa9261871e03ab9dc058a550d240f25859b05b67fd70f80d4b3/multidict-6.7.1-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:ce1bbd7d780bb5a0da032e095c951f7014d6b0a205f8318308140f1a6aba159e", size = 240337, upload-time = "2026-01-26T02:45:58.698Z" }, + { url = "https://files.pythonhosted.org/packages/a1/29/fdc6a43c203890dc2ae9249971ecd0c41deaedfe00d25cb6564b2edd99eb/multidict-6.7.1-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:b26684587228afed0d50cf804cc71062cc9c1cdf55051c4c6345d372947b268c", size = 248788, upload-time = "2026-01-26T02:46:00.862Z" }, + { url = "https://files.pythonhosted.org/packages/a9/14/a153a06101323e4cf086ecee3faadba52ff71633d471f9685c42e3736163/multidict-6.7.1-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:9f9af11306994335398293f9958071019e3ab95e9a707dc1383a35613f6abcb9", size = 242842, upload-time = "2026-01-26T02:46:02.824Z" }, + { url = "https://files.pythonhosted.org/packages/41/5f/604ae839e64a4a6efc80db94465348d3b328ee955e37acb24badbcd24d83/multidict-6.7.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:b4938326284c4f1224178a560987b6cf8b4d38458b113d9b8c1db1a836e640a2", size = 240237, upload-time = "2026-01-26T02:46:05.898Z" }, + { url = "https://files.pythonhosted.org/packages/5f/60/c3a5187bf66f6fb546ff4ab8fb5a077cbdd832d7b1908d4365c7f74a1917/multidict-6.7.1-cp314-cp314t-win32.whl", hash = "sha256:98655c737850c064a65e006a3df7c997cd3b220be4ec8fe26215760b9697d4d7", size = 48008, upload-time = "2026-01-26T02:46:07.468Z" }, + { url = "https://files.pythonhosted.org/packages/0c/f7/addf1087b860ac60e6f382240f64fb99f8bfb532bb06f7c542b83c29ca61/multidict-6.7.1-cp314-cp314t-win_amd64.whl", hash = "sha256:497bde6223c212ba11d462853cfa4f0ae6ef97465033e7dc9940cdb3ab5b48e5", size = 53542, upload-time = "2026-01-26T02:46:08.809Z" }, + { url = "https://files.pythonhosted.org/packages/4c/81/4629d0aa32302ef7b2ec65c75a728cc5ff4fa410c50096174c1632e70b3e/multidict-6.7.1-cp314-cp314t-win_arm64.whl", hash = "sha256:2bbd113e0d4af5db41d5ebfe9ccaff89de2120578164f86a5d17d5a576d1e5b2", size = 44719, upload-time = "2026-01-26T02:46:11.146Z" }, + { url = "https://files.pythonhosted.org/packages/81/08/7036c080d7117f28a4af526d794aab6a84463126db031b007717c1a6676e/multidict-6.7.1-py3-none-any.whl", hash = "sha256:55d97cc6dae627efa6a6e548885712d4864b81110ac76fa4e534c03819fa4a56", size = 12319, upload-time = "2026-01-26T02:46:44.004Z" }, +] + +[[package]] +name = "multiprocess" +version = "0.70.19" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "dill" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a2/f2/e783ac7f2aeeed14e9e12801f22529cc7e6b7ab80928d6dcce4e9f00922d/multiprocess-0.70.19.tar.gz", hash = "sha256:952021e0e6c55a4a9fe4cd787895b86e239a40e76802a789d6305398d3975897", size = 2079989, upload-time = "2026-01-19T06:47:39.744Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7e/aa/714635c727dbfc251139226fa4eaf1b07f00dc12d9cd2eb25f931adaf873/multiprocess-0.70.19-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:1bbf1b69af1cf64cd05f65337d9215b88079ec819cd0ea7bac4dab84e162efe7", size = 144743, upload-time = "2026-01-19T06:47:24.562Z" }, + { url = "https://files.pythonhosted.org/packages/0f/e1/155f6abf5e6b5d9cef29b6d0167c180846157a4aca9b9bee1a217f67c959/multiprocess-0.70.19-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:5be9ec7f0c1c49a4f4a6fd20d5dda4aeabc2d39a50f4ad53720f1cd02b3a7c2e", size = 144738, upload-time = "2026-01-19T06:47:26.636Z" }, + { url = "https://files.pythonhosted.org/packages/af/cb/f421c2869d75750a4f32301cc20c4b63fab6376e9a75c8e5e655bdeb3d9b/multiprocess-0.70.19-pp311-pypy311_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:1c3dce098845a0db43b32a0b76a228ca059a668071cfeaa0f40c36c0b1585d45", size = 144741, upload-time = "2026-01-19T06:47:27.985Z" }, + { url = "https://files.pythonhosted.org/packages/e3/45/8004d1e6b9185c1a444d6b55ac5682acf9d98035e54386d967366035a03a/multiprocess-0.70.19-py310-none-any.whl", hash = "sha256:97404393419dcb2a8385910864eedf47a3cadf82c66345b44f036420eb0b5d87", size = 134948, upload-time = "2026-01-19T06:47:32.325Z" }, + { url = "https://files.pythonhosted.org/packages/86/c2/dec9722dc3474c164a0b6bcd9a7ed7da542c98af8cabce05374abab35edd/multiprocess-0.70.19-py311-none-any.whl", hash = "sha256:928851ae7973aea4ce0eaf330bbdafb2e01398a91518d5c8818802845564f45c", size = 144457, upload-time = "2026-01-19T06:47:33.711Z" }, + { url = "https://files.pythonhosted.org/packages/71/70/38998b950a97ea279e6bd657575d22d1a2047256caf707d9a10fbce4f065/multiprocess-0.70.19-py312-none-any.whl", hash = "sha256:3a56c0e85dd5025161bac5ce138dcac1e49174c7d8e74596537e729fd5c53c28", size = 150281, upload-time = "2026-01-19T06:47:35.037Z" }, + { url = "https://files.pythonhosted.org/packages/7f/74/d2c27e03cb84251dfe7249b8e82923643c6d48fa4883b9476b025e7dc7eb/multiprocess-0.70.19-py313-none-any.whl", hash = "sha256:8d5eb4ec5017ba2fab4e34a747c6d2c2b6fecfe9e7236e77988db91580ada952", size = 156414, upload-time = "2026-01-19T06:47:35.915Z" }, + { url = "https://files.pythonhosted.org/packages/a0/61/af9115673a5870fd885247e2f1b68c4f1197737da315b520a91c757a861a/multiprocess-0.70.19-py314-none-any.whl", hash = "sha256:e8cc7fbdff15c0613f0a1f1f8744bef961b0a164c0ca29bdff53e9d2d93c5e5f", size = 160318, upload-time = "2026-01-19T06:47:37.497Z" }, + { url = "https://files.pythonhosted.org/packages/7e/82/69e539c4c2027f1e1697e09aaa2449243085a0edf81ae2c6341e84d769b6/multiprocess-0.70.19-py39-none-any.whl", hash = "sha256:0d4b4397ed669d371c81dcd1ef33fd384a44d6c3de1bd0ca7ac06d837720d3c5", size = 133477, upload-time = "2026-01-19T06:47:38.619Z" }, +] + [[package]] name = "myst-parser" version = "5.0.0" @@ -718,22 +1154,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d3/ac/686789b9145413f1a61878c407210e41bfdb097976864e0913078b24098c/myst_parser-5.0.0-py3-none-any.whl", hash = "sha256:ab31e516024918296e169139072b81592336f2fef55b8986aa31c9f04b5f7211", size = 84533, upload-time = "2026-01-15T09:08:16.788Z" }, ] -[[package]] -name = "nodejs-wheel-binaries" -version = "22.20.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/0f/54/02f58c8119e2f1984e2572cc77a7b469dbaf4f8d171ad376e305749ef48e/nodejs_wheel_binaries-22.20.0.tar.gz", hash = "sha256:a62d47c9fd9c32191dff65bbe60261504f26992a0a19fe8b4d523256a84bd351", size = 8058, upload-time = "2025-09-26T09:48:00.906Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/24/6d/333e5458422f12318e3c3e6e7f194353aa68b0d633217c7e89833427ca01/nodejs_wheel_binaries-22.20.0-py2.py3-none-macosx_11_0_arm64.whl", hash = "sha256:455add5ac4f01c9c830ab6771dbfad0fdf373f9b040d3aabe8cca9b6c56654fb", size = 53246314, upload-time = "2025-09-26T09:47:32.536Z" }, - { url = "https://files.pythonhosted.org/packages/56/30/dcd6879d286a35b3c4c8f9e5e0e1bcf4f9e25fe35310fc77ecf97f915a23/nodejs_wheel_binaries-22.20.0-py2.py3-none-macosx_11_0_x86_64.whl", hash = "sha256:5d8c12f97eea7028b34a84446eb5ca81829d0c428dfb4e647e09ac617f4e21fa", size = 53644391, upload-time = "2025-09-26T09:47:36.093Z" }, - { url = "https://files.pythonhosted.org/packages/58/be/c7b2e7aa3bb281d380a1c531f84d0ccfe225832dfc3bed1ca171753b9630/nodejs_wheel_binaries-22.20.0-py2.py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7a2b0989194148f66e9295d8f11bc463bde02cbe276517f4d20a310fb84780ae", size = 60282516, upload-time = "2025-09-26T09:47:39.88Z" }, - { url = "https://files.pythonhosted.org/packages/3e/c5/8befacf4190e03babbae54cb0809fb1a76e1600ec3967ab8ee9f8fc85b65/nodejs_wheel_binaries-22.20.0-py2.py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b5c500aa4dc046333ecb0a80f183e069e5c30ce637f1c1a37166b2c0b642dc21", size = 60347290, upload-time = "2025-09-26T09:47:43.712Z" }, - { url = "https://files.pythonhosted.org/packages/c0/bd/cfffd1e334277afa0714962c6ec432b5fe339340a6bca2e5fa8e678e7590/nodejs_wheel_binaries-22.20.0-py2.py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:3279eb1b99521f0d20a850bbfc0159a658e0e85b843b3cf31b090d7da9f10dfc", size = 62178798, upload-time = "2025-09-26T09:47:47.752Z" }, - { url = "https://files.pythonhosted.org/packages/08/14/10b83a9c02faac985b3e9f5e65d63a34fc0f46b48d8a2c3e4caa3e1e7318/nodejs_wheel_binaries-22.20.0-py2.py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:d29705797b33bade62d79d8f106c2453c8a26442a9b2a5576610c0f7e7c351ed", size = 62772957, upload-time = "2025-09-26T09:47:51.266Z" }, - { url = "https://files.pythonhosted.org/packages/b4/a9/c6a480259aa0d6b270aac2c6ba73a97444b9267adde983a5b7e34f17e45a/nodejs_wheel_binaries-22.20.0-py2.py3-none-win_amd64.whl", hash = "sha256:4bd658962f24958503541963e5a6f2cc512a8cb301e48a69dc03c879f40a28ae", size = 40120431, upload-time = "2025-09-26T09:47:54.363Z" }, - { url = "https://files.pythonhosted.org/packages/42/b1/6a4eb2c6e9efa028074b0001b61008c9d202b6b46caee9e5d1b18c088216/nodejs_wheel_binaries-22.20.0-py2.py3-none-win_arm64.whl", hash = "sha256:1fccac931faa210d22b6962bcdbc99269d16221d831b9a118bbb80fe434a60b8", size = 38844133, upload-time = "2025-09-26T09:47:57.357Z" }, -] - [[package]] name = "numpy" version = "2.3.3" @@ -952,11 +1372,11 @@ wheels = [ [[package]] name = "pip" -version = "26.1.2" +version = "26.2.1" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/01/91/47e7d486260f618783899587af63ccf7980fb60245c3e63dd4571c6b57ad/pip-26.1.2.tar.gz", hash = "sha256:f49cd134c61cf2fd75e0ce2676db03e4054504a5a4986d00f8299ae632dc4605", size = 1840799, upload-time = "2026-05-31T17:33:58.56Z" } +sdist = { url = "https://files.pythonhosted.org/packages/ae/15/4500e320e6b101ec3b719ae85b697d9940b6cda672bc555bd6016fc60c6f/pip-26.2.1.tar.gz", hash = "sha256:f6ad667e89a1fe78046c8f13232b247200f5258d7828f3f7883d660878e0813f", size = 1848877, upload-time = "2026-08-04T22:51:14.148Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5d/95/6b5cb3461ea5673ba0995989746db58eb18b91b54dbf331e72f569540946/pip-26.1.2-py3-none-any.whl", hash = "sha256:382ff9f685ee3bc25864f820aa50505825f10f5458ffff07e30a6d96e5715cab", size = 1813144, upload-time = "2026-05-31T17:33:56.772Z" }, + { url = "https://files.pythonhosted.org/packages/f3/6e/1736e5b4ae2b778ef2f81c47d797de9f891d4d8acb047a24ca37a60294dd/pip-26.2.1-py3-none-any.whl", hash = "sha256:71138adf1f4ca900cdb7d289c21b7494329f2332b6d85f0e1c42108c0384ed3e", size = 1816632, upload-time = "2026-08-04T22:51:12.472Z" }, ] [[package]] @@ -1006,6 +1426,117 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/84/03/0d3ce49e2505ae70cf43bc5bb3033955d2fc9f932163e84dc0779cc47f48/prompt_toolkit-3.0.52-py3-none-any.whl", hash = "sha256:9aac639a3bbd33284347de5ad8d68ecc044b91a762dc39b7c21095fcd6a19955", size = 391431, upload-time = "2025-08-27T15:23:59.498Z" }, ] +[[package]] +name = "propcache" +version = "0.5.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ec/44/c87281c333769159c50594f22610f77398a47ccbfbbf23074e744e86f87c/propcache-0.5.2.tar.gz", hash = "sha256:01c4fc7480cd0598bb4b57022df55b9ca296da7fc5a8760bd8451a7e63a7d427", size = 50208, upload-time = "2026-05-08T21:02:12.199Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/f1/8a8cc1c2c7e7934ab77e0163414f736fadbc0f5e8dd9673b952355ac175b/propcache-0.5.2-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:74b70780220e2dd89175ca24b81b68b67c83db499ae611e7f2313cb329801c78", size = 90744, upload-time = "2026-05-08T20:59:45.799Z" }, + { url = "https://files.pythonhosted.org/packages/c2/f4/651b1225e976bd1a2ba5cfba0c29d096581c2636b437e3a9a7ab6276270a/propcache-0.5.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:a4840ab0ae0216d952f4b53dc6d0b992bfc2bedbfe360bdd9b548bc184c08959", size = 52033, upload-time = "2026-05-08T20:59:47.408Z" }, + { url = "https://files.pythonhosted.org/packages/15/a8/8ede85d6aa1f79fc7dc2f8fd2c8d65920b8272c3892903c8a1affde48cfb/propcache-0.5.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:c6844ba6364fb12f403928a82cfd295ab103a2b315c77c747b2dbe4a41894ea7", size = 52754, upload-time = "2026-05-08T20:59:49.202Z" }, + { url = "https://files.pythonhosted.org/packages/7d/fe/b3551b41bbc2f5b5bb088fc6920567cd43101253e68fbaa261339eb96fe1/propcache-0.5.2-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2293949b855ce597f2826452d17c2d545fb5622379c4ea6fdf525e9b8e8a2511", size = 57573, upload-time = "2026-05-08T20:59:50.778Z" }, + { url = "https://files.pythonhosted.org/packages/83/27/ab851ebd1b7172e3e161f5f8d39e315d54a91bea246f01f4d872d3376aef/propcache-0.5.2-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:0fd59b5af35f74da48d905dcbad55449ba13be91823cb05a9bd590bbf5b61660", size = 60645, upload-time = "2026-05-08T20:59:52.227Z" }, + { url = "https://files.pythonhosted.org/packages/95/7d/466b3d18022e9897cbda9c735c493c5bd747d7a4c6f5ea1480b4cec434b6/propcache-0.5.2-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:29f9309a2e42b0d273be006fdb4be2d6c39a47f6f57d8fb1cf9f81481df81b66", size = 61563, upload-time = "2026-05-08T20:59:53.866Z" }, + { url = "https://files.pythonhosted.org/packages/27/1b/16ab7f2cf2041da2f60d156ba64c2484eadf9168075b4ff43c3ef60045af/propcache-0.5.2-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5aaa2b923c1944ac8febd6609cb373540a5563e7cbcb0fd770f75dace2eb817b", size = 58888, upload-time = "2026-05-08T20:59:55.457Z" }, + { url = "https://files.pythonhosted.org/packages/0a/67/bb777ffd907633563bf35fd859c4ce97b0512c32f4633cf5d1eb7c33512b/propcache-0.5.2-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:66ea454f095ddf5b6b14f56c064c0941c4788be11e18d2464cf643bf7203ff67", size = 59253, upload-time = "2026-05-08T20:59:57.075Z" }, + { url = "https://files.pythonhosted.org/packages/b9/42/64f8d90b73fd9cdc1499b48057ff6d9cd2a98a25734c9bb62ecf07e87061/propcache-0.5.2-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:95f1e3f4760d404b13c9976c0229b2b49a3c8e2c62a9ce92efdd2b11ada75e3f", size = 57558, upload-time = "2026-05-08T20:59:58.602Z" }, + { url = "https://files.pythonhosted.org/packages/eb/02/dba5bc03c9041f2092ea55a449caf5dfe68352c6654511b29ba0654ddb69/propcache-0.5.2-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:85341b12b9d55bad0bded24cac341bb34289469e03a11f3f583ea1cc1db0326c", size = 55007, upload-time = "2026-05-08T20:59:59.837Z" }, + { url = "https://files.pythonhosted.org/packages/14/c0/43f649c7aa2a77a3b100d84e9dea3a483120ecb608bfe36ce49eaff517fe/propcache-0.5.2-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:26a4dca084132874e639895c3135dfad5eb20bae209f62d1aeb31b03e601c3c0", size = 60355, upload-time = "2026-05-08T21:00:01.144Z" }, + { url = "https://files.pythonhosted.org/packages/83/c0/435dafd27f1cb4a495381dae60e25883ccfe4020bb72818e8184c1678092/propcache-0.5.2-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:3b199b9b2b3d6a7edf3183ba8a9a137a22b97f7df525feb5ae1eccf026d2a9c6", size = 59057, upload-time = "2026-05-08T21:00:02.401Z" }, + { url = "https://files.pythonhosted.org/packages/53/ae/6e292df9135d659944e96cb3389258e4a663e5b2b5f6c217ef0ddc8d2f73/propcache-0.5.2-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:e59bc9e66329185b93dab73f210f1a37f81cb40f321501db8017c9aea15dba27", size = 61938, upload-time = "2026-05-08T21:00:03.638Z" }, + { url = "https://files.pythonhosted.org/packages/0b/42/314ebc50d8159055411fd6b0bda322ff510e4b1f7d2e4927940ad0f6af20/propcache-0.5.2-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:552ffadf6ad409844bc5919c42a0a83d88314cedddaea0e41e80a8b8fffe881f", size = 59731, upload-time = "2026-05-08T21:00:04.881Z" }, + { url = "https://files.pythonhosted.org/packages/b8/9b/2da6dee38871c3c8772fabc2758325a5c9077d6d18c597737dc04dd884cd/propcache-0.5.2-cp311-cp311-win32.whl", hash = "sha256:cd416c1de191973c52ff1a12a57446bfc7642797b282d7caf2162d7d1b8aa9a0", size = 38966, upload-time = "2026-05-08T21:00:06.511Z" }, + { url = "https://files.pythonhosted.org/packages/42/4e/f17363fb58c0afe05b067361cb6d86ed2d29de6506779a27547c4d183075/propcache-0.5.2-cp311-cp311-win_amd64.whl", hash = "sha256:44e488ef40dbb452700b2b1f8188934121f6648f52c295055662d2191959ff82", size = 42135, upload-time = "2026-05-08T21:00:08.088Z" }, + { url = "https://files.pythonhosted.org/packages/c6/eb/6af6685077d22e8b33358d3c548e3282706a0b3cd85044ffba4e5dd08e3b/propcache-0.5.2-cp311-cp311-win_arm64.whl", hash = "sha256:54adaa85a22078d1e306304a40984dc5be99d599bf3dc0a24dc98f7daeab89ab", size = 38381, upload-time = "2026-05-08T21:00:09.692Z" }, + { url = "https://files.pythonhosted.org/packages/4a/cb/e27bc2b2737a0bb49962b275efa051e8f1c35a936df7d5139b6b658b7dc9/propcache-0.5.2-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:806719138ecd720339a12410fb9614ac9b2b2d3a5fdf8235d56981c36f4039ba", size = 95887, upload-time = "2026-05-08T21:00:11.277Z" }, + { url = "https://files.pythonhosted.org/packages/e6/13/b8ae04c59392f8d11c6cd9fb4011d1dc7c86b81225c770280300e259ffe1/propcache-0.5.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:db2b80ea58eab4f86b2beec3cc8b39e8ff9276ac20e96b7cce43c8ae84cd6b5a", size = 54654, upload-time = "2026-05-08T21:00:12.604Z" }, + { url = "https://files.pythonhosted.org/packages/2c/7d/49777a3e20b55863d4794384a38acd460c04157b0a00f8602b0d508b8431/propcache-0.5.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:e5cbfac9f61484f7e9f3597775500cd3ebe8274e9b050c38f9525c77c97520bf", size = 55190, upload-time = "2026-05-08T21:00:13.935Z" }, + { url = "https://files.pythonhosted.org/packages/44/c7/085d0cd63062e84044e3f05797749c3f8e3938ff3aeb0eb2f69d43fafc91/propcache-0.5.2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5dbc581d2814337da56222fab8dc5f161cd798a434e49bac27930aaef798e144", size = 59995, upload-time = "2026-05-08T21:00:15.526Z" }, + { url = "https://files.pythonhosted.org/packages/9c/42/32cf8e3009e92b2645cf1e944f701e8ea4e924dffde1ee26db860bcbf7e4/propcache-0.5.2-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:857187f381f88c8e2fa2fe56ab94879d011b883d5a2ee5a1b60a8cd2a06846d9", size = 63422, upload-time = "2026-05-08T21:00:16.824Z" }, + { url = "https://files.pythonhosted.org/packages/9e/1b/f112433f99fc979431b87a39ef169e3f8df070d99a72792c56d6937ac48b/propcache-0.5.2-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:178b4a2cdaac1818e2bf1c5a99b94383fa73ea5382e032a48dec07dc5668dc42", size = 64342, upload-time = "2026-05-08T21:00:18.362Z" }, + { url = "https://files.pythonhosted.org/packages/14/15/5574111ae50dd6e879456888c0eadd4c5a869959775854e18e18a6b345f3/propcache-0.5.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6f328175a2cde1f0ff2c4ed8ce968b9dcfb55f3a7153f39e2957ed994da13476", size = 61639, upload-time = "2026-05-08T21:00:19.692Z" }, + { url = "https://files.pythonhosted.org/packages/cc/da/4d775080b1490c0ae604acda868bd71aabe3a89ed16f2aa4339eb8a283e7/propcache-0.5.2-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:5671d09a36b06d0fd4a3da0fccbcae360e9b1570924171a15e9e0997f0249fba", size = 61588, upload-time = "2026-05-08T21:00:21.155Z" }, + { url = "https://files.pythonhosted.org/packages/04/ac/f076982cbe2195ee9cf32de5a1e46951d9fb399fc207f390562dd0fd8fb2/propcache-0.5.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:80168e2ebe4d3ec6599d10ad8f520304ae1cad9b6c5a95372aef1b66b7bfb53a", size = 60029, upload-time = "2026-05-08T21:00:22.713Z" }, + { url = "https://files.pythonhosted.org/packages/70/60/189be62e0dd898dce3b331e1b8c7a543cd3a405ac0c81fe8ee8a9d5d77e1/propcache-0.5.2-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:45f11346f884bc47444f6e6647131055844134c3175b629f84952e2b5cd62b64", size = 56774, upload-time = "2026-05-08T21:00:24.001Z" }, + { url = "https://files.pythonhosted.org/packages/ea/9e/93377b9c7939c1ffae98f878dee955efadfd638078bc86dbc21f9d52f651/propcache-0.5.2-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:8e778ebd44ef4f66ed60a0416b06b489687db264a9c0b3620362f26489492913", size = 63532, upload-time = "2026-05-08T21:00:25.545Z" }, + { url = "https://files.pythonhosted.org/packages/14/f9/590ef6cfb9b8028d516d287812ece32bb0bc5f11fbb9c8bf6b2e6313fec8/propcache-0.5.2-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:c0cb9ed24c8964e172768d455a38254c2dd8a552905729ce006cad3d3dda59b1", size = 61592, upload-time = "2026-05-08T21:00:27.186Z" }, + { url = "https://files.pythonhosted.org/packages/b4/5e/70958b3034c297a630bba2f17ca7abc2d5f39a803ad7e370ab79d1ecd022/propcache-0.5.2-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:1d1ad32d9d4355e2be65574fd0bfd3677e7066b009cd5b9b2dee8aa6a6393b33", size = 64788, upload-time = "2026-05-08T21:00:28.8Z" }, + { url = "https://files.pythonhosted.org/packages/12/fd/77fe5936d8c3086ca9048f7f415f122ed82e53884a9ec193646b42deef06/propcache-0.5.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:c80f4ba3e8f00189165999a742ee526ebeccedf6c3f7beb0c7df821e9772435a", size = 62514, upload-time = "2026-05-08T21:00:30.098Z" }, + { url = "https://files.pythonhosted.org/packages/cf/74/66bd798b5b3be70aa1b391f5cc9d6a0a5532d7fd3b19ec0b213e72e6ad9d/propcache-0.5.2-cp312-cp312-win32.whl", hash = "sha256:8c7972d8f193740d9175f0998ab38717e6cd322d5935c5b0fef8c0d323fd9031", size = 39018, upload-time = "2026-05-08T21:00:31.622Z" }, + { url = "https://files.pythonhosted.org/packages/61/7c/5c0d34aa3024694d6dcb9271cdbdd08c4e47c1c0ad95ec7e7bc74cdea145/propcache-0.5.2-cp312-cp312-win_amd64.whl", hash = "sha256:d9ee8826a7d47863a08ac44e1a5f611a462eefc3a194b492da242128bec75b42", size = 42322, upload-time = "2026-05-08T21:00:32.918Z" }, + { url = "https://files.pythonhosted.org/packages/4d/91/875812f1a3feb20ceba818ef39fbe4d92f1081e04ac815c822496d0d038b/propcache-0.5.2-cp312-cp312-win_arm64.whl", hash = "sha256:2800a4a8ead6b28cccd1ec54b59346f0def7922ee1c7598e8499c733cfbb7c84", size = 38172, upload-time = "2026-05-08T21:00:35.124Z" }, + { url = "https://files.pythonhosted.org/packages/c5/09/f049e45385503fe67db75a6b6186a7b9f0c3930366dc960522c312a825b1/propcache-0.5.2-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:099aaf4b4d1a02265b92a977edf00b5c4f63b3b17ac6de39b0d637c9cac0188a", size = 94457, upload-time = "2026-05-08T21:00:36.355Z" }, + { url = "https://files.pythonhosted.org/packages/6b/65/83d1d05655baf63113731bd5a1008435e14f8d1e5a06cbe4ec5b23ad7a31/propcache-0.5.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:68ce1c44c7a813a7f71ea04315a8c7b330b63db99d059a797a4651bb6f69f117", size = 53835, upload-time = "2026-05-08T21:00:38.072Z" }, + { url = "https://files.pythonhosted.org/packages/a9/12/a6ba6482bb5ea3260c000c9b20881c95fa11c6b30173715668259f844ed7/propcache-0.5.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:fc299c129490f55f254cd90be0deca4764e36e9a7c08b4aa588479a3bbed3098", size = 54545, upload-time = "2026-05-08T21:00:39.319Z" }, + { url = "https://files.pythonhosted.org/packages/a9/19/7fa086f5764c59ec8a8e157cd93aa8497acc00aba9dcdec56bfffb32602d/propcache-0.5.2-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a6ae2198be502c10f09b2516e7b5d019816924bc3183a43ce792a7bd6625e6f4", size = 59886, upload-time = "2026-05-08T21:00:40.621Z" }, + { url = "https://files.pythonhosted.org/packages/a1/e4/5d7663dc8235956c8f5281698a3af1d351d8820341ddd890f59d9a9127f2/propcache-0.5.2-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:6041d31504dc1779d700e1edcfb08eea334b357620b06681a4eabb57a74e574e", size = 63261, upload-time = "2026-05-08T21:00:41.775Z" }, + { url = "https://files.pythonhosted.org/packages/4a/4a/15a03adee24d6350da4292caeac44c34c033d2afe5e87eb370f38854560f/propcache-0.5.2-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:f7eabc04151c78a9f4d5bbb5f1faf571e4defeb4b585e0fe95b60ff2dbe4d3d7", size = 64184, upload-time = "2026-05-08T21:00:43.018Z" }, + { url = "https://files.pythonhosted.org/packages/8b/c6/979176efdaa3d239e36d503d5af63a0a773b36662ed8f52e5b6a6d9fd40e/propcache-0.5.2-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4db0ba63d693afd40d249bd93f842b5f144f8fcbb83de05660373bcf30517b1d", size = 61534, upload-time = "2026-05-08T21:00:44.507Z" }, + { url = "https://files.pythonhosted.org/packages/c8/22/63e8cd1bae4c2d2be6493b6b7d10566ddafad88137cfbc99964a1119853c/propcache-0.5.2-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:1dbcf7675229b35d31abb6547d8ebc8c27a830ac3f9a794edff6254873ec7c0a", size = 61500, upload-time = "2026-05-08T21:00:45.796Z" }, + { url = "https://files.pythonhosted.org/packages/60/5a/28e5d9acbac1cc9ccb67045e8c1b943aa8d79fdf39c93bd73cacd68008ea/propcache-0.5.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:d310c013aad2c72f1c3f2f8dd3279d460a858c551f97aeb8c63e4693cca7b4d2", size = 59994, upload-time = "2026-05-08T21:00:47.093Z" }, + { url = "https://files.pythonhosted.org/packages/f3/40/db650677f554a95b9c01a7c9d93d629e93a15562f5deb4573c9ee136fed2/propcache-0.5.2-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:06187263ddad280d05b4d8a8b3bb7d164cbebd469236544a42e6d9b28ac6a4fa", size = 56884, upload-time = "2026-05-08T21:00:48.376Z" }, + { url = "https://files.pythonhosted.org/packages/80/45/70b39b89516ff8b96bf732fa6fded8cef20f293cb1508690101c3c07ec51/propcache-0.5.2-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:3115559b8effafd63b142ea5ed53d63a16ea6469cbc63dce4ee194b42db5d853", size = 63464, upload-time = "2026-05-08T21:00:49.954Z" }, + { url = "https://files.pythonhosted.org/packages/f9/e2/fa59d3a89eac5534293124af4f1d0d0ada091ce4a0ab4610ce03fd2bdd8d/propcache-0.5.2-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:c60462af8e6dc30c35407c7237ea908d777b22862bbee27bc4699c0d8bcdc45a", size = 61588, upload-time = "2026-05-08T21:00:51.281Z" }, + { url = "https://files.pythonhosted.org/packages/0b/97/efb547a55c4bc7381cfb202d6a2239ac621045277bc1ea5dfd3a7f0516c0/propcache-0.5.2-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:40314bca9ac559716fe374094fc81c11dcc34b64fd6c585360f5775690505704", size = 64667, upload-time = "2026-05-08T21:00:52.602Z" }, + { url = "https://files.pythonhosted.org/packages/92/56/f5c7d9b4b7595d5127da38974d791b2153f3d1eae6c674af3583ace92ad3/propcache-0.5.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:cfa21e036ce1e1db2be04ba3b85d2df1bb1702fa01932d984c5464c665228ff4", size = 62463, upload-time = "2026-05-08T21:00:54.303Z" }, + { url = "https://files.pythonhosted.org/packages/bd/3b/484a3a65fc9f9f60c41dcd17b428bace5389544e2c680994534a20755066/propcache-0.5.2-cp313-cp313-win32.whl", hash = "sha256:f156a3529f38063b6dbaf356e15602a7f95f8055b1295a438433a6386f10463d", size = 38621, upload-time = "2026-05-08T21:00:55.808Z" }, + { url = "https://files.pythonhosted.org/packages/1c/fd/3f0f10dba4dabad3bf53102be007abf55481067952bde0fdddff439e7c61/propcache-0.5.2-cp313-cp313-win_amd64.whl", hash = "sha256:dfed59d0a5aeb01e242e66ff0300bc4a265a7c05f612d30016f0b60b1017d757", size = 41649, upload-time = "2026-05-08T21:00:57.061Z" }, + { url = "https://files.pythonhosted.org/packages/90/ec/6ce619cc32bb500a482f811f9cd509368b4e58e638d13f2c68f370d6b475/propcache-0.5.2-cp313-cp313-win_arm64.whl", hash = "sha256:ba338430e87ceb9c8f0cf754de38a9860560261e56c00376debd628698a7364f", size = 37636, upload-time = "2026-05-08T21:00:58.646Z" }, + { url = "https://files.pythonhosted.org/packages/1b/82/c1d268bbbf2ef981c5bf0fbbe746db617c66e3bcefe431a1aa8943fbe23a/propcache-0.5.2-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:a592f5f3da71c8691c788c13cb6734b6d17663d2e1cb8caddf0673d01ef8847d", size = 98872, upload-time = "2026-05-08T21:00:59.889Z" }, + { url = "https://files.pythonhosted.org/packages/f4/d4/52c871e73e864e6b34c0e2d58ac1ec5ccd149497ddc7ad2137ae98323a35/propcache-0.5.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:6a997d0489e9668a384fcfd5061b857aa5361de73191cac204d04b889cfbbafa", size = 56257, upload-time = "2026-05-08T21:01:01.195Z" }, + { url = "https://files.pythonhosted.org/packages/67/f0/9b90ca2a210b3d09bcfcd96ecd0f55545c091535abce2a45de2775cfd357/propcache-0.5.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:10734b5484ea113152ee25a91dccedf81631791805d2c9ccb054958e51842c94", size = 56696, upload-time = "2026-05-08T21:01:02.941Z" }, + { url = "https://files.pythonhosted.org/packages/9d/0e/6e9d4ba07c8e56e21ddec1e75f12148142b21ca83a51871babce095334f4/propcache-0.5.2-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cafca7e56c12bb02ae16d283742bef25a61122e9dab2b5b3f2ccbe589ce32164", size = 62378, upload-time = "2026-05-08T21:01:04.475Z" }, + { url = "https://files.pythonhosted.org/packages/65/19/c10badaa463dde8a27ce884f8ee2ec37e6035b7c9f5ff0c8f74f06f08dac/propcache-0.5.2-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:f064f8d2b59177878b7615df1735cd8fe3462ed6be8c7b217d17a276489c2b7f", size = 65283, upload-time = "2026-05-08T21:01:05.959Z" }, + { url = "https://files.pythonhosted.org/packages/b0/b6/93bea99ca80e19cef6512a8580e5b7857bbe09422d9daa7fd4ef5723306c/propcache-0.5.2-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:f78abfa8dfc32376fd1aacf597b2f2fbbe0ea751419aee718af5d4f82537ef8c", size = 66616, upload-time = "2026-05-08T21:01:07.228Z" }, + { url = "https://files.pythonhosted.org/packages/83/e4/5c7462e50625f051f37fb38b8224f7639f667184bbd34424ec83819bb1b7/propcache-0.5.2-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f7467da8a9822bf1a55336f877340c5bcbd3c482afc43a99771169f74a26dedc", size = 63773, upload-time = "2026-05-08T21:01:08.514Z" }, + { url = "https://files.pythonhosted.org/packages/ca/b6/99238894047b13c823be25027e736626cd414a52a5e30d2c3347c2733529/propcache-0.5.2-cp313-cp313t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:a6ddc6ac9e25de626c1f129c1b467d7ecd33ce2237d3fd0c4e429feef0a7ee1f", size = 63664, upload-time = "2026-05-08T21:01:09.874Z" }, + { url = "https://files.pythonhosted.org/packages/85/1e/a3a1a63116a2b8edb415a8bb9a6f0c34bd03830b1e18e8ce2904e1dc1cf4/propcache-0.5.2-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:2f22cbbac9e26a8e864c0985ff1268d5d939d53d9d9411a9824279097e03a2cb", size = 62643, upload-time = "2026-05-08T21:01:11.132Z" }, + { url = "https://files.pythonhosted.org/packages/e4/03/893cf147de2fc6543c5eaa07ad833170e7e2a2385725bbebe8c0503723bb/propcache-0.5.2-cp313-cp313t-musllinux_1_2_armv7l.whl", hash = "sha256:fc76378c62a0f04d0cd82fbb1a2cd2d7e28fcb40d5873f28a6c44e388aaa2751", size = 59595, upload-time = "2026-05-08T21:01:12.387Z" }, + { url = "https://files.pythonhosted.org/packages/86/3b/04c1a2e12c57766568ba75ba72b3bf2042818d4c1425fab6fc07155c7cff/propcache-0.5.2-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:acd2c8edba48e31e58a363b8cf4e5c7db3b04b3f9e371f601df30d9b0d244836", size = 65711, upload-time = "2026-05-08T21:01:13.676Z" }, + { url = "https://files.pythonhosted.org/packages/1c/34/80f8d0099f8d6bacc4de1624c85672681c8cd1149ca2da0e38fd120b817f/propcache-0.5.2-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:452b5065457eb9991ec5eb38ff41d6cd4c991c9ac7c531c4d5849ae473a9a13f", size = 64247, upload-time = "2026-05-08T21:01:14.936Z" }, + { url = "https://files.pythonhosted.org/packages/f3/1a/8b08f3a5f1037e9e370c55883ceeeee0f6dd0416fb2d2d67b8bfc91f2a79/propcache-0.5.2-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:3430bb2bfe1331885c427745a751e774ee679fd4344f80b97bf879815fe8fa55", size = 67102, upload-time = "2026-05-08T21:01:16.281Z" }, + { url = "https://files.pythonhosted.org/packages/34/68/8bdb7bb7756d76e005490649d10e4a8369e610c74d619f71e1aedf889e9c/propcache-0.5.2-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:cef6cea3922890dd6c9654971001fa797b526c16ab5e1e46c05fd6f877be7568", size = 64964, upload-time = "2026-05-08T21:01:17.57Z" }, + { url = "https://files.pythonhosted.org/packages/0a/aa/50fb0b5d3968b61a510926ff8b8465f1d6e976b3ab74496d7a4b9fc42515/propcache-0.5.2-cp313-cp313t-win32.whl", hash = "sha256:72d61e16dd78228b58c5d47be830ff3da7e5f139abdf0aef9d86cde1c5cf2191", size = 42546, upload-time = "2026-05-08T21:01:18.946Z" }, + { url = "https://files.pythonhosted.org/packages/ae/4c/0ddbae64321bd4a95bcbfc19307238016b5b1fee645c84626c8d539e5b74/propcache-0.5.2-cp313-cp313t-win_amd64.whl", hash = "sha256:0958834041a0166d343b8d2cedcd8bcbaeb4fdbe0cf08320c5379f143c3be6e7", size = 46330, upload-time = "2026-05-08T21:01:20.162Z" }, + { url = "https://files.pythonhosted.org/packages/00/d9/9cddc8efb78d8af264c5ec9f6d10b62f57c515feda8d321595f56010fb23/propcache-0.5.2-cp313-cp313t-win_arm64.whl", hash = "sha256:6de8bd93ddde9b992cf2b2e0d796d501a19026b5b9fd87356d7d0779531a8d96", size = 40521, upload-time = "2026-05-08T21:01:21.399Z" }, + { url = "https://files.pythonhosted.org/packages/e2/ea/23ee535d90ce8bcc465a3028eb3cc0ce3bd1005f4bb27710b30587de798d/propcache-0.5.2-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:46088abff4cba581dea21ae0467a480526cb25aa5f3c269e909f800328bc3999", size = 94662, upload-time = "2026-05-08T21:01:22.683Z" }, + { url = "https://files.pythonhosted.org/packages/b5/06/c5a52f419b5d8972f8d46a7577476090d8e3263ff589ce40b5ca4968d5be/propcache-0.5.2-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:fc88b26f08d634f7bc819a7852e5214f5802641ab8d9fd5326892292eee1993e", size = 53928, upload-time = "2026-05-08T21:01:23.986Z" }, + { url = "https://files.pythonhosted.org/packages/63/b1/4260d67d6bd85e58a66b72d54ce15d5de789b6f3870cc6bedf8ff9667401/propcache-0.5.2-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:97797ebb098e670a2f92dd66f32897e30d7615b14e7f59711de23e30a9072539", size = 54650, upload-time = "2026-05-08T21:01:25.305Z" }, + { url = "https://files.pythonhosted.org/packages/70/06/2f46c318e3307cd7a6a7481def374ce838c0fe20084b39dd54b0879d0e99/propcache-0.5.2-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ba57fffe4ac99c5d30076161b5866336d97600769bad35cc68f7774b15298a4e", size = 59912, upload-time = "2026-05-08T21:01:26.545Z" }, + { url = "https://files.pythonhosted.org/packages/4c/29/fe1aebec2ce57ab985a9c382bded1124431f85078113aa222c5d278430d4/propcache-0.5.2-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:583c19759d9eec1e5b69e2fbef36a7d9c326041be9746cb822d335c8cedc2979", size = 63300, upload-time = "2026-05-08T21:01:27.937Z" }, + { url = "https://files.pythonhosted.org/packages/b4/18/2334b26768b6c82be8c69e83671b767d5ef426aa09b0cba6c2ea47816774/propcache-0.5.2-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:d0326e2e5e1f3163fa306c834e48e8d490e5fae607a097a40c0648109b47ba80", size = 64208, upload-time = "2026-05-08T21:01:29.484Z" }, + { url = "https://files.pythonhosted.org/packages/2b/76/7f1bfd6afff4c5e38e36a3c6d68eb5f4b7311ea80baf693db78d95b603c4/propcache-0.5.2-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e00820e192c8dbebcafb383ebbf99030895f09905e7a0eb2e0340a0bcc2bc825", size = 61633, upload-time = "2026-05-08T21:01:31.068Z" }, + { url = "https://files.pythonhosted.org/packages/c4/46/b3ff8aba2b4953a3e50de2cf72f1b5748b8eca93b15f3dc2c84339084c09/propcache-0.5.2-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:c66afea89b1e43725731d2004732a046fe6fe955d51f952c3e95a7314a284a39", size = 61724, upload-time = "2026-05-08T21:01:32.374Z" }, + { url = "https://files.pythonhosted.org/packages/c5/01/814cfcafbcff954f94c01cf30e097ddc88a076b5440fbcf4570753437d40/propcache-0.5.2-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:d4dc37dec6c6cdad0b57881a5658fd14fbf53e333b1a86cf86559f190e1d9ec4", size = 60069, upload-time = "2026-05-08T21:01:33.67Z" }, + { url = "https://files.pythonhosted.org/packages/da/68/5c6f7622d510cc666a300687e06fd060c1a43361c0c9b20d284f06d8096a/propcache-0.5.2-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:5570dbcc97571c15f68068e529c92715a12f8d54030e272d264b377e22bd17a5", size = 57099, upload-time = "2026-05-08T21:01:34.915Z" }, + { url = "https://files.pythonhosted.org/packages/55/27/9cb0b4c679124085327957d42521c99dba04c88c90c3e55a6f0b633ebccc/propcache-0.5.2-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:f814362777a9f841adddb200ecdf8f5cb1e5a3c4b7a86378edbd6ccb26edd702", size = 63391, upload-time = "2026-05-08T21:01:36.231Z" }, + { url = "https://files.pythonhosted.org/packages/f0/9d/7258aaa5bdf60fc6f27591eef6fe52768cb0beda7140be477c8b12c9794a/propcache-0.5.2-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:196913dea116aeb5a2ba95af4ddcb7ea85559ae07d8eee8751688310d09168c3", size = 61626, upload-time = "2026-05-08T21:01:37.545Z" }, + { url = "https://files.pythonhosted.org/packages/8e/0d/41c602003e8a9b16fe1e7eadf62c7bfba9d5474370b24200bf48b315f45f/propcache-0.5.2-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:6e7b8719005dd1175be4ab1cd25e9b98659a5e0347331506ec6760d2773a7fb5", size = 64781, upload-time = "2026-05-08T21:01:38.83Z" }, + { url = "https://files.pythonhosted.org/packages/8b/f3/38e66b1856e9bd079deea015bc4a55f7767c0e4db2f7dcf69e7e680ba4ce/propcache-0.5.2-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:51f96d685ab16e88cab128cd37a52c5da540809c8b879fa047731bfcb4ad35a4", size = 62570, upload-time = "2026-05-08T21:01:40.415Z" }, + { url = "https://files.pythonhosted.org/packages/95/ca/bbfe9b910ce57dde8bb4876b4520fc02a4e89497c10de26be936758a3aaa/propcache-0.5.2-cp314-cp314-win32.whl", hash = "sha256:cc6fc3cc62e8501d3ed62894425040d2728ecddb1ed072737a5c70bd537aa9f0", size = 39436, upload-time = "2026-05-08T21:01:41.654Z" }, + { url = "https://files.pythonhosted.org/packages/61/d2/45c9defbaa1ea297035d9d4cce9e8f80daafbf19319c6007f157c6256ea9/propcache-0.5.2-cp314-cp314-win_amd64.whl", hash = "sha256:81e3a30b0bb60caa22033dd0f8a3618d1d67356212514f62c57db75cb0ef410c", size = 42373, upload-time = "2026-05-08T21:01:43.041Z" }, + { url = "https://files.pythonhosted.org/packages/44/68/9ea5103f41d5217d7d6ec24db90018e23aebec070c3f9a6e54d12b841fd8/propcache-0.5.2-cp314-cp314-win_arm64.whl", hash = "sha256:0d2c9bf8528f135dbb805ce027567e09164f7efa51a2be07458a2c0420f292d0", size = 38554, upload-time = "2026-05-08T21:01:44.336Z" }, + { url = "https://files.pythonhosted.org/packages/8a/81/fadf555f42d3b762eea8a53950b0489fdc0aa9da5f8ed9e10ce0a4e01b48/propcache-0.5.2-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:4bc8ff1feffc6a61c7002ffe84634c41b822e104990ae009f44a0834430070bb", size = 99395, upload-time = "2026-05-08T21:01:45.883Z" }, + { url = "https://files.pythonhosted.org/packages/f5/c9/c61e134a686949cf7971af3a390148b1156f7be81c73bc0cd12c873e2d48/propcache-0.5.2-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:79aa3ff0a9b566633b642fa9caf7e21ed1c13d6feca718187873f199e1514078", size = 56653, upload-time = "2026-05-08T21:01:47.307Z" }, + { url = "https://files.pythonhosted.org/packages/cb/73/daf935ea7048ddd7ec8eec5345b4a40b619d2d178b3c0a0900796bc3c794/propcache-0.5.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:1b31822f4474c4036bae62de9402710051d431a606d6a0f907fec79935a071aa", size = 56914, upload-time = "2026-05-08T21:01:48.573Z" }, + { url = "https://files.pythonhosted.org/packages/79/9f/aba959b435ea18617edd7cf0a7ad0b9c574b8fc7e3d2cd55fb59cb255d33/propcache-0.5.2-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:13fef48778b5a2a756523fdb781326b028ca75e32858b04f2cdd19f394564917", size = 62567, upload-time = "2026-05-08T21:01:49.903Z" }, + { url = "https://files.pythonhosted.org/packages/6c/a1/859942de9a791ff42f6141736f5b37749b8f53e65edfa49638c67dd67e6a/propcache-0.5.2-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:8b73ab70f1a3351fbc71f663b3e645af6dd0329100c353081cf69c37433fc6fe", size = 65542, upload-time = "2026-05-08T21:01:51.204Z" }, + { url = "https://files.pythonhosted.org/packages/b5/61/315bc0fd6c0fc7f80a528b8afd209e5fc4a875ea79571b91b8f50f442907/propcache-0.5.2-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5538d2c13d93e4698af7e092b57bc7298fd35d1d58e656ae18f23ee0d0378e03", size = 66845, upload-time = "2026-05-08T21:01:52.539Z" }, + { url = "https://files.pythonhosted.org/packages/47/f7/9f8122e3132e8e354ac41975ef8f1099be7d5a16bc7ae562734e993665c0/propcache-0.5.2-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:cd645f03898405cabe694fb8bc35241e3a9c332ec85627584fe3de201452b335", size = 63985, upload-time = "2026-05-08T21:01:53.847Z" }, + { url = "https://files.pythonhosted.org/packages/c8/54/c317819ec157cbf6f35df9df9657a6f82daf34d5faf15948b2f639c2192e/propcache-0.5.2-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:a473b3440261e0c60706e732b2ed2f517857344fc21bf48fdfe211e2d98eb285", size = 63999, upload-time = "2026-05-08T21:01:55.179Z" }, + { url = "https://files.pythonhosted.org/packages/5a/56/387e3f7dfce0a9233df41fb888aa1c30222cb4bbbf09537c02dd9bd85fe2/propcache-0.5.2-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:7afa37062e6650640e932e4cc9297d81f9f42d9944029cc386b8247dea4da837", size = 62779, upload-time = "2026-05-08T21:01:57.489Z" }, + { url = "https://files.pythonhosted.org/packages/a1/9c/596784cb5824ed61ee960d3f8655a3f0993e107c6e98ab6c818b7fb92ccb/propcache-0.5.2-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:8a90efd5777e996e42d568db9ac740b944d691e565cbfd31b2f7832f9184b2b8", size = 59796, upload-time = "2026-05-08T21:01:58.736Z" }, + { url = "https://files.pythonhosted.org/packages/c2/3d/1a6cfa1726a48542c1e8784a0761421476a5b68e09b7f36bf95eb954aaba/propcache-0.5.2-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:f19bb891234d72535764d703bfed1153cc34f4214d5bd7150aee1eec9e8f4366", size = 66023, upload-time = "2026-05-08T21:02:00.228Z" }, + { url = "https://files.pythonhosted.org/packages/e4/0e/05fd6990369477076e4e280bcb970de760fddf0161a46e988bc95f7940ec/propcache-0.5.2-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:32775082acd2d807ee3db715c7770d38767b817870acfa08c29e057f3c4d5b56", size = 64448, upload-time = "2026-05-08T21:02:01.888Z" }, + { url = "https://files.pythonhosted.org/packages/cd/86/5f8da315a4309c62c10c0b2516b17492d5d3bbe1bb862b96604db67e2a37/propcache-0.5.2-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:9282fb1a3bccd038da9f768b927b24a0c753e466c086b7c4f3c6982851eefb2d", size = 67329, upload-time = "2026-05-08T21:02:03.484Z" }, + { url = "https://files.pythonhosted.org/packages/da/d3/3368efe79ab21f0cdf86ef49895811c9cc933131d4cde1f28a624e22e712/propcache-0.5.2-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:cc49723e2f60d6b32a0f0b08a3fd6d13203c07f1cd9566cfce0f12a917c967a2", size = 65172, upload-time = "2026-05-08T21:02:04.745Z" }, + { url = "https://files.pythonhosted.org/packages/d5/07/127e8b0bacfb325396196f9d976a22453049b89b9b2b08477cc3145faa44/propcache-0.5.2-cp314-cp314t-win32.whl", hash = "sha256:2d7aa89ebca5acc98cba9d1472d976e394782f587bad6661003602a619fd1821", size = 43813, upload-time = "2026-05-08T21:02:06.025Z" }, + { url = "https://files.pythonhosted.org/packages/88/fb/46dad6c0ae49ed230ab1b16c890c2b6314e2403e6c412976f4a72d64a527/propcache-0.5.2-cp314-cp314t-win_amd64.whl", hash = "sha256:d447bb0b3054be5818458fbb171208b1d9ff11eba14e18ca18b90cbb45767370", size = 47764, upload-time = "2026-05-08T21:02:07.353Z" }, + { url = "https://files.pythonhosted.org/packages/e7/c4/a47d0a63aa309d10d59ede6e9d4cff03a344a79d1f0f4cd0cd74997b53e0/propcache-0.5.2-cp314-cp314t-win_arm64.whl", hash = "sha256:fe67a3d11cd9b4efabfa45c3d00ffba2b26811442a73a581a94b67c2b5faccf6", size = 41140, upload-time = "2026-05-08T21:02:09.065Z" }, + { url = "https://files.pythonhosted.org/packages/3a/ed/1cdcab6ba3d6ab7feca11fc14f0eeea80755bb53ef4e892079f31b10a25f/propcache-0.5.2-py3-none-any.whl", hash = "sha256:be1ddfcbb376e3de5d2e2db1d58d6d67463e6b4f9f040c000de8e300295465fe", size = 14036, upload-time = "2026-05-08T21:02:10.673Z" }, +] + [[package]] name = "protobuf" version = "6.33.5" @@ -1050,38 +1581,45 @@ wheels = [ [[package]] name = "pyarrow" -version = "21.0.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/ef/c2/ea068b8f00905c06329a3dfcd40d0fcc2b7d0f2e355bdb25b65e0a0e4cd4/pyarrow-21.0.0.tar.gz", hash = "sha256:5051f2dccf0e283ff56335760cbc8622cf52264d67e359d5569541ac11b6d5bc", size = 1133487, upload-time = "2025-07-18T00:57:31.761Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/94/dc/80564a3071a57c20b7c32575e4a0120e8a330ef487c319b122942d665960/pyarrow-21.0.0-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:c077f48aab61738c237802836fc3844f85409a46015635198761b0d6a688f87b", size = 31243234, upload-time = "2025-07-18T00:55:03.812Z" }, - { url = "https://files.pythonhosted.org/packages/ea/cc/3b51cb2db26fe535d14f74cab4c79b191ed9a8cd4cbba45e2379b5ca2746/pyarrow-21.0.0-cp311-cp311-macosx_12_0_x86_64.whl", hash = "sha256:689f448066781856237eca8d1975b98cace19b8dd2ab6145bf49475478bcaa10", size = 32714370, upload-time = "2025-07-18T00:55:07.495Z" }, - { url = "https://files.pythonhosted.org/packages/24/11/a4431f36d5ad7d83b87146f515c063e4d07ef0b7240876ddb885e6b44f2e/pyarrow-21.0.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:479ee41399fcddc46159a551705b89c05f11e8b8cb8e968f7fec64f62d91985e", size = 41135424, upload-time = "2025-07-18T00:55:11.461Z" }, - { url = "https://files.pythonhosted.org/packages/74/dc/035d54638fc5d2971cbf1e987ccd45f1091c83bcf747281cf6cc25e72c88/pyarrow-21.0.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:40ebfcb54a4f11bcde86bc586cbd0272bac0d516cfa539c799c2453768477569", size = 42823810, upload-time = "2025-07-18T00:55:16.301Z" }, - { url = "https://files.pythonhosted.org/packages/2e/3b/89fced102448a9e3e0d4dded1f37fa3ce4700f02cdb8665457fcc8015f5b/pyarrow-21.0.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:8d58d8497814274d3d20214fbb24abcad2f7e351474357d552a8d53bce70c70e", size = 43391538, upload-time = "2025-07-18T00:55:23.82Z" }, - { url = "https://files.pythonhosted.org/packages/fb/bb/ea7f1bd08978d39debd3b23611c293f64a642557e8141c80635d501e6d53/pyarrow-21.0.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:585e7224f21124dd57836b1530ac8f2df2afc43c861d7bf3d58a4870c42ae36c", size = 45120056, upload-time = "2025-07-18T00:55:28.231Z" }, - { url = "https://files.pythonhosted.org/packages/6e/0b/77ea0600009842b30ceebc3337639a7380cd946061b620ac1a2f3cb541e2/pyarrow-21.0.0-cp311-cp311-win_amd64.whl", hash = "sha256:555ca6935b2cbca2c0e932bedd853e9bc523098c39636de9ad4693b5b1df86d6", size = 26220568, upload-time = "2025-07-18T00:55:32.122Z" }, - { url = "https://files.pythonhosted.org/packages/ca/d4/d4f817b21aacc30195cf6a46ba041dd1be827efa4a623cc8bf39a1c2a0c0/pyarrow-21.0.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:3a302f0e0963db37e0a24a70c56cf91a4faa0bca51c23812279ca2e23481fccd", size = 31160305, upload-time = "2025-07-18T00:55:35.373Z" }, - { url = "https://files.pythonhosted.org/packages/a2/9c/dcd38ce6e4b4d9a19e1d36914cb8e2b1da4e6003dd075474c4cfcdfe0601/pyarrow-21.0.0-cp312-cp312-macosx_12_0_x86_64.whl", hash = "sha256:b6b27cf01e243871390474a211a7922bfbe3bda21e39bc9160daf0da3fe48876", size = 32684264, upload-time = "2025-07-18T00:55:39.303Z" }, - { url = "https://files.pythonhosted.org/packages/4f/74/2a2d9f8d7a59b639523454bec12dba35ae3d0a07d8ab529dc0809f74b23c/pyarrow-21.0.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:e72a8ec6b868e258a2cd2672d91f2860ad532d590ce94cdf7d5e7ec674ccf03d", size = 41108099, upload-time = "2025-07-18T00:55:42.889Z" }, - { url = "https://files.pythonhosted.org/packages/ad/90/2660332eeb31303c13b653ea566a9918484b6e4d6b9d2d46879a33ab0622/pyarrow-21.0.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:b7ae0bbdc8c6674259b25bef5d2a1d6af5d39d7200c819cf99e07f7dfef1c51e", size = 42829529, upload-time = "2025-07-18T00:55:47.069Z" }, - { url = "https://files.pythonhosted.org/packages/33/27/1a93a25c92717f6aa0fca06eb4700860577d016cd3ae51aad0e0488ac899/pyarrow-21.0.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:58c30a1729f82d201627c173d91bd431db88ea74dcaa3885855bc6203e433b82", size = 43367883, upload-time = "2025-07-18T00:55:53.069Z" }, - { url = "https://files.pythonhosted.org/packages/05/d9/4d09d919f35d599bc05c6950095e358c3e15148ead26292dfca1fb659b0c/pyarrow-21.0.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:072116f65604b822a7f22945a7a6e581cfa28e3454fdcc6939d4ff6090126623", size = 45133802, upload-time = "2025-07-18T00:55:57.714Z" }, - { url = "https://files.pythonhosted.org/packages/71/30/f3795b6e192c3ab881325ffe172e526499eb3780e306a15103a2764916a2/pyarrow-21.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:cf56ec8b0a5c8c9d7021d6fd754e688104f9ebebf1bf4449613c9531f5346a18", size = 26203175, upload-time = "2025-07-18T00:56:01.364Z" }, - { url = "https://files.pythonhosted.org/packages/16/ca/c7eaa8e62db8fb37ce942b1ea0c6d7abfe3786ca193957afa25e71b81b66/pyarrow-21.0.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:e99310a4ebd4479bcd1964dff9e14af33746300cb014aa4a3781738ac63baf4a", size = 31154306, upload-time = "2025-07-18T00:56:04.42Z" }, - { url = "https://files.pythonhosted.org/packages/ce/e8/e87d9e3b2489302b3a1aea709aaca4b781c5252fcb812a17ab6275a9a484/pyarrow-21.0.0-cp313-cp313-macosx_12_0_x86_64.whl", hash = "sha256:d2fe8e7f3ce329a71b7ddd7498b3cfac0eeb200c2789bd840234f0dc271a8efe", size = 32680622, upload-time = "2025-07-18T00:56:07.505Z" }, - { url = "https://files.pythonhosted.org/packages/84/52/79095d73a742aa0aba370c7942b1b655f598069489ab387fe47261a849e1/pyarrow-21.0.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:f522e5709379d72fb3da7785aa489ff0bb87448a9dc5a75f45763a795a089ebd", size = 41104094, upload-time = "2025-07-18T00:56:10.994Z" }, - { url = "https://files.pythonhosted.org/packages/89/4b/7782438b551dbb0468892a276b8c789b8bbdb25ea5c5eb27faadd753e037/pyarrow-21.0.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:69cbbdf0631396e9925e048cfa5bce4e8c3d3b41562bbd70c685a8eb53a91e61", size = 42825576, upload-time = "2025-07-18T00:56:15.569Z" }, - { url = "https://files.pythonhosted.org/packages/b3/62/0f29de6e0a1e33518dec92c65be0351d32d7ca351e51ec5f4f837a9aab91/pyarrow-21.0.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:731c7022587006b755d0bdb27626a1a3bb004bb56b11fb30d98b6c1b4718579d", size = 43368342, upload-time = "2025-07-18T00:56:19.531Z" }, - { url = "https://files.pythonhosted.org/packages/90/c7/0fa1f3f29cf75f339768cc698c8ad4ddd2481c1742e9741459911c9ac477/pyarrow-21.0.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:dc56bc708f2d8ac71bd1dcb927e458c93cec10b98eb4120206a4091db7b67b99", size = 45131218, upload-time = "2025-07-18T00:56:23.347Z" }, - { url = "https://files.pythonhosted.org/packages/01/63/581f2076465e67b23bc5a37d4a2abff8362d389d29d8105832e82c9c811c/pyarrow-21.0.0-cp313-cp313-win_amd64.whl", hash = "sha256:186aa00bca62139f75b7de8420f745f2af12941595bbbfa7ed3870ff63e25636", size = 26087551, upload-time = "2025-07-18T00:56:26.758Z" }, - { url = "https://files.pythonhosted.org/packages/c9/ab/357d0d9648bb8241ee7348e564f2479d206ebe6e1c47ac5027c2e31ecd39/pyarrow-21.0.0-cp313-cp313t-macosx_12_0_arm64.whl", hash = "sha256:a7a102574faa3f421141a64c10216e078df467ab9576684d5cd696952546e2da", size = 31290064, upload-time = "2025-07-18T00:56:30.214Z" }, - { url = "https://files.pythonhosted.org/packages/3f/8a/5685d62a990e4cac2043fc76b4661bf38d06efed55cf45a334b455bd2759/pyarrow-21.0.0-cp313-cp313t-macosx_12_0_x86_64.whl", hash = "sha256:1e005378c4a2c6db3ada3ad4c217b381f6c886f0a80d6a316fe586b90f77efd7", size = 32727837, upload-time = "2025-07-18T00:56:33.935Z" }, - { url = "https://files.pythonhosted.org/packages/fc/de/c0828ee09525c2bafefd3e736a248ebe764d07d0fd762d4f0929dbc516c9/pyarrow-21.0.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:65f8e85f79031449ec8706b74504a316805217b35b6099155dd7e227eef0d4b6", size = 41014158, upload-time = "2025-07-18T00:56:37.528Z" }, - { url = "https://files.pythonhosted.org/packages/6e/26/a2865c420c50b7a3748320b614f3484bfcde8347b2639b2b903b21ce6a72/pyarrow-21.0.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:3a81486adc665c7eb1a2bde0224cfca6ceaba344a82a971ef059678417880eb8", size = 42667885, upload-time = "2025-07-18T00:56:41.483Z" }, - { url = "https://files.pythonhosted.org/packages/0a/f9/4ee798dc902533159250fb4321267730bc0a107d8c6889e07c3add4fe3a5/pyarrow-21.0.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:fc0d2f88b81dcf3ccf9a6ae17f89183762c8a94a5bdcfa09e05cfe413acf0503", size = 43276625, upload-time = "2025-07-18T00:56:48.002Z" }, - { url = "https://files.pythonhosted.org/packages/5a/da/e02544d6997037a4b0d22d8e5f66bc9315c3671371a8b18c79ade1cefe14/pyarrow-21.0.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:6299449adf89df38537837487a4f8d3bd91ec94354fdd2a7d30bc11c48ef6e79", size = 44951890, upload-time = "2025-07-18T00:56:52.568Z" }, - { url = "https://files.pythonhosted.org/packages/e5/4e/519c1bc1876625fe6b71e9a28287c43ec2f20f73c658b9ae1d485c0c206e/pyarrow-21.0.0-cp313-cp313t-win_amd64.whl", hash = "sha256:222c39e2c70113543982c6b34f3077962b44fca38c0bd9e68bb6781534425c10", size = 26371006, upload-time = "2025-07-18T00:56:56.379Z" }, +version = "25.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/27/f3/95428098d1fa7d04432fb750eed06b41304c2f6a5d3319985e64db2d9d41/pyarrow-25.0.0.tar.gz", hash = "sha256:d2d697008b5ec06d75952ef260c2e9a8a0f6ccfce24266c04c9c8ade927cb3b4", size = 1199181, upload-time = "2026-07-10T08:29:50.116Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a7/98/ae2b5acf9876dbeffa6f320776242c52caab062df55c8ac5501ed2679e74/pyarrow-25.0.0-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:2e3b6544e26e393fe2cd530f523e36c1c8d3c345bbbb60cca3fd866be8322517", size = 35939080, upload-time = "2026-07-10T08:26:04.53Z" }, + { url = "https://files.pythonhosted.org/packages/80/09/3de2a968edbd496c86cb8b932cdbee2d4b08c4a28e9884a15e5c705a646b/pyarrow-25.0.0-cp311-cp311-macosx_12_0_x86_64.whl", hash = "sha256:b724d127783b4c19f088fcdfc844cbc318809246a30307bcabd5ed02045e890e", size = 37633420, upload-time = "2026-07-10T08:26:10.354Z" }, + { url = "https://files.pythonhosted.org/packages/19/86/8399243a4ce080426ec37db18d5e29148b7ec960a8a8c7f9059a7bf6ef0a/pyarrow-25.0.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:244f98a595f70fa4fd35faa7508c4ae67e14a173397a4b3b49d2b3c360fb0062", size = 46861050, upload-time = "2026-07-10T08:26:16.397Z" }, + { url = "https://files.pythonhosted.org/packages/7b/79/72d704b02bc5fc6d06954d76a0208c1e79cad3ab370f6d6a91ffe5078870/pyarrow-25.0.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:0222f0071d13313962a88d21bf28b80d355ac39d81bfa6ff3fe00eeaf748e4be", size = 50056458, upload-time = "2026-07-10T08:26:23.271Z" }, + { url = "https://files.pythonhosted.org/packages/06/5d/3c31a60b6403d63cad2e0f829096f5fc5763a129ead4207a5d4690b96448/pyarrow-25.0.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:b58726f118c079f9d4ed7e904975d4f15fd69d0741ba511a4e2dcaa4ef16354f", size = 49957793, upload-time = "2026-07-10T08:26:30.232Z" }, + { url = "https://files.pythonhosted.org/packages/34/f7/8f8a019061f9863a831915329264372a87ed25eaf9109ce56eb0e84012c5/pyarrow-25.0.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:38a2c887cb3883e241b70201688db34133b6dfadd04f03c8f9213df53770c18e", size = 53100544, upload-time = "2026-07-10T08:26:36.414Z" }, + { url = "https://files.pythonhosted.org/packages/f1/e2/738071e95c5ddad7b3dfc12f569ffa992db89d7d7b4a95258fd184191249/pyarrow-25.0.0-cp311-cp311-win_amd64.whl", hash = "sha256:161649d60a7a46c613a19fd795763ea8a88c36ba997dd99d9bc66e6794ee36e8", size = 27848311, upload-time = "2026-07-10T08:26:41.429Z" }, + { url = "https://files.pythonhosted.org/packages/73/44/fdd3a4377807b7dcabe2d4b5aa99dbbc98e2e5df3f1ca4e7f0aec492d987/pyarrow-25.0.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:149730a3d1f0fb59d663a0b8aa210adfd9c17c27cd94a0d143e60daea8320d4e", size = 35850884, upload-time = "2026-07-10T08:26:47.357Z" }, + { url = "https://files.pythonhosted.org/packages/bf/71/9f053177a7709b8c90abb00a2375b916286f9f0d6cfb21a5cadd4ef811e8/pyarrow-25.0.0-cp312-cp312-macosx_12_0_x86_64.whl", hash = "sha256:0721332c30fdd453fdd1fc203b2ac1f4c9db5aea28fa38d41f2574c4b068b9ec", size = 37616197, upload-time = "2026-07-10T08:26:53.564Z" }, + { url = "https://files.pythonhosted.org/packages/95/1a/22bfb6597dcdc861fa83c39c06e1457cb56f698940eff42fbb25de30e8e5/pyarrow-25.0.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:fa1482b3da10cac2d4db6e26b81da543e237616af2ef6d466018b31ca586496f", size = 46841966, upload-time = "2026-07-10T08:27:07.685Z" }, + { url = "https://files.pythonhosted.org/packages/55/0e/cd705c042bc4fe7022478db577fcab4abdcfabb9bc37ab7a75556b3fcb2b/pyarrow-25.0.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:5d1dbf24e151042f2fa3c129563f65d66674128868496fb008c4272b16bdf778", size = 50088993, upload-time = "2026-07-10T08:27:14.268Z" }, + { url = "https://files.pythonhosted.org/packages/98/ee/d822e1ee31fe31ec5d057210e0605c950b975dcd8d9a332976cc859a9df8/pyarrow-25.0.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:20887a762dd61dcc530f93a140840ab1f6aa7836b33270e42d627ab3cf11e537", size = 49941005, upload-time = "2026-07-10T08:27:21.274Z" }, + { url = "https://files.pythonhosted.org/packages/33/1b/207a90cc64619a095eb75a263ae069735f2810056d43c667befd573ec083/pyarrow-25.0.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:58d1ab556b0cea1c93fdb799b24ad58adb2f2a2788dbce782a94f64ae1a5cc9b", size = 53112355, upload-time = "2026-07-10T08:27:27.911Z" }, + { url = "https://files.pythonhosted.org/packages/7e/fe/81d1e5f8beed15c01e98649d5c6e2167b67fd395884a2488f18bf1cf0dba/pyarrow-25.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:3f356afe61186395c861d5cd63dc21ff7d5fa335012a4668d979257df7fea0f5", size = 27945954, upload-time = "2026-07-10T08:27:32.903Z" }, + { url = "https://files.pythonhosted.org/packages/6c/c8/098ce17d778fd9d29e40bb8c5f19a40cc90c3f0b46c9057b0d7993f42f54/pyarrow-25.0.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:8831a3ba52fa7cdb78d368d968b1dcd06171e6dff5461e16d90de91d371e47bc", size = 35844549, upload-time = "2026-07-10T08:27:37.956Z" }, + { url = "https://files.pythonhosted.org/packages/bc/66/24c28877219abf6263d909b1592c97ff82c59f13a59acbed11fc87c0654f/pyarrow-25.0.0-cp313-cp313-macosx_12_0_x86_64.whl", hash = "sha256:5f4bacb60f91dd2fca6c52f1b9a0012cd090e0294f1f781dc1881a247a352f8e", size = 37610397, upload-time = "2026-07-10T08:27:43.803Z" }, + { url = "https://files.pythonhosted.org/packages/53/55/6d1d5f5aff317ec5de9421594679ed51ed828fe7e2ce209327f819d801e4/pyarrow-25.0.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:59516c822d5fd8e544aaa0dfe72f36fed5d4c24ea8390aab1bcd31d7e959c6be", size = 46841701, upload-time = "2026-07-10T08:27:49.741Z" }, + { url = "https://files.pythonhosted.org/packages/b5/5d/f790fb6965ab54c9da0dda7856abc75fd0d7648d865f8d603c111d203a64/pyarrow-25.0.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:6f9dbd83e91c239a1f5ee7ce13f108b5f6c0efbe40a4375260d8f08b43ad05e9", size = 50090118, upload-time = "2026-07-10T08:27:56.051Z" }, + { url = "https://files.pythonhosted.org/packages/0c/8c/faf025357ebf31bc96777f234277aa31e2aeca6dd4ecaa391f29085473c2/pyarrow-25.0.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:18dcc8cc50b5e72eae6fcbfc6c8776c21a007176b27a3cdec5c2f5bcf126708d", size = 49945559, upload-time = "2026-07-10T08:28:01.927Z" }, + { url = "https://files.pythonhosted.org/packages/07/a1/bd051871708ea99a5e0fc711926c26c6f2c6d0130c7aaac8093e34998af6/pyarrow-25.0.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:4ec1895a87aa834c3b99b7a1e758747eb8bb57f922b32c0e0fa04afb8d6998b1", size = 53114238, upload-time = "2026-07-10T08:28:08.594Z" }, + { url = "https://files.pythonhosted.org/packages/7c/31/737f0c3cffcd6af647849477d1dd68045deac2e3963c3f9f211bedc48540/pyarrow-25.0.0-cp313-cp313-win_amd64.whl", hash = "sha256:77c8d1ae46a44b4006e8db1cc977bbcc6ce4873c92f74137d68e45503b97fb18", size = 27861162, upload-time = "2026-07-10T08:28:12.975Z" }, + { url = "https://files.pythonhosted.org/packages/55/c7/581ccbcdb3d897eb2893328d68db3d52eca373bf2a7e964d0a6276b8e85b/pyarrow-25.0.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:72132b9a8a0a1840197794d4dea26080069b6b0981c116bc078762dc9691b21b", size = 35878945, upload-time = "2026-07-10T08:28:18.222Z" }, + { url = "https://files.pythonhosted.org/packages/64/d1/ccb01db7329ea0411ef4fbd9b62a04d3268b36777d4e758d5e39b91ddeab/pyarrow-25.0.0-cp314-cp314-macosx_12_0_x86_64.whl", hash = "sha256:e009ef945e498dca2f050ea10d2e9764cb44017254826fc4574fdb8d2530173b", size = 37630854, upload-time = "2026-07-10T08:28:23.452Z" }, + { url = "https://files.pythonhosted.org/packages/af/9f/2d81ba89d1e4198d0cb25fe7529de936830fdaec0db926bb52a1ef7080d4/pyarrow-25.0.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:f57a39dbcb416345401c2e77a4373669b45fd111a1768e6cf267a7a0607ff0ec", size = 46905617, upload-time = "2026-07-10T08:28:29.376Z" }, + { url = "https://files.pythonhosted.org/packages/6a/29/0ed312ec800fb536f93783215126cee4b8977dcfeccba6f0f44df0cc87d7/pyarrow-25.0.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:447df764beb07c544f0178a5f6b70ef44b9ecf382b3cdfad4c2d7867353c3887", size = 50119765, upload-time = "2026-07-10T08:28:35.826Z" }, + { url = "https://files.pythonhosted.org/packages/ca/88/cab5063ba0c4d46a9f6b4b7eb1c9029dc0302d65cd5ab3510c949a386568/pyarrow-25.0.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ac5dfeee59f9ceb4d45ba76e83b026c38c24334135bb329d8274baa49cec3c62", size = 50027563, upload-time = "2026-07-10T08:28:43.848Z" }, + { url = "https://files.pythonhosted.org/packages/7b/fb/4d24f1b7fe2e042dc4ef315ef75e4e702d8e46fe10c37e63caff00502b03/pyarrow-25.0.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:f0f100dacf2c0f400601664a79d1a907ced4740514bb2b00917341038e2ce76f", size = 53162437, upload-time = "2026-07-10T08:28:52.819Z" }, + { url = "https://files.pythonhosted.org/packages/fa/65/da20806de93ca6ee91e72cb6a9b08b3ac890b46efc8d94a7326c651c4c81/pyarrow-25.0.0-cp314-cp314-win_amd64.whl", hash = "sha256:2e093efbecb5317372f819228fa4b4e6157eee48d3f0a7b0303705ebf81a7104", size = 28613262, upload-time = "2026-07-10T08:29:47.544Z" }, + { url = "https://files.pythonhosted.org/packages/86/9f/c632afb1d3ef4a7814cee236718235f3a47eac46e97eb87df40f550b6b48/pyarrow-25.0.0-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:26be35b80780d2d21f4bae3d568b1666337c3a89722cc1794c956a77017cb24e", size = 36120702, upload-time = "2026-07-10T08:28:59.577Z" }, + { url = "https://files.pythonhosted.org/packages/36/0a/093d53a0e72ad06e45d6443e00651bbc2d21af4211295086cbf4d873d3b9/pyarrow-25.0.0-cp314-cp314t-macosx_12_0_x86_64.whl", hash = "sha256:6f4812bfbf11ca7d8faf59eb8fff8bf4dd25ce3a38b62baa010cc17a0926d1b2", size = 37750674, upload-time = "2026-07-10T08:29:06.916Z" }, + { url = "https://files.pythonhosted.org/packages/8a/18/b37fc31a69cff4bdfb8842683def5612f551b93fff6f44375e4a4a6a5535/pyarrow-25.0.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:b8af8ceedf0c9c160fd2b63440f2d205b9404db85866c1217bfea601de7cfb50", size = 46912304, upload-time = "2026-07-10T08:29:14.656Z" }, + { url = "https://files.pythonhosted.org/packages/32/35/5cae19ba72493e5598022468b56f6a5571f399f485bf412f157356476caa/pyarrow-25.0.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:c70a5fd9a82bd1a702fd482bdc62d38dcb672fb2b449b1d7c0d7d1f4be7b7bfe", size = 50073652, upload-time = "2026-07-10T08:29:22.467Z" }, + { url = "https://files.pythonhosted.org/packages/2e/a5/ddd508424bdfd5e6945765e9e2ffc687e2f6115972badc8ecf423076c407/pyarrow-25.0.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:0490a7f8b38ffe11cc26526b50c65d111cb54ddac3717cec781806793f1244dc", size = 50058654, upload-time = "2026-07-10T08:29:29.689Z" }, + { url = "https://files.pythonhosted.org/packages/5f/a4/324d0db203ff5eebe8694ec2d6ec5a23f9aaa5d02e5b8c692914c518c33c/pyarrow-25.0.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:e83916bbcf380866b4e14255850b33323ff678dc9758411d0409cdd2523880b0", size = 53140153, upload-time = "2026-07-10T08:29:36.041Z" }, + { url = "https://files.pythonhosted.org/packages/bd/8d/d236e9c82fe315f9128885c8be3ec719f41965a1eb6b6f4b42470904cd41/pyarrow-25.0.0-cp314-cp314t-win_amd64.whl", hash = "sha256:13240f0d3dc5932ccd0bfa90cd76d835680b9d94a7661c635df4b703d40ce849", size = 28743657, upload-time = "2026-07-10T08:29:42.742Z" }, ] [[package]] @@ -1096,51 +1634,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/83/13/75c86a8ef61ea2c758c924318cf894dced2436b0f7aeb3c5f0fe9e4305b4/pyarrow_stubs-20.0.0.20250928-py3-none-any.whl", hash = "sha256:5389057a55db3c2662c05f22685a52e15e5effaf4345f41f12fb9b6b348647b9", size = 235745, upload-time = "2025-09-28T02:50:03.205Z" }, ] -[[package]] -name = "pybtex" -version = "0.25.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "latexcodec" }, - { name = "pyyaml" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/5f/bc/c2be05ca72f8c103670e983df8be26d1e288bc6556f487fa8cccaa27779f/pybtex-0.25.1.tar.gz", hash = "sha256:9eaf90267c7e83e225af89fea65c370afbf65f458220d3946a9e3049e1eca491", size = 406157, upload-time = "2025-06-26T13:27:41.903Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/25/68/ceb5d6679baa326261f5d3e5113d9cfed6efef2810afd9f18bffb8ed312b/pybtex-0.25.1-py2.py3-none-any.whl", hash = "sha256:9053b0d619409a0a83f38abad5d9921de5f7b3ede00742beafcd9f10ad0d8c5c", size = 127437, upload-time = "2025-06-26T13:27:43.585Z" }, -] - -[[package]] -name = "pybtex-docutils" -version = "1.0.3" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "docutils" }, - { name = "pybtex" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/7e/84/796ea94d26188a853660f81bded39f8de4cfe595130aef0dea1088705a11/pybtex-docutils-1.0.3.tar.gz", hash = "sha256:3a7ebdf92b593e00e8c1c538aa9a20bca5d92d84231124715acc964d51d93c6b", size = 18348, upload-time = "2023-08-22T18:47:54.833Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/11/b1/ce1f4596211efb5410e178a803f08e59b20bedb66837dcf41e21c54f9ec1/pybtex_docutils-1.0.3-py3-none-any.whl", hash = "sha256:8fd290d2ae48e32fcb54d86b0efb8d573198653c7e2447d5bec5847095f430b9", size = 6385, upload-time = "2023-08-22T06:43:20.513Z" }, -] - -[[package]] -name = "pydata-sphinx-theme" -version = "0.16.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "accessible-pygments" }, - { name = "babel" }, - { name = "beautifulsoup4" }, - { name = "docutils" }, - { name = "pygments" }, - { name = "sphinx", version = "9.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, - { name = "sphinx", version = "9.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/00/20/bb50f9de3a6de69e6abd6b087b52fa2418a0418b19597601605f855ad044/pydata_sphinx_theme-0.16.1.tar.gz", hash = "sha256:a08b7f0b7f70387219dc659bff0893a7554d5eb39b59d3b8ef37b8401b7642d7", size = 2412693, upload-time = "2024-12-17T10:53:39.537Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/e2/0d/8ba33fa83a7dcde13eb3c1c2a0c1cc29950a048bfed6d9b0d8b6bd710b4c/pydata_sphinx_theme-0.16.1-py3-none-any.whl", hash = "sha256:225331e8ac4b32682c18fcac5a57a6f717c4e632cea5dd0e247b55155faeccde", size = 6723264, upload-time = "2024-12-17T10:53:35.645Z" }, -] - [[package]] name = "pygments" version = "2.19.2" @@ -1299,6 +1792,50 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d5/5b/33ac5616706d1bc0bd40e22baf12acccd687b505f85d107f549355379d6f/ray-2.56.0-cp314-cp314-manylinux2014_x86_64.whl", hash = "sha256:73edb6fb5fd05481b1f358ac2e8a4c7f6a031d89f9b823d25a587ddb7529070f", size = 74099117, upload-time = "2026-06-29T20:51:48.19Z" }, ] +[[package]] +name = "rcssmin" +version = "1.2.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/81/af/c9654b4f9b054ec163ed7cb20d8db0e5ae05e2e9ce99a4c11d91a2180b3f/rcssmin-1.2.2.tar.gz", hash = "sha256:806986eaf7414545edc28a1d29523e9560e49e151ff4a337d9d1f0271d6e1cc4", size = 587012, upload-time = "2025-10-12T10:48:08.932Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9c/3b/d4db4a2fb0d0033d222b56526bb1935e892e2560516b0378b4ccffec8d9c/rcssmin-1.2.2-cp311-cp311-manylinux1_i686.whl", hash = "sha256:da4801f4f429d66f9922871a7c71dee54c87f0ea5666cae6f1eb84c3fbc4e1f4", size = 49090, upload-time = "2025-10-12T10:48:24.827Z" }, + { url = "https://files.pythonhosted.org/packages/bd/f3/aeed5758339ccba61a82de12897762bad8f4317883a20de2dcc78842afda/rcssmin-1.2.2-cp311-cp311-manylinux1_x86_64.whl", hash = "sha256:e6b5913f3e8cb249044e916bc6ddcb5158815121548686a0fc8e2b8a5961a62e", size = 49368, upload-time = "2025-10-12T10:48:26.237Z" }, + { url = "https://files.pythonhosted.org/packages/2a/7b/e4206002c8c1bdcac6905ad7b200d62662d20d2b23f3d3e7df4e89447fdd/rcssmin-1.2.2-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:1472a98142d10d6c6772d96424ddcaf99d7e1d3217475f7f28f7d40dd84f24a2", size = 50714, upload-time = "2025-10-12T10:48:27.305Z" }, + { url = "https://files.pythonhosted.org/packages/d7/7c/e1cd335ce659af50a2c16dad37eee4b166536d73a463cdfeab5bb8e0833b/rcssmin-1.2.2-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:b6c91878a7e6f708f90c1bbc1a02729f45b2e5dee89045b395e997aa71744ee4", size = 53376, upload-time = "2025-10-12T10:48:28.374Z" }, + { url = "https://files.pythonhosted.org/packages/98/0b/9071882a74df398bf40e668a89cf2dd7eb95ff2e02c111c4c156aaad745a/rcssmin-1.2.2-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:844227668a235451eb544455b911067ba5495d680857d4bad2b0b78878f30a5c", size = 53194, upload-time = "2025-10-12T10:48:29.331Z" }, + { url = "https://files.pythonhosted.org/packages/03/48/295e57fbb4767226b5231ee99c056eab5447845259e4172e2db76f07c26d/rcssmin-1.2.2-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:cc866e23121bc4e29014b588fb67c8242a80ce053196f511c4c806b30ca6a393", size = 52998, upload-time = "2025-10-12T10:48:30.266Z" }, + { url = "https://files.pythonhosted.org/packages/86/5c/29af37ffb21a3069d108902868262b25fbbf731821cf5c7de76bba986dd1/rcssmin-1.2.2-cp312-cp312-manylinux1_i686.whl", hash = "sha256:78249189d39344a1e9d813c51362831537500e104c5bdce4ff24fe59010e9ee1", size = 48819, upload-time = "2025-10-12T10:48:31.3Z" }, + { url = "https://files.pythonhosted.org/packages/89/dc/b522a5e1a0a8ef8af50adbb3bdd9f5a059a9890b9fc5ce3f44a37a996a74/rcssmin-1.2.2-cp312-cp312-manylinux1_x86_64.whl", hash = "sha256:217efed0dff304d503bf481068ddb13ae72176ed5970f1011fb1a1e379308d9c", size = 49248, upload-time = "2025-10-12T10:48:32.254Z" }, + { url = "https://files.pythonhosted.org/packages/66/0b/7c0018793080ed26939d9beaf09591cf58fb9cda3253a891137b841a902a/rcssmin-1.2.2-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:c51aa47b1752ae55ad4cf4332e7316c5206a6a686d65bc15431a6bfea393e665", size = 50723, upload-time = "2025-10-12T10:48:33.398Z" }, + { url = "https://files.pythonhosted.org/packages/52/b4/c8b0d2588719af2b9c454e6e95f18bd60b0c474da4b65af61bb6457a7555/rcssmin-1.2.2-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:354215283f32413ced87b358934ebbb7c5529f51f5d316e80bc2889486d388b3", size = 53302, upload-time = "2025-10-12T10:48:34.452Z" }, + { url = "https://files.pythonhosted.org/packages/8d/cb/1a8d6ace1ff845d57fa31087510a88dc10ac01cea41317e976bbf2413f91/rcssmin-1.2.2-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:265b57de87949b505bcd658f4f5bbfc1f077390108cd12e288ba2f7824bee52c", size = 53001, upload-time = "2025-10-12T10:48:35.389Z" }, + { url = "https://files.pythonhosted.org/packages/5e/25/2d97155edb351a28e5a46a35f7b4e54bfe3933847bd2ba6674216a817d9e/rcssmin-1.2.2-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:83ecd3093640d69f7582839788e012ecf9a85faeb95760032626977a7d3904b2", size = 53285, upload-time = "2025-10-12T10:48:36.513Z" }, + { url = "https://files.pythonhosted.org/packages/8b/ae/3a7911e1c773f3deb039a42588ae6cee59d6bcec07b5081db376677b293a/rcssmin-1.2.2-cp313-cp313-manylinux1_i686.whl", hash = "sha256:e91449b612a08e5e80df3487e941c86e2c73c5088169588c31c382eb94da0521", size = 48764, upload-time = "2025-10-12T10:48:37.45Z" }, + { url = "https://files.pythonhosted.org/packages/30/a7/6d311986d76da0a538bae3f584d2b7579dd11648e74f539b177d8af51f6b/rcssmin-1.2.2-cp313-cp313-manylinux1_x86_64.whl", hash = "sha256:20a32c49d65b65c3ac80305d8a31b98f3d92b1b052dd63b57fbebc7003f9ae38", size = 49175, upload-time = "2025-10-12T10:48:39.601Z" }, + { url = "https://files.pythonhosted.org/packages/f2/9b/ceb12f3397695d075a1f3e12e295d84f021562540a6579144cb985d80ccb/rcssmin-1.2.2-cp313-cp313-manylinux2014_aarch64.whl", hash = "sha256:cd0a5ca4a0fc3b193ab0dcd251bd2463900558108cc4306a5cc4ab77c6bfffde", size = 50675, upload-time = "2025-10-12T10:48:40.664Z" }, + { url = "https://files.pythonhosted.org/packages/3b/8c/efb41baaea20567fa0c335705bff1f187e35301b684891d26282a16aff1b/rcssmin-1.2.2-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:2b008aa77a92f9db2d88f7e7ab45b81f37253cb0baafda59dd5b857c2de9b09f", size = 52999, upload-time = "2025-10-12T10:48:41.749Z" }, + { url = "https://files.pythonhosted.org/packages/4d/f6/cf692cca8837375fd21bf31cd134e10684fc11283a68c04160619aa826dc/rcssmin-1.2.2-cp313-cp313-musllinux_1_1_i686.whl", hash = "sha256:b3843e0501fa45d7c911dd7b3b78fd5f51c8159dd36d780ee12060da2d526aa0", size = 52752, upload-time = "2025-10-12T10:48:42.823Z" }, + { url = "https://files.pythonhosted.org/packages/8d/fa/5b0f1df380f598a397414dfaba74b05901379918f4d6b1746462190ae011/rcssmin-1.2.2-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:bb75b0e412a5419d62f39d89d0b3920a6697d2b12c8dad57f8bde1c76332c640", size = 52994, upload-time = "2025-10-12T10:48:44.023Z" }, + { url = "https://files.pythonhosted.org/packages/16/df/7157985ff9e2f3fecb15b03370ee0c8de42fd5a07c4b54e0be2c0f3f8133/rcssmin-1.2.2-cp313-cp313t-manylinux1_i686.whl", hash = "sha256:dd192a876a7af9a14628ff20818df80187294db96d86ddccf72371a6ae3e7ce7", size = 51314, upload-time = "2025-10-12T10:48:45.569Z" }, + { url = "https://files.pythonhosted.org/packages/56/d5/e6c176b8d39faf0fac3a6896022febb00d0ac5c4b99d4924572c579af210/rcssmin-1.2.2-cp313-cp313t-manylinux1_x86_64.whl", hash = "sha256:5724ed426c1444c35584f0bcda43c81ac47da769228722207aea7b8eedf31224", size = 51509, upload-time = "2025-10-12T10:48:47.03Z" }, + { url = "https://files.pythonhosted.org/packages/b7/3e/493ef8b7ce621b45f1be4505295fe604280918f67a01a28c82f9d0621a3f/rcssmin-1.2.2-cp313-cp313t-manylinux2014_aarch64.whl", hash = "sha256:228cc8d192ba4bd82305c085cbb5594d45d8dc6605d4eddc319543fb9f47b319", size = 52730, upload-time = "2025-10-12T10:48:48.13Z" }, + { url = "https://files.pythonhosted.org/packages/b2/90/77cf149fac7f247dac530a96beac7c52cea9fab928b5c3e2a45c6da86147/rcssmin-1.2.2-cp313-cp313t-musllinux_1_1_aarch64.whl", hash = "sha256:714390aac7c4cb611eecc845a5d9bb01495a3c9fccf9d8a2d6aa75a109276f7b", size = 54939, upload-time = "2025-10-12T10:48:49.423Z" }, + { url = "https://files.pythonhosted.org/packages/35/37/b8347b3a817b99eab9cc987f1090c7192d9d5f077fdc84c04d12f5186b87/rcssmin-1.2.2-cp313-cp313t-musllinux_1_1_i686.whl", hash = "sha256:64ec506fce7a3f1e993f4c4b55c7b3d9ad8259191cf20d986aa1d1a13e920fe8", size = 55108, upload-time = "2025-10-12T10:48:50.463Z" }, + { url = "https://files.pythonhosted.org/packages/5e/e1/fb555f831d5e5674a91444007434ba85b4ae98cfd97dd7bc9c2962c0f56b/rcssmin-1.2.2-cp313-cp313t-musllinux_1_1_x86_64.whl", hash = "sha256:742bb522d1efe0f1d362d81e00b5dc93ca2ddd1e435ed2d921cfa84fbb9f6887", size = 54858, upload-time = "2025-10-12T10:48:51.941Z" }, + { url = "https://files.pythonhosted.org/packages/a0/40/9c4cb3133f6d4ddfbeada76988a10ff2a974706fd6fcbb97edd8c0f4cc76/rcssmin-1.2.2-cp314-cp314-manylinux1_i686.whl", hash = "sha256:540dd3aa586b5f8f4c4b90db37e6a31c04718cdf90dbe9bec43c3b4dd50519e7", size = 49032, upload-time = "2025-10-12T10:48:53.014Z" }, + { url = "https://files.pythonhosted.org/packages/07/84/a411a48fd4179a88c68a2ad3649b408fa7887a421d3435c10ae6f5724e3a/rcssmin-1.2.2-cp314-cp314-manylinux1_x86_64.whl", hash = "sha256:6ea38a38eec263858b70bed6715478dcfed7fbc5d63333a8c512631ee22baad9", size = 49497, upload-time = "2025-10-12T10:48:54.009Z" }, + { url = "https://files.pythonhosted.org/packages/a1/32/5663a71a9304e0c9f33b765264508229d026359cfff746e1d0a593d809ea/rcssmin-1.2.2-cp314-cp314-manylinux2014_aarch64.whl", hash = "sha256:07dc7d352e8eb08de82fc4c545dd04f9f487466c8370051e0bee4eb1e4dc85d0", size = 50382, upload-time = "2025-10-12T10:48:55.079Z" }, + { url = "https://files.pythonhosted.org/packages/d7/28/e411eb191ffff7bd712f2eb0f691cb7ca514b1876d6bff2f5ae61359b8db/rcssmin-1.2.2-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:cdccb0e08281f0dd5d463c16ec61a06bd1534de50206dc72918be3c10dcb82e5", size = 50962, upload-time = "2025-10-12T10:48:56.494Z" }, + { url = "https://files.pythonhosted.org/packages/fb/3f/cdb99526d294c5dd4b919dc4ef492b7bd11e08b585d15ec641dfb9423493/rcssmin-1.2.2-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:2b6d5e2e2fd65738d57ef65aaaed2cff2288eccff7f704bf3d579e6f451cb60a", size = 52504, upload-time = "2025-10-12T10:48:57.886Z" }, + { url = "https://files.pythonhosted.org/packages/e8/60/a8183401fa64e93e1d52b2cdf275a2c11e0993f5f3162c573a67872b535d/rcssmin-1.2.2-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:7018d4197713c7797d1a67ed47ab53d4706c2e9ed134123c30a47d389dda5386", size = 50561, upload-time = "2025-10-12T10:48:58.935Z" }, + { url = "https://files.pythonhosted.org/packages/47/5e/496d6c9c309e2fe79e6a69f25f7a6d18f545edb4ea3584f461b9f84b0d60/rcssmin-1.2.2-cp314-cp314t-manylinux1_i686.whl", hash = "sha256:0162c32ce946978edc834d4fba705ac5f9422d7f556f3264cc4fc67c7ee39171", size = 51214, upload-time = "2025-10-12T10:49:00.021Z" }, + { url = "https://files.pythonhosted.org/packages/5e/78/87da6706d5856ceee71421ba831d2f5d93c3e6865acfbb56ace8d54587cc/rcssmin-1.2.2-cp314-cp314t-manylinux1_x86_64.whl", hash = "sha256:f17dc92553a46412c49f972f0ab31088032b9482a9c421bc2d39691a5d8842aa", size = 51608, upload-time = "2025-10-12T10:49:01.422Z" }, + { url = "https://files.pythonhosted.org/packages/cd/6c/204b0262c11ac2da2b8df2d8fed76f1959273fbc8376450d0ac022d754b7/rcssmin-1.2.2-cp314-cp314t-manylinux2014_aarch64.whl", hash = "sha256:40c7dfba098bbd129d8c35dd8b604275585f9dc0496e5d17dbe7fd6b873b0233", size = 53349, upload-time = "2025-10-12T10:49:02.512Z" }, + { url = "https://files.pythonhosted.org/packages/c3/7b/9aae16756d3f33cbc512760ba3e69c3856a51aa293e463f2ca97760d1b1b/rcssmin-1.2.2-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:d0197fab78ebbe33f5df9caf2572ef2d44bbe243a9130881a0c5c53ba03641fa", size = 53066, upload-time = "2025-10-12T10:49:03.589Z" }, + { url = "https://files.pythonhosted.org/packages/4e/18/b06fadfa9b85e486bb1571050217cb539c062d1ae4cd32b1a31c36f67fd4/rcssmin-1.2.2-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:19e53c58768369366fdaef00da59f275f724f229994ea885309df6ca368ff3c8", size = 54271, upload-time = "2025-10-12T10:49:04.735Z" }, + { url = "https://files.pythonhosted.org/packages/79/55/f29ce21f8e5a1f3c19d43b67b907268d227b7edcda2ca200ca0028734a5e/rcssmin-1.2.2-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:8d3de1a870e00d157f3a7b1797498fdc09a3774629079572350f75783bb94b9a", size = 52423, upload-time = "2025-10-12T10:49:06.04Z" }, +] + [[package]] name = "referencing" version = "0.36.2" @@ -1484,15 +2021,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c6/2a/65880dfd0e13f7f13a775998f34703674a4554906167dce02daf7865b954/ruff-0.14.0-py3-none-win_arm64.whl", hash = "sha256:f42c9495f5c13ff841b1da4cb3c2a42075409592825dada7c5885c2c844ac730", size = 12565142, upload-time = "2025-10-07T18:21:53.577Z" }, ] -[[package]] -name = "setuptools" -version = "83.0.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/34/26/f5d29e25ffdb535afef2d35cdb55b325298f96debd670da4c325e08d70f4/setuptools-83.0.0.tar.gz", hash = "sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef", size = 1154254, upload-time = "2026-07-04T15:31:22.699Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/5d/40/e1e72872c6354b306daef1703549e8e83b4d43cfea356311bf722a043752/setuptools-83.0.0-py3-none-any.whl", hash = "sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3", size = 1008090, upload-time = "2026-07-04T15:31:20.885Z" }, -] - [[package]] name = "shellingham" version = "1.5.4" @@ -1511,15 +2039,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, ] -[[package]] -name = "sniffio" -version = "1.3.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/a2/87/a6771e1546d97e7e041b6ae58d80074f81b7d5121207425c964ddf5cfdbd/sniffio-1.3.1.tar.gz", hash = "sha256:f4324edc670a0f49750a81b895f35c3adb843cca46f0530f79fc1babb23789dc", size = 20372, upload-time = "2024-02-25T23:20:04.057Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl", hash = "sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2", size = 10235, upload-time = "2024-02-25T23:20:01.196Z" }, -] - [[package]] name = "snowballstemmer" version = "3.0.1" @@ -1529,15 +2048,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c8/78/3565d011c61f5a43488987ee32b6f3f656e7f107ac2782dd57bdd7d91d9a/snowballstemmer-3.0.1-py3-none-any.whl", hash = "sha256:6cd7b3897da8d6c9ffb968a6781fa6532dce9c3618a4b127d920dab764a19064", size = 103274, upload-time = "2025-05-09T16:34:50.371Z" }, ] -[[package]] -name = "soupsieve" -version = "2.8" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/6d/e6/21ccce3262dd4889aa3332e5a119a3491a95e8f60939870a3a035aabac0d/soupsieve-2.8.tar.gz", hash = "sha256:e2dd4a40a628cb5f28f6d4b0db8800b8f581b65bb380b97de22ba5ca8d72572f", size = 103472, upload-time = "2025-08-27T15:39:51.78Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/14/a0/bb38d3b76b8cae341dad93a2dd83ab7462e6dbcdd84d43f54ee60a8dc167/soupsieve-2.8-py3-none-any.whl", hash = "sha256:0cc76456a30e20f5d7f2e14a98a4ae2ee4e5abdc7c5ea0aafe795f344bc7984c", size = 36679, upload-time = "2025-08-27T15:39:50.179Z" }, -] - [[package]] name = "sphinx" version = "9.0.4" @@ -1546,23 +2056,23 @@ resolution-markers = [ "python_full_version < '3.12'", ] dependencies = [ - { name = "alabaster" }, - { name = "babel" }, - { name = "colorama", marker = "sys_platform == 'win32'" }, - { name = "docutils" }, - { name = "imagesize" }, - { name = "jinja2" }, - { name = "packaging" }, - { name = "pygments" }, - { name = "requests" }, - { name = "roman-numerals" }, - { name = "snowballstemmer" }, - { name = "sphinxcontrib-applehelp" }, - { name = "sphinxcontrib-devhelp" }, - { name = "sphinxcontrib-htmlhelp" }, - { name = "sphinxcontrib-jsmath" }, - { name = "sphinxcontrib-qthelp" }, - { name = "sphinxcontrib-serializinghtml" }, + { name = "alabaster", marker = "python_full_version < '3.12'" }, + { name = "babel", marker = "python_full_version < '3.12'" }, + { name = "colorama", marker = "python_full_version < '3.12' and sys_platform == 'win32'" }, + { name = "docutils", marker = "python_full_version < '3.12'" }, + { name = "imagesize", marker = "python_full_version < '3.12'" }, + { name = "jinja2", marker = "python_full_version < '3.12'" }, + { name = "packaging", marker = "python_full_version < '3.12'" }, + { name = "pygments", marker = "python_full_version < '3.12'" }, + { name = "requests", marker = "python_full_version < '3.12'" }, + { name = "roman-numerals", marker = "python_full_version < '3.12'" }, + { name = "snowballstemmer", marker = "python_full_version < '3.12'" }, + { name = "sphinxcontrib-applehelp", marker = "python_full_version < '3.12'" }, + { name = "sphinxcontrib-devhelp", marker = "python_full_version < '3.12'" }, + { name = "sphinxcontrib-htmlhelp", marker = "python_full_version < '3.12'" }, + { name = "sphinxcontrib-jsmath", marker = "python_full_version < '3.12'" }, + { name = "sphinxcontrib-qthelp", marker = "python_full_version < '3.12'" }, + { name = "sphinxcontrib-serializinghtml", marker = "python_full_version < '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/42/50/a8c6ccc36d5eacdfd7913ddccd15a9cee03ecafc5ee2bc40e1f168d85022/sphinx-9.0.4.tar.gz", hash = "sha256:594ef59d042972abbc581d8baa577404abe4e6c3b04ef61bd7fc2acbd51f3fa3", size = 8710502, upload-time = "2025-12-04T07:45:27.343Z" } wheels = [ @@ -1577,86 +2087,29 @@ resolution-markers = [ "python_full_version >= '3.12'", ] dependencies = [ - { name = "alabaster" }, - { name = "babel" }, - { name = "colorama", marker = "sys_platform == 'win32'" }, - { name = "docutils" }, - { name = "imagesize" }, - { name = "jinja2" }, - { name = "packaging" }, - { name = "pygments" }, - { name = "requests" }, - { name = "roman-numerals" }, - { name = "snowballstemmer" }, - { name = "sphinxcontrib-applehelp" }, - { name = "sphinxcontrib-devhelp" }, - { name = "sphinxcontrib-htmlhelp" }, - { name = "sphinxcontrib-jsmath" }, - { name = "sphinxcontrib-qthelp" }, - { name = "sphinxcontrib-serializinghtml" }, + { name = "alabaster", marker = "python_full_version >= '3.12'" }, + { name = "babel", marker = "python_full_version >= '3.12'" }, + { name = "colorama", marker = "python_full_version >= '3.12' and sys_platform == 'win32'" }, + { name = "docutils", marker = "python_full_version >= '3.12'" }, + { name = "imagesize", marker = "python_full_version >= '3.12'" }, + { name = "jinja2", marker = "python_full_version >= '3.12'" }, + { name = "packaging", marker = "python_full_version >= '3.12'" }, + { name = "pygments", marker = "python_full_version >= '3.12'" }, + { name = "requests", marker = "python_full_version >= '3.12'" }, + { name = "roman-numerals", marker = "python_full_version >= '3.12'" }, + { name = "snowballstemmer", marker = "python_full_version >= '3.12'" }, + { name = "sphinxcontrib-applehelp", marker = "python_full_version >= '3.12'" }, + { name = "sphinxcontrib-devhelp", marker = "python_full_version >= '3.12'" }, + { name = "sphinxcontrib-htmlhelp", marker = "python_full_version >= '3.12'" }, + { name = "sphinxcontrib-jsmath", marker = "python_full_version >= '3.12'" }, + { name = "sphinxcontrib-qthelp", marker = "python_full_version >= '3.12'" }, + { name = "sphinxcontrib-serializinghtml", marker = "python_full_version >= '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cd/bd/f08eb0f4eed5c83f1ba2a3bd18f7745a2b1525fad70660a1c00224ec468a/sphinx-9.1.0.tar.gz", hash = "sha256:7741722357dd75f8190766926071fed3bdc211c74dd2d7d4df5404da95930ddb", size = 8718324, upload-time = "2025-12-31T15:09:27.646Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/73/f7/b1884cb3188ab181fc81fa00c266699dab600f927a964df02ec3d5d1916a/sphinx-9.1.0-py3-none-any.whl", hash = "sha256:c84fdd4e782504495fe4f2c0b3413d6c2bf388589bb352d439b2a3bb99991978", size = 3921742, upload-time = "2025-12-31T15:09:25.561Z" }, ] -[[package]] -name = "sphinx-autobuild" -version = "2025.8.25" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "colorama" }, - { name = "sphinx", version = "9.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, - { name = "sphinx", version = "9.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, - { name = "starlette" }, - { name = "uvicorn" }, - { name = "watchfiles" }, - { name = "websockets" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/e0/3c/a59a3a453d4133777f7ed2e83c80b7dc817d43c74b74298ca0af869662ad/sphinx_autobuild-2025.8.25.tar.gz", hash = "sha256:9cf5aab32853c8c31af572e4fecdc09c997e2b8be5a07daf2a389e270e85b213", size = 15200, upload-time = "2025-08-25T18:44:55.436Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/d7/20/56411b52f917696995f5ad27d2ea7e9492c84a043c5b49a3a3173573cd93/sphinx_autobuild-2025.8.25-py3-none-any.whl", hash = "sha256:b750ac7d5a18603e4665294323fd20f6dcc0a984117026d1986704fa68f0379a", size = 12535, upload-time = "2025-08-25T18:44:54.164Z" }, -] - -[[package]] -name = "sphinx-copybutton" -version = "0.5.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "sphinx", version = "9.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, - { name = "sphinx", version = "9.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/fc/2b/a964715e7f5295f77509e59309959f4125122d648f86b4fe7d70ca1d882c/sphinx-copybutton-0.5.2.tar.gz", hash = "sha256:4cf17c82fb9646d1bc9ca92ac280813a3b605d8c421225fd9913154103ee1fbd", size = 23039, upload-time = "2023-04-14T08:10:22.998Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/9e/48/1ea60e74949eecb12cdd6ac43987f9fd331156388dcc2319b45e2ebb81bf/sphinx_copybutton-0.5.2-py3-none-any.whl", hash = "sha256:fb543fd386d917746c9a2c50360c7905b605726b9355cd26e9974857afeae06e", size = 13343, upload-time = "2023-04-14T08:10:20.844Z" }, -] - -[[package]] -name = "sphinx-design" -version = "0.7.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "sphinx", version = "9.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, - { name = "sphinx", version = "9.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/13/7b/804f311da4663a4aecc6cf7abd83443f3d4ded970826d0c958edc77d4527/sphinx_design-0.7.0.tar.gz", hash = "sha256:d2a3f5b19c24b916adb52f97c5f00efab4009ca337812001109084a740ec9b7a", size = 2203582, upload-time = "2026-01-19T13:12:53.297Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/30/cf/45dd359f6ca0c3762ce0490f681da242f0530c49c81050c035c016bfdd3a/sphinx_design-0.7.0-py3-none-any.whl", hash = "sha256:f82bf179951d58f55dca78ab3706aeafa496b741a91b1911d371441127d64282", size = 2220350, upload-time = "2026-01-19T13:12:51.077Z" }, -] - -[[package]] -name = "sphinx-inline-tabs" -version = "2023.4.21" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "sphinx", version = "9.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, - { name = "sphinx", version = "9.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/48/f5/f8a2be63ed7be9f91a4c2bea0e25bcb56aa4c5cc37ec4d8ead8065f926b1/sphinx_inline_tabs-2023.4.21.tar.gz", hash = "sha256:5df2f13f602c158f3f5f6c509e008aeada199a8c76d97ba3aa2822206683bebc", size = 42664, upload-time = "2023-04-21T20:25:30.578Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/9d/60/1e4c9017d722b9c7731abc11f39ac8b083b479fbcefe12015b57e457a296/sphinx_inline_tabs-2023.4.21-py3-none-any.whl", hash = "sha256:06809ac613f7c48ddd6e2fa588413e3fe92cff2397b56e2ccf0b0218f9ef6a78", size = 6850, upload-time = "2023-04-21T20:25:28.778Z" }, -] - [[package]] name = "sphinxcontrib-applehelp" version = "2.0.0" @@ -1666,22 +2119,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5d/85/9ebeae2f76e9e77b952f4b274c27238156eae7979c5421fba91a28f4970d/sphinxcontrib_applehelp-2.0.0-py3-none-any.whl", hash = "sha256:4cd3f0ec4ac5dd9c17ec65e9ab272c9b867ea77425228e68ecf08d6b28ddbdb5", size = 119300, upload-time = "2024-07-29T01:08:58.99Z" }, ] -[[package]] -name = "sphinxcontrib-bibtex" -version = "2.6.5" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "docutils" }, - { name = "pybtex" }, - { name = "pybtex-docutils" }, - { name = "sphinx", version = "9.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, - { name = "sphinx", version = "9.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/de/83/1488c9879f2fa3c2cbd6f666c7a3a42a1fa9e08462bec73281fa6c092cba/sphinxcontrib_bibtex-2.6.5.tar.gz", hash = "sha256:9b3224dd6fece9268ebd8c905dc0a83ff2f6c54148a9235fe70e9d1e9ff149c0", size = 118462, upload-time = "2025-06-27T10:40:14.061Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/9e/a0/3a612da94f828f26cabb247817393e79472c32b12c49222bf85fb6d7b6c8/sphinxcontrib_bibtex-2.6.5-py3-none-any.whl", hash = "sha256:455ea4509642ea0b28ede3721550273626f85af65af01f161bfd8e19dc1edd7d", size = 40410, upload-time = "2025-06-27T10:40:12.274Z" }, -] - [[package]] name = "sphinxcontrib-devhelp" version = "2.0.0" @@ -1709,21 +2146,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c2/42/4c8646762ee83602e3fb3fbe774c2fac12f317deb0b5dbeeedd2d3ba4b77/sphinxcontrib_jsmath-1.0.1-py2.py3-none-any.whl", hash = "sha256:2ec2eaebfb78f3f2078e73666b1415417a116cc848b72e5172e596c871103178", size = 5071, upload-time = "2019-01-21T16:10:14.333Z" }, ] -[[package]] -name = "sphinxcontrib-mermaid" -version = "2.0.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "jinja2" }, - { name = "pyyaml" }, - { name = "sphinx", version = "9.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, - { name = "sphinx", version = "9.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/96/a5/65a5c439cc14ba80483b9891e9350f11efb80cd3bdccb222f0c738068c78/sphinxcontrib_mermaid-2.0.0.tar.gz", hash = "sha256:cf4f7d453d001132eaba5d1fdf53d42049f02e913213cf8337427483bfca26f4", size = 18194, upload-time = "2026-01-13T17:13:42.563Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f9/de/bd96c69b62e967bffd02c6d89dfca9471b04e761c466725fc39746abf41d/sphinxcontrib_mermaid-2.0.0-py3-none-any.whl", hash = "sha256:59a73249bbee2c74b1a4db036f8e8899ade65982bdda6712cf22b4f4e9874bb5", size = 14055, upload-time = "2026-01-13T17:13:41.481Z" }, -] - [[package]] name = "sphinxcontrib-qthelp" version = "2.0.0" @@ -1742,19 +2164,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/52/a7/d2782e4e3f77c8450f727ba74a8f12756d5ba823d81b941f1b04da9d033a/sphinxcontrib_serializinghtml-2.0.0-py3-none-any.whl", hash = "sha256:6e2cb0eef194e10c27ec0023bfeb25badbbb5868244cf5bc5bdc04e4464bf331", size = 92072, upload-time = "2024-07-29T01:10:08.203Z" }, ] -[[package]] -name = "sphinxext-opengraph" -version = "0.13.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "sphinx", version = "9.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, - { name = "sphinx", version = "9.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/f6/c0/eb6838e3bae624ce6c8b90b245d17e84252863150e95efdb88f92c8aa3fb/sphinxext_opengraph-0.13.0.tar.gz", hash = "sha256:103335d08567ad8468faf1425f575e3b698e9621f9323949a6c8b96d9793e80b", size = 1026875, upload-time = "2025-08-29T12:20:31.066Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/bf/a4/66c1fd4f8fab88faf71cee04a945f9806ba0fef753f2cfc8be6353f64508/sphinxext_opengraph-0.13.0-py3-none-any.whl", hash = "sha256:936c07828edc9ad9a7b07908b29596dc84ed0b3ceaa77acdf51282d232d4d80e", size = 1004152, upload-time = "2025-08-29T12:20:29.072Z" }, -] - [[package]] name = "stack-data" version = "0.6.3" @@ -1769,19 +2178,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f1/7b/ce1eafaf1a76852e2ec9b22edecf1daa58175c090266e9f6c64afcd81d91/stack_data-0.6.3-py3-none-any.whl", hash = "sha256:d5558e0c25a4cb0853cddad3d77da9891a08cb85dd9f9f91b9f8cd66e511e695", size = 24521, upload-time = "2023-09-30T13:58:03.53Z" }, ] -[[package]] -name = "starlette" -version = "1.3.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "anyio" }, - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/eb/e3/7c1dc7381d9f8ab7d854328ebfa884e62cb3f3d8549ddfd37c7814f42afa/starlette-1.3.1.tar.gz", hash = "sha256:05d0213193f2fbaae60e2ecb593b4add4262ad4e46536b54abe36f11a71724e0", size = 2703240, upload-time = "2026-06-12T09:23:11.602Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ec/bb/2799cc2ede3ed41131f8975621e7213dfc7ef4acbbaadfa440f32500c370/starlette-1.3.1-py3-none-any.whl", hash = "sha256:c7372aae11c3c3f26a42df7bd626cec2f47d03483d261d369516a615a53714c6", size = 73632, upload-time = "2026-06-12T09:23:10.017Z" }, -] - [[package]] name = "substrait" version = "0.27.0" @@ -1816,6 +2212,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/40/44/4a5f08c96eb108af5cb50b41f76142f0afa346dfa99d5296fe7202a11854/tabulate-0.9.0-py3-none-any.whl", hash = "sha256:024ca478df22e9340661486f85298cff5f6dcdba14f3813e8830015b9ed1948f", size = 35252, upload-time = "2022-10-06T17:21:44.262Z" }, ] +[[package]] +name = "tqdm" +version = "4.70.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/21/3b/6c24bec5be5e743ffd99576daa5cc077722fc7d5bbc00bd133fa0c698dc6/tqdm-4.70.0.tar.gz", hash = "sha256:55b0b0dbd97462d06ebee91e4dac24ed4d4702be82b24f07e6c1d27e08cea220", size = 795438, upload-time = "2026-07-27T11:33:15.271Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f9/1c/01bfd571a64e7f270e6bab5e33777debe0edc56759233ce84f27dec92d14/tqdm-4.70.0-py3-none-any.whl", hash = "sha256:7f585706bfddbdebf89daac705b2dfcc16890130727d3197ca62c732b4310953", size = 80184, upload-time = "2026-07-27T11:33:13.167Z" }, +] + [[package]] name = "traitlets" version = "5.14.3" @@ -1876,19 +2284,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" }, ] -[[package]] -name = "uvicorn" -version = "0.37.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "click" }, - { name = "h11" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/71/57/1616c8274c3442d802621abf5deb230771c7a0fec9414cb6763900eb3868/uvicorn-0.37.0.tar.gz", hash = "sha256:4115c8add6d3fd536c8ee77f0e14a7fd2ebba939fed9b02583a97f80648f9e13", size = 80367, upload-time = "2025-09-23T13:33:47.486Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/85/cd/584a2ceb5532af99dd09e50919e3615ba99aa127e9850eafe5f31ddfdb9a/uvicorn-0.37.0-py3-none-any.whl", hash = "sha256:913b2b88672343739927ce381ff9e2ad62541f9f8289664fa1d1d3803fa2ce6c", size = 67976, upload-time = "2025-09-23T13:33:45.842Z" }, -] - [[package]] name = "vortex-data" source = { editable = "vortex-python" } @@ -1905,6 +2300,10 @@ cuda = [ duckdb = [ { name = "duckdb" }, ] +hf = [ + { name = "datasets" }, + { name = "huggingface-hub" }, +] numpy = [ { name = "numpy" }, ] @@ -1920,7 +2319,9 @@ ray = [ [package.dev-dependencies] dev = [ + { name = "datasets" }, { name = "duckdb" }, + { name = "huggingface-hub" }, { name = "maturin" }, { name = "numpy" }, { name = "pandas", extra = ["output-formatting"] }, @@ -1936,7 +2337,9 @@ dev = [ [package.metadata] requires-dist = [ + { name = "datasets", marker = "extra == 'hf'", specifier = ">=4.0.0,<6" }, { name = "duckdb", marker = "extra == 'duckdb'", specifier = ">=1.1.2" }, + { name = "huggingface-hub", marker = "extra == 'hf'", specifier = ">=0.34.0" }, { name = "numpy", marker = "extra == 'numpy'", specifier = ">=1.26.0" }, { name = "pandas", marker = "extra == 'pandas'", specifier = ">=2.2.0" }, { name = "polars", marker = "extra == 'polars'", specifier = ">=1.31.0" }, @@ -1946,12 +2349,14 @@ requires-dist = [ { name = "typing-extensions", specifier = ">=4.5.0" }, { name = "vortex-data-cuda", marker = "extra == 'cuda'", directory = "vortex-python-cuda" }, ] -provides-extras = ["cuda", "duckdb", "numpy", "pandas", "polars", "ray"] +provides-extras = ["cuda", "duckdb", "hf", "numpy", "pandas", "polars", "ray"] [package.metadata.requires-dev] dev = [ + { name = "datasets", specifier = ">=4.0.0,<6" }, { name = "duckdb", specifier = ">=1.1.2" }, - { name = "maturin", specifier = ">=1.10.2" }, + { name = "huggingface-hub", specifier = ">=0.34.0" }, + { name = "maturin", specifier = ">=1.15.0" }, { name = "numpy", specifier = ">=2.2.2" }, { name = "pandas", extras = ["output-formatting"], specifier = ">=2.2.3" }, { name = "pandas-stubs", specifier = ">=2.2.3.241126" }, @@ -1974,90 +2379,6 @@ dependencies = [ [package.metadata] requires-dist = [{ name = "vortex-data", specifier = "==0.1.0" }] -[[package]] -name = "watchfiles" -version = "1.1.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "anyio" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/2a/9a/d451fcc97d029f5812e898fd30a53fd8c15c7bbd058fd75cfc6beb9bd761/watchfiles-1.1.0.tar.gz", hash = "sha256:693ed7ec72cbfcee399e92c895362b6e66d63dac6b91e2c11ae03d10d503e575", size = 94406, upload-time = "2025-06-15T19:06:59.42Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/8b/78/7401154b78ab484ccaaeef970dc2af0cb88b5ba8a1b415383da444cdd8d3/watchfiles-1.1.0-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:c9649dfc57cc1f9835551deb17689e8d44666315f2e82d337b9f07bd76ae3aa2", size = 405751, upload-time = "2025-06-15T19:05:07.679Z" }, - { url = "https://files.pythonhosted.org/packages/76/63/e6c3dbc1f78d001589b75e56a288c47723de28c580ad715eb116639152b5/watchfiles-1.1.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:406520216186b99374cdb58bc48e34bb74535adec160c8459894884c983a149c", size = 397313, upload-time = "2025-06-15T19:05:08.764Z" }, - { url = "https://files.pythonhosted.org/packages/6c/a2/8afa359ff52e99af1632f90cbf359da46184207e893a5f179301b0c8d6df/watchfiles-1.1.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:cb45350fd1dc75cd68d3d72c47f5b513cb0578da716df5fba02fff31c69d5f2d", size = 450792, upload-time = "2025-06-15T19:05:09.869Z" }, - { url = "https://files.pythonhosted.org/packages/1d/bf/7446b401667f5c64972a57a0233be1104157fc3abf72c4ef2666c1bd09b2/watchfiles-1.1.0-cp311-cp311-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:11ee4444250fcbeb47459a877e5e80ed994ce8e8d20283857fc128be1715dac7", size = 458196, upload-time = "2025-06-15T19:05:11.91Z" }, - { url = "https://files.pythonhosted.org/packages/58/2f/501ddbdfa3fa874ea5597c77eeea3d413579c29af26c1091b08d0c792280/watchfiles-1.1.0-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:bda8136e6a80bdea23e5e74e09df0362744d24ffb8cd59c4a95a6ce3d142f79c", size = 484788, upload-time = "2025-06-15T19:05:13.373Z" }, - { url = "https://files.pythonhosted.org/packages/61/1e/9c18eb2eb5c953c96bc0e5f626f0e53cfef4bd19bd50d71d1a049c63a575/watchfiles-1.1.0-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b915daeb2d8c1f5cee4b970f2e2c988ce6514aace3c9296e58dd64dc9aa5d575", size = 597879, upload-time = "2025-06-15T19:05:14.725Z" }, - { url = "https://files.pythonhosted.org/packages/8b/6c/1467402e5185d89388b4486745af1e0325007af0017c3384cc786fff0542/watchfiles-1.1.0-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:ed8fc66786de8d0376f9f913c09e963c66e90ced9aa11997f93bdb30f7c872a8", size = 477447, upload-time = "2025-06-15T19:05:15.775Z" }, - { url = "https://files.pythonhosted.org/packages/2b/a1/ec0a606bde4853d6c4a578f9391eeb3684a9aea736a8eb217e3e00aa89a1/watchfiles-1.1.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fe4371595edf78c41ef8ac8df20df3943e13defd0efcb732b2e393b5a8a7a71f", size = 453145, upload-time = "2025-06-15T19:05:17.17Z" }, - { url = "https://files.pythonhosted.org/packages/90/b9/ef6f0c247a6a35d689fc970dc7f6734f9257451aefb30def5d100d6246a5/watchfiles-1.1.0-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:b7c5f6fe273291f4d414d55b2c80d33c457b8a42677ad14b4b47ff025d0893e4", size = 626539, upload-time = "2025-06-15T19:05:18.557Z" }, - { url = "https://files.pythonhosted.org/packages/34/44/6ffda5537085106ff5aaa762b0d130ac6c75a08015dd1621376f708c94de/watchfiles-1.1.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:7738027989881e70e3723c75921f1efa45225084228788fc59ea8c6d732eb30d", size = 624472, upload-time = "2025-06-15T19:05:19.588Z" }, - { url = "https://files.pythonhosted.org/packages/c3/e3/71170985c48028fa3f0a50946916a14055e741db11c2e7bc2f3b61f4d0e3/watchfiles-1.1.0-cp311-cp311-win32.whl", hash = "sha256:622d6b2c06be19f6e89b1d951485a232e3b59618def88dbeda575ed8f0d8dbf2", size = 279348, upload-time = "2025-06-15T19:05:20.856Z" }, - { url = "https://files.pythonhosted.org/packages/89/1b/3e39c68b68a7a171070f81fc2561d23ce8d6859659406842a0e4bebf3bba/watchfiles-1.1.0-cp311-cp311-win_amd64.whl", hash = "sha256:48aa25e5992b61debc908a61ab4d3f216b64f44fdaa71eb082d8b2de846b7d12", size = 292607, upload-time = "2025-06-15T19:05:21.937Z" }, - { url = "https://files.pythonhosted.org/packages/61/9f/2973b7539f2bdb6ea86d2c87f70f615a71a1fc2dba2911795cea25968aea/watchfiles-1.1.0-cp311-cp311-win_arm64.whl", hash = "sha256:00645eb79a3faa70d9cb15c8d4187bb72970b2470e938670240c7998dad9f13a", size = 285056, upload-time = "2025-06-15T19:05:23.12Z" }, - { url = "https://files.pythonhosted.org/packages/f6/b8/858957045a38a4079203a33aaa7d23ea9269ca7761c8a074af3524fbb240/watchfiles-1.1.0-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:9dc001c3e10de4725c749d4c2f2bdc6ae24de5a88a339c4bce32300a31ede179", size = 402339, upload-time = "2025-06-15T19:05:24.516Z" }, - { url = "https://files.pythonhosted.org/packages/80/28/98b222cca751ba68e88521fabd79a4fab64005fc5976ea49b53fa205d1fa/watchfiles-1.1.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d9ba68ec283153dead62cbe81872d28e053745f12335d037de9cbd14bd1877f5", size = 394409, upload-time = "2025-06-15T19:05:25.469Z" }, - { url = "https://files.pythonhosted.org/packages/86/50/dee79968566c03190677c26f7f47960aff738d32087087bdf63a5473e7df/watchfiles-1.1.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:130fc497b8ee68dce163e4254d9b0356411d1490e868bd8790028bc46c5cc297", size = 450939, upload-time = "2025-06-15T19:05:26.494Z" }, - { url = "https://files.pythonhosted.org/packages/40/45/a7b56fb129700f3cfe2594a01aa38d033b92a33dddce86c8dfdfc1247b72/watchfiles-1.1.0-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:50a51a90610d0845a5931a780d8e51d7bd7f309ebc25132ba975aca016b576a0", size = 457270, upload-time = "2025-06-15T19:05:27.466Z" }, - { url = "https://files.pythonhosted.org/packages/b5/c8/fa5ef9476b1d02dc6b5e258f515fcaaecf559037edf8b6feffcbc097c4b8/watchfiles-1.1.0-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:dc44678a72ac0910bac46fa6a0de6af9ba1355669b3dfaf1ce5f05ca7a74364e", size = 483370, upload-time = "2025-06-15T19:05:28.548Z" }, - { url = "https://files.pythonhosted.org/packages/98/68/42cfcdd6533ec94f0a7aab83f759ec11280f70b11bfba0b0f885e298f9bd/watchfiles-1.1.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:a543492513a93b001975ae283a51f4b67973662a375a403ae82f420d2c7205ee", size = 598654, upload-time = "2025-06-15T19:05:29.997Z" }, - { url = "https://files.pythonhosted.org/packages/d3/74/b2a1544224118cc28df7e59008a929e711f9c68ce7d554e171b2dc531352/watchfiles-1.1.0-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:8ac164e20d17cc285f2b94dc31c384bc3aa3dd5e7490473b3db043dd70fbccfd", size = 478667, upload-time = "2025-06-15T19:05:31.172Z" }, - { url = "https://files.pythonhosted.org/packages/8c/77/e3362fe308358dc9f8588102481e599c83e1b91c2ae843780a7ded939a35/watchfiles-1.1.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f7590d5a455321e53857892ab8879dce62d1f4b04748769f5adf2e707afb9d4f", size = 452213, upload-time = "2025-06-15T19:05:32.299Z" }, - { url = "https://files.pythonhosted.org/packages/6e/17/c8f1a36540c9a1558d4faf08e909399e8133599fa359bf52ec8fcee5be6f/watchfiles-1.1.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:37d3d3f7defb13f62ece99e9be912afe9dd8a0077b7c45ee5a57c74811d581a4", size = 626718, upload-time = "2025-06-15T19:05:33.415Z" }, - { url = "https://files.pythonhosted.org/packages/26/45/fb599be38b4bd38032643783d7496a26a6f9ae05dea1a42e58229a20ac13/watchfiles-1.1.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:7080c4bb3efd70a07b1cc2df99a7aa51d98685be56be6038c3169199d0a1c69f", size = 623098, upload-time = "2025-06-15T19:05:34.534Z" }, - { url = "https://files.pythonhosted.org/packages/a1/e7/fdf40e038475498e160cd167333c946e45d8563ae4dd65caf757e9ffe6b4/watchfiles-1.1.0-cp312-cp312-win32.whl", hash = "sha256:cbcf8630ef4afb05dc30107bfa17f16c0896bb30ee48fc24bf64c1f970f3b1fd", size = 279209, upload-time = "2025-06-15T19:05:35.577Z" }, - { url = "https://files.pythonhosted.org/packages/3f/d3/3ae9d5124ec75143bdf088d436cba39812122edc47709cd2caafeac3266f/watchfiles-1.1.0-cp312-cp312-win_amd64.whl", hash = "sha256:cbd949bdd87567b0ad183d7676feb98136cde5bb9025403794a4c0db28ed3a47", size = 292786, upload-time = "2025-06-15T19:05:36.559Z" }, - { url = "https://files.pythonhosted.org/packages/26/2f/7dd4fc8b5f2b34b545e19629b4a018bfb1de23b3a496766a2c1165ca890d/watchfiles-1.1.0-cp312-cp312-win_arm64.whl", hash = "sha256:0a7d40b77f07be87c6faa93d0951a0fcd8cbca1ddff60a1b65d741bac6f3a9f6", size = 284343, upload-time = "2025-06-15T19:05:37.5Z" }, - { url = "https://files.pythonhosted.org/packages/d3/42/fae874df96595556a9089ade83be34a2e04f0f11eb53a8dbf8a8a5e562b4/watchfiles-1.1.0-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:5007f860c7f1f8df471e4e04aaa8c43673429047d63205d1630880f7637bca30", size = 402004, upload-time = "2025-06-15T19:05:38.499Z" }, - { url = "https://files.pythonhosted.org/packages/fa/55/a77e533e59c3003d9803c09c44c3651224067cbe7fb5d574ddbaa31e11ca/watchfiles-1.1.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:20ecc8abbd957046f1fe9562757903f5eaf57c3bce70929fda6c7711bb58074a", size = 393671, upload-time = "2025-06-15T19:05:39.52Z" }, - { url = "https://files.pythonhosted.org/packages/05/68/b0afb3f79c8e832e6571022611adbdc36e35a44e14f129ba09709aa4bb7a/watchfiles-1.1.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f2f0498b7d2a3c072766dba3274fe22a183dbea1f99d188f1c6c72209a1063dc", size = 449772, upload-time = "2025-06-15T19:05:40.897Z" }, - { url = "https://files.pythonhosted.org/packages/ff/05/46dd1f6879bc40e1e74c6c39a1b9ab9e790bf1f5a2fe6c08b463d9a807f4/watchfiles-1.1.0-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:239736577e848678e13b201bba14e89718f5c2133dfd6b1f7846fa1b58a8532b", size = 456789, upload-time = "2025-06-15T19:05:42.045Z" }, - { url = "https://files.pythonhosted.org/packages/8b/ca/0eeb2c06227ca7f12e50a47a3679df0cd1ba487ea19cf844a905920f8e95/watchfiles-1.1.0-cp313-cp313-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:eff4b8d89f444f7e49136dc695599a591ff769300734446c0a86cba2eb2f9895", size = 482551, upload-time = "2025-06-15T19:05:43.781Z" }, - { url = "https://files.pythonhosted.org/packages/31/47/2cecbd8694095647406645f822781008cc524320466ea393f55fe70eed3b/watchfiles-1.1.0-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:12b0a02a91762c08f7264e2e79542f76870c3040bbc847fb67410ab81474932a", size = 597420, upload-time = "2025-06-15T19:05:45.244Z" }, - { url = "https://files.pythonhosted.org/packages/d9/7e/82abc4240e0806846548559d70f0b1a6dfdca75c1b4f9fa62b504ae9b083/watchfiles-1.1.0-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:29e7bc2eee15cbb339c68445959108803dc14ee0c7b4eea556400131a8de462b", size = 477950, upload-time = "2025-06-15T19:05:46.332Z" }, - { url = "https://files.pythonhosted.org/packages/25/0d/4d564798a49bf5482a4fa9416dea6b6c0733a3b5700cb8a5a503c4b15853/watchfiles-1.1.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d9481174d3ed982e269c090f780122fb59cee6c3796f74efe74e70f7780ed94c", size = 451706, upload-time = "2025-06-15T19:05:47.459Z" }, - { url = "https://files.pythonhosted.org/packages/81/b5/5516cf46b033192d544102ea07c65b6f770f10ed1d0a6d388f5d3874f6e4/watchfiles-1.1.0-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:80f811146831c8c86ab17b640801c25dc0a88c630e855e2bef3568f30434d52b", size = 625814, upload-time = "2025-06-15T19:05:48.654Z" }, - { url = "https://files.pythonhosted.org/packages/0c/dd/7c1331f902f30669ac3e754680b6edb9a0dd06dea5438e61128111fadd2c/watchfiles-1.1.0-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:60022527e71d1d1fda67a33150ee42869042bce3d0fcc9cc49be009a9cded3fb", size = 622820, upload-time = "2025-06-15T19:05:50.088Z" }, - { url = "https://files.pythonhosted.org/packages/1b/14/36d7a8e27cd128d7b1009e7715a7c02f6c131be9d4ce1e5c3b73d0e342d8/watchfiles-1.1.0-cp313-cp313-win32.whl", hash = "sha256:32d6d4e583593cb8576e129879ea0991660b935177c0f93c6681359b3654bfa9", size = 279194, upload-time = "2025-06-15T19:05:51.186Z" }, - { url = "https://files.pythonhosted.org/packages/25/41/2dd88054b849aa546dbeef5696019c58f8e0774f4d1c42123273304cdb2e/watchfiles-1.1.0-cp313-cp313-win_amd64.whl", hash = "sha256:f21af781a4a6fbad54f03c598ab620e3a77032c5878f3d780448421a6e1818c7", size = 292349, upload-time = "2025-06-15T19:05:52.201Z" }, - { url = "https://files.pythonhosted.org/packages/c8/cf/421d659de88285eb13941cf11a81f875c176f76a6d99342599be88e08d03/watchfiles-1.1.0-cp313-cp313-win_arm64.whl", hash = "sha256:5366164391873ed76bfdf618818c82084c9db7fac82b64a20c44d335eec9ced5", size = 283836, upload-time = "2025-06-15T19:05:53.265Z" }, - { url = "https://files.pythonhosted.org/packages/45/10/6faf6858d527e3599cc50ec9fcae73590fbddc1420bd4fdccfebffeedbc6/watchfiles-1.1.0-cp313-cp313t-macosx_10_12_x86_64.whl", hash = "sha256:17ab167cca6339c2b830b744eaf10803d2a5b6683be4d79d8475d88b4a8a4be1", size = 400343, upload-time = "2025-06-15T19:05:54.252Z" }, - { url = "https://files.pythonhosted.org/packages/03/20/5cb7d3966f5e8c718006d0e97dfe379a82f16fecd3caa7810f634412047a/watchfiles-1.1.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:328dbc9bff7205c215a7807da7c18dce37da7da718e798356212d22696404339", size = 392916, upload-time = "2025-06-15T19:05:55.264Z" }, - { url = "https://files.pythonhosted.org/packages/8c/07/d8f1176328fa9e9581b6f120b017e286d2a2d22ae3f554efd9515c8e1b49/watchfiles-1.1.0-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f7208ab6e009c627b7557ce55c465c98967e8caa8b11833531fdf95799372633", size = 449582, upload-time = "2025-06-15T19:05:56.317Z" }, - { url = "https://files.pythonhosted.org/packages/66/e8/80a14a453cf6038e81d072a86c05276692a1826471fef91df7537dba8b46/watchfiles-1.1.0-cp313-cp313t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:a8f6f72974a19efead54195bc9bed4d850fc047bb7aa971268fd9a8387c89011", size = 456752, upload-time = "2025-06-15T19:05:57.359Z" }, - { url = "https://files.pythonhosted.org/packages/5a/25/0853b3fe0e3c2f5af9ea60eb2e781eade939760239a72c2d38fc4cc335f6/watchfiles-1.1.0-cp313-cp313t-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d181ef50923c29cf0450c3cd47e2f0557b62218c50b2ab8ce2ecaa02bd97e670", size = 481436, upload-time = "2025-06-15T19:05:58.447Z" }, - { url = "https://files.pythonhosted.org/packages/fe/9e/4af0056c258b861fbb29dcb36258de1e2b857be4a9509e6298abcf31e5c9/watchfiles-1.1.0-cp313-cp313t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:adb4167043d3a78280d5d05ce0ba22055c266cf8655ce942f2fb881262ff3cdf", size = 596016, upload-time = "2025-06-15T19:05:59.59Z" }, - { url = "https://files.pythonhosted.org/packages/c5/fa/95d604b58aa375e781daf350897aaaa089cff59d84147e9ccff2447c8294/watchfiles-1.1.0-cp313-cp313t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:8c5701dc474b041e2934a26d31d39f90fac8a3dee2322b39f7729867f932b1d4", size = 476727, upload-time = "2025-06-15T19:06:01.086Z" }, - { url = "https://files.pythonhosted.org/packages/65/95/fe479b2664f19be4cf5ceeb21be05afd491d95f142e72d26a42f41b7c4f8/watchfiles-1.1.0-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b067915e3c3936966a8607f6fe5487df0c9c4afb85226613b520890049deea20", size = 451864, upload-time = "2025-06-15T19:06:02.144Z" }, - { url = "https://files.pythonhosted.org/packages/d3/8a/3c4af14b93a15ce55901cd7a92e1a4701910f1768c78fb30f61d2b79785b/watchfiles-1.1.0-cp313-cp313t-musllinux_1_1_aarch64.whl", hash = "sha256:9c733cda03b6d636b4219625a4acb5c6ffb10803338e437fb614fef9516825ef", size = 625626, upload-time = "2025-06-15T19:06:03.578Z" }, - { url = "https://files.pythonhosted.org/packages/da/f5/cf6aa047d4d9e128f4b7cde615236a915673775ef171ff85971d698f3c2c/watchfiles-1.1.0-cp313-cp313t-musllinux_1_1_x86_64.whl", hash = "sha256:cc08ef8b90d78bfac66f0def80240b0197008e4852c9f285907377b2947ffdcb", size = 622744, upload-time = "2025-06-15T19:06:05.066Z" }, - { url = "https://files.pythonhosted.org/packages/2c/00/70f75c47f05dea6fd30df90f047765f6fc2d6eb8b5a3921379b0b04defa2/watchfiles-1.1.0-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:9974d2f7dc561cce3bb88dfa8eb309dab64c729de85fba32e98d75cf24b66297", size = 402114, upload-time = "2025-06-15T19:06:06.186Z" }, - { url = "https://files.pythonhosted.org/packages/53/03/acd69c48db4a1ed1de26b349d94077cca2238ff98fd64393f3e97484cae6/watchfiles-1.1.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c68e9f1fcb4d43798ad8814c4c1b61547b014b667216cb754e606bfade587018", size = 393879, upload-time = "2025-06-15T19:06:07.369Z" }, - { url = "https://files.pythonhosted.org/packages/2f/c8/a9a2a6f9c8baa4eceae5887fecd421e1b7ce86802bcfc8b6a942e2add834/watchfiles-1.1.0-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:95ab1594377effac17110e1352989bdd7bdfca9ff0e5eeccd8c69c5389b826d0", size = 450026, upload-time = "2025-06-15T19:06:08.476Z" }, - { url = "https://files.pythonhosted.org/packages/fe/51/d572260d98388e6e2b967425c985e07d47ee6f62e6455cefb46a6e06eda5/watchfiles-1.1.0-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:fba9b62da882c1be1280a7584ec4515d0a6006a94d6e5819730ec2eab60ffe12", size = 457917, upload-time = "2025-06-15T19:06:09.988Z" }, - { url = "https://files.pythonhosted.org/packages/c6/2d/4258e52917bf9f12909b6ec314ff9636276f3542f9d3807d143f27309104/watchfiles-1.1.0-cp314-cp314-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3434e401f3ce0ed6b42569128b3d1e3af773d7ec18751b918b89cd49c14eaafb", size = 483602, upload-time = "2025-06-15T19:06:11.088Z" }, - { url = "https://files.pythonhosted.org/packages/84/99/bee17a5f341a4345fe7b7972a475809af9e528deba056f8963d61ea49f75/watchfiles-1.1.0-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:fa257a4d0d21fcbca5b5fcba9dca5a78011cb93c0323fb8855c6d2dfbc76eb77", size = 596758, upload-time = "2025-06-15T19:06:12.197Z" }, - { url = "https://files.pythonhosted.org/packages/40/76/e4bec1d59b25b89d2b0716b41b461ed655a9a53c60dc78ad5771fda5b3e6/watchfiles-1.1.0-cp314-cp314-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:7fd1b3879a578a8ec2076c7961076df540b9af317123f84569f5a9ddee64ce92", size = 477601, upload-time = "2025-06-15T19:06:13.391Z" }, - { url = "https://files.pythonhosted.org/packages/1f/fa/a514292956f4a9ce3c567ec0c13cce427c158e9f272062685a8a727d08fc/watchfiles-1.1.0-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:62cc7a30eeb0e20ecc5f4bd113cd69dcdb745a07c68c0370cea919f373f65d9e", size = 451936, upload-time = "2025-06-15T19:06:14.656Z" }, - { url = "https://files.pythonhosted.org/packages/32/5d/c3bf927ec3bbeb4566984eba8dd7a8eb69569400f5509904545576741f88/watchfiles-1.1.0-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:891c69e027748b4a73847335d208e374ce54ca3c335907d381fde4e41661b13b", size = 626243, upload-time = "2025-06-15T19:06:16.232Z" }, - { url = "https://files.pythonhosted.org/packages/e6/65/6e12c042f1a68c556802a84d54bb06d35577c81e29fba14019562479159c/watchfiles-1.1.0-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:12fe8eaffaf0faa7906895b4f8bb88264035b3f0243275e0bf24af0436b27259", size = 623073, upload-time = "2025-06-15T19:06:17.457Z" }, - { url = "https://files.pythonhosted.org/packages/89/ab/7f79d9bf57329e7cbb0a6fd4c7bd7d0cee1e4a8ef0041459f5409da3506c/watchfiles-1.1.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:bfe3c517c283e484843cb2e357dd57ba009cff351edf45fb455b5fbd1f45b15f", size = 400872, upload-time = "2025-06-15T19:06:18.57Z" }, - { url = "https://files.pythonhosted.org/packages/df/d5/3f7bf9912798e9e6c516094db6b8932df53b223660c781ee37607030b6d3/watchfiles-1.1.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:a9ccbf1f129480ed3044f540c0fdbc4ee556f7175e5ab40fe077ff6baf286d4e", size = 392877, upload-time = "2025-06-15T19:06:19.55Z" }, - { url = "https://files.pythonhosted.org/packages/0d/c5/54ec7601a2798604e01c75294770dbee8150e81c6e471445d7601610b495/watchfiles-1.1.0-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ba0e3255b0396cac3cc7bbace76404dd72b5438bf0d8e7cefa2f79a7f3649caa", size = 449645, upload-time = "2025-06-15T19:06:20.66Z" }, - { url = "https://files.pythonhosted.org/packages/0a/04/c2f44afc3b2fce21ca0b7802cbd37ed90a29874f96069ed30a36dfe57c2b/watchfiles-1.1.0-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:4281cd9fce9fc0a9dbf0fc1217f39bf9cf2b4d315d9626ef1d4e87b84699e7e8", size = 457424, upload-time = "2025-06-15T19:06:21.712Z" }, - { url = "https://files.pythonhosted.org/packages/9f/b0/eec32cb6c14d248095261a04f290636da3df3119d4040ef91a4a50b29fa5/watchfiles-1.1.0-cp314-cp314t-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:6d2404af8db1329f9a3c9b79ff63e0ae7131986446901582067d9304ae8aaf7f", size = 481584, upload-time = "2025-06-15T19:06:22.777Z" }, - { url = "https://files.pythonhosted.org/packages/d1/e2/ca4bb71c68a937d7145aa25709e4f5d68eb7698a25ce266e84b55d591bbd/watchfiles-1.1.0-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:e78b6ed8165996013165eeabd875c5dfc19d41b54f94b40e9fff0eb3193e5e8e", size = 596675, upload-time = "2025-06-15T19:06:24.226Z" }, - { url = "https://files.pythonhosted.org/packages/a1/dd/b0e4b7fb5acf783816bc950180a6cd7c6c1d2cf7e9372c0ea634e722712b/watchfiles-1.1.0-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:249590eb75ccc117f488e2fabd1bfa33c580e24b96f00658ad88e38844a040bb", size = 477363, upload-time = "2025-06-15T19:06:25.42Z" }, - { url = "https://files.pythonhosted.org/packages/69/c4/088825b75489cb5b6a761a4542645718893d395d8c530b38734f19da44d2/watchfiles-1.1.0-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d05686b5487cfa2e2c28ff1aa370ea3e6c5accfe6435944ddea1e10d93872147", size = 452240, upload-time = "2025-06-15T19:06:26.552Z" }, - { url = "https://files.pythonhosted.org/packages/10/8c/22b074814970eeef43b7c44df98c3e9667c1f7bf5b83e0ff0201b0bd43f9/watchfiles-1.1.0-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:d0e10e6f8f6dc5762adee7dece33b722282e1f59aa6a55da5d493a97282fedd8", size = 625607, upload-time = "2025-06-15T19:06:27.606Z" }, - { url = "https://files.pythonhosted.org/packages/32/fa/a4f5c2046385492b2273213ef815bf71a0d4c1943b784fb904e184e30201/watchfiles-1.1.0-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:af06c863f152005c7592df1d6a7009c836a247c9d8adb78fef8575a5a98699db", size = 623315, upload-time = "2025-06-15T19:06:29.076Z" }, - { url = "https://files.pythonhosted.org/packages/8c/6b/686dcf5d3525ad17b384fd94708e95193529b460a1b7bf40851f1328ec6e/watchfiles-1.1.0-pp311-pypy311_pp73-macosx_10_12_x86_64.whl", hash = "sha256:0ece16b563b17ab26eaa2d52230c9a7ae46cf01759621f4fbbca280e438267b3", size = 406910, upload-time = "2025-06-15T19:06:49.335Z" }, - { url = "https://files.pythonhosted.org/packages/f3/d3/71c2dcf81dc1edcf8af9f4d8d63b1316fb0a2dd90cbfd427e8d9dd584a90/watchfiles-1.1.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:51b81e55d40c4b4aa8658427a3ee7ea847c591ae9e8b81ef94a90b668999353c", size = 398816, upload-time = "2025-06-15T19:06:50.433Z" }, - { url = "https://files.pythonhosted.org/packages/b8/fa/12269467b2fc006f8fce4cd6c3acfa77491dd0777d2a747415f28ccc8c60/watchfiles-1.1.0-pp311-pypy311_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f2bcdc54ea267fe72bfc7d83c041e4eb58d7d8dc6f578dfddb52f037ce62f432", size = 451584, upload-time = "2025-06-15T19:06:51.834Z" }, - { url = "https://files.pythonhosted.org/packages/bd/d3/254cea30f918f489db09d6a8435a7de7047f8cb68584477a515f160541d6/watchfiles-1.1.0-pp311-pypy311_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:923fec6e5461c42bd7e3fd5ec37492c6f3468be0499bc0707b4bbbc16ac21792", size = 454009, upload-time = "2025-06-15T19:06:52.896Z" }, -] - [[package]] name = "wcwidth" version = "0.2.14" @@ -2067,48 +2388,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/af/b5/123f13c975e9f27ab9c0770f514345bd406d0e8d3b7a0723af9d43f710af/wcwidth-0.2.14-py2.py3-none-any.whl", hash = "sha256:a7bb560c8aee30f9957e5f9895805edd20602f2d7f720186dfd906e82b4982e1", size = 37286, upload-time = "2025-09-22T16:29:51.641Z" }, ] -[[package]] -name = "websockets" -version = "15.0.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/21/e6/26d09fab466b7ca9c7737474c52be4f76a40301b08362eb2dbc19dcc16c1/websockets-15.0.1.tar.gz", hash = "sha256:82544de02076bafba038ce055ee6412d68da13ab47f0c60cab827346de828dee", size = 177016, upload-time = "2025-03-05T20:03:41.606Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/9f/32/18fcd5919c293a398db67443acd33fde142f283853076049824fc58e6f75/websockets-15.0.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:823c248b690b2fd9303ba00c4f66cd5e2d8c3ba4aa968b2779be9532a4dad431", size = 175423, upload-time = "2025-03-05T20:01:56.276Z" }, - { url = "https://files.pythonhosted.org/packages/76/70/ba1ad96b07869275ef42e2ce21f07a5b0148936688c2baf7e4a1f60d5058/websockets-15.0.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:678999709e68425ae2593acf2e3ebcbcf2e69885a5ee78f9eb80e6e371f1bf57", size = 173082, upload-time = "2025-03-05T20:01:57.563Z" }, - { url = "https://files.pythonhosted.org/packages/86/f2/10b55821dd40eb696ce4704a87d57774696f9451108cff0d2824c97e0f97/websockets-15.0.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:d50fd1ee42388dcfb2b3676132c78116490976f1300da28eb629272d5d93e905", size = 173330, upload-time = "2025-03-05T20:01:59.063Z" }, - { url = "https://files.pythonhosted.org/packages/a5/90/1c37ae8b8a113d3daf1065222b6af61cc44102da95388ac0018fcb7d93d9/websockets-15.0.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d99e5546bf73dbad5bf3547174cd6cb8ba7273062a23808ffea025ecb1cf8562", size = 182878, upload-time = "2025-03-05T20:02:00.305Z" }, - { url = "https://files.pythonhosted.org/packages/8e/8d/96e8e288b2a41dffafb78e8904ea7367ee4f891dafc2ab8d87e2124cb3d3/websockets-15.0.1-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:66dd88c918e3287efc22409d426c8f729688d89a0c587c88971a0faa2c2f3792", size = 181883, upload-time = "2025-03-05T20:02:03.148Z" }, - { url = "https://files.pythonhosted.org/packages/93/1f/5d6dbf551766308f6f50f8baf8e9860be6182911e8106da7a7f73785f4c4/websockets-15.0.1-cp311-cp311-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8dd8327c795b3e3f219760fa603dcae1dcc148172290a8ab15158cf85a953413", size = 182252, upload-time = "2025-03-05T20:02:05.29Z" }, - { url = "https://files.pythonhosted.org/packages/d4/78/2d4fed9123e6620cbf1706c0de8a1632e1a28e7774d94346d7de1bba2ca3/websockets-15.0.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:8fdc51055e6ff4adeb88d58a11042ec9a5eae317a0a53d12c062c8a8865909e8", size = 182521, upload-time = "2025-03-05T20:02:07.458Z" }, - { url = "https://files.pythonhosted.org/packages/e7/3b/66d4c1b444dd1a9823c4a81f50231b921bab54eee2f69e70319b4e21f1ca/websockets-15.0.1-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:693f0192126df6c2327cce3baa7c06f2a117575e32ab2308f7f8216c29d9e2e3", size = 181958, upload-time = "2025-03-05T20:02:09.842Z" }, - { url = "https://files.pythonhosted.org/packages/08/ff/e9eed2ee5fed6f76fdd6032ca5cd38c57ca9661430bb3d5fb2872dc8703c/websockets-15.0.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:54479983bd5fb469c38f2f5c7e3a24f9a4e70594cd68cd1fa6b9340dadaff7cf", size = 181918, upload-time = "2025-03-05T20:02:11.968Z" }, - { url = "https://files.pythonhosted.org/packages/d8/75/994634a49b7e12532be6a42103597b71098fd25900f7437d6055ed39930a/websockets-15.0.1-cp311-cp311-win32.whl", hash = "sha256:16b6c1b3e57799b9d38427dda63edcbe4926352c47cf88588c0be4ace18dac85", size = 176388, upload-time = "2025-03-05T20:02:13.32Z" }, - { url = "https://files.pythonhosted.org/packages/98/93/e36c73f78400a65f5e236cd376713c34182e6663f6889cd45a4a04d8f203/websockets-15.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:27ccee0071a0e75d22cb35849b1db43f2ecd3e161041ac1ee9d2352ddf72f065", size = 176828, upload-time = "2025-03-05T20:02:14.585Z" }, - { url = "https://files.pythonhosted.org/packages/51/6b/4545a0d843594f5d0771e86463606a3988b5a09ca5123136f8a76580dd63/websockets-15.0.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:3e90baa811a5d73f3ca0bcbf32064d663ed81318ab225ee4f427ad4e26e5aff3", size = 175437, upload-time = "2025-03-05T20:02:16.706Z" }, - { url = "https://files.pythonhosted.org/packages/f4/71/809a0f5f6a06522af902e0f2ea2757f71ead94610010cf570ab5c98e99ed/websockets-15.0.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:592f1a9fe869c778694f0aa806ba0374e97648ab57936f092fd9d87f8bc03665", size = 173096, upload-time = "2025-03-05T20:02:18.832Z" }, - { url = "https://files.pythonhosted.org/packages/3d/69/1a681dd6f02180916f116894181eab8b2e25b31e484c5d0eae637ec01f7c/websockets-15.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:0701bc3cfcb9164d04a14b149fd74be7347a530ad3bbf15ab2c678a2cd3dd9a2", size = 173332, upload-time = "2025-03-05T20:02:20.187Z" }, - { url = "https://files.pythonhosted.org/packages/a6/02/0073b3952f5bce97eafbb35757f8d0d54812b6174ed8dd952aa08429bcc3/websockets-15.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e8b56bdcdb4505c8078cb6c7157d9811a85790f2f2b3632c7d1462ab5783d215", size = 183152, upload-time = "2025-03-05T20:02:22.286Z" }, - { url = "https://files.pythonhosted.org/packages/74/45/c205c8480eafd114b428284840da0b1be9ffd0e4f87338dc95dc6ff961a1/websockets-15.0.1-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0af68c55afbd5f07986df82831c7bff04846928ea8d1fd7f30052638788bc9b5", size = 182096, upload-time = "2025-03-05T20:02:24.368Z" }, - { url = "https://files.pythonhosted.org/packages/14/8f/aa61f528fba38578ec553c145857a181384c72b98156f858ca5c8e82d9d3/websockets-15.0.1-cp312-cp312-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:64dee438fed052b52e4f98f76c5790513235efaa1ef7f3f2192c392cd7c91b65", size = 182523, upload-time = "2025-03-05T20:02:25.669Z" }, - { url = "https://files.pythonhosted.org/packages/ec/6d/0267396610add5bc0d0d3e77f546d4cd287200804fe02323797de77dbce9/websockets-15.0.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:d5f6b181bb38171a8ad1d6aa58a67a6aa9d4b38d0f8c5f496b9e42561dfc62fe", size = 182790, upload-time = "2025-03-05T20:02:26.99Z" }, - { url = "https://files.pythonhosted.org/packages/02/05/c68c5adbf679cf610ae2f74a9b871ae84564462955d991178f95a1ddb7dd/websockets-15.0.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:5d54b09eba2bada6011aea5375542a157637b91029687eb4fdb2dab11059c1b4", size = 182165, upload-time = "2025-03-05T20:02:30.291Z" }, - { url = "https://files.pythonhosted.org/packages/29/93/bb672df7b2f5faac89761cb5fa34f5cec45a4026c383a4b5761c6cea5c16/websockets-15.0.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:3be571a8b5afed347da347bfcf27ba12b069d9d7f42cb8c7028b5e98bbb12597", size = 182160, upload-time = "2025-03-05T20:02:31.634Z" }, - { url = "https://files.pythonhosted.org/packages/ff/83/de1f7709376dc3ca9b7eeb4b9a07b4526b14876b6d372a4dc62312bebee0/websockets-15.0.1-cp312-cp312-win32.whl", hash = "sha256:c338ffa0520bdb12fbc527265235639fb76e7bc7faafbb93f6ba80d9c06578a9", size = 176395, upload-time = "2025-03-05T20:02:33.017Z" }, - { url = "https://files.pythonhosted.org/packages/7d/71/abf2ebc3bbfa40f391ce1428c7168fb20582d0ff57019b69ea20fa698043/websockets-15.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:fcd5cf9e305d7b8338754470cf69cf81f420459dbae8a3b40cee57417f4614a7", size = 176841, upload-time = "2025-03-05T20:02:34.498Z" }, - { url = "https://files.pythonhosted.org/packages/cb/9f/51f0cf64471a9d2b4d0fc6c534f323b664e7095640c34562f5182e5a7195/websockets-15.0.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:ee443ef070bb3b6ed74514f5efaa37a252af57c90eb33b956d35c8e9c10a1931", size = 175440, upload-time = "2025-03-05T20:02:36.695Z" }, - { url = "https://files.pythonhosted.org/packages/8a/05/aa116ec9943c718905997412c5989f7ed671bc0188ee2ba89520e8765d7b/websockets-15.0.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:5a939de6b7b4e18ca683218320fc67ea886038265fd1ed30173f5ce3f8e85675", size = 173098, upload-time = "2025-03-05T20:02:37.985Z" }, - { url = "https://files.pythonhosted.org/packages/ff/0b/33cef55ff24f2d92924923c99926dcce78e7bd922d649467f0eda8368923/websockets-15.0.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:746ee8dba912cd6fc889a8147168991d50ed70447bf18bcda7039f7d2e3d9151", size = 173329, upload-time = "2025-03-05T20:02:39.298Z" }, - { url = "https://files.pythonhosted.org/packages/31/1d/063b25dcc01faa8fada1469bdf769de3768b7044eac9d41f734fd7b6ad6d/websockets-15.0.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:595b6c3969023ecf9041b2936ac3827e4623bfa3ccf007575f04c5a6aa318c22", size = 183111, upload-time = "2025-03-05T20:02:40.595Z" }, - { url = "https://files.pythonhosted.org/packages/93/53/9a87ee494a51bf63e4ec9241c1ccc4f7c2f45fff85d5bde2ff74fcb68b9e/websockets-15.0.1-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3c714d2fc58b5ca3e285461a4cc0c9a66bd0e24c5da9911e30158286c9b5be7f", size = 182054, upload-time = "2025-03-05T20:02:41.926Z" }, - { url = "https://files.pythonhosted.org/packages/ff/b2/83a6ddf56cdcbad4e3d841fcc55d6ba7d19aeb89c50f24dd7e859ec0805f/websockets-15.0.1-cp313-cp313-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0f3c1e2ab208db911594ae5b4f79addeb3501604a165019dd221c0bdcabe4db8", size = 182496, upload-time = "2025-03-05T20:02:43.304Z" }, - { url = "https://files.pythonhosted.org/packages/98/41/e7038944ed0abf34c45aa4635ba28136f06052e08fc2168520bb8b25149f/websockets-15.0.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:229cf1d3ca6c1804400b0a9790dc66528e08a6a1feec0d5040e8b9eb14422375", size = 182829, upload-time = "2025-03-05T20:02:48.812Z" }, - { url = "https://files.pythonhosted.org/packages/e0/17/de15b6158680c7623c6ef0db361da965ab25d813ae54fcfeae2e5b9ef910/websockets-15.0.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:756c56e867a90fb00177d530dca4b097dd753cde348448a1012ed6c5131f8b7d", size = 182217, upload-time = "2025-03-05T20:02:50.14Z" }, - { url = "https://files.pythonhosted.org/packages/33/2b/1f168cb6041853eef0362fb9554c3824367c5560cbdaad89ac40f8c2edfc/websockets-15.0.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:558d023b3df0bffe50a04e710bc87742de35060580a293c2a984299ed83bc4e4", size = 182195, upload-time = "2025-03-05T20:02:51.561Z" }, - { url = "https://files.pythonhosted.org/packages/86/eb/20b6cdf273913d0ad05a6a14aed4b9a85591c18a987a3d47f20fa13dcc47/websockets-15.0.1-cp313-cp313-win32.whl", hash = "sha256:ba9e56e8ceeeedb2e080147ba85ffcd5cd0711b89576b83784d8605a7df455fa", size = 176393, upload-time = "2025-03-05T20:02:53.814Z" }, - { url = "https://files.pythonhosted.org/packages/1b/6c/c65773d6cab416a64d191d6ee8a8b1c68a09970ea6909d16965d26bfed1e/websockets-15.0.1-cp313-cp313-win_amd64.whl", hash = "sha256:e09473f095a819042ecb2ab9465aee615bd9c2028e4ef7d933600a8401c79561", size = 176837, upload-time = "2025-03-05T20:02:55.237Z" }, - { url = "https://files.pythonhosted.org/packages/fa/a8/5b41e0da817d64113292ab1f8247140aac61cbf6cfd085d6a0fa77f4984f/websockets-15.0.1-py3-none-any.whl", hash = "sha256:f7a866fbc1e97b5c617ee4116daaa09b722101d4a3c170c787450ba409f9736f", size = 169743, upload-time = "2025-03-05T20:03:39.41Z" }, -] - [[package]] name = "workspace" version = "0.1.0" @@ -2121,7 +2400,6 @@ dependencies = [ [package.dev-dependencies] dev = [ - { name = "basedpyright" }, { name = "filelock" }, { name = "ipython" }, { name = "pip" }, @@ -2139,7 +2417,6 @@ requires-dist = [ [package.metadata.requires-dev] dev = [ - { name = "basedpyright", specifier = ">=1.31" }, { name = "filelock", specifier = ">=3.20.3" }, { name = "ipython", specifier = ">=8.26.0" }, { name = "pip", specifier = ">=23.3.2" }, @@ -2147,3 +2424,240 @@ dev = [ { name = "ruff", specifier = ">=0.7.1" }, { name = "urllib3", specifier = ">=2.6.3" }, ] + +[[package]] +name = "xxhash" +version = "3.8.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/8e/63/71aa56b151a1b28770037a61bd4e461c2619cfc8866a4fcaf1548605e325/xxhash-3.8.1.tar.gz", hash = "sha256:b0de4bf3aa66363552d52c6a89003c479911f12098cd48a53d44a0f7a25f7c46", size = 86223, upload-time = "2026-07-06T10:49:58.937Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8a/5a/05eaa129555f85476a3e16ff869e95f81a78bbe4647eef9d0229f515a317/xxhash-3.8.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:602efcad4a42c184e81d43a2b7e6e4f524d619878f2b6ee2ba469011f47c8147", size = 34699, upload-time = "2026-07-06T10:44:10.14Z" }, + { url = "https://files.pythonhosted.org/packages/80/59/0df1133958b2228929355e022aab1e958c7b2c43e27bf7f59bc9edfa8a54/xxhash-3.8.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:131324f719957b988861714de7d6ddf57b47abec3b0cc691302ffeaba0e05e10", size = 32373, upload-time = "2026-07-06T10:44:11.353Z" }, + { url = "https://files.pythonhosted.org/packages/3e/bf/1cfda5b5e6bf26617812b4a31662ef2220d2ad04e0a55b8ff9eb36e56a5c/xxhash-3.8.1-cp311-cp311-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:db77278a6eddadbf44ce5aae2fee5ebb4d061f026b1ce2130d058cd4d7a7b670", size = 220284, upload-time = "2026-07-06T10:44:12.683Z" }, + { url = "https://files.pythonhosted.org/packages/70/93/45dc0ad7913b69e5b08bd039236cf628380e4c9cc76a8a4c6625a328e058/xxhash-3.8.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1c332dd48b8cb050da2bb2a3c96d72b1664168650a250ef9718e423df7989e05", size = 240980, upload-time = "2026-07-06T10:44:14.297Z" }, + { url = "https://files.pythonhosted.org/packages/e9/02/f28ba7d17f2c1410ee397982c817ab1bd5b2701070c2d2c373539aad000a/xxhash-3.8.1-cp311-cp311-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:a5cd96f6dcdf4fa657b2d95668d71d58455248f98712ecffaa9c528edf40ccae", size = 264526, upload-time = "2026-07-06T10:44:16.017Z" }, + { url = "https://files.pythonhosted.org/packages/5c/d0/f10651cec2c7981b20d693deae6bdfc438427d92be2db4ccabb6181f0021/xxhash-3.8.1-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c959f88160b13b4e730b0d75b459b7929fc0d2225c284c9683ac95d6feeeac6a", size = 241369, upload-time = "2026-07-06T10:44:17.698Z" }, + { url = "https://files.pythonhosted.org/packages/ff/40/136e0cbaf5db51e191423b1c98643593189f02b6cd90837bf64b19113d70/xxhash-3.8.1-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:027dee4355f3fcc41481650d846cf6cfc895c85a1ab7acd063063821a0df5b4c", size = 473186, upload-time = "2026-07-06T10:44:19.354Z" }, + { url = "https://files.pythonhosted.org/packages/4b/3f/6aa808a96bdc43dba9a740dec56c744526ee3c0019e32c75e810fa90ae4d/xxhash-3.8.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ad52a0e4bcc0ba956a953a169d1feec2734a64981d689e4fc8f490f7bf91af60", size = 220092, upload-time = "2026-07-06T10:44:20.956Z" }, + { url = "https://files.pythonhosted.org/packages/47/28/a8675e78a9ced96dab853416162268e10e05b452e95db7888cf69f58ac5f/xxhash-3.8.1-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:5d3dfb1f0ff146da7952867a9414f0c7a29762f8825a84879592612fd6139342", size = 309846, upload-time = "2026-07-06T10:44:22.543Z" }, + { url = "https://files.pythonhosted.org/packages/89/0f/7fe4d4ef4e69f0033e012396ee2a115886bca7b10b7e45ce398626436bfc/xxhash-3.8.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:4482380b462ca9e59994d072a877ecadd1cf51102daeeab2db696f96ab763723", size = 237659, upload-time = "2026-07-06T10:44:24.135Z" }, + { url = "https://files.pythonhosted.org/packages/38/8f/83e9e31d4ed57fe963b99cb5b13a23e3e0f0dad1885aa0ebd2a7819dd423/xxhash-3.8.1-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:950ac754d16daea42038f38e7465eb84cda4d08d7343c1c915771b29470f065a", size = 268737, upload-time = "2026-07-06T10:44:25.875Z" }, + { url = "https://files.pythonhosted.org/packages/57/79/7e7de46dbe5d1f49afc96a0bc42e6b8df24eae3d6bad6007b99e42f48430/xxhash-3.8.1-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:0418ec8b2331b9d4d575fc9284427e8e69449d7172e99e1a86fcdd1f51a0a937", size = 224955, upload-time = "2026-07-06T10:44:27.777Z" }, + { url = "https://files.pythonhosted.org/packages/ec/34/b8540839e958d5ef5c6101af6f16032109e7099698ae8edbc8dcefe4d8f4/xxhash-3.8.1-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:32a94ad2763e0263d9102037d349002c3d3c401e42770542c3eeb4801f311661", size = 239653, upload-time = "2026-07-06T10:44:29.422Z" }, + { url = "https://files.pythonhosted.org/packages/ce/87/a735d05f7f859354acadabe470ff40e2c46672275f96dcf096a761904def/xxhash-3.8.1-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:89b11a5cdd441aa463f6d34ca0241602bc09b001a76994b6059828494108c673", size = 300213, upload-time = "2026-07-06T10:44:31.401Z" }, + { url = "https://files.pythonhosted.org/packages/98/31/3e1cb020237b68117fc212dc5f9753b87f865b4dfee7c1ce62d0836955b5/xxhash-3.8.1-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:09a204dd4bb0823daf938cdd0dc8057d5f1e14fe3cbde929424255f23f9de872", size = 442508, upload-time = "2026-07-06T10:44:33.023Z" }, + { url = "https://files.pythonhosted.org/packages/23/bf/f80090622141cc734b039ce1d15ce3ff6dced375e9680249bf5b9b8c6bf9/xxhash-3.8.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:e710ad822c493fb80a4fbc1e3d0a807b1422cb90adbe64378f98291b7fa48fef", size = 216853, upload-time = "2026-07-06T10:44:34.983Z" }, + { url = "https://files.pythonhosted.org/packages/a6/a3/60157acecc307b238d3651c2483168e224b48b23a36ae6d6903588341d80/xxhash-3.8.1-cp311-cp311-win32.whl", hash = "sha256:5013be3bea7612852c62a7437f3302c1cfb91ca7e703b194459db0b2b2e0d792", size = 31936, upload-time = "2026-07-06T10:44:36.542Z" }, + { url = "https://files.pythonhosted.org/packages/59/5c/ef70c418d878d187b8da56d4cdc06aea6cf5e456b301e96e51e1d2cc8625/xxhash-3.8.1-cp311-cp311-win_amd64.whl", hash = "sha256:f377012b86c0a23a1df0cf5a1b05aa7187649e472f71c7892e5f2c2815bbe74f", size = 32724, upload-time = "2026-07-06T10:44:38.177Z" }, + { url = "https://files.pythonhosted.org/packages/2c/25/f008db952cec6b2a26445b456eeed2ebebd65e08e848ebe09ed6ac0634e6/xxhash-3.8.1-cp311-cp311-win_arm64.whl", hash = "sha256:836f11d4474d3228e9909d97216faa4f7505df41cfaf3927eb29809de785a78d", size = 29212, upload-time = "2026-07-06T10:44:39.577Z" }, + { url = "https://files.pythonhosted.org/packages/42/91/f65c34a7aa7b4e7cf4854f8e6ef3f7ee32ceac41d4f008da0780db0612f6/xxhash-3.8.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:e6e49370822c1f4d8d90e678b06dbcb08b51a026a7c4b55479e7d467f2e813bc", size = 34680, upload-time = "2026-07-06T10:44:40.932Z" }, + { url = "https://files.pythonhosted.org/packages/57/04/b10a245a4c09a9cfa88f8e9ae755029413ad1ac17047f9a61906e5ae0799/xxhash-3.8.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:220d68130f83f7cc86d6edfdeab176adc73d7200bf3a8ec10c629e8cf605c215", size = 32397, upload-time = "2026-07-06T10:44:42.196Z" }, + { url = "https://files.pythonhosted.org/packages/3a/75/45ab795b5945b6388583bd75202106af505537935566c15a1577797a0e08/xxhash-3.8.1-cp312-cp312-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:4d365ee1892c1fa803536f8c6ce21d24b29c9718ec75eb856095c07830f8c478", size = 220549, upload-time = "2026-07-06T10:44:43.603Z" }, + { url = "https://files.pythonhosted.org/packages/13/44/5ba2bd0a14ddf4193fc7d8ec29625f659f22c06d60b28f04bf46305d8330/xxhash-3.8.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:852bfe059720632e2f16a6a4745e41d20937b2bf2a42a401e2412046bb6971cc", size = 241186, upload-time = "2026-07-06T10:44:45.534Z" }, + { url = "https://files.pythonhosted.org/packages/23/32/c4147def4d1e4538b906f82731e0ba23424377fc50a7cddd03cd284c8f63/xxhash-3.8.1-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:2f8c25a7061d952de589bd0ea0eaadee32378ff83dd6a677b267f9cd86f401f8", size = 264852, upload-time = "2026-07-06T10:44:47.199Z" }, + { url = "https://files.pythonhosted.org/packages/6c/bd/71ed14f4f0318bb7fd7b2ec51999413487fa8da8d41208e84d50d1ef0f98/xxhash-3.8.1-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:868a8dcaff1a84ba78038e1cef14fc88ccf84d9b4d12ea604696e0693296aa56", size = 242663, upload-time = "2026-07-06T10:44:48.846Z" }, + { url = "https://files.pythonhosted.org/packages/91/09/70af22c565a8473b3f2ae73f88e7721af281bc4a575236dbd1970c9f76f6/xxhash-3.8.1-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:6536d8677d2fff7e64cd0b98b976df9de7aee0e69590044c2af5f51b76b7a170", size = 473510, upload-time = "2026-07-06T10:44:50.695Z" }, + { url = "https://files.pythonhosted.org/packages/18/96/34db781c8f0cf99c544ca1f2bc2e5bf55426e1eb4ca6de8ea5da56a9f352/xxhash-3.8.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:82c0cedd280eab2e8291270e6c04894dbc096f8159a39dcf1807429f026ca3cc", size = 220469, upload-time = "2026-07-06T10:44:52.422Z" }, + { url = "https://files.pythonhosted.org/packages/93/5f/9a184f615fa5a4dce30c01534f62946ce5a11ce40f73785cbd356ccabaa9/xxhash-3.8.1-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:daa86e4b68221d38e669bb236ba112d0335353829fb627c82e5909e4bbe8694c", size = 310290, upload-time = "2026-07-06T10:44:54.142Z" }, + { url = "https://files.pythonhosted.org/packages/a9/dc/9b9a9789011ee153723a5eb9e7dd7fcbae2ba9b3fe7a729249ca7c252056/xxhash-3.8.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:2bc7113e6f2b6b3922dd61796ca9f36af09da3773898e7003038dc992fc83b8d", size = 238173, upload-time = "2026-07-06T10:44:55.693Z" }, + { url = "https://files.pythonhosted.org/packages/ec/4d/71c6005ada9dcb608a4e1902e8475ecadb5f3fbfa04e1e244d276a2d0c43/xxhash-3.8.1-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:5eed32dad81d6ba8e62dc7b9ffa0500199385d7810a8dd9d4eafaceb8c6e20bb", size = 269026, upload-time = "2026-07-06T10:44:57.424Z" }, + { url = "https://files.pythonhosted.org/packages/2f/87/d6c036ba25dfbd9c8633be5aa86fc9474bbb9e2c68212a841d090abe7344/xxhash-3.8.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:83697b0ea1f10e7f5d8b26a4906fa851393c61546c63839643a2b7fe2d868061", size = 224970, upload-time = "2026-07-06T10:44:59.085Z" }, + { url = "https://files.pythonhosted.org/packages/48/62/4c1f035a41c5752aa05e195b6c904c07b94fe9061a16de61e72a6e6b135f/xxhash-3.8.1-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:36fc69160465ae75c6ec4ac9f781bb2aa16ae7ff869e73c26fee85fbb11b9887", size = 240820, upload-time = "2026-07-06T10:45:00.746Z" }, + { url = "https://files.pythonhosted.org/packages/da/14/d39d565069b87e86d21a2af2a31d04db79249d25aa8d5b62959056a89857/xxhash-3.8.1-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:445e0f5a31f2f3546ae0895d4811e159518cdc9d824c11419898d40cfadb677e", size = 300619, upload-time = "2026-07-06T10:45:02.716Z" }, + { url = "https://files.pythonhosted.org/packages/13/22/75467acc887edc8cf71c97ab1708feb3df7a88bda589b9f399765c6387d2/xxhash-3.8.1-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:dfe0580fbfd5e4af87d0cc52d2044f155d55ebd8c8a93568758a2ea7d8e15975", size = 443267, upload-time = "2026-07-06T10:45:04.653Z" }, + { url = "https://files.pythonhosted.org/packages/a4/b6/1da3baa5fa6ef705e3425fddd382be7dfc4dfba2686df90a20f16e9c7b1b/xxhash-3.8.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:095e1323fa108be1292c54c86da3ef3c7a7dc015b105a52133973bc07a6ad11a", size = 217338, upload-time = "2026-07-06T10:45:06.304Z" }, + { url = "https://files.pythonhosted.org/packages/78/dd/b5295a9f97484e7a1c2b283a742ca45e3104991c55a1ef670dde161829ba/xxhash-3.8.1-cp312-cp312-win32.whl", hash = "sha256:bf28f55e427e0483acb1f666bd0d869b6d5e5a716680c216ad7befe3d4cfba2e", size = 31970, upload-time = "2026-07-06T10:45:07.823Z" }, + { url = "https://files.pythonhosted.org/packages/ec/31/3fa0b807d7e21515cd975e7fe5c039d52ac3e9401a96d6ad68dae6305215/xxhash-3.8.1-cp312-cp312-win_amd64.whl", hash = "sha256:2256e80e4960ee282f63428adb349cb7f8bd8efe4db770d88eb815f4b9860724", size = 32741, upload-time = "2026-07-06T10:45:09.42Z" }, + { url = "https://files.pythonhosted.org/packages/b8/05/86feada74e239600e6875aa507afb40482a89b92700aa74a92da83bdcb77/xxhash-3.8.1-cp312-cp312-win_arm64.whl", hash = "sha256:9df56e6df96a60590935e22373041cccc91fd55858763dcffb55bf63b3a2b396", size = 29234, upload-time = "2026-07-06T10:45:10.809Z" }, + { url = "https://files.pythonhosted.org/packages/6b/8c/446bb782cd0d27007a917b5569a08dd73219c3e8d6e459014db104b27bdb/xxhash-3.8.1-cp313-cp313-android_21_arm64_v8a.whl", hash = "sha256:3c682fcd96eb4bf64be32a4d95f96107e1588005831bd8a741b324fdda01b913", size = 38562, upload-time = "2026-07-06T10:45:12.425Z" }, + { url = "https://files.pythonhosted.org/packages/d7/ec/c0c45627eaa6be7a5d6117423adf8f7a15b17ee74b4b17072cca5959a225/xxhash-3.8.1-cp313-cp313-android_21_x86_64.whl", hash = "sha256:036a024d8b9c01f70782e09ed98d532e76fd23f950ae7154bd950fe94e90ebec", size = 36656, upload-time = "2026-07-06T10:45:13.932Z" }, + { url = "https://files.pythonhosted.org/packages/f6/94/8324c04cc7597154caaeba6c094e01fbd2e7601d01e7a13eea9f5420e77b/xxhash-3.8.1-cp313-cp313-ios_13_0_arm64_iphoneos.whl", hash = "sha256:d6a5c0bce213b23b0166fe0d35bcbbe23ce4b968f257cc7eb6fd57cb8e1e6297", size = 31169, upload-time = "2026-07-06T10:45:15.687Z" }, + { url = "https://files.pythonhosted.org/packages/40/a4/beb6bb26e1184e126dbe7a5682330214ef54dcfbf882078aa9f4b5428d42/xxhash-3.8.1-cp313-cp313-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:5177aa44eddaa97c6ef0cc00c6d540edb64d51781d2f8fb941612ec61a92c9ed", size = 32177, upload-time = "2026-07-06T10:45:17.035Z" }, + { url = "https://files.pythonhosted.org/packages/56/0f/fc4c92a5a528f839b34b6419b2e53c8597f2a629d5a1f5d721f65bfa1fd6/xxhash-3.8.1-cp313-cp313-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:7801b7223db017b9c0c9ccf37e44524edb35a1544a1c032add22c061c6af0276", size = 34642, upload-time = "2026-07-06T10:45:18.39Z" }, + { url = "https://files.pythonhosted.org/packages/d4/58/edbfb141d4000767ac6a9694f8ac0763e2c2e983e65c9e31620ba56e2667/xxhash-3.8.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:9e80238259655bf69d7bcd08226a970d7f42605f3157786bfa76dd13472d7fa0", size = 34684, upload-time = "2026-07-06T10:45:20.033Z" }, + { url = "https://files.pythonhosted.org/packages/07/3f/5072f1f0f5714186f0ac2a0b5a4929ce30d4b845e94886b6c01b6ebda0be/xxhash-3.8.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:bcab50a389cc04d87f90092af78a6adba2ab3deca63175a3344ca83514045315", size = 32401, upload-time = "2026-07-06T10:45:21.414Z" }, + { url = "https://files.pythonhosted.org/packages/49/c7/802ea2f9c2ed59219934d6d65c470d502b1788043eae277a52af8658bda6/xxhash-3.8.1-cp313-cp313-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:a2489d3a776fa380cb8e71f54c7fda268a9baf3de9b1395093fd280f95735907", size = 220617, upload-time = "2026-07-06T10:45:23.234Z" }, + { url = "https://files.pythonhosted.org/packages/99/a8/e10488efd31fcb13fcd6acbc6e788f10c6f8e3a0cc4ae3eb89dc19c55a12/xxhash-3.8.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:32ab1e5432690276e71192be7401b55f96db2d0eedea5d44eb1f164505669cc0", size = 241295, upload-time = "2026-07-06T10:45:25.364Z" }, + { url = "https://files.pythonhosted.org/packages/18/cc/14180b17d44892a631f8ae7323c30bfbb1328efc8209e528a480293528ac/xxhash-3.8.1-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:b30e01a0b97a4bc3f519a4d7a82da3dc53251fb0de5eeea8660dcd4ff094c0c2", size = 264688, upload-time = "2026-07-06T10:45:27.09Z" }, + { url = "https://files.pythonhosted.org/packages/a9/72/a14019d0c5f6c41ee407a503036ae32787c91325ca218a96a9b5627be651/xxhash-3.8.1-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:1f44275ddb0978b67a58a951501903f04d49335a91f7681c9ce122ecb8ccb329", size = 242740, upload-time = "2026-07-06T10:45:28.753Z" }, + { url = "https://files.pythonhosted.org/packages/68/08/92550e556c6fcfcb96c6a336945eb53a431ed43120ed749636debb16c5cf/xxhash-3.8.1-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:e3b87cbd974512c0c5fc7b469c36b2cdc9ee6d76e4ec78bccb2c7184611c49b0", size = 473599, upload-time = "2026-07-06T10:45:30.524Z" }, + { url = "https://files.pythonhosted.org/packages/29/83/e361d3c1acd1b21e1d489616de6fa4aaf843365d8179f612e3743eac20a9/xxhash-3.8.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:98ee81b4b7f3023c9cb04a78cc67610baffcb5812d92f2096cb5a5efc6f19437", size = 220559, upload-time = "2026-07-06T10:45:32.979Z" }, + { url = "https://files.pythonhosted.org/packages/05/01/006a4243c2c2a6831827f9999f6d1c23feeef100eb023c1f886022a00bf3/xxhash-3.8.1-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:2666f059a1588a99267e33605365ed89cea92f424b3522806a9f4bd8ad2e3d62", size = 310383, upload-time = "2026-07-06T10:45:35.875Z" }, + { url = "https://files.pythonhosted.org/packages/d8/20/af388e8bf9f9a0f89eeef7d2a1935d176ee1c20bc6adeda05035879379cf/xxhash-3.8.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:b0093cf7eeb91b84776e8742113afa4bdf47533d36cf719179aaaf1f56f6f8bf", size = 238228, upload-time = "2026-07-06T10:45:38.02Z" }, + { url = "https://files.pythonhosted.org/packages/63/6b/4666579a87eebd1744663c404297355fa0658617b015cedfa58810ee7036/xxhash-3.8.1-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:3a800912a2e5e975d4128969d645c4a2a80aa886ccd6c9b1c6f44529e327e8cf", size = 269137, upload-time = "2026-07-06T10:45:39.954Z" }, + { url = "https://files.pythonhosted.org/packages/de/d3/e963a8a46f900a137d91b02144d8ea07a8f812971b138204a3b2f8b8e55c/xxhash-3.8.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:0fe37f72a207223d22a4eddc3149d4298993385aa9daef25c039246ca5a309f3", size = 225068, upload-time = "2026-07-06T10:45:41.718Z" }, + { url = "https://files.pythonhosted.org/packages/aa/80/9d181dbcde4b0fe48375f48833a5832d4b8cd2b349b15110c92ee472d874/xxhash-3.8.1-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:5db43f249b4be9f99ef4b967863f37094fb40e67effafb78ba4f0356b6396104", size = 240874, upload-time = "2026-07-06T10:45:43.414Z" }, + { url = "https://files.pythonhosted.org/packages/39/15/ce3ab5a1cd27ead25a5196e55a7284220f6ad6e316da494ffd900b2b600f/xxhash-3.8.1-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:c4ed42965c2cd9081f011be22f69d0e65d3b6165fe7734072fd0c232840bbd4e", size = 300702, upload-time = "2026-07-06T10:45:45.135Z" }, + { url = "https://files.pythonhosted.org/packages/96/c0/2281a8ab5f2a62dbf57a23c58a01ccc1d98abf40f71193c8a81f59e759b5/xxhash-3.8.1-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:3557bec8fcb11738a8920eeb68974bc76b75262f6947998d3147954ce0a4b893", size = 443351, upload-time = "2026-07-06T10:45:47.188Z" }, + { url = "https://files.pythonhosted.org/packages/81/2e/071a58c1a53a52d4f7a3aa0987be0c396dffd40da8204805fe1b130a81f4/xxhash-3.8.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:00de40f3b42240db23a82a5c682b55d7263d84a26a953240c1aee463409660e3", size = 217396, upload-time = "2026-07-06T10:45:48.925Z" }, + { url = "https://files.pythonhosted.org/packages/68/44/36ab58134badd9d3433fc7b53c4ca8d113d8e807782885628640f8297a4d/xxhash-3.8.1-cp313-cp313-win32.whl", hash = "sha256:b5196cc2574cfec572a5f3fb7cfa5ade27305ae3d06516a082132441aff4c83a", size = 31974, upload-time = "2026-07-06T10:45:50.591Z" }, + { url = "https://files.pythonhosted.org/packages/96/2a/2a0b84798448e766f7b89ceed073cb0cb5a43fc9ebbacbdea74a38de18e3/xxhash-3.8.1-cp313-cp313-win_amd64.whl", hash = "sha256:538f5f865df6cd8c32dd63158a0e5b4f5dd08d732a7da8b7228a5a0776c8ce55", size = 32739, upload-time = "2026-07-06T10:45:52.221Z" }, + { url = "https://files.pythonhosted.org/packages/d4/60/bb51dbf7c363ff88a7cbd50b7959718219577ef44d7cf255929ffc4a2194/xxhash-3.8.1-cp313-cp313-win_arm64.whl", hash = "sha256:a6617f30641ba0d8baa1635fbefb1dffc5165ec36d26921bd5cee13497cd937a", size = 29239, upload-time = "2026-07-06T10:45:53.714Z" }, + { url = "https://files.pythonhosted.org/packages/56/d3/827ca123c2ee5443a6aaed3c5dd199237dc2f010e2bebd7ec09ef36f3a5f/xxhash-3.8.1-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:bfcd82852c62a60e314670a9602de354c4460f8adad916e2e42a20860c7870bc", size = 34964, upload-time = "2026-07-06T10:45:55.535Z" }, + { url = "https://files.pythonhosted.org/packages/05/67/67ae2a3ccdeb8b8ef025d35aee9edd1d26c3abe5051d47da9286232afbf8/xxhash-3.8.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:08ea2081f5e88615fec8622a9f87fbe21b8ea58d88cfc02163ca11026ee62a92", size = 32697, upload-time = "2026-07-06T10:45:57.288Z" }, + { url = "https://files.pythonhosted.org/packages/38/5a/3d3994346e1f45493679cb5c1ffc2bf454e410e9d1e8a662d253becee91e/xxhash-3.8.1-cp313-cp313t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:2e32855b6f9e5b18f449e59d45e3d5778bdeb660632ef2693cca267a11246c75", size = 225954, upload-time = "2026-07-06T10:45:58.897Z" }, + { url = "https://files.pythonhosted.org/packages/3f/2c/53169270309b7cd8e05504e07fe123bac053b89d00ac63617faacf0a2ec0/xxhash-3.8.1-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a6e088bd7870775624256a0d84c2a6714afd223b2eeb56b0ca58398e52a32fda", size = 249776, upload-time = "2026-07-06T10:46:00.977Z" }, + { url = "https://files.pythonhosted.org/packages/70/e0/5c551d8d592f944506f7c5185e210255c15e672a3c6008c156a1bd9b775e/xxhash-3.8.1-cp313-cp313t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:72eb5ae575cc7ae2b23f6f8064a8b10f638c7149819ae9cc6d20ebd4d37a1629", size = 274776, upload-time = "2026-07-06T10:46:02.869Z" }, + { url = "https://files.pythonhosted.org/packages/a0/2a/d3a762270cee2d7bcd0e25e28c623e5f3f5c0dc637b66e3e47dd5b0bb3f0/xxhash-3.8.1-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:d0b48cdf690a64cedf7258c3dc9506cc41fc86edd7739c40e3098952265dc068", size = 252056, upload-time = "2026-07-06T10:46:04.688Z" }, + { url = "https://files.pythonhosted.org/packages/c1/8f/b78e4373b2cb6d1c42af60ea2d7e9146ad0710b239ac7f706d5d31d5bb98/xxhash-3.8.1-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:fb9e256a357dfcede7818c6d34e70db2d6b664394803d1de4b6984d2de76c0f1", size = 482108, upload-time = "2026-07-06T10:46:06.498Z" }, + { url = "https://files.pythonhosted.org/packages/e6/0d/642d923336ea61a15f8ce64fc7e078729e6e06c3a026e517fa79b2c23b7a/xxhash-3.8.1-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:51f71a6e2ad071e70c937e41fcb6c19f82c3f9f49831eba850ed4a106ffbb647", size = 226739, upload-time = "2026-07-06T10:46:08.598Z" }, + { url = "https://files.pythonhosted.org/packages/a6/0a/a37d6da6427d45a8d23e3ee3a0ca9c9d4a90364849c6637fe2963a755f9b/xxhash-3.8.1-cp313-cp313t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e4a6443968c4e8dc69967e12776776a5952c119cc1bd94168ad1c5ad667c2be1", size = 319658, upload-time = "2026-07-06T10:46:10.504Z" }, + { url = "https://files.pythonhosted.org/packages/4a/51/ebbd40da8a3f1bc53b4b7a9a87f8e28bd95c5f21bc14b8a57860cf367d1b/xxhash-3.8.1-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:714503083a1f2065c9ad15340dd49ac8a8e948a505a705ffa1750cb951519113", size = 246059, upload-time = "2026-07-06T10:46:12.634Z" }, + { url = "https://files.pythonhosted.org/packages/24/4c/d9014030147e1f0bb26e7da47aa240dd9ec61c763c573e558111d869f8e1/xxhash-3.8.1-cp313-cp313t-musllinux_1_2_armv7l.whl", hash = "sha256:77f74e45a1e5574bbbf80181c8027b3a4c65c2248fffbd557bd596fff13102f9", size = 275535, upload-time = "2026-07-06T10:46:14.614Z" }, + { url = "https://files.pythonhosted.org/packages/84/86/caee2db41fadcd5a25aa4323213f9afec5a8586d4e419241e3d659362bd7/xxhash-3.8.1-cp313-cp313t-musllinux_1_2_i686.whl", hash = "sha256:4e0e1b0fb0259c1b75d1251ac0bb4d7ab675d36f7a6bf4ba6aa630dae94f9ffa", size = 231292, upload-time = "2026-07-06T10:46:16.452Z" }, + { url = "https://files.pythonhosted.org/packages/0b/60/f52f08bcdc904c4514ea5c25caa19e9f3214144434a6ff96dc82dc1cbddd/xxhash-3.8.1-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:10e4393ec33633c2f05ad01869e546ad080b1a18f2650503731f153774608b31", size = 250490, upload-time = "2026-07-06T10:46:18.318Z" }, + { url = "https://files.pythonhosted.org/packages/24/a0/94dc7ae310838f250669c6ad7168e6d6fca17d49dac1053f06dc232c4a56/xxhash-3.8.1-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:b3ba794c3d885803db6c3116686923f1ec13bc86e621e169a375282b63ea1cc6", size = 309861, upload-time = "2026-07-06T10:46:20.503Z" }, + { url = "https://files.pythonhosted.org/packages/8b/f9/adeead7d0eb28cdfc2832544ea639ffbc6749ccde47a8e228d667459182e/xxhash-3.8.1-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:57189a69c0891e4818853feaa521c972d22c880a001453addea015f48e3c3398", size = 448739, upload-time = "2026-07-06T10:46:22.79Z" }, + { url = "https://files.pythonhosted.org/packages/04/a4/22ec0e07db57d901c9298ae98aa3cf2be45bafded6f07c13131e85b89032/xxhash-3.8.1-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:d59e71153fe9ff85648d00e18649b07e9b22c797291abb7e27274fa06df8b838", size = 223657, upload-time = "2026-07-06T10:46:24.831Z" }, + { url = "https://files.pythonhosted.org/packages/94/32/8a9531f37b59e5a013003db7cb7414baf4ce7e0e1268e0d5947cd3d6a2df/xxhash-3.8.1-cp313-cp313t-win32.whl", hash = "sha256:5b96f0024e9840f449bd91b2d005c921a4b666055a0d1b6492463799f32aae22", size = 32377, upload-time = "2026-07-06T10:46:26.86Z" }, + { url = "https://files.pythonhosted.org/packages/e7/ab/2ca45fd7f671de5f81fc297ef1c95080b40c86ec6be0cc6034b8f7707ac8/xxhash-3.8.1-cp313-cp313t-win_amd64.whl", hash = "sha256:37d5a56c36dcc0b9a87b814cd992598d33863ff683749de6c86081f278d5e629", size = 33274, upload-time = "2026-07-06T10:46:28.39Z" }, + { url = "https://files.pythonhosted.org/packages/5a/54/20d7163463ddb6438b73a427d1655a77a502cf9b9b0c3ada3599629d9c0a/xxhash-3.8.1-cp313-cp313t-win_arm64.whl", hash = "sha256:6696c8752aded28ff3b16f33ef28ce28fb5d209b80c206746f943199fcf5fd65", size = 29375, upload-time = "2026-07-06T10:46:29.962Z" }, + { url = "https://files.pythonhosted.org/packages/c2/8b/df2ba04f22a6cd6b39f96a6577329a8471a55c90ef8d8e2f7c102363613f/xxhash-3.8.1-cp314-cp314-android_24_arm64_v8a.whl", hash = "sha256:9db455cb649dcfe4504d6d68a6d83a7315a99a3ca59871dc3ff840671f99adba", size = 38430, upload-time = "2026-07-06T10:46:31.496Z" }, + { url = "https://files.pythonhosted.org/packages/b2/4f/6a059e8ad3ca8deedc91dfe335b211204900895152212c03ebbe721de68b/xxhash-3.8.1-cp314-cp314-android_24_x86_64.whl", hash = "sha256:affb37f152e55b5e4494bb9d0107f7bb08515c6704fbed82d9f61214d74adc17", size = 36558, upload-time = "2026-07-06T10:46:33.078Z" }, + { url = "https://files.pythonhosted.org/packages/cb/95/40be178205acce092ae418feb20ac737b32a02c7b864926ed0717354c9f8/xxhash-3.8.1-cp314-cp314-ios_13_0_arm64_iphoneos.whl", hash = "sha256:460261045936975193bfd20549a0de1cd52a33b405cbb972f0d80940c42266cd", size = 31181, upload-time = "2026-07-06T10:46:34.793Z" }, + { url = "https://files.pythonhosted.org/packages/3f/89/2da4dbf051bafa156c0e3f12012db2b0ac3b84ff37ca1f021f6bfffcdfbb/xxhash-3.8.1-cp314-cp314-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:38c887aedb696ef8bca19983206d270848558cfae4a91afa6a2fb05dde58ffc5", size = 32192, upload-time = "2026-07-06T10:46:36.393Z" }, + { url = "https://files.pythonhosted.org/packages/7c/4e/e000bbae3566bc8e0be771a8a0f294aa99075e3f0bc4ef43922ebffdebc8/xxhash-3.8.1-cp314-cp314-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:594131ce1aad18db3689781f806db1b065cdaa04f4df36b4c038d2013aefd0bf", size = 34691, upload-time = "2026-07-06T10:46:38.1Z" }, + { url = "https://files.pythonhosted.org/packages/b4/4a/ea954aacc7d1c8711880ac2b55da94429a9b4296b151c4fc0966549ca1ee/xxhash-3.8.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:78c794b643d214f1522e7a288bcf5a2de120d26cd170516749a4009dc92722c9", size = 34807, upload-time = "2026-07-06T10:46:39.647Z" }, + { url = "https://files.pythonhosted.org/packages/ca/29/df598e738ff37558ac627264deb2e560902d9bf7f46d3bd5175c9eee593e/xxhash-3.8.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:af0c9fedc4a2c24e8664953882fe8185f3790b8338c9c700f76f5ad660817711", size = 32410, upload-time = "2026-07-06T10:46:41.359Z" }, + { url = "https://files.pythonhosted.org/packages/59/9c/81ab40e7d33ada0b3df5d1bc884894d15dbf4f805cd645b685e4606bb8e0/xxhash-3.8.1-cp314-cp314-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:115772daeb71b2f3b9381177017f53e6cf3f3439c840737fdabd21aba6e54920", size = 220564, upload-time = "2026-07-06T10:46:43.463Z" }, + { url = "https://files.pythonhosted.org/packages/fd/6f/62ae6f5c8606320a0e2a41c2dc8c6d91cc5d63d0f84dd9582e9543779dd8/xxhash-3.8.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:000435984a0469b0f822fe76f35bddea0f96a4d6521b3339a60a6428cdee1edc", size = 241462, upload-time = "2026-07-06T10:46:45.509Z" }, + { url = "https://files.pythonhosted.org/packages/15/a1/9c3a0ec6cb524396f551eddd102a76690a795494eb9784fc67542b0daa37/xxhash-3.8.1-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:2f1c68394818e0595569c2ff3cbc1e6d5a36a434e796f5c526b987b80c8a8c62", size = 264491, upload-time = "2026-07-06T10:46:47.655Z" }, + { url = "https://files.pythonhosted.org/packages/64/f2/700a4674e4308eb59d2fdb973977e82eae231bea5044753fee5c9eec0e0c/xxhash-3.8.1-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:46b39976d008e2a845758650f0ff7136bca004f40da0c8798bd37ac37860154f", size = 242905, upload-time = "2026-07-06T10:46:49.857Z" }, + { url = "https://files.pythonhosted.org/packages/f3/8a/72d9874375c8d4cbc64a8cd1d659d5695a8765c3db82efa82dc5bd9f14d0/xxhash-3.8.1-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:d5006c65ec507a333479e76e00e2c368781f16c24ededa764763956b32a0e93e", size = 473873, upload-time = "2026-07-06T10:46:51.953Z" }, + { url = "https://files.pythonhosted.org/packages/03/f0/6db07590ed7e0a77f186ef0bcea8d52553bf1ba57833e09467a2411f0f2d/xxhash-3.8.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c31a2649bcf1fe97cf11c79848d761df33ac46b3896942d31b640557b486ff6b", size = 220765, upload-time = "2026-07-06T10:46:55.41Z" }, + { url = "https://files.pythonhosted.org/packages/8f/10/00d12d8b8beabbf49a8bbc626fb9f40445145a8887eb41a6acfb69149ac4/xxhash-3.8.1-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:8f759eed402448c2bdbb492e4fba1f20668ffe29688605ea61f0f67f9e4e386d", size = 310478, upload-time = "2026-07-06T10:46:57.729Z" }, + { url = "https://files.pythonhosted.org/packages/1f/f9/12a82394eefb0f185d15a7f7b9f627c61c475a72dd83718436a5b84b42ac/xxhash-3.8.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:7b5f97ecfede10d5b2870383620e2d25c8561e217c7bf9081073802b54248d2b", size = 238393, upload-time = "2026-07-06T10:46:59.87Z" }, + { url = "https://files.pythonhosted.org/packages/20/f3/53f963e320b9ce678337aa7273f39ce692ded8b99e3d22a866ec722159ab/xxhash-3.8.1-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:1da930bbcac3e8fbe2191850e2abb57977a99348c12c4b385e1058ac1b0a9ecc", size = 268704, upload-time = "2026-07-06T10:47:01.806Z" }, + { url = "https://files.pythonhosted.org/packages/0a/50/5b5badbd87c82d9f9b5f58ac74a3f29ef08f6fc387b324b8fd482450b862/xxhash-3.8.1-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:747476436f6891b9773374ce8d48edcc8b12cb5b61b67c6fb6289633747d088f", size = 225015, upload-time = "2026-07-06T10:47:03.784Z" }, + { url = "https://files.pythonhosted.org/packages/30/93/3ca68265afe7b4e69435e08a7b6a1d9d0f2a071e889da1f8041ed00fe878/xxhash-3.8.1-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:4ef09bbc2519a93cd0f95f2ceb5f7b85919dffea643278e02362bf40e3c4bed1", size = 240951, upload-time = "2026-07-06T10:47:05.816Z" }, + { url = "https://files.pythonhosted.org/packages/dd/a6/27e19670c40f46b5e76e11f2f4713d21054804568425d870670e757172ad/xxhash-3.8.1-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:a5eed9d41995a83f3332b4e3396abb7f433cac584222bd7e305b606d8353861e", size = 300751, upload-time = "2026-07-06T10:47:07.95Z" }, + { url = "https://files.pythonhosted.org/packages/bc/fb/b33e27689959fe7ed2ae0b830af41560d65213943983afa9db3a8d481bce/xxhash-3.8.1-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:53f3ed9118397074ff63a79b66b7fec1c84c782eecde35c5bc94e420a971c231", size = 443480, upload-time = "2026-07-06T10:47:10Z" }, + { url = "https://files.pythonhosted.org/packages/26/60/0e0d973be5fe280753ef02fbc89349492ad6e903bf1dcb870b668f94b662/xxhash-3.8.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:d247b34bf433c92b41689318fd25d246313cab2275a6a47e2efac178b80d6efe", size = 217657, upload-time = "2026-07-06T10:47:12.196Z" }, + { url = "https://files.pythonhosted.org/packages/ad/68/c9e3ecef4a9a417d464cb5bd200aa12f73192dee677901b9e08e0ad0d1bb/xxhash-3.8.1-cp314-cp314-win32.whl", hash = "sha256:d58ce8b6cfa9c4d2f230557f69caf7c06369e318015d0b19485095bc2c5963ab", size = 32690, upload-time = "2026-07-06T10:47:14.204Z" }, + { url = "https://files.pythonhosted.org/packages/d7/99/e9e44588c0b62837bbec5ba7927816de0afa03406b1a0b6c7a7e1d1a30a0/xxhash-3.8.1-cp314-cp314-win_amd64.whl", hash = "sha256:6cee733fe4ccb1737e0997135283c82341e5cfa9cf214b165f9087fb663aaf4f", size = 33460, upload-time = "2026-07-06T10:47:16.021Z" }, + { url = "https://files.pythonhosted.org/packages/45/2b/64f36d86380b3657ad9031967ab814f3ef31307174650853f69c18932ebc/xxhash-3.8.1-cp314-cp314-win_arm64.whl", hash = "sha256:58346024d47e84f7d8b3e7f5d6faa1d58acbbe49a8771497872059f58c1d8ea5", size = 30092, upload-time = "2026-07-06T10:47:17.81Z" }, + { url = "https://files.pythonhosted.org/packages/92/cb/18b64bff88c58a0ca209dc533e63cf02d7ae5aa6b1b9a9fd14e81b5dbd60/xxhash-3.8.1-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:01cab782f8a0a05ecad2c63d7ef10f7ab475f660e0d6419d069418c14d88de7c", size = 35024, upload-time = "2026-07-06T10:47:19.821Z" }, + { url = "https://files.pythonhosted.org/packages/af/1d/72d8a70520e5dcddb472ea0486d299da3240745a10658290cd7b5690ede2/xxhash-3.8.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:717b12fdc51819833704e85e6926d76981ffa3f780ef92e33ebb8b26d46bb230", size = 32697, upload-time = "2026-07-06T10:47:21.649Z" }, + { url = "https://files.pythonhosted.org/packages/9c/b8/e041f555903c56db3d0a731b3d72a6575d75e0ed868b1bd2e5176111ca44/xxhash-3.8.1-cp314-cp314t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:ec55d80e9b8a519d742669e0b49e8ce9e6747be42bf3c138158b6543a9c8e489", size = 226044, upload-time = "2026-07-06T10:47:23.612Z" }, + { url = "https://files.pythonhosted.org/packages/3a/7e/5cdcf06bf6ec4b5d2ac073feb23432ec1d603fd438864cbd2c09c7cb45e1/xxhash-3.8.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:98d8ac1129b4dd39098cffed94d1284aceb61c3aa396757ccc736ac392e4cee5", size = 249899, upload-time = "2026-07-06T10:47:25.812Z" }, + { url = "https://files.pythonhosted.org/packages/c0/c0/eb7e059cb5e1dba11fd30d2fdf882f56e5a417a3eaa43669d43623767f45/xxhash-3.8.1-cp314-cp314t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:3bc0fa90830df1e1277f33cc6e55de9990b83c0319fd8c7412866cfde38b025e", size = 274892, upload-time = "2026-07-06T10:47:27.931Z" }, + { url = "https://files.pythonhosted.org/packages/66/74/a600aaf7cd39957fd1510adeedb1749c1e7eb82bd632a1153d9c664c3135/xxhash-3.8.1-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c73b6f652f0745425aa6378319c331293b5341756262e9408ed3d45f183375e6", size = 252243, upload-time = "2026-07-06T10:47:30.288Z" }, + { url = "https://files.pythonhosted.org/packages/ad/04/78d88fa75a6763e5d09bf1b947a392a27988903381b219006f92f3c68fc8/xxhash-3.8.1-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:f6114692261eff4266386cdec0f7d87eee24e317ab397c218b7ae6a76b4c6339", size = 482191, upload-time = "2026-07-06T10:47:32.45Z" }, + { url = "https://files.pythonhosted.org/packages/7f/06/07a8aea1108d682de8791ce608cdf367d75ff4e7e57cd3c154bdc6f47b23/xxhash-3.8.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4df57c0b161ec1b3ed0526a67b0db0914b557e86ee8aae51887aec941b261542", size = 226877, upload-time = "2026-07-06T10:47:34.705Z" }, + { url = "https://files.pythonhosted.org/packages/ed/b5/86bade5618a524d2c06c4041aa2fe8e5749ce16e88afba60d67c1684a21f/xxhash-3.8.1-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:9043877a917be88ccf230aa5667c1bd059bce80f4c2727e4defa1b29b7f48b08", size = 319794, upload-time = "2026-07-06T10:47:37.08Z" }, + { url = "https://files.pythonhosted.org/packages/23/69/9b1a2b89b1621bb740fbcb7beb512f60f99480c1bdc680c0c90e1f56ff75/xxhash-3.8.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:559e3cabe522231909f9de98ef06929edbd53782046bd21aae0c72db6f2a0775", size = 246202, upload-time = "2026-07-06T10:47:39.676Z" }, + { url = "https://files.pythonhosted.org/packages/08/ea/662ed6cb49f1d34078b6a3a3e0f3d29ff93fd7b5a03c0bc9ecfd9b2159c3/xxhash-3.8.1-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:264710bd335016f303763ce1275c6486df30bb57c2245c91b224c983d7ac39b8", size = 275628, upload-time = "2026-07-06T10:47:41.99Z" }, + { url = "https://files.pythonhosted.org/packages/13/f5/49fc9e4c6728a5a3bd8fe639199d2fa67609b3a84f938aff6e8568dd3e4f/xxhash-3.8.1-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:e14800b9b10bb39d7a60ad4a310e403164d7b8988a27ae933d4e40618a44088e", size = 231390, upload-time = "2026-07-06T10:47:44.233Z" }, + { url = "https://files.pythonhosted.org/packages/64/9d/3acaf8f599c0e0b30e910a3a11ba32929da53c86dc73c7c55fe6a010b4e9/xxhash-3.8.1-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:ea6a3e734b0fd41b82784a400be946821900daebe610c050a5e0760838a34f99", size = 250600, upload-time = "2026-07-06T10:47:47.611Z" }, + { url = "https://files.pythonhosted.org/packages/23/64/8acab4c5ec60dbe664b5b9858fd44c2413b07e535b09556a0a5022e78aa6/xxhash-3.8.1-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:cf399fac542a1c7a4734a435b93df2c55e858c7d31abf6c1bdf46f9ae67fbfd0", size = 310032, upload-time = "2026-07-06T10:47:49.88Z" }, + { url = "https://files.pythonhosted.org/packages/56/47/a0288d7329b1fe63e2734a32d19d444a96ae2b4810f545bc61e561224917/xxhash-3.8.1-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:44c89d915a75c11d2547eaee9098fcd80398987c4bff2974a0497a925bf92c07", size = 448882, upload-time = "2026-07-06T10:47:52.631Z" }, + { url = "https://files.pythonhosted.org/packages/01/e7/3071dfd3beb5c38204ce1cf56bf7749fce08de900fa92714b81d1d8ca1f2/xxhash-3.8.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:358650d5bda9c635da699c53adf4e8134af492ecc79c960f917eebf088bb6799", size = 223728, upload-time = "2026-07-06T10:47:55.093Z" }, + { url = "https://files.pythonhosted.org/packages/12/11/b99949f0ba2b07e9f9ffe83b9c86faa685f9080725dc21a916a607313be5/xxhash-3.8.1-cp314-cp314t-win32.whl", hash = "sha256:c240939e963653054fc7e4a17c382829cda4aa88a7daf0af841715dbded1b497", size = 33150, upload-time = "2026-07-06T10:47:57.274Z" }, + { url = "https://files.pythonhosted.org/packages/54/1c/09703eb341f8416e74e58d6c6732d4b5c46de59c942363203cb237cc95b0/xxhash-3.8.1-cp314-cp314t-win_amd64.whl", hash = "sha256:7258ee276e8772599bc19e14b36f6260306e21b637190cd7cb489a2449d48684", size = 34005, upload-time = "2026-07-06T10:47:59.434Z" }, + { url = "https://files.pythonhosted.org/packages/d6/f9/6ed7251bb6a8af10ac73b1821c60583d2826e5b2064e45a979c935287c98/xxhash-3.8.1-cp314-cp314t-win_arm64.whl", hash = "sha256:8f454166c2ffed45636c8d501741e649851ba2f346c4eb73a64c07ac00428f20", size = 30239, upload-time = "2026-07-06T10:48:01.874Z" }, + { url = "https://files.pythonhosted.org/packages/99/e4/4d8040435aeac814fc69ba63621565fbeb19229a138e2568324a26b2a45c/xxhash-3.8.1-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:39c9d5b61508b0bb68f29e54546de0ed2a74943c6a18585535a7e37356f1dd12", size = 32687, upload-time = "2026-07-06T10:49:42.803Z" }, + { url = "https://files.pythonhosted.org/packages/da/6a/975f1f2318c760e5bcec109ed379713ae645d8d856c2a3b9ec5d26857087/xxhash-3.8.1-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:83b9130b80b216d56fdf9e87131946b353c9627930c061955a101ea82b09fed9", size = 29879, upload-time = "2026-07-06T10:49:45.172Z" }, + { url = "https://files.pythonhosted.org/packages/08/0b/40a2a55ff52cf635bfdc5eae67a772bec85b4f44c6c737f73f6f528d51d1/xxhash-3.8.1-pp311-pypy311_pp73-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:8304be0982130954b7fd3aad18e2c6f8ee40254bc3d2e635991c16d77c91e2bd", size = 43246, upload-time = "2026-07-06T10:49:47.905Z" }, + { url = "https://files.pythonhosted.org/packages/9c/6d/56ed2b6b200f26fb474f3fd387d95d0601efcd5bb33430c90c68924bdd77/xxhash-3.8.1-pp311-pypy311_pp73-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4b512261801b1e5fde7b6ebf2fef7977339c620cbbca88a0040ad9ad134f4d02", size = 38202, upload-time = "2026-07-06T10:49:50.59Z" }, + { url = "https://files.pythonhosted.org/packages/0d/a3/56864d895d1161a9f17502088e9c1fb7c06bde2c2efdde620d22bb7a9c43/xxhash-3.8.1-pp311-pypy311_pp73-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:49aa8692507835dcc1e8ad8021f20c74c2dc13d83b5112e87877faa2a0035b20", size = 34448, upload-time = "2026-07-06T10:49:53.242Z" }, + { url = "https://files.pythonhosted.org/packages/6b/57/5c6e0908a47f61dca96d01c8ee6fce01ed1050611eb779083ba8758fed81/xxhash-3.8.1-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:345b07b78e2bf583d71682aa34ae5b5fab575f7a1cb31e10263ebbc6f89f8c42", size = 32869, upload-time = "2026-07-06T10:49:55.972Z" }, +] + +[[package]] +name = "yarl" +version = "1.24.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "idna" }, + { name = "multidict" }, + { name = "propcache" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/31/33/ebe9e3d1f86c7a0b51094c0a146392045ca1631d2664889539dec8088a33/yarl-1.24.5.tar.gz", hash = "sha256:e81b83143bee16329c23db3c1b2d82b29892fcbcb849186d2f6e98a5abe9a57f", size = 228679, upload-time = "2026-07-20T02:07:45.435Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fe/db/3cb5df059756a45761cc3dee8fd25ec82b83a6585ea3542b969fda850f99/yarl-1.24.5-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:2c1fe720934a16ea8e7146175cba2126f87f54912c8c5435e7f7c7a51ef808d3", size = 135043, upload-time = "2026-07-20T02:04:52.39Z" }, + { url = "https://files.pythonhosted.org/packages/44/f8/767d6bd5a03db63bc467df2fb56d6fafeae9667d74aea92cd6af399f828b/yarl-1.24.5-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:c687ed078e145f5fd53a14854beff320e1d2ab76df03e2009c98f39a0f68f39a", size = 96942, upload-time = "2026-07-20T02:04:54.26Z" }, + { url = "https://files.pythonhosted.org/packages/ce/97/10b939c44d7b28d1dbc389cfc7012306d1ea8dba01eaef44b39fffaee52a/yarl-1.24.5-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:709f1efed56c4a145793c046cd4939f9959bcd818979a787b77d8e09c57a0840", size = 97046, upload-time = "2026-07-20T02:04:56.638Z" }, + { url = "https://files.pythonhosted.org/packages/5b/7a/b410dbe39b6255c55fb2a2bcee96eb844d0789235ddc381a889a90dc72d6/yarl-1.24.5-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:874019bd513008b009f58657134e5d0c5e030b3559bd0553976837adf52fe966", size = 110512, upload-time = "2026-07-20T02:04:58.955Z" }, + { url = "https://files.pythonhosted.org/packages/83/c7/da591971f78a5617e1f21f5699858ebccd836fe181a6493788ffc91ba69b/yarl-1.24.5-cp311-cp311-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:a4582acf7ef76482f6f511ebaf1946dae7f2e85ec4728b81a678c01df63bd723", size = 102454, upload-time = "2026-07-20T02:05:00.623Z" }, + { url = "https://files.pythonhosted.org/packages/c4/8e/73b0ed4de47289a78a96045d76d1cfe5e41848bf0da59ce25b2ec87ee05d/yarl-1.24.5-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:2cabe6546e41dabe439999a23fcb5246e0c3b595b4315b96ef755252be90caeb", size = 117617, upload-time = "2026-07-20T02:05:02.325Z" }, + { url = "https://files.pythonhosted.org/packages/cf/14/b744747bc4f57a8d55bd744df463457524583e1e9f7538b5ace0346ab92e/yarl-1.24.5-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:17f57620f5475b3c69109376cc87e42a7af5db13c9398e4292772a706ff10780", size = 116135, upload-time = "2026-07-20T02:05:04.05Z" }, + { url = "https://files.pythonhosted.org/packages/66/ca/95aa4d0e5b7ea4f20e4d577c42d001ed9df207569fdb063cc5ed4ebb496b/yarl-1.24.5-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:570fec8fbd22b032733625f03f10b7ff023bc399213db15e72a7acaef28c2f4e", size = 111935, upload-time = "2026-07-20T02:05:05.738Z" }, + { url = "https://files.pythonhosted.org/packages/72/0d/d2ad8d6b147832d177a4e720ba1962fe686eb0913b74503b3eca094b8bba/yarl-1.24.5-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:5fede79c6f73ff2c3ef822864cb1ada23196e62756df53bc6231d351a49516a2", size = 110010, upload-time = "2026-07-20T02:05:07.471Z" }, + { url = "https://files.pythonhosted.org/packages/50/18/eb335e4120903903f4865041355ae46256a2406eb2865bc24827f4f27b61/yarl-1.24.5-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:8ccf9aca873b767977c73df497a85dbedee4ee086ae9ae49dc461333b9b79f58", size = 110058, upload-time = "2026-07-20T02:05:09.246Z" }, + { url = "https://files.pythonhosted.org/packages/44/70/97353add32c62ad6f206d948ac5a5ee84398225e534dc6ed6433d1b335b6/yarl-1.24.5-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:ad5d8201d310b031e6cd839d9bac2d4e5a01533ce5d3d5b50b7de1ef3af1de61", size = 103308, upload-time = "2026-07-20T02:05:11.31Z" }, + { url = "https://files.pythonhosted.org/packages/68/39/5e7398d4b6f6b3c9062823ebc60802df5b272e3fe9e788f9734c6ee46c85/yarl-1.24.5-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:841f0852f48fefea3b12c9dfec00704dfa3aef5215d0e3ce564bb3d7cd8d57c6", size = 116898, upload-time = "2026-07-20T02:05:13.099Z" }, + { url = "https://files.pythonhosted.org/packages/e4/c9/09e52f2239e8b96357eccca05915382e4ba5405ebfb623b6036040d99654/yarl-1.24.5-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:9baafc71b04f8f4bb0703b21d6fc9f0c30b346c636a532ff16ec8491a5ea4b1f", size = 109400, upload-time = "2026-07-20T02:05:14.821Z" }, + { url = "https://files.pythonhosted.org/packages/4b/6a/e94133d4c2d1a14d2384310bf3e79d9cf32c9d1eae1c6f034fb80d098fa1/yarl-1.24.5-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:d897129df1a22b12aeed2c2c98df0785a2e8e6e0bde87b389491d0025c187077", size = 115934, upload-time = "2026-07-20T02:05:17.78Z" }, + { url = "https://files.pythonhosted.org/packages/4e/3c/34955ed967b976fc38edcbb6d538dee79dbda4cb7fc7f72a0907a7c78e0f/yarl-1.24.5-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:dd625535328fd9882374356269227670189adfcc6a2d90284f323c05862eecbd", size = 112178, upload-time = "2026-07-20T02:05:19.675Z" }, + { url = "https://files.pythonhosted.org/packages/f5/46/d7bd3a8859d47dcfaffd7127af7076032a7da278a9a02e17b5f37bfb6712/yarl-1.24.5-cp311-cp311-win_amd64.whl", hash = "sha256:f4239bbec5a3577ddb49e4b50aeb32d8e5792098262ae2f63723f916a29b1a25", size = 97544, upload-time = "2026-07-20T02:05:21.523Z" }, + { url = "https://files.pythonhosted.org/packages/01/69/c1bfd21e32c638974ea2c542a0b8c53ef1fa9eff336020f5d014f9503ff2/yarl-1.24.5-cp311-cp311-win_arm64.whl", hash = "sha256:3ac6aff147deb9c09461b2d4bbdf6256831198f5d8a23f5d37138213090b6d8a", size = 93359, upload-time = "2026-07-20T02:05:23.493Z" }, + { url = "https://files.pythonhosted.org/packages/1b/84/71d051c850b5af41d168c679d9eb67eb7c55283ac4ee131673edf134bc4e/yarl-1.24.5-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:d693396e5aea78db03decd60aec9ece16c9b40ba00a587f089615ff4e718a81d", size = 136035, upload-time = "2026-07-20T02:05:25.489Z" }, + { url = "https://files.pythonhosted.org/packages/03/4d/8ad27f9a1b7e69313cca5d695b925b48efe51208d3490e0844bae97cabc0/yarl-1.24.5-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:3363fcc96e665878946ad7a106b9a13eac0541766a690ef287c0232ac768b6ec", size = 97642, upload-time = "2026-07-20T02:05:27.429Z" }, + { url = "https://files.pythonhosted.org/packages/ea/b4/05b4131c407006cd1e410e9c6539f16a0945724677e5364447313c15ea3e/yarl-1.24.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:9d399bdcfb4a0f659b9b3788bbc89babe63d9a6a65aacdf4d4e7065ff2e6316c", size = 97323, upload-time = "2026-07-20T02:05:29.441Z" }, + { url = "https://files.pythonhosted.org/packages/20/16/e618c875c73e0e39611f20a581b3d5e8d59b8857bf001bee3263044c6deb/yarl-1.24.5-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:90333fd89b43c0d08ac85f3f1447593fc2c66de18c3d6378d7125ea118dc7a54", size = 107741, upload-time = "2026-07-20T02:05:31.367Z" }, + { url = "https://files.pythonhosted.org/packages/d9/9a/c4defeaf3ed33fcb346aacf9c6e971a8d4e2bde04a0310e79abb208e7965/yarl-1.24.5-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:665b0a2c463cc9423dd647e0bfd9f4ccc9b50f768c55304d5e9f80b177c1de12", size = 103570, upload-time = "2026-07-20T02:05:33.303Z" }, + { url = "https://files.pythonhosted.org/packages/5f/e7/0e0e0de5865ebd5914537ef486f36c727a59865c3ac0cf5ff1b32aececbf/yarl-1.24.5-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:e006d3a974c4ee19512e5f058abedb6eef36a5e553c14812bdeba1758d812e6d", size = 115815, upload-time = "2026-07-20T02:05:35.292Z" }, + { url = "https://files.pythonhosted.org/packages/2b/27/ca56b700cb170aba25a3893b75355b213935657dc5714d2383354a270e62/yarl-1.24.5-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:e7d42c531243450ef0d4d9c172e7ed6ef052640f195629065041b5add4e058d1", size = 116025, upload-time = "2026-07-20T02:05:37.503Z" }, + { url = "https://files.pythonhosted.org/packages/d6/d0/d56c859b8222116f5d68459199f48359e0bf121b6f65a69bf329b3602ba0/yarl-1.24.5-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f08c7513ecef5aad65687bfdf6bc601ae9fccd04a42904501f8f7141abad9eb9", size = 109835, upload-time = "2026-07-20T02:05:39.506Z" }, + { url = "https://files.pythonhosted.org/packages/70/a2/3a35557e4d1a79425040eba202ccaf08bdc8717680fc77e2498a1ad2e0a5/yarl-1.24.5-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:6c95b17fe34ed802f17e205112e6e10db92275c34fee290aa9bdc55a9c724027", size = 108884, upload-time = "2026-07-20T02:05:41.584Z" }, + { url = "https://files.pythonhosted.org/packages/e4/35/ef4c26356b7913c68983bac2d72a4212b3347af551cb8d250b99b5ed7b7f/yarl-1.24.5-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:56b149b22de33b23b0c6077ab9518c6dcb538ad462e1830e68d06591ccf6e38b", size = 107308, upload-time = "2026-07-20T02:05:43.697Z" }, + { url = "https://files.pythonhosted.org/packages/d5/91/ff0dc66c2ccf3e0153ab97ff61eabab4400e6a5264af427ab30cd69f1857/yarl-1.24.5-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:a8fe66b8f300da93798025a785a5b90b42f3810dc2b72283ff84a41aaaebc293", size = 103646, upload-time = "2026-07-20T02:05:45.895Z" }, + { url = "https://files.pythonhosted.org/packages/74/f0/33b9271c7f881766359d58266fa0811d2e5210ed860e28da7dc6d7786344/yarl-1.24.5-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:377fe3732edbaf78ee74efdf2c9f49f6e99f20e7f9d2649fda3eb4badd77d76e", size = 115305, upload-time = "2026-07-20T02:05:47.832Z" }, + { url = "https://files.pythonhosted.org/packages/ef/65/fd79fb1868c4a80db8661091de525bf430f63c3bea1b20e8b6a84fc7d359/yarl-1.24.5-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:e8ffa78582120024f476a611d7befc123cee59e47e8309d470cf667d806e613b", size = 108404, upload-time = "2026-07-20T02:05:49.604Z" }, + { url = "https://files.pythonhosted.org/packages/ff/ba/dbabe6b262f17a816c70cfc09558dbf03ece3ec76684d02f911a3d3a189c/yarl-1.24.5-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:daba5e594f06114e37db186efd2dd916609071e59daca901a0a2e71f02b142ce", size = 115940, upload-time = "2026-07-20T02:05:51.741Z" }, + { url = "https://files.pythonhosted.org/packages/a5/43/fab2d1dad9d340a268cdde63756a123d069723efff6a372d123fa74a9517/yarl-1.24.5-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:65be18ec59496c13908f02a2472751d9ef840b4f3fb5726f129306bf6a2a7bba", size = 110006, upload-time = "2026-07-20T02:05:53.554Z" }, + { url = "https://files.pythonhosted.org/packages/c4/27/41eb51bbd1b8d89546b83897cfb0164f1e109304fd408dbb151b639eec0f/yarl-1.24.5-cp312-cp312-win_amd64.whl", hash = "sha256:a929d878fec099030c292803b31e5d5540a7b6a31e6a3cc76cb4685fc2a2f51b", size = 97618, upload-time = "2026-07-20T02:05:55.57Z" }, + { url = "https://files.pythonhosted.org/packages/3c/25/b2553764b3d65db711d8f45416351ec4f420847558eb669edcbcaadf5780/yarl-1.24.5-cp312-cp312-win_arm64.whl", hash = "sha256:7ce27823052e2013b597e0c738b13e7e36b8ccb9400df8959417b052ab0fd92c", size = 93018, upload-time = "2026-07-20T02:05:57.554Z" }, + { url = "https://files.pythonhosted.org/packages/e1/63/64ef361967cc983573149dc1515d531db5da8a4c92d22bb833d59e01b313/yarl-1.24.5-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:79af890482fc94648e8cde4c68620378f7fef60932710fa17a66abc039244da2", size = 135075, upload-time = "2026-07-20T02:05:59.671Z" }, + { url = "https://files.pythonhosted.org/packages/bb/89/55920fd853ce43e608adbc3962456f0d649d6bb15250dc2988321da0fe1c/yarl-1.24.5-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:46c2f213e23a04b93a392942d782eb9e413e6ef6bf7c8c53884e599a5c174dcb", size = 97225, upload-time = "2026-07-20T02:06:01.769Z" }, + { url = "https://files.pythonhosted.org/packages/15/f0/7688d3f2cfff7590df2af38ec46d969f4281a4dddb08a9ad2eafbcdddf98/yarl-1.24.5-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:92ab3e11448f2ff7bf53c5a26eff0edc086898ec8b21fb154b85839ce1d88075", size = 96751, upload-time = "2026-07-20T02:06:03.676Z" }, + { url = "https://files.pythonhosted.org/packages/05/1a/a851a0f94aaaf379dd4f901bfc80f634280bec51eb260b47363e2a4cd62e/yarl-1.24.5-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ebb0ec7f17803063d5aeb982f3b1bd2b2f4e4fae6751226cbd6ba1fcfe9e63ff", size = 107960, upload-time = "2026-07-20T02:06:05.699Z" }, + { url = "https://files.pythonhosted.org/packages/6c/a8/faea066c12f9c77ca0de90641f1655f9dd7b412477bf28c76d692f3aecff/yarl-1.24.5-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:82632daed195dcc8ea664e8556dc9bdbd671960fb3776bd92806ce05792c2448", size = 103500, upload-time = "2026-07-20T02:06:07.556Z" }, + { url = "https://files.pythonhosted.org/packages/fb/9c/1e67084c2a6e2f2db0e3be798328cb3be42c0119b621d25461479a224d21/yarl-1.24.5-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:53e549287ef628fecba270045c9701b0c564563a9b0577d24a4ec75b8ab8040f", size = 115780, upload-time = "2026-07-20T02:06:09.599Z" }, + { url = "https://files.pythonhosted.org/packages/58/86/1f94664e147474337e3359f52012cf3d02f825f694317b178bfba1078c62/yarl-1.24.5-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:fcd3b77e2f17bbe4ca56ec7bcb07992647d19d0b9c05d84886dcd6f9eb810afd", size = 115308, upload-time = "2026-07-20T02:06:11.352Z" }, + { url = "https://files.pythonhosted.org/packages/0a/43/8e55ae7538ba5f28ccb3c845c6dd4549cf7016d5992e5326512519107cdd/yarl-1.24.5-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d46b86567dd4e248c6c159fcbcdcce01e0a5c8a7cd2334a0fff759d0fa075b16", size = 110574, upload-time = "2026-07-20T02:06:13.129Z" }, + { url = "https://files.pythonhosted.org/packages/ce/ba/a889ec8765cedcf2ac44dcb02d6a21e4861399b243b263c5f2dde27ee740/yarl-1.24.5-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:7f72c74aa99359e27a2ee8d6613fefa28b5f76a983c083074dfc2aaa4ab46213", size = 109914, upload-time = "2026-07-20T02:06:15.243Z" }, + { url = "https://files.pythonhosted.org/packages/9c/c3/e45f821af67b791c2dbbe4a9f4137a1d33f8d386654a05a0c3f47bdfa25d/yarl-1.24.5-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:3f45789ce415a7ec0820dc4f82925f9b5f7732070be1dec1f5f23ec381435a24", size = 107712, upload-time = "2026-07-20T02:06:17.443Z" }, + { url = "https://files.pythonhosted.org/packages/02/00/2ab0f42c9857fcb490bfaa6647b14540b53d241ab209f23220b958cc5832/yarl-1.24.5-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:6e73e7fe93f17a7b191f52ec9da9dd8c06a8fe735a1ecbd13b97d1c723bff385", size = 104251, upload-time = "2026-07-20T02:06:19.259Z" }, + { url = "https://files.pythonhosted.org/packages/7a/70/709d9a286e98af2c7fd8e4e6cada658b5c0e30d87dd7e2a63c2fb5767217/yarl-1.24.5-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:4a36f9becdd4c5c52a20c3e9484128b070b1dcfc8944c006f3a528295a359a9c", size = 115319, upload-time = "2026-07-20T02:06:21.207Z" }, + { url = "https://files.pythonhosted.org/packages/5c/6c/3eaa515142991fe84cfc483ff986492211f1978f90161ccefdbec919d09b/yarl-1.24.5-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:7bcbe0fcf850eae67b6b01749815a4f7161c560a844c769ad7b48fcd99f791c4", size = 109163, upload-time = "2026-07-20T02:06:23.006Z" }, + { url = "https://files.pythonhosted.org/packages/bb/64/711dafce66c323a3144d470547a71c5384c57623308ac8bb5e4b903ac148/yarl-1.24.5-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:24e861e9630e0daddcb9191fb187f60f034e17a4426f8101279f0c475cd74144", size = 115435, upload-time = "2026-07-20T02:06:24.923Z" }, + { url = "https://files.pythonhosted.org/packages/cf/f3/9b9d0e6d84bea851eb1ba99e4bdc755b86fd813e49ec86dfe42f26befdef/yarl-1.24.5-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:9335a099ad87287c37fe5d1a982ff392fa5efe5d14b40a730b1ec1d6a41382b4", size = 110691, upload-time = "2026-07-20T02:06:26.973Z" }, + { url = "https://files.pythonhosted.org/packages/86/e4/62a06b7e87c4246ac76b7c2da136f972eb4a3a1fc94abb07e7022d6fdb0a/yarl-1.24.5-cp313-cp313-win_amd64.whl", hash = "sha256:2dbe06fc16bc91502bca713704022182e5729861ae00277c3a23354b40929740", size = 97454, upload-time = "2026-07-20T02:06:29.163Z" }, + { url = "https://files.pythonhosted.org/packages/9e/c9/5fc8025b318ab10db413b61056bd0d95c557a70e8df4210c7511f866329c/yarl-1.24.5-cp313-cp313-win_arm64.whl", hash = "sha256:6b8536851f9f65e7f00c7a1d49ba7f2be0ffe2c11555367fc9f50d9f842410a1", size = 92813, upload-time = "2026-07-20T02:06:31.113Z" }, + { url = "https://files.pythonhosted.org/packages/a9/08/5f3085fef9564217074db9dd8573de1795bc82cde61a7ad10b6a7234a569/yarl-1.24.5-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:2729fcfc4f6a596fb0c50f32090400aa9367774ac296a00387e65098c0befa76", size = 135680, upload-time = "2026-07-20T02:06:33.273Z" }, + { url = "https://files.pythonhosted.org/packages/98/35/ba9436e579bd48a8801f2021d842d9ab4994c26e4c7dd3a4c1f1bcb57a9e/yarl-1.24.5-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:ff330d3c30db4eb6b01d79e29d2d0b407a7ecad39cfd9ec993ece57396a2ec0d", size = 97395, upload-time = "2026-07-20T02:06:35.259Z" }, + { url = "https://files.pythonhosted.org/packages/18/a9/a07f76f3c44e02b25cc743af5ef93eef27f7013eadca770451b6a6ccb5db/yarl-1.24.5-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:e42d75862735da90e7fc5a7b23db0c976f737113a54b3c9777a9b665e9cbff75", size = 97223, upload-time = "2026-07-20T02:06:37.216Z" }, + { url = "https://files.pythonhosted.org/packages/77/f7/a9a1d6fa7dd9e388f95b30f6ad3ec4e285f6c8f61f44ce16070c3fcfe414/yarl-1.24.5-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a3732e66413163e72508da9eff9ce9d2846fde51fae45d3605393d3e6cd303e9", size = 108777, upload-time = "2026-07-20T02:06:39.292Z" }, + { url = "https://files.pythonhosted.org/packages/2f/44/e0b86c302471fabd6f02808ecf2ac52b8412b624787849d4bf2cdb466f6f/yarl-1.24.5-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:5b8ee53be440a0cffc991a27be3057e0530122548dbe7c0892df08822fce5ede", size = 103119, upload-time = "2026-07-20T02:06:41.456Z" }, + { url = "https://files.pythonhosted.org/packages/d1/16/9c16d180bf8faaf223225eb50e1245870ff1ae0e302a27153988e65c51fd/yarl-1.24.5-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:af3aefa655adb5869491fa907e652290386800ae99cc50095cba71e2c6aefdca", size = 116471, upload-time = "2026-07-20T02:06:43.696Z" }, + { url = "https://files.pythonhosted.org/packages/d2/8d/b219b9df28a02ce95cfbdd41d2f7caa5669d0ff979c1c9975697145e33c5/yarl-1.24.5-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:2120b96872df4a117cde97d270bac96aea7cc52205d305cf4611df694a487027", size = 115974, upload-time = "2026-07-20T02:06:45.874Z" }, + { url = "https://files.pythonhosted.org/packages/9b/e8/f20557aca240d88e69850ad1ee91756821d094bb1310565c04d25c6682a2/yarl-1.24.5-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:66410eb6345d467151934b49bfa70fb32f5b35a6140baa40ad97d6436abea2e9", size = 110830, upload-time = "2026-07-20T02:06:47.852Z" }, + { url = "https://files.pythonhosted.org/packages/db/18/199b85109a53eeca64ee19c9cca228287e8e4ab0cc1a09b28f530e65cce0/yarl-1.24.5-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:4af7b7e1be0a69bee8210735fe6dcfc38879adfac6d62e789d53ba432d1ffa41", size = 110054, upload-time = "2026-07-20T02:06:49.84Z" }, + { url = "https://files.pythonhosted.org/packages/aa/2f/ed28147f8cd7f48c49367c90713b30a555284b6105a6a56f3a05568da795/yarl-1.24.5-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:fa139875ff98ab97da323cfadfaff08900d1ad42f1b5087b0b812a55c5a06373", size = 108312, upload-time = "2026-07-20T02:06:51.835Z" }, + { url = "https://files.pythonhosted.org/packages/c5/c5/55e16ae0a5c227cea8df1c6871ba57d614a34243146c05729caf2a1bd9c5/yarl-1.24.5-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:0055afc45e864b92729ac7600e2d102c17bef060647e74bca75fa84d66b9ff36", size = 103662, upload-time = "2026-07-20T02:06:54.061Z" }, + { url = "https://files.pythonhosted.org/packages/8d/ea/dbd7c2caec459c9a426f18b02688ecbfb58620d0f6a3422d24769fbaf8ab/yarl-1.24.5-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:f0e466ed7511fe9d459a819edbc6c2585c0b6eabde9fa8a8947552468a7a6ef0", size = 116090, upload-time = "2026-07-20T02:06:56.015Z" }, + { url = "https://files.pythonhosted.org/packages/06/84/39ce4ce3059e07fece5fbdbee8c4053406af9aca911ce9fa5f8548aab6af/yarl-1.24.5-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:f141474e85b7e54998ec5180530a7cda99ab29e282fa50e0756d89981a9b43c5", size = 109523, upload-time = "2026-07-20T02:06:57.926Z" }, + { url = "https://files.pythonhosted.org/packages/a9/8b/71ff44137b405c64a7788075669c24010019f57a7464b78c3a6cbee539d9/yarl-1.24.5-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:e2935f8c39e3b03e83519292d78f075189978f3f4adc15a78144c7c8e2a1cba5", size = 116084, upload-time = "2026-07-20T02:06:59.868Z" }, + { url = "https://files.pythonhosted.org/packages/62/c0/423078fdd4042e1862c11f0ffd977a0ffa393783c12bee94685923bc189e/yarl-1.24.5-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:9d1216a7f6f77836617dba35687c5b78a4170afc3c3f18fc788f785ba26565c4", size = 111006, upload-time = "2026-07-20T02:07:01.907Z" }, + { url = "https://files.pythonhosted.org/packages/cf/52/6daa2ee9d95e5c98b8128f8df91eb692eb423ab274b8cf08db52152fad26/yarl-1.24.5-cp314-cp314-win_amd64.whl", hash = "sha256:5ba4f78df2bcc19f764a4b26a8a4f5049c110090ad5825993aacb052bf8003ad", size = 99215, upload-time = "2026-07-20T02:07:03.852Z" }, + { url = "https://files.pythonhosted.org/packages/ec/0e/464a847d7359e0da75dd9fc5c1d1aa35d0159ea31e5f8e66a3c1c29ff3d0/yarl-1.24.5-cp314-cp314-win_arm64.whl", hash = "sha256:9e4e16c73d717c5cf27626c524d0a2e261ad20e46932b2670f64ad5dde23e26f", size = 94566, upload-time = "2026-07-20T02:07:06.074Z" }, + { url = "https://files.pythonhosted.org/packages/e2/55/e03acc4446772660bc335e86e41ef31e4d0d838fd641531a11a5ee33b493/yarl-1.24.5-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:e1ae548a9d901adca07899a4147a7c826bbcc06239d3ce9a59f57886a28a4c88", size = 142533, upload-time = "2026-07-20T02:07:08.284Z" }, + { url = "https://files.pythonhosted.org/packages/ae/71/4acd3a1fc7cf14345cdb302665ecd2097f62c365b4f14ca17d4f37775cf9/yarl-1.24.5-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:ff405d91509d88e8d44129cd87b18d70acd1f0c1aeabd7bc3c46792b1fe2acba", size = 100776, upload-time = "2026-07-20T02:07:10.197Z" }, + { url = "https://files.pythonhosted.org/packages/ff/0b/cfb76b7fe99686db264bff829779a539d923e7564ffd7ef18da6c54c3774/yarl-1.24.5-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:47e98aab9d8d82ff682e7b0b5dded33bf138a32b817fcf7fa3b27b2d7c412928", size = 100913, upload-time = "2026-07-20T02:07:12.357Z" }, + { url = "https://files.pythonhosted.org/packages/8b/3f/7116e782992abbd4fb6948488aec72078895e929a23078290739e8396fce/yarl-1.24.5-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f0a658a6d3fafee5c6f63c58f3e785c8c43c93fbc02bf9f2b6663f8185e0971f", size = 106507, upload-time = "2026-07-20T02:07:14.173Z" }, + { url = "https://files.pythonhosted.org/packages/33/90/d4d2d73ee78229cc889872eb8e085d8f5c6f51abdb178409fd9b23cf74fd/yarl-1.24.5-cp314-cp314t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:4377407001ca3c057773f44d8ddd6358fa5f691407c1ba92210bd3cf8d9e4c95", size = 99219, upload-time = "2026-07-20T02:07:16.019Z" }, + { url = "https://files.pythonhosted.org/packages/3e/fa/a6df1a9bccd644eec00abee0dff4277416222cec435330fd1f2858523ec1/yarl-1.24.5-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:7c0494a31a1ac5461a226e7947a9c9b78c44e1dc7185164fa7e9651557a5d9bc", size = 111804, upload-time = "2026-07-20T02:07:18.141Z" }, + { url = "https://files.pythonhosted.org/packages/8a/9e/7b2a1f4bcc20e9447156dd2b1c4d01f70d9df0759025ee7d09a84ffae134/yarl-1.24.5-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:a7cff474ab7cd149765bb784cf6d78b32e18e20473fb7bda860bce98ab58e9da", size = 110943, upload-time = "2026-07-20T02:07:20.06Z" }, + { url = "https://files.pythonhosted.org/packages/08/ff/22c92affb0f9b623ca753d27d968b5625b868f12c6378d049d55ae247643/yarl-1.24.5-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:cbb833ccacdb5519eff9b8b71ee618cc2801c878e77e288775d77c3a2ced858a", size = 108251, upload-time = "2026-07-20T02:07:22.217Z" }, + { url = "https://files.pythonhosted.org/packages/45/44/5769b96298c1e195fb412997b6090af2a84105cf59c17613558a2d011d1f/yarl-1.24.5-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:82f75e05912e84b7a0fe57075d9c59de3cb352b928330f2eb69b2e1f54c3e1f0", size = 106025, upload-time = "2026-07-20T02:07:24.083Z" }, + { url = "https://files.pythonhosted.org/packages/4c/40/009e8e791fd9762c0e1567e69248acb4f49064597e1680874c16dd8bb798/yarl-1.24.5-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:16a2f5010280020e90f5330257e6944bc33e73593b136cc5a241e6c1dc292498", size = 106573, upload-time = "2026-07-20T02:07:26.248Z" }, + { url = "https://files.pythonhosted.org/packages/20/c6/b7480578f8a0a80946f36ad6df547ecec704f9ba69d2de60f8aa6f1c1cbf/yarl-1.24.5-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:ffcd54362564dc1a30fb74d8b8a6e5a6b11ebd5e27266adc3b7427a21a6c9104", size = 100751, upload-time = "2026-07-20T02:07:28.098Z" }, + { url = "https://files.pythonhosted.org/packages/d4/27/4476f3360b91a48c5cf125e91f59a3bd35299d84a431a258d57f5977bb11/yarl-1.24.5-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:0465ec8cedc2349b97a6b595ace64084a50c6e839eca40aa0626f38b8350e331", size = 111643, upload-time = "2026-07-20T02:07:30.88Z" }, + { url = "https://files.pythonhosted.org/packages/4c/4b/5cdd3e5ee944e8af31e52f6cd3d3af5fd7b937e036ccbbba2c9ffebede95/yarl-1.24.5-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:4db9aecb141cb7a5447171b57aa1ed3a8fee06af40b992ffc31206c0b0121550", size = 106312, upload-time = "2026-07-20T02:07:33.06Z" }, + { url = "https://files.pythonhosted.org/packages/18/86/f406b0c2a6f99575de2da671ef47aa06f89a5be83a27a46971c3b86cecdb/yarl-1.24.5-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:f540c013589084679a6c7fac07096b10159737918174f5dfc5e11bf5bca4dfe6", size = 110379, upload-time = "2026-07-20T02:07:35.155Z" }, + { url = "https://files.pythonhosted.org/packages/f0/6c/9f3adfbd3b30b4fa0f7ccb3a83eba2c1152d3fff554d535e640ba0f7ba2b/yarl-1.24.5-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:a61834fb15d81322d872eaafd333838ae7c9cea84067f232656f75965933d047", size = 108497, upload-time = "2026-07-20T02:07:37.35Z" }, + { url = "https://files.pythonhosted.org/packages/dd/37/91eb2e5ca883a529c1b390348a74cd9fc0512171727f547ce70bfe02be5c/yarl-1.24.5-cp314-cp314t-win_amd64.whl", hash = "sha256:5c88e5815a49d289e599f3513aa7fde0bc2092ff188f99c940f007f90f53d104", size = 102450, upload-time = "2026-07-20T02:07:39.578Z" }, + { url = "https://files.pythonhosted.org/packages/bf/f4/ed5c402ac8fde4403ed3366c2716bfddc8a6677ebd59f3d62772cc7fe468/yarl-1.24.5-cp314-cp314t-win_arm64.whl", hash = "sha256:cf139c02f5f23ef6532040a30ff662c00a318c952334f211046b8e60b7f17688", size = 97222, upload-time = "2026-07-20T02:07:41.55Z" }, + { url = "https://files.pythonhosted.org/packages/61/02/962c1cbfc401a30c1d034dc67ff395f64b52302c6d62de556c1fca99acc0/yarl-1.24.5-py3-none-any.whl", hash = "sha256:a33700d13d9b7d84fd10947b09ff69fb9a792e519c8cb9764a3ca70baa6c23a7", size = 58612, upload-time = "2026-07-20T02:07:43.461Z" }, +] diff --git a/vortex-array-macros/src/lib.rs b/vortex-array-macros/src/lib.rs index 3105dac4add..5270e8b8e00 100644 --- a/vortex-array-macros/src/lib.rs +++ b/vortex-array-macros/src/lib.rs @@ -4,32 +4,49 @@ //! Proc macros for `vortex-array`. use proc_macro::TokenStream; +use proc_macro2::Span; use quote::format_ident; use quote::quote; use syn::Field; use syn::Fields; use syn::Ident; use syn::ItemStruct; +use syn::LitInt; use syn::Path; +use syn::Token; use syn::Type; use syn::Visibility; +use syn::parse::Parse; +use syn::parse::ParseStream; use syn::parse_macro_input; use syn::spanned::Spanned; +/// Name of the per-field attribute that pins a field to a slot index. +const SLOT_ATTR: &str = "slot"; + /// Generate slot index constants, a borrowed view struct, and a typed ext trait /// from a slot struct definition. /// -/// Fields must be `ArrayRef` (required slot) or `Option` (optional slot). -/// Field declaration order determines slot indices. +/// Fields must be `ArrayRef` (required slot), `Option` (optional slot), or +/// `Vec` (variadic tail of required slots). +/// +/// Every field must carry a `#[slot(..)]` attribute naming the exact slot index it maps +/// to. The attribute — not the declaration order — defines the storage layout, so fields +/// may be reordered, grouped, or documented in any order without changing the slot +/// indices an array is built from or read back with. /// /// # Example /// /// ```ignore /// #[array_slots(Patched)] /// pub struct PatchedSlots { +/// #[slot(0)] /// pub inner: ArrayRef, +/// #[slot(1)] /// pub lane_offsets: ArrayRef, +/// #[slot(2)] /// pub patch_indices: ArrayRef, +/// #[slot(3)] /// pub patch_values: ArrayRef, /// } /// ``` @@ -39,7 +56,7 @@ use syn::spanned::Spanned; /// Given the above, the macro generates: /// /// ```ignore -/// // --- The original struct is preserved as-is --- +/// // --- The original struct, minus the consumed `#[slot(..)]` attributes --- /// pub struct PatchedSlots { ... } /// /// // --- Slot index constants and conversion methods on the struct --- @@ -83,6 +100,16 @@ use syn::spanned::Spanned; /// impl> PatchedArraySlotsExt for T {} /// ``` /// +/// # Slot index annotations +/// +/// - Fixed fields use `#[slot(N)]`, where `N` is the exact index of the slot. +/// - A variadic tail uses `#[slot(N..)]`, where `N` is the index its first slot occupies. +/// +/// The annotations are validated at compile time: the fixed indices must cover +/// `0..FIXED_COUNT` exactly — no duplicates and no gaps — and a variadic tail must start +/// immediately after the last fixed slot. A field without a `#[slot(..)]` attribute is a +/// compile error, so the layout can never silently fall back to declaration order. +/// /// # Required vs optional slots /// /// - `ArrayRef` — the slot must be present. `from_slots()` panics if `None`. @@ -94,6 +121,61 @@ use syn::spanned::Spanned; /// /// The underlying storage is always `ArraySlots` — the field type only /// controls whether the macro inserts a `.vortex_expect()` unwrap or not. +/// +/// # Variadic tail slots +/// +/// One field may be `Vec`, declaring that every slot from its index onward +/// belongs to a homogeneous, variable-length run of required slots. This supports +/// encodings like `Chunked` (`[chunk_offsets, chunks...]`), `Struct` +/// (`[validity?, fields...]`), and `Union` (`[type_ids, children...]`). +/// +/// ```ignore +/// #[array_slots(Chunked)] +/// pub struct ChunkedSlots { +/// #[slot(0)] +/// pub chunk_offsets: ArrayRef, +/// #[slot(1..)] +/// pub chunks: Vec, +/// } +/// ``` +/// +/// For a struct with a variadic tail, the macro generates a different set of +/// constants — slot count is no longer a compile-time constant: +/// +/// ```ignore +/// impl ChunkedSlots { +/// pub const CHUNK_OFFSETS: usize = 0; +/// /// Offset at which the `chunks` slots begin. +/// pub const CHUNKS_OFFSET: usize = 1; +/// /// Number of fixed (non-variadic) slots. +/// pub const FIXED_COUNT: usize = 1; +/// /// Names of the fixed slots in storage order. +/// pub const FIXED_NAMES: [&'static str; 1] = ["chunk_offsets"]; +/// +/// /// Name of the slot at `idx`, e.g. "chunk_offsets" or "chunks[3]". +/// pub fn slot_name(idx: usize) -> String { ... } +/// +/// pub fn from_slots(slots: ArraySlots) -> Self { ... } +/// pub fn into_slots(self) -> ArraySlots { ... } +/// } +/// ``` +/// +/// The view field and ext trait accessor for the tail are a `vortex_array::SlotSlice`, +/// a borrowed run of required slots supporting `len()`, `get()`, `iter()`, and +/// indexing: +/// +/// ```ignore +/// pub struct ChunkedSlotsView<'a> { +/// pub chunk_offsets: &'a ArrayRef, +/// pub chunks: SlotSlice<'a>, +/// } +/// +/// pub trait ChunkedArraySlotsExt: TypedArrayRef { +/// fn chunk_offsets(&self) -> &ArrayRef { ... } +/// fn chunks(&self) -> SlotSlice<'_> { ... } +/// fn slots_view(&self) -> ChunkedSlotsView<'_> { ... } +/// } +/// ``` #[proc_macro_attribute] pub fn array_slots(attr: TokenStream, item: TokenStream) -> TokenStream { let encoding = parse_macro_input!(attr as Path); @@ -132,26 +214,31 @@ fn expand_array_slots( .map(|segment| &segment.ident) .ok_or_else(|| syn::Error::new(encoding.span(), "missing encoding type"))?; - let struct_ident = &item_struct.ident; - let struct_vis = &item_struct.vis; - let view_ident = format_ident!("{}View", ident_name(struct_ident)); + let struct_ident = item_struct.ident.clone(); + let struct_vis = item_struct.vis.clone(); + let view_ident = format_ident!("{}View", ident_name(&struct_ident)); let ext_ident = format_ident!("{}ArraySlotsExt", ident_name(encoding_ident)); let field_specs = fields .iter() - .enumerate() - .map(|(index, field)| SlotField::new(field, index, struct_ident)) + .map(|field| SlotField::new(field, &struct_ident)) .collect::>>()?; - let idx_consts = field_specs.iter().map(SlotField::idx_const); + // The `#[slot(..)]` annotations, not the declaration order, define storage order. + let (fixed_specs, tail_spec) = partition_by_slot_index(&field_specs)?; + + let idx_consts = fixed_specs.iter().copied().map(SlotField::idx_const); let view_fields = field_specs.iter().map(SlotField::view_field); let view_from_slots = field_specs.iter().map(SlotField::view_from_slots); let view_to_owned = field_specs.iter().map(SlotField::view_to_owned); - let owned_from_slots = field_specs.iter().map(SlotField::owned_from_slots); - let into_slots = field_specs.iter().map(SlotField::storage_slot); let ext_methods = field_specs.iter().map(SlotField::ext_method); - let slot_names = field_specs.iter().map(|field| field.slot_name.as_str()); - let slot_count = field_specs.len(); + + let counts = gen_counts(&fixed_specs, tail_spec); + let from_slots = gen_from_slots(&fixed_specs, tail_spec); + let into_slots = gen_into_slots(&fixed_specs, tail_spec); + + // `#[slot(..)]` is inert helper syntax consumed here; strip it from the emitted struct. + let item_struct = strip_slot_attrs(item_struct); Ok(quote! { #item_struct @@ -159,23 +246,13 @@ fn expand_array_slots( impl #struct_ident { #(#idx_consts)* - #[doc = "Total number of slots."] - pub const COUNT: usize = #slot_count; - - #[doc = "Slot names in storage order."] - pub const NAMES: [&'static str; #slot_count] = [#(#slot_names),*]; + #counts #[doc = "Convert owned slot storage into an owned slot struct."] - pub fn from_slots(mut slots: ::vortex_array::ArraySlots) -> Self { - Self { - #(#owned_from_slots,)* - } - } + #from_slots #[doc = "Convert this slot struct into storage order."] - pub fn into_slots(self) -> ::vortex_array::ArraySlots { - ::vortex_array::smallvec::smallvec![#(#into_slots),*] - } + #into_slots } #[derive(Clone, Copy, Debug)] @@ -214,6 +291,195 @@ fn expand_array_slots( }) } +/// Split the fields into the fixed slots (sorted by slot index) and the optional variadic +/// tail, validating that the annotated indices describe a complete, gap-free layout. +fn partition_by_slot_index( + field_specs: &[SlotField], +) -> syn::Result<(Vec<&SlotField>, Option<&SlotField>)> { + let mut fixed_specs = Vec::with_capacity(field_specs.len()); + let mut tail_spec: Option<&SlotField> = None; + + for spec in field_specs { + if matches!(spec.slot_type, SlotFieldType::VariadicTail) { + if let Some(previous) = tail_spec { + return Err(syn::Error::new( + spec.index_span, + format!( + "#[array_slots] allows at most one variadic tail, but `{}` is already \ + declared as one", + previous.slot_name + ), + )); + } + tail_spec = Some(spec); + } else { + fixed_specs.push(spec); + } + } + + fixed_specs.sort_by_key(|spec| spec.index); + + for (expected, spec) in fixed_specs.iter().enumerate() { + if spec.index == expected { + continue; + } + // `fixed_specs` is sorted, so a smaller index than expected means the previous + // field claimed the same slot, and a larger one means nothing claimed `expected`. + return Err(if spec.index < expected { + syn::Error::new( + spec.index_span, + format!( + "#[array_slots] slot index {} is claimed by both `{}` and `{}`", + spec.index, + fixed_specs[expected - 1].slot_name, + spec.slot_name + ), + ) + } else { + syn::Error::new( + spec.index_span, + format!( + "#[array_slots] no field claims slot index {expected}; fixed slot indices \ + must cover 0..{} without gaps", + fixed_specs.len() + ), + ) + }); + } + + if let Some(tail) = tail_spec + && tail.index != fixed_specs.len() + { + return Err(syn::Error::new( + tail.index_span, + format!( + "#[array_slots] variadic tail `{}` must start at slot index {}, immediately after \ + the {} fixed slot(s), but is annotated `#[slot({}..)]`", + tail.slot_name, + fixed_specs.len(), + fixed_specs.len(), + tail.index + ), + )); + } + + Ok((fixed_specs, tail_spec)) +} + +/// Remove the inert `#[slot(..)]` helper attributes so the struct can be re-emitted. +fn strip_slot_attrs(mut item_struct: ItemStruct) -> ItemStruct { + if let Fields::Named(fields) = &mut item_struct.fields { + for field in &mut fields.named { + field.attrs.retain(|attr| !attr.path().is_ident(SLOT_ATTR)); + } + } + item_struct +} + +fn gen_counts( + fixed_specs: &[&SlotField], + tail_spec: Option<&SlotField>, +) -> proc_macro2::TokenStream { + let names = fixed_specs.iter().map(|field| field.slot_name.as_str()); + let fixed_count = fixed_specs.len(); + + match tail_spec { + None => quote! { + #[doc = "Total number of slots."] + pub const COUNT: usize = #fixed_count; + + #[doc = "Slot names in storage order."] + pub const NAMES: [&'static str; #fixed_count] = [#(#names),*]; + }, + Some(tail) => { + let offset_const = &tail.const_ident; + let tail_name = &tail.slot_name; + quote! { + #[doc = concat!("Offset at which the `", #tail_name, "` slots begin.")] + pub const #offset_const: usize = #fixed_count; + + #[doc = "Number of fixed (non-variadic) slots."] + pub const FIXED_COUNT: usize = #fixed_count; + + #[doc = "Names of the fixed slots in storage order."] + pub const FIXED_NAMES: [&'static str; #fixed_count] = [#(#names),*]; + + #[doc = "Name of the slot at the given index."] + pub fn slot_name(idx: usize) -> String { + if idx < Self::FIXED_COUNT { + Self::FIXED_NAMES[idx].to_string() + } else { + format!(concat!(#tail_name, "[{}]"), idx - Self::#offset_const) + } + } + } + } + } +} + +fn gen_from_slots( + fixed_specs: &[&SlotField], + tail_spec: Option<&SlotField>, +) -> proc_macro2::TokenStream { + let owned_from_slots = fixed_specs.iter().copied().map(SlotField::owned_from_slots); + + match tail_spec { + None => quote! { + pub fn from_slots(mut slots: ::vortex_array::ArraySlots) -> Self { + Self { + #(#owned_from_slots,)* + } + } + }, + Some(tail) => { + let tail_ident = &tail.field_ident; + let offset_const = &tail.const_ident; + let expect_message = &tail.expect_message; + quote! { + pub fn from_slots(mut slots: ::vortex_array::ArraySlots) -> Self { + let __variadic_tail: ::std::vec::Vec<::vortex_array::ArrayRef> = slots + .drain(Self::#offset_const..) + .map(|slot| ::vortex_error::VortexExpect::vortex_expect( + slot, + #expect_message, + )) + .collect(); + Self { + #(#owned_from_slots,)* + #tail_ident: __variadic_tail, + } + } + } + } + } +} + +fn gen_into_slots( + fixed_specs: &[&SlotField], + tail_spec: Option<&SlotField>, +) -> proc_macro2::TokenStream { + let fixed_into_slots = fixed_specs.iter().copied().map(SlotField::storage_slot); + + match tail_spec { + None => quote! { + pub fn into_slots(self) -> ::vortex_array::ArraySlots { + ::vortex_array::smallvec::smallvec![#(#fixed_into_slots),*] + } + }, + Some(tail) => { + let tail_ident = &tail.field_ident; + quote! { + pub fn into_slots(self) -> ::vortex_array::ArraySlots { + let mut slots: ::vortex_array::ArraySlots = + ::vortex_array::smallvec::smallvec![#(#fixed_into_slots),*]; + slots.extend(self.#tail_ident.into_iter().map(Some)); + slots + } + } + } + } +} + struct SlotField { field_ident: Ident, field_vis: Visibility, @@ -221,19 +487,51 @@ struct SlotField { slot_name: String, slot_type: SlotFieldType, index: usize, + index_span: Span, expect_message: syn::LitStr, struct_ident: Ident, } impl SlotField { - fn new(field: &Field, index: usize, struct_ident: &Ident) -> syn::Result { + fn new(field: &Field, struct_ident: &Ident) -> syn::Result { let field_ident = field .ident .clone() .ok_or_else(|| syn::Error::new(field.span(), "slot fields must be named"))?; let field_name = ident_name(&field_ident); - let const_ident = format_ident!("{}", to_screaming_snake_case(&field_name)); let slot_type = SlotFieldType::from_syn_type(&field.ty)?; + let annotation = SlotIndexAttr::from_field(field, &field_name)?; + + match (slot_type, annotation.variadic) { + (SlotFieldType::VariadicTail, false) => { + return Err(syn::Error::new( + annotation.span, + format!( + "`{field_name}` is a variadic `Vec` tail, so it must be \ + annotated `#[slot({}..)]`", + annotation.index + ), + )); + } + (SlotFieldType::Required | SlotFieldType::Optional, true) => { + return Err(syn::Error::new( + annotation.span, + format!( + "`#[slot(N..)]` declares a variadic `Vec` tail; `{field_name}` \ + occupies a single slot, so annotate it `#[slot({})]`", + annotation.index + ), + )); + } + _ => {} + } + + let const_ident = match slot_type { + SlotFieldType::VariadicTail => { + format_ident!("{}_OFFSET", to_screaming_snake_case(&field_name)) + } + _ => format_ident!("{}", to_screaming_snake_case(&field_name)), + }; let expect_message = syn::LitStr::new( &format!("{} {} slot", ident_name(struct_ident), field_name), field.span(), @@ -245,7 +543,8 @@ impl SlotField { const_ident, slot_name: field_name, slot_type, - index, + index: annotation.index, + index_span: annotation.span, expect_message, struct_ident: struct_ident.clone(), }) @@ -288,6 +587,12 @@ impl SlotField { SlotFieldType::Optional => quote! { #field_ident: slots[#struct_ident::#const_ident].as_ref() }, + SlotFieldType::VariadicTail => quote! { + #field_ident: ::vortex_array::SlotSlice::new( + &slots[#struct_ident::#const_ident..], + #expect_message, + ) + }, } } @@ -301,6 +606,9 @@ impl SlotField { SlotFieldType::Optional => quote! { #field_ident: self.#field_ident.cloned() }, + SlotFieldType::VariadicTail => quote! { + #field_ident: self.#field_ident.to_vec() + }, } } @@ -320,6 +628,9 @@ impl SlotField { SlotFieldType::Optional => quote! { #field_ident: slots[#struct_ident::#const_ident].take() }, + SlotFieldType::VariadicTail => { + unreachable!("variadic tail is drained before fixed fields") + } } } @@ -333,6 +644,9 @@ impl SlotField { SlotFieldType::Optional => quote! { self.#field_ident }, + SlotFieldType::VariadicTail => { + unreachable!("variadic tail is appended after fixed fields") + } } } @@ -358,7 +672,73 @@ impl SlotField { self.as_ref().slots()[#struct_ident::#const_ident].as_ref() } }, + SlotFieldType::VariadicTail => quote! { + #[inline] + fn #field_ident(&self) -> ::vortex_array::SlotSlice<'_> { + ::vortex_array::SlotSlice::new( + &self.as_ref().slots()[#struct_ident::#const_ident..], + #expect_message, + ) + } + }, + } + } +} + +/// A parsed `#[slot(N)]` or `#[slot(N..)]` field annotation. +struct SlotIndexAttr { + index: usize, + /// Whether the trailing `..` was present, marking a variadic tail. + variadic: bool, + span: Span, +} + +impl SlotIndexAttr { + fn from_field(field: &Field, field_name: &str) -> syn::Result { + let mut annotation = None; + for attr in &field.attrs { + if !attr.path().is_ident(SLOT_ATTR) { + continue; + } + if annotation.is_some() { + return Err(syn::Error::new( + attr.span(), + format!("`{field_name}` has more than one `#[slot(..)]` attribute"), + )); + } + annotation = Some(attr.parse_args::()?); + } + + annotation.ok_or_else(|| { + syn::Error::new( + field.span(), + format!( + "`{field_name}` is missing a `#[slot(N)]` attribute; every field of an \ + `#[array_slots]` struct must pin itself to a slot index so that reordering \ + field declarations cannot change the slot layout" + ), + ) + }) + } +} + +impl Parse for SlotIndexAttr { + fn parse(input: ParseStream) -> syn::Result { + let literal: LitInt = input.parse()?; + let index = literal.base10_parse::()?; + let variadic = input.peek(Token![..]); + if variadic { + input.parse::()?; + } + if !input.is_empty() { + return Err(input.error("expected `#[slot(N)]` or `#[slot(N..)]`")); } + + Ok(Self { + index, + variadic, + span: literal.span(), + }) } } @@ -366,6 +746,7 @@ impl SlotField { enum SlotFieldType { Required, Optional, + VariadicTail, } impl SlotFieldType { @@ -374,15 +755,21 @@ impl SlotFieldType { return Ok(Self::Required); } - if let Some(inner_ty) = option_inner_type(ty) + if let Some(inner_ty) = wrapper_inner_type(ty, "Option") && is_array_ref_type(inner_ty) { return Ok(Self::Optional); } + if let Some(inner_ty) = wrapper_inner_type(ty, "Vec") + && is_array_ref_type(inner_ty) + { + return Ok(Self::VariadicTail); + } + Err(syn::Error::new( ty.span(), - "#[array_slots] fields must be ArrayRef or Option", + "#[array_slots] fields must be ArrayRef, Option, or Vec", )) } @@ -390,6 +777,7 @@ impl SlotFieldType { match self { Self::Required => quote! { &'a ::vortex_array::ArrayRef }, Self::Optional => quote! { Option<&'a ::vortex_array::ArrayRef> }, + Self::VariadicTail => quote! { ::vortex_array::SlotSlice<'a> }, } } } @@ -407,12 +795,12 @@ fn is_array_ref_type(ty: &Type) -> bool { ) } -fn option_inner_type(ty: &Type) -> Option<&Type> { +fn wrapper_inner_type<'a>(ty: &'a Type, wrapper: &str) -> Option<&'a Type> { let Type::Path(type_path) = ty else { return None; }; let segment = type_path.path.segments.last()?; - if segment.ident != "Option" { + if segment.ident != wrapper { return None; } diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index a1aa5b047f8..f6b06544baf 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -60,7 +60,6 @@ static_assertions = { workspace = true } tabled = { workspace = true, optional = true, default-features = false, features = [ "std", ] } -termtree = { workspace = true } tracing = { workspace = true } uuid = { workspace = true } vortex-array-macros = { workspace = true } @@ -80,8 +79,11 @@ cudarc = ["dep:cudarc"] table-display = ["dep:tabled"] _test-harness = ["dep:goldenfile", "dep:rstest", "dep:rstest_reuse"] serde = ["dep:serde", "vortex-buffer/serde", "vortex-mask/serde"] +# Exposes experimental row-function APIs without compatibility guarantees. +unstable_row_fns = [] [dev-dependencies] +allocator-api2 = { workspace = true } divan = { workspace = true } futures = { workspace = true, features = ["executor"] } insta = { workspace = true } @@ -91,7 +93,12 @@ rstest = { workspace = true } serde_json = { workspace = true } serde_test = { workspace = true } test-with = { workspace = true } -vortex-array = { path = ".", features = ["_test-harness", "table-display"] } +vortex-array = { path = ".", features = [ + "_test-harness", + "table-display", + "unstable_row_fns", +] } +vortex-bench-support = { workspace = true } [[bench]] name = "aggregate_max" @@ -129,6 +136,10 @@ harness = false name = "binary_ops" harness = false +[[bench]] +name = "row_fn_output" +harness = false + [[bench]] name = "kleene_bool" harness = false @@ -180,6 +191,11 @@ name = "expr_optimize" path = "benches/expr/optimize_bench.rs" harness = false +[[bench]] +name = "expr_optimize_predicate" +path = "benches/expr/optimize_predicate.rs" +harness = false + [[bench]] name = "chunked_dict_builder" harness = false @@ -211,6 +227,10 @@ harness = false name = "varbinview_zip" harness = false +[[bench]] +name = "varbinview_try_new" +harness = false + [[bench]] name = "bool_zip" harness = false @@ -227,6 +247,18 @@ harness = false name = "take_primitive" harness = false +[[bench]] +name = "take_chunked" +harness = false + +[[bench]] +name = "piecewise_sequence_take_primitive" +harness = false + +[[bench]] +name = "take_slices_to_buffer_matrix" +harness = false + [[bench]] name = "take_struct" harness = false @@ -235,6 +267,10 @@ harness = false name = "take_fsl" harness = false +[[bench]] +name = "take_varbin" +harness = false + [[bench]] name = "take_filter" harness = false @@ -247,6 +283,10 @@ harness = false name = "filter_fixed_width" harness = false +[[bench]] +name = "filter_list" +harness = false + [[bench]] name = "list_length" harness = false diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index 11477e57503..64820997778 100644 --- a/vortex-array/benches/aggregate_grouped.rs +++ b/vortex-array/benches/aggregate_grouped.rs @@ -11,6 +11,7 @@ use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; use vortex_array::ArrayRef; +use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnVTable; @@ -19,6 +20,7 @@ use vortex_array::aggregate_fn::GroupedAccumulator; use vortex_array::aggregate_fn::NumericalAggregateOpts; use vortex_array::aggregate_fn::fns::count::Count; use vortex_array::aggregate_fn::fns::sum::Sum; +use vortex_array::aggregate_fn::fns::sum_v2::SumV2; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; @@ -199,6 +201,125 @@ fn sum_f64_clustered_nulls(bencher: Bencher) { .bench_refs(|input| grouped_accumulator(input, Sum)); } +#[divan::bench] +fn sum_v2_i32_nullable_all_valid(bencher: Bencher) { + let input = i32_nullable_all_valid_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator(input, SumV2)); +} + +#[divan::bench] +fn sum_v2_i32_clustered_nulls(bencher: Bencher) { + let input = i32_clustered_nulls_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator(input, SumV2)); +} + +#[divan::bench] +fn sum_v2_f64_all_valid(bencher: Bencher) { + let input = f64_all_valid_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator(input, SumV2)); +} + +#[divan::bench] +fn sum_v2_f64_clustered_nulls(bencher: Bencher) { + let input = f64_clustered_nulls_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator(input, SumV2)); +} + +/// Execute the lazy finalize result to canonical form so the benchmark includes the full cost of +/// producing usable sums. +fn grouped_accumulator_canonical(list_view: &ArrayRef, vtable: V) -> ArrayRef +where + V: AggregateFnVTable + Clone, +{ + let mut acc = GroupedAccumulator::try_new( + vtable, + NumericalAggregateOpts::default(), + list_element_dtype(list_view), + ) + .unwrap(); + let mut ctx = SESSION.create_execution_ctx(); + acc.accumulate_list(list_view, &mut ctx).unwrap(); + let result = acc + .finish() + .unwrap() + .execute::(&mut ctx) + .unwrap() + .into_array(); + divan::black_box(result) +} + +#[divan::bench] +fn canonical_sum_i32_nullable_all_valid(bencher: Bencher) { + let input = i32_nullable_all_valid_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, Sum)); +} + +#[divan::bench] +fn canonical_sum_i32_clustered_nulls(bencher: Bencher) { + let input = i32_clustered_nulls_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, Sum)); +} + +#[divan::bench] +fn canonical_sum_f64_all_valid(bencher: Bencher) { + let input = f64_all_valid_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, Sum)); +} + +#[divan::bench] +fn canonical_sum_f64_clustered_nulls(bencher: Bencher) { + let input = f64_clustered_nulls_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, Sum)); +} + +#[divan::bench] +fn canonical_sum_v2_i32_nullable_all_valid(bencher: Bencher) { + let input = i32_nullable_all_valid_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, SumV2)); +} + +#[divan::bench] +fn canonical_sum_v2_i32_clustered_nulls(bencher: Bencher) { + let input = i32_clustered_nulls_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, SumV2)); +} + +#[divan::bench] +fn canonical_sum_v2_f64_all_valid(bencher: Bencher) { + let input = f64_all_valid_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, SumV2)); +} + +#[divan::bench] +fn canonical_sum_v2_f64_clustered_nulls(bencher: Bencher) { + let input = f64_clustered_nulls_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator_canonical(input, SumV2)); +} + #[divan::bench] fn count_i32_clustered_nulls(bencher: Bencher) { let input = i32_clustered_nulls_input(); diff --git a/vortex-array/benches/aggregate_max.rs b/vortex-array/benches/aggregate_max.rs index a47bdb28164..20677522b10 100644 --- a/vortex-array/benches/aggregate_max.rs +++ b/vortex-array/benches/aggregate_max.rs @@ -16,7 +16,8 @@ fn main() { divan::main(); } -const N: usize = 100_000; +// Sized to keep the CodSpeed simulation under 1ms per benchmark. +const N: usize = 15_000; static SESSION: LazyLock = LazyLock::new(array_session); diff --git a/vortex-array/benches/aggregate_sum.rs b/vortex-array/benches/aggregate_sum.rs index 98e6845edb7..31856e22f9f 100644 --- a/vortex-array/benches/aggregate_sum.rs +++ b/vortex-array/benches/aggregate_sum.rs @@ -7,6 +7,7 @@ use divan::Bencher; use rand::prelude::*; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::fns::sum_v2::sum_v2; use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::expr::stats::Stat; @@ -17,7 +18,8 @@ fn main() { divan::main(); } -const N: usize = 100_000; +// Sized to keep the CodSpeed simulation under 1ms per benchmark. +const N: usize = 15_000; static SESSION: LazyLock = LazyLock::new(array_session); @@ -35,6 +37,20 @@ fn sum_i32(bencher: Bencher) { .bench_refs(|(a, ctx)| a.statistics().compute_as::(Stat::Sum, ctx)); } +#[divan::bench] +fn sum_v2_i32(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(1); + let data: Vec = (0..N).map(|_| rng.random_range(-1000..1000)).collect(); + bencher + .with_inputs(|| { + ( + PrimitiveArray::from_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| sum_v2(a, ctx)); +} + #[divan::bench] fn sum_u32(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(2); @@ -49,6 +65,20 @@ fn sum_u32(bencher: Bencher) { .bench_refs(|(a, ctx)| a.statistics().compute_as::(Stat::Sum, ctx)); } +#[divan::bench] +fn sum_v2_u32(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(2); + let data: Vec = (0..N).map(|_| rng.random_range(0..2000)).collect(); + bencher + .with_inputs(|| { + ( + PrimitiveArray::from_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| sum_v2(a, ctx)); +} + #[divan::bench] fn sum_i64(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(3); @@ -63,6 +93,20 @@ fn sum_i64(bencher: Bencher) { .bench_refs(|(a, ctx)| a.statistics().compute_as::(Stat::Sum, ctx)); } +#[divan::bench] +fn sum_v2_i64(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(3); + let data: Vec = (0..N).map(|_| rng.random_range(-1000..1000)).collect(); + bencher + .with_inputs(|| { + ( + PrimitiveArray::from_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| sum_v2(a, ctx)); +} + #[divan::bench] fn sum_f64(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(6); @@ -77,6 +121,20 @@ fn sum_f64(bencher: Bencher) { .bench_refs(|(a, ctx)| a.statistics().compute_as::(Stat::Sum, ctx)); } +#[divan::bench] +fn sum_v2_f64(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(6); + let data: Vec = (0..N).map(|_| rng.random_range(-1000.0..1000.0)).collect(); + bencher + .with_inputs(|| { + ( + PrimitiveArray::from_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| sum_v2(a, ctx)); +} + #[divan::bench] fn sum_f64_nulls_clustered(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(7); @@ -99,6 +157,28 @@ fn sum_f64_nulls_clustered(bencher: Bencher) { .bench_refs(|(a, ctx)| a.statistics().compute_as::(Stat::Sum, ctx)); } +#[divan::bench] +fn sum_v2_f64_nulls_clustered(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(7); + let data: Vec> = (0..N) + .map(|i| { + if (i / 64) % 10 == 0 { + None + } else { + Some(rng.random_range(-1000.0..1000.0)) + } + }) + .collect(); + bencher + .with_inputs(|| { + ( + PrimitiveArray::from_option_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| sum_v2(a, ctx)); +} + // Clustered nulls: long runs of valid values broken up by occasional null blocks. This is the // case the run-based valid path is expected to accelerate. #[divan::bench] @@ -123,6 +203,28 @@ fn sum_i32_nulls_clustered(bencher: Bencher) { .bench_refs(|(a, ctx)| a.statistics().compute_as::(Stat::Sum, ctx)); } +#[divan::bench] +fn sum_v2_i32_nulls_clustered(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(4); + let data: Vec> = (0..N) + .map(|i| { + if (i / 64) % 10 == 0 { + None + } else { + Some(rng.random_range(-1000..1000)) + } + }) + .collect(); + bencher + .with_inputs(|| { + ( + PrimitiveArray::from_option_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| sum_v2(a, ctx)); +} + // Scattered nulls: ~50% nulls placed at random, producing many short runs. This is the worst case // for a run-based valid path, used to guard against regressions versus a per-element loop. #[divan::bench] @@ -140,3 +242,19 @@ fn sum_i32_nulls_scattered(bencher: Bencher) { }) .bench_refs(|(a, ctx)| a.statistics().compute_as::(Stat::Sum, ctx)); } + +#[divan::bench] +fn sum_v2_i32_nulls_scattered(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(5); + let data: Vec> = (0..N) + .map(|_| rng.random_bool(0.5).then(|| rng.random_range(-1000..1000))) + .collect(); + bencher + .with_inputs(|| { + ( + PrimitiveArray::from_option_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| sum_v2(a, ctx)); +} diff --git a/vortex-array/benches/binary_ops.rs b/vortex-array/benches/binary_ops.rs index 0740eac43f0..a1f14ad9bb1 100644 --- a/vortex-array/benches/binary_ops.rs +++ b/vortex-array/benches/binary_ops.rs @@ -1,12 +1,24 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Benchmarks for the binary `Operator` path, over primitive, decimal, and boolean inputs. +//! +//! The primitive arithmetic and comparison cases carry `#[cpu_features]`, so they are +//! measured on every walltime CPU-feature leg rather than in simulation. Each is written +//! once and compiled differently per leg: the kernel underneath is a portable lane loop, and +//! how wide the compiler vectorizes it is decided by the build flags, not by the source. +//! +//! The decimal and boolean cases are not tagged. Decimal arithmetic is `i128` widening and +//! per-lane rescaling, and the boolean kernels are already word-at-a-time over a bitmap; +//! neither is where a wider vector register shows up. They stay in simulation. + #![expect(clippy::unwrap_used)] #![expect( clippy::cast_possible_truncation, reason = "benchmark fixtures use indices that fit in the chosen widths" )] +use std::mem::size_of; use std::sync::LazyLock; use divan::Bencher; @@ -19,8 +31,10 @@ use vortex_array::VortexSessionExecute; use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DecimalDType; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_session::VortexSession; @@ -34,124 +48,277 @@ fn main() { static SESSION: LazyLock = LazyLock::new(array_session); -const LEN: usize = 65_536; +/// Sized to keep CodSpeed simulation under 1ms per benchmark. +const LEN: usize = 4_096; + +/// Primitive cases process at least this many rows and this many value bytes per varying input. +/// This lengthens narrow integer cases while keeping the current CodSpeed simulations below 1 ms. +const MIN_PRIMITIVE_LEN: usize = 16_384; +const MIN_PRIMITIVE_INPUT_BYTES: usize = 96 * 1_024; + +const I8_LEN: usize = primitive_len::(); +const I16_LEN: usize = primitive_len::(); +const I32_LEN: usize = primitive_len::(); +const I64_LEN: usize = primitive_len::(); + +/// Per-row against per-row, short and long. This is the shape the operators are tuned for, so it +/// is the one every operator is measured on. +const BINARY_SHAPE_CASES: &[(usize, BinaryShape)] = &[ + (128, BinaryShape::PerRowPerRow), + (I64_LEN, BinaryShape::PerRowPerRow), +]; + +/// Constant operands are measured on `Add` alone, at the long length. +/// +/// Decoding a constant operand, orienting it, and carrying its validity are shared by every +/// operator, so repeating all three shapes under `Sub` and `Mul` costs three walltime legs apiece +/// and measures the same code again. These cases stay as the regression guard on that path. +const CONSTANT_SHAPE_CASES: &[(usize, BinaryShape)] = &[ + (I64_LEN, BinaryShape::PerRowConstant), + (I64_LEN, BinaryShape::ConstantPerRow), + (I64_LEN, BinaryShape::PerRowNullableConstant), +]; + +#[derive(Clone, Copy, Debug)] +enum BinaryShape { + PerRowPerRow, + PerRowConstant, + ConstantPerRow, + PerRowNullableConstant, +} + +/// Decimal Mul and Div cost far more per lane than Add, so they run over a shorter array to keep +/// the instrumented CodSpeed runs quick. +const DECIMAL_MUL_DIV_LEN: usize = 1_024; + +#[vortex_bench_support::cpu_features] +#[divan::bench(args = BINARY_SHAPE_CASES)] +fn add_shapes(bencher: Bencher, &(len, shape): &(usize, BinaryShape)) { + bench_binary_shape(bencher, len, shape, Operator::Add); +} + +#[vortex_bench_support::cpu_features] +#[divan::bench(args = CONSTANT_SHAPE_CASES)] +fn add_constant_shapes(bencher: Bencher, &(len, shape): &(usize, BinaryShape)) { + bench_binary_shape(bencher, len, shape, Operator::Add); +} + +#[vortex_bench_support::cpu_features] +#[divan::bench(args = BINARY_SHAPE_CASES)] +fn subtract_shapes(bencher: Bencher, &(len, shape): &(usize, BinaryShape)) { + bench_binary_shape(bencher, len, shape, Operator::Sub); +} + +#[vortex_bench_support::cpu_features] +#[divan::bench(args = BINARY_SHAPE_CASES)] +fn multiply_shapes(bencher: Bencher, &(len, shape): &(usize, BinaryShape)) { + bench_binary_shape(bencher, len, shape, Operator::Mul); +} + +fn bench_binary_shape(bencher: Bencher, len: usize, shape: BinaryShape, operator: Operator) { + let per_row = + || PrimitiveArray::from_iter((0..len).map(|index| (index % 1_024) as i64 + 1)).into_array(); + let constant = || ConstantArray::new(17_i64, len).into_array(); + let nullable_constant = || ConstantArray::new(Some(17_i64), len).into_array(); + + let (lhs, rhs) = match shape { + BinaryShape::PerRowPerRow => (per_row(), per_row()), + BinaryShape::PerRowConstant => (per_row(), constant()), + BinaryShape::ConstantPerRow => (constant(), per_row()), + BinaryShape::PerRowNullableConstant => (per_row(), nullable_constant()), + }; + bench_primitive(bencher, lhs, rhs, operator); +} + +#[vortex_bench_support::cpu_features] #[divan::bench] fn add_i64_nonnull(bencher: Bencher) { - let lhs = primitive_nonnull(0).into_array(); - let rhs = primitive_nonnull(1_000_000).into_array(); + let lhs = primitive_nonnull(0, I64_LEN).into_array(); + let rhs = primitive_nonnull(1_000_000, I64_LEN).into_array(); bench_primitive(bencher, lhs, rhs, Operator::Add); } +#[vortex_bench_support::cpu_features] #[divan::bench] fn add_i64_nullable(bencher: Bencher) { - let lhs = primitive_nullable(0, 7).into_array(); - let rhs = primitive_nullable(1_000_000, 5).into_array(); + let lhs = primitive_nullable(0, 7, I64_LEN).into_array(); + let rhs = primitive_nullable(1_000_000, 5, I64_LEN).into_array(); bench_primitive(bencher, lhs, rhs, Operator::Add); } +#[vortex_bench_support::cpu_features] +#[divan::bench] +fn add_i32_nonnull(bencher: Bencher) { + let lhs = primitive_i32_small_nonnull(1, I32_LEN).into_array(); + let rhs = primitive_i32_small_nonnull(17, I32_LEN).into_array(); + + bench_primitive(bencher, lhs, rhs, Operator::Add); +} + +#[vortex_bench_support::cpu_features] +#[divan::bench] +fn add_u32_nonnull(bencher: Bencher) { + let lhs = primitive_u32_small_nonnull(1, I32_LEN).into_array(); + let rhs = primitive_u32_small_nonnull(17, I32_LEN).into_array(); + + bench_primitive(bencher, lhs, rhs, Operator::Add); +} + +#[vortex_bench_support::cpu_features] #[divan::bench] fn mul_i64_nonnull(bencher: Bencher) { - let lhs = primitive_small_nonnull(1).into_array(); - let rhs = primitive_small_nonnull(17).into_array(); + let lhs = primitive_small_nonnull(1, I64_LEN).into_array(); + let rhs = primitive_small_nonnull(17, I64_LEN).into_array(); bench_primitive(bencher, lhs, rhs, Operator::Mul); } +#[vortex_bench_support::cpu_features] #[divan::bench] fn mul_i8_nonnull(bencher: Bencher) { - let lhs = primitive_i8_small_nonnull(1).into_array(); - let rhs = primitive_i8_small_nonnull(7).into_array(); + let lhs = primitive_i8_small_nonnull(1, I8_LEN).into_array(); + let rhs = primitive_i8_small_nonnull(7, I8_LEN).into_array(); bench_primitive(bencher, lhs, rhs, Operator::Mul); } +#[vortex_bench_support::cpu_features] #[divan::bench] fn mul_u8_nonnull(bencher: Bencher) { - let lhs = primitive_u8_small_nonnull(1).into_array(); - let rhs = primitive_u8_small_nonnull(7).into_array(); + let lhs = primitive_u8_small_nonnull(1, I8_LEN).into_array(); + let rhs = primitive_u8_small_nonnull(7, I8_LEN).into_array(); bench_primitive(bencher, lhs, rhs, Operator::Mul); } +#[vortex_bench_support::cpu_features] #[divan::bench] fn mul_i16_nonnull(bencher: Bencher) { - let lhs = primitive_i16_small_nonnull(1).into_array(); - let rhs = primitive_i16_small_nonnull(17).into_array(); + let lhs = primitive_i16_small_nonnull(1, I16_LEN).into_array(); + let rhs = primitive_i16_small_nonnull(17, I16_LEN).into_array(); bench_primitive(bencher, lhs, rhs, Operator::Mul); } +#[vortex_bench_support::cpu_features] #[divan::bench] fn mul_u16_nonnull(bencher: Bencher) { - let lhs = primitive_u16_small_nonnull(1).into_array(); - let rhs = primitive_u16_small_nonnull(17).into_array(); + let lhs = primitive_u16_small_nonnull(1, I16_LEN).into_array(); + let rhs = primitive_u16_small_nonnull(17, I16_LEN).into_array(); bench_primitive(bencher, lhs, rhs, Operator::Mul); } +#[vortex_bench_support::cpu_features] #[divan::bench] fn mul_i32_nonnull(bencher: Bencher) { - let lhs = primitive_i32_small_nonnull(1).into_array(); - let rhs = primitive_i32_small_nonnull(17).into_array(); + let lhs = primitive_i32_small_nonnull(1, I32_LEN).into_array(); + let rhs = primitive_i32_small_nonnull(17, I32_LEN).into_array(); bench_primitive(bencher, lhs, rhs, Operator::Mul); } +#[vortex_bench_support::cpu_features] #[divan::bench] fn mul_u32_nonnull(bencher: Bencher) { - let lhs = primitive_u32_small_nonnull(1).into_array(); - let rhs = primitive_u32_small_nonnull(17).into_array(); + let lhs = primitive_u32_small_nonnull(1, I32_LEN).into_array(); + let rhs = primitive_u32_small_nonnull(17, I32_LEN).into_array(); bench_primitive(bencher, lhs, rhs, Operator::Mul); } +#[vortex_bench_support::cpu_features] #[divan::bench] -fn mul_i32_nullable(bencher: Bencher) { - let lhs = primitive_i32_small_nullable(1, 7).into_array(); - let rhs = primitive_i32_small_nullable(17, 5).into_array(); +fn mul_u64_nonnull(bencher: Bencher) { + let lhs = primitive_u64_small_nonnull(1, I64_LEN).into_array(); + let rhs = primitive_u64_small_nonnull(17, I64_LEN).into_array(); bench_primitive(bencher, lhs, rhs, Operator::Mul); } +#[vortex_bench_support::cpu_features] #[divan::bench] -fn mul_i32_constant(bencher: Bencher) { - let lhs = primitive_i32_small_nonnull(1).into_array(); - let rhs = ConstantArray::new(31i32, LEN).into_array(); +fn mul_i32_nullable(bencher: Bencher) { + let lhs = primitive_i32_small_nullable(1, 7, I32_LEN).into_array(); + let rhs = primitive_i32_small_nullable(17, 5, I32_LEN).into_array(); bench_primitive(bencher, lhs, rhs, Operator::Mul); } +#[vortex_bench_support::cpu_features] #[divan::bench] fn div_i64_nonnull(bencher: Bencher) { - let lhs = primitive_nonnull(1_000_000).into_array(); - let rhs = primitive_nonzero().into_array(); + let lhs = primitive_nonnull(1_000_000, I64_LEN).into_array(); + let rhs = primitive_nonzero(I64_LEN).into_array(); bench_primitive(bencher, lhs, rhs, Operator::Div); } +#[vortex_bench_support::cpu_features] #[divan::bench] -fn sub_i64_constant(bencher: Bencher) { - let lhs = primitive_nonnull(0).into_array(); - let rhs = ConstantArray::new(37i64, LEN).into_array(); +fn div_i64_nullable(bencher: Bencher) { + let lhs = primitive_nullable(1_000_000, 7, I64_LEN).into_array(); + let rhs = primitive_nullable(17, 5, I64_LEN).into_array(); - bench_primitive(bencher, lhs, rhs, Operator::Sub); + bench_primitive(bencher, lhs, rhs, Operator::Div); } #[divan::bench] -fn eq_i64_constant(bencher: Bencher) { - let lhs = primitive_nonnull(0).into_array(); - let rhs = ConstantArray::new(1024i64, LEN).into_array(); +fn add_decimal_i64_nonnull(bencher: Bencher) { + let lhs = decimal_i64_nonnull(0, LEN).into_array(); + let rhs = decimal_i64_nonnull(1_000_000, LEN).into_array(); - bench_bool(bencher, lhs, rhs, Operator::Eq); + bench_decimal(bencher, lhs, rhs, Operator::Add); } +#[divan::bench] +fn add_decimal_i128_nullable(bencher: Bencher) { + let lhs = decimal_i128_nullable(0, 7, LEN).into_array(); + let rhs = decimal_i128_nullable(1_000_000, 5, LEN).into_array(); + + bench_decimal(bencher, lhs, rhs, Operator::Add); +} + +#[divan::bench] +fn mul_decimal_i64_nonnull(bencher: Bencher) { + let lhs = decimal_i64_nonnull(0, DECIMAL_MUL_DIV_LEN).into_array(); + let rhs = decimal_i64_nonnull(1_000_000, DECIMAL_MUL_DIV_LEN).into_array(); + + bench_decimal(bencher, lhs, rhs, Operator::Mul); +} + +#[divan::bench] +fn mul_decimal_i128_nullable(bencher: Bencher) { + let lhs = decimal_i128_nullable(0, 7, DECIMAL_MUL_DIV_LEN).into_array(); + let rhs = decimal_i128_nullable(1_000_000, 5, DECIMAL_MUL_DIV_LEN).into_array(); + + bench_decimal(bencher, lhs, rhs, Operator::Mul); +} + +#[divan::bench] +fn div_decimal_i64_nonnull(bencher: Bencher) { + let lhs = decimal_i64_nonnull(0, DECIMAL_MUL_DIV_LEN).into_array(); + let rhs = decimal_i64_nonnull(1_000_000, DECIMAL_MUL_DIV_LEN).into_array(); + + bench_decimal(bencher, lhs, rhs, Operator::Div); +} + +#[divan::bench] +fn div_decimal_i128_nullable(bencher: Bencher) { + let lhs = decimal_i128_nullable(0, 7, DECIMAL_MUL_DIV_LEN).into_array(); + let rhs = decimal_i128_nullable(1_000_000, 5, DECIMAL_MUL_DIV_LEN).into_array(); + + bench_decimal(bencher, lhs, rhs, Operator::Div); +} + +#[vortex_bench_support::cpu_features] #[divan::bench] fn lt_i64_nullable(bencher: Bencher) { - let lhs = primitive_nullable(0, 7).into_array(); - let rhs = primitive_nullable(1_000_000, 5).into_array(); + let lhs = primitive_nullable(0, 7, LEN).into_array(); + let rhs = primitive_nullable(1_000_000, 5, LEN).into_array(); bench_bool(bencher, lhs, rhs, Operator::Lt); } @@ -164,18 +331,14 @@ fn and_bool_nullable(bencher: Bencher) { bench_bool(bencher, lhs, rhs, Operator::And); } -#[divan::bench] -fn or_bool_constant(bencher: Bencher) { - let lhs = bool_nullable(2, 7).into_array(); - let rhs = ConstantArray::new(true, LEN).into_array(); - - bench_bool(bencher, lhs, rhs, Operator::Or); -} - fn bench_primitive(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, operator: Operator) { bench_binary::(bencher, lhs, rhs, operator); } +fn bench_decimal(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, operator: Operator) { + bench_binary::(bencher, lhs, rhs, operator); +} + fn bench_bool(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, operator: Operator) { bench_binary::(bencher, lhs, rhs, operator); } @@ -187,8 +350,9 @@ fn bench_binary( operator: Operator, ) { let mut ctx = SESSION.create_execution_ctx(); + let len = lhs.len(); - bencher.counter(ItemsCount::new(LEN)).bench_local(|| { + bencher.counter(ItemsCount::new(len)).bench_local(|| { lhs.clone() .binary(rhs.clone(), operator) .unwrap() @@ -197,50 +361,74 @@ fn bench_binary( }); } -fn primitive_nonnull(base: i64) -> PrimitiveArray { - PrimitiveArray::from_iter((0..LEN as i64).map(|i| base + i)) +const fn primitive_len() -> usize { + let input_len = MIN_PRIMITIVE_INPUT_BYTES / size_of::(); + if input_len > MIN_PRIMITIVE_LEN { + input_len + } else { + MIN_PRIMITIVE_LEN + } +} + +fn primitive_nonnull(base: i64, len: usize) -> PrimitiveArray { + PrimitiveArray::from_iter((0..len as i64).map(|i| base + i)) +} + +fn decimal_i64_nonnull(base: i64, len: usize) -> DecimalArray { + DecimalArray::from_iter::((0..len as i64).map(|i| base + i), DecimalDType::new(18, 2)) +} + +fn decimal_i128_nullable(base: i128, null_every: usize, len: usize) -> DecimalArray { + DecimalArray::from_option_iter::( + (0..len as i128).map(|i| (!(i as usize).is_multiple_of(null_every)).then_some(base + i)), + DecimalDType::new(38, 2), + ) +} + +fn primitive_small_nonnull(offset: i64, len: usize) -> PrimitiveArray { + PrimitiveArray::from_iter((0..len as i64).map(|i| ((i + offset) % 1024) + 1)) } -fn primitive_small_nonnull(offset: i64) -> PrimitiveArray { - PrimitiveArray::from_iter((0..LEN as i64).map(|i| ((i + offset) % 1024) + 1)) +fn primitive_i8_small_nonnull(offset: i8, len: usize) -> PrimitiveArray { + PrimitiveArray::from_iter((0..len).map(|i| (((i as i32 + offset as i32) % 21) - 10) as i8)) } -fn primitive_i8_small_nonnull(offset: i8) -> PrimitiveArray { - PrimitiveArray::from_iter((0..LEN).map(|i| (((i as i32 + offset as i32) % 21) - 10) as i8)) +fn primitive_u8_small_nonnull(offset: u8, len: usize) -> PrimitiveArray { + PrimitiveArray::from_iter((0..len).map(|i| ((i + offset as usize) % 15 + 1) as u8)) } -fn primitive_u8_small_nonnull(offset: u8) -> PrimitiveArray { - PrimitiveArray::from_iter((0..LEN).map(|i| ((i + offset as usize) % 15 + 1) as u8)) +fn primitive_i16_small_nonnull(offset: i16, len: usize) -> PrimitiveArray { + PrimitiveArray::from_iter((0..len).map(|i| (((i as i32 + offset as i32) % 255) - 127) as i16)) } -fn primitive_i16_small_nonnull(offset: i16) -> PrimitiveArray { - PrimitiveArray::from_iter((0..LEN).map(|i| (((i as i32 + offset as i32) % 255) - 127) as i16)) +fn primitive_u16_small_nonnull(offset: u16, len: usize) -> PrimitiveArray { + PrimitiveArray::from_iter((0..len).map(|i| ((i + offset as usize) % 251 + 1) as u16)) } -fn primitive_u16_small_nonnull(offset: u16) -> PrimitiveArray { - PrimitiveArray::from_iter((0..LEN).map(|i| ((i + offset as usize) % 251 + 1) as u16)) +fn primitive_i32_small_nonnull(offset: i32, len: usize) -> PrimitiveArray { + PrimitiveArray::from_iter((0..len).map(|i| (((i as i64 + offset as i64) % 4096) - 2048) as i32)) } -fn primitive_i32_small_nonnull(offset: i32) -> PrimitiveArray { - PrimitiveArray::from_iter((0..LEN).map(|i| (((i as i64 + offset as i64) % 4096) - 2048) as i32)) +fn primitive_u32_small_nonnull(offset: u32, len: usize) -> PrimitiveArray { + PrimitiveArray::from_iter((0..len).map(|i| ((i + offset as usize) % 4096 + 1) as u32)) } -fn primitive_u32_small_nonnull(offset: u32) -> PrimitiveArray { - PrimitiveArray::from_iter((0..LEN).map(|i| ((i + offset as usize) % 4096 + 1) as u32)) +fn primitive_u64_small_nonnull(offset: u64, len: usize) -> PrimitiveArray { + PrimitiveArray::from_iter((0..len).map(|i| ((i + offset as usize) % 4096 + 1) as u64)) } -fn primitive_nonzero() -> PrimitiveArray { - PrimitiveArray::from_iter((0..LEN as i64).map(|i| (i % 255) + 1)) +fn primitive_nonzero(len: usize) -> PrimitiveArray { + PrimitiveArray::from_iter((0..len as i64).map(|i| (i % 255) + 1)) } -fn primitive_nullable(base: i64, null_every: usize) -> PrimitiveArray { +fn primitive_nullable(base: i64, null_every: usize, len: usize) -> PrimitiveArray { PrimitiveArray::from_option_iter( - (0..LEN as i64).map(|i| (!(i as usize).is_multiple_of(null_every)).then_some(base + i)), + (0..len as i64).map(|i| (!(i as usize).is_multiple_of(null_every)).then_some(base + i)), ) } -fn primitive_i32_small_nullable(offset: i32, null_every: usize) -> PrimitiveArray { - PrimitiveArray::from_option_iter((0..LEN).map(|i| { +fn primitive_i32_small_nullable(offset: i32, null_every: usize, len: usize) -> PrimitiveArray { + PrimitiveArray::from_option_iter((0..len).map(|i| { (!i.is_multiple_of(null_every)) .then_some((((i as i64 + offset as i64) % 4096) - 2048) as i32) })) diff --git a/vortex-array/benches/cast_decimal.rs b/vortex-array/benches/cast_decimal.rs index 8e0ddb19b6b..c1c32147059 100644 --- a/vortex-array/benches/cast_decimal.rs +++ b/vortex-array/benches/cast_decimal.rs @@ -14,6 +14,7 @@ use std::sync::LazyLock; use divan::Bencher; +use mimalloc::MiMalloc; use rand::prelude::*; use vortex_array::ArrayRef; use vortex_array::Canonical; @@ -28,12 +29,19 @@ use vortex_array::validity::Validity; use vortex_buffer::BufferMut; use vortex_session::VortexSession; +// The copy casts allocate their output buffer inside the timed region, so route allocation +// through vendored mimalloc to keep glibc malloc (which varies across runner images) out of +// the measured trace. +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + fn main() { divan::main(); } -// Kept small enough to stay in L2 so the kernel cost shows up rather than DRAM bandwidth. -const SIZES: &[usize] = &[65_536]; +// Kept small enough to stay in L2 so the kernel cost shows up rather than DRAM bandwidth, +// and to keep the CodSpeed simulation under 1ms per benchmark. +const SIZES: &[usize] = &[16_384]; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); diff --git a/vortex-array/benches/cast_primitive.rs b/vortex-array/benches/cast_primitive.rs index 69a4c5d2cc5..5ffc7ba3f39 100644 --- a/vortex-array/benches/cast_primitive.rs +++ b/vortex-array/benches/cast_primitive.rs @@ -14,6 +14,7 @@ use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::expr::stats::Stat; @@ -25,8 +26,9 @@ fn main() { } // Sizes used for the fallible-path benches below. Kept small enough to fit in L2 so -// the kernel cost shows up clearly rather than being hidden by DRAM bandwidth. -const SIZES: &[usize] = &[65_536]; +// the kernel cost shows up clearly rather than being hidden by DRAM bandwidth, and +// to keep the CodSpeed simulation under 1ms per benchmark. +const SIZES: &[usize] = &[8_192]; static SESSION: LazyLock = LazyLock::new(array_session); @@ -87,3 +89,55 @@ fn cast_i32_to_u32(bencher: Bencher, n: usize) { .execute::(ctx) }); } + +/// Integer-to-decimal cast at the largest precision that uses an i128 backing buffer. This +/// exercises the common rescaling path without paying for i256 arithmetic per input value. +#[divan::bench(args = SIZES)] +fn cast_i64_to_decimal38_scale2(bencher: Bencher, n: usize) { + let arr = + PrimitiveArray::from_iter((0..n).map(|value| i64::try_from(value).unwrap())).into_array(); + bencher + .with_inputs(|| (arr.clone(), SESSION.create_execution_ctx())) + .bench_refs(|(a, ctx)| { + a.cast(DType::Decimal( + DecimalDType::new(38, 2), + Nullability::NonNullable, + )) + .unwrap() + .execute::(ctx) + }); +} + +/// Common integer-to-decimal cast that rescales directly in its i32 output buffer. +#[divan::bench(args = SIZES)] +fn cast_i32_to_decimal9_scale2(bencher: Bencher, n: usize) { + let arr = + PrimitiveArray::from_iter((0..n).map(|value| i32::try_from(value).unwrap())).into_array(); + bencher + .with_inputs(|| (arr.clone(), SESSION.create_execution_ctx())) + .bench_refs(|(a, ctx)| { + a.cast(DType::Decimal( + DecimalDType::new(9, 2), + Nullability::NonNullable, + )) + .unwrap() + .execute::(ctx) + }); +} + +/// Same-width scale-zero cast that validates then reuses the source values buffer. +#[divan::bench(args = SIZES)] +fn cast_i32_to_decimal9_scale0(bencher: Bencher, n: usize) { + let arr = + PrimitiveArray::from_iter((0..n).map(|value| i32::try_from(value).unwrap())).into_array(); + bencher + .with_inputs(|| (arr.clone(), SESSION.create_execution_ctx())) + .bench_refs(|(a, ctx)| { + a.cast(DType::Decimal( + DecimalDType::new(9, 0), + Nullability::NonNullable, + )) + .unwrap() + .execute::(ctx) + }); +} diff --git a/vortex-array/benches/chunk_array_builder.rs b/vortex-array/benches/chunk_array_builder.rs index c2dedc6ef46..954c6a677e1 100644 --- a/vortex-array/benches/chunk_array_builder.rs +++ b/vortex-array/benches/chunk_array_builder.rs @@ -17,8 +17,9 @@ use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::ConstantArray; use vortex_array::builders::ArrayBuilder; +use vortex_array::builders::VarBinBuilder; use vortex_array::builders::VarBinViewBuilder; -use vortex_array::builders::builder_with_capacity; +use vortex_array::builders::builder_with_capacity_in; use vortex_array::dtype::DType; use vortex_error::VortexExpect; use vortex_session::VortexSession; @@ -33,8 +34,9 @@ fn main() { const BENCH_ARGS: &[(usize, usize)] = &[ // length, chunk_count - (10, 1000), - (100, 100), + // Chunk counts sized to keep CodSpeed simulation under 1ms per benchmark. + (10, 100), + (100, 50), (1000, 10), ]; @@ -47,7 +49,8 @@ fn chunked_bool_canonical_into(bencher: Bencher, (len, chunk_count): (usize, usi bencher .with_inputs(|| (&chunk, SESSION.create_execution_ctx())) .bench_refs(|(chunk, ctx)| { - let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count); + let mut builder = + builder_with_capacity_in(chunk.dtype(), len * chunk_count, ctx.allocator()); chunk .append_to_builder(builder.as_mut(), ctx) .vortex_expect("append failed"); @@ -62,7 +65,8 @@ fn chunked_opt_bool_canonical_into(bencher: Bencher, (len, chunk_count): (usize, bencher .with_inputs(|| (&chunk, SESSION.create_execution_ctx())) .bench_refs(|(chunk, ctx)| { - let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count); + let mut builder = + builder_with_capacity_in(chunk.dtype(), len * chunk_count, ctx.allocator()); chunk .append_to_builder(builder.as_mut(), ctx) .vortex_expect("append failed"); @@ -86,9 +90,10 @@ fn chunked_varbinview_canonical_into(bencher: Bencher, (len, chunk_count): (usiz bencher .with_inputs(|| (&chunks, SESSION.create_execution_ctx())) .bench_refs(|(chunk, ctx)| { - let mut builder = VarBinViewBuilder::with_capacity( + let mut builder = VarBinViewBuilder::with_capacity_in( DType::Utf8(chunk.dtype().nullability()), len * chunk_count, + ctx.allocator().clone(), ); chunk .append_to_builder(&mut builder, ctx) @@ -113,9 +118,10 @@ fn chunked_varbinview_opt_canonical_into(bencher: Bencher, (len, chunk_count): ( bencher .with_inputs(|| (&chunks, SESSION.create_execution_ctx())) .bench_refs(|(chunk, ctx)| { - let mut builder = VarBinViewBuilder::with_capacity( + let mut builder = VarBinViewBuilder::with_capacity_in( DType::Utf8(chunk.dtype().nullability()), len * chunk_count, + ctx.allocator().clone(), ); chunk .append_to_builder(&mut builder, ctx) @@ -133,6 +139,60 @@ fn chunked_varbinview_opt_into_canonical(bencher: Bencher, (len, chunk_count): ( .bench_refs(|(chunk, ctx)| (**chunk).clone().execute::(ctx)) } +// Fewer rows than BENCH_ARGS: decoding VarBin values is the most expensive work in this file. +const VARBIN_BENCH_ARGS: &[(usize, usize)] = &[ + // length, chunk_count + (10, 100), + (500, 2), +]; + +#[divan::bench(args = VARBIN_BENCH_ARGS)] +fn chunked_varbin_to_varbinview_builder(bencher: Bencher, (len, chunk_count): (usize, usize)) { + let chunks = make_varbin_chunks(false, len, chunk_count); + + bencher + .with_inputs(|| (&chunks, SESSION.create_execution_ctx())) + .bench_refs(|(chunk, ctx)| { + let mut builder = VarBinViewBuilder::with_capacity_in( + chunk.dtype().clone(), + len * chunk_count, + ctx.allocator().clone(), + ); + chunk + .append_to_builder(&mut builder, ctx) + .vortex_expect("append failed"); + builder.finish() + }) +} + +#[divan::bench(args = VARBIN_BENCH_ARGS)] +fn chunked_varbin_opt_to_varbinview_builder(bencher: Bencher, (len, chunk_count): (usize, usize)) { + let chunks = make_varbin_chunks(true, len, chunk_count); + + bencher + .with_inputs(|| (&chunks, SESSION.create_execution_ctx())) + .bench_refs(|(chunk, ctx)| { + let mut builder = VarBinViewBuilder::with_capacity_in( + chunk.dtype().clone(), + len * chunk_count, + ctx.allocator().clone(), + ); + chunk + .append_to_builder(&mut builder, ctx) + .vortex_expect("append failed"); + builder.finish() + }) +} + +#[divan::bench(args = VARBIN_BENCH_ARGS)] +fn chunked_varbin_into_canonical(bencher: Bencher, (len, chunk_count): (usize, usize)) { + let chunks = make_varbin_chunks(false, len, chunk_count); + + bencher + .with_inputs(|| (&chunks, SESSION.create_execution_ctx())) + .bench_refs(|(chunk, ctx)| (**chunk).clone().execute::(ctx)) +} + #[divan::bench(args = BENCH_ARGS)] fn chunked_constant_i32_append_to_builder(bencher: Bencher, (len, chunk_count): (usize, usize)) { let chunk = make_constant_i32_chunks(len, chunk_count); @@ -140,7 +200,8 @@ fn chunked_constant_i32_append_to_builder(bencher: Bencher, (len, chunk_count): bencher .with_inputs(|| (&chunk, SESSION.create_execution_ctx())) .bench_refs(|(chunk, ctx)| { - let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count); + let mut builder = + builder_with_capacity_in(chunk.dtype(), len * chunk_count, ctx.allocator()); chunk .append_to_builder(builder.as_mut(), ctx) .vortex_expect("append failed"); @@ -164,7 +225,8 @@ fn chunked_constant_utf8_append_to_builder( bencher .with_inputs(|| (&chunk, SESSION.create_execution_ctx())) .bench_refs(|(chunk, ctx)| { - let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count); + let mut builder = + builder_with_capacity_in(chunk.dtype(), len * chunk_count, ctx.allocator()); chunk .append_to_builder(builder.as_mut(), ctx) .vortex_expect("append failed"); @@ -226,12 +288,44 @@ fn make_bool_chunks(len: usize, chunk_count: usize) -> ArrayRef { .into_array() } +fn make_varbin_chunks(nullable: bool, len: usize, chunk_count: usize) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(123); + let dtype = DType::Utf8(nullable.into()); + + (0..chunk_count) + .map(|_| { + let mut builder = VarBinBuilder::::with_capacity_in( + dtype.clone(), + len, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + (0..len).for_each(|_| { + if nullable && rng.random_bool(0.2) { + builder.push_null() + } else { + builder.append_value( + (0..rng.random_range(0..=20)) + .map(|_| rng.random_range(b'a'..=b'z')) + .collect::>(), + ) + } + }); + builder.finish() + }) + .collect::() + .into_array() +} + fn make_string_chunks(nullable: bool, len: usize, chunk_count: usize) -> ArrayRef { let mut rng = StdRng::seed_from_u64(123); (0..chunk_count) .map(|_| { - let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(nullable.into()), len); + let mut builder = VarBinViewBuilder::with_capacity_in( + DType::Utf8(nullable.into()), + len, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); (0..len).for_each(|_| { if nullable && rng.random_bool(0.2) { builder.append_null() diff --git a/vortex-array/benches/chunked_dict_builder.rs b/vortex-array/benches/chunked_dict_builder.rs index de2f40dcbf5..41fc16b19e8 100644 --- a/vortex-array/benches/chunked_dict_builder.rs +++ b/vortex-array/benches/chunked_dict_builder.rs @@ -9,7 +9,7 @@ use rand::distr::StandardUniform; use vortex_array::Canonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::dict_test::gen_dict_primitive_chunks; -use vortex_array::builders::builder_with_capacity; +use vortex_array::builders::builder_with_capacity_in; use vortex_array::dtype::NativePType; use vortex_error::VortexExpect; use vortex_session::VortexSession; @@ -21,13 +21,14 @@ fn main() { static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +// Chunk counts sized to keep CodSpeed simulation under 1ms per benchmark. const BENCH_ARGS: &[(usize, usize, usize)] = &[ (1000, 10, 10), (1000, 100, 10), (1000, 1000, 10), - (1000, 10, 100), - (1000, 100, 100), - (1000, 1000, 100), + (1000, 10, 20), + (1000, 100, 20), + (1000, 1000, 20), ]; #[divan::bench(types = [u32, u64, f32, f64], args = BENCH_ARGS)] @@ -42,7 +43,8 @@ fn chunked_dict_primitive_canonical_into( bencher .with_inputs(|| (&chunk, SESSION.create_execution_ctx())) .bench_refs(|(chunk, ctx)| { - let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count); + let mut builder = + builder_with_capacity_in(chunk.dtype(), len * chunk_count, ctx.allocator()); chunk .append_to_builder(builder.as_mut(), ctx) .vortex_expect("append failed"); diff --git a/vortex-array/benches/chunked_fsl_canonicalize.rs b/vortex-array/benches/chunked_fsl_canonicalize.rs index 9be8433ce9d..3718e756456 100644 --- a/vortex-array/benches/chunked_fsl_canonicalize.rs +++ b/vortex-array/benches/chunked_fsl_canonicalize.rs @@ -55,9 +55,7 @@ fn create_fsl(list_size: usize, num_lists: usize) -> FixedSizeListArray { fn create_chunked_fsl(list_size: usize, num_chunks: usize) -> ChunkedArray { let chunk = create_fsl(list_size, LISTS_PER_CHUNK); let dtype = chunk.dtype().clone(); - let chunks = (0..num_chunks) - .map(|_| chunk.clone().into_array()) - .collect(); + let chunks = (0..num_chunks).map(|_| chunk.clone().into_array()); ChunkedArray::try_new(chunks, dtype).unwrap() } diff --git a/vortex-array/benches/compare.rs b/vortex-array/benches/compare.rs index 4a399760dc2..8e040ce1b91 100644 --- a/vortex-array/benches/compare.rs +++ b/vortex-array/benches/compare.rs @@ -1,9 +1,29 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Benchmarks for the binary comparison path, over every array kind it accepts. +//! +//! The primitive cases carry `#[cpu_features]`, so they are measured on every walltime +//! CPU-feature leg rather than in simulation. Each is written once and compiled differently +//! per leg: today the primitive comparison path is a portable lane kernel, and how well it +//! auto-vectorizes is decided by the build. That is the baseline a hand-written kernel +//! selected through `cfg(target_feature)` has to beat, measured on the silicon it would run +//! on. +//! +//! Every case here compares one array against another, which is the shape the path is tuned for. +//! The three constant cases that remain — boolean, integer, and string — are a regression guard on +//! constant handling rather than coverage of it: a constant operand is decoded the same way +//! whatever it holds. Constant orientation is measured once, in `binary_ops.rs`. +//! +//! The boolean, decimal, string, and struct cases are not tagged. A wider vector register is +//! not what decides them: booleans are already word-at-a-time over a bitmap, decimals are +//! `i128`, and the string and struct cases are dominated by view chasing and per-field +//! dispatch. They stay in simulation. + #![expect(clippy::unwrap_used)] use divan::Bencher; +use divan::counter::ItemsCount; use mimalloc::MiMalloc; use rand::RngExt; use rand::SeedableRng; @@ -33,11 +53,15 @@ fn main() { divan::main(); } -const ARRAY_SIZE: usize = 65_536; +// Sized to keep CodSpeed simulation under 1ms per benchmark. +const ARRAY_SIZE: usize = 8_192; fn bench_compare(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, op: Operator) { let session = vortex_array::array_session(); + let len = lhs.len(); + bencher + .counter(ItemsCount::new(len)) .with_inputs(|| (&lhs, &rhs, session.create_execution_ctx())) .bench_refs(|input| { input @@ -87,6 +111,27 @@ fn float_array(rng: &mut StdRng) -> ArrayRef { .into_array() } +fn u8_array(rng: &mut StdRng) -> ArrayRef { + (0..ARRAY_SIZE) + .map(|_| rng.random::()) + .collect::>() + .into_array() +} + +fn u64_array(rng: &mut StdRng) -> ArrayRef { + (0..ARRAY_SIZE) + .map(|_| rng.random::()) + .collect::>() + .into_array() +} + +fn f32_array(rng: &mut StdRng) -> ArrayRef { + (0..ARRAY_SIZE) + .map(|_| rng.random_range(0.0f32..1.0)) + .collect::>() + .into_array() +} + fn string_array(rng: &mut StdRng) -> ArrayRef { VarBinViewArray::from_iter_str((0..ARRAY_SIZE).map(|_| { let len = rng.random_range(1usize..24); @@ -129,6 +174,7 @@ fn compare_bool_constant(bencher: Bencher) { bench_compare(bencher, arr, constant, Operator::Eq); } +#[vortex_bench_support::cpu_features] #[divan::bench] fn compare_int(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); @@ -137,6 +183,7 @@ fn compare_int(bencher: Bencher) { bench_compare(bencher, arr1, arr2, Operator::Gte); } +#[vortex_bench_support::cpu_features] #[divan::bench] fn compare_int_nullable(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); @@ -145,6 +192,7 @@ fn compare_int_nullable(bencher: Bencher) { bench_compare(bencher, arr1, arr2, Operator::Gte); } +#[vortex_bench_support::cpu_features] #[divan::bench] fn compare_int_constant(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); @@ -153,6 +201,34 @@ fn compare_int_constant(bencher: Bencher) { bench_compare(bencher, arr, constant, Operator::Gte); } +#[vortex_bench_support::cpu_features] +#[divan::bench] +fn compare_u8(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = u8_array(&mut rng); + let arr2 = u8_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Gte); +} + +#[vortex_bench_support::cpu_features] +#[divan::bench] +fn compare_u64(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = u64_array(&mut rng); + let arr2 = u64_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Gte); +} + +#[vortex_bench_support::cpu_features] +#[divan::bench] +fn compare_f32(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = f32_array(&mut rng); + let arr2 = f32_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Gte); +} + +#[vortex_bench_support::cpu_features] #[divan::bench] fn compare_int_eq(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); @@ -161,6 +237,7 @@ fn compare_int_eq(bencher: Bencher) { bench_compare(bencher, arr1, arr2, Operator::Eq); } +#[vortex_bench_support::cpu_features] #[divan::bench] fn compare_float(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); diff --git a/vortex-array/benches/dict_compare.rs b/vortex-array/benches/dict_compare.rs index 5b31cf51831..adf7e3cfb49 100644 --- a/vortex-array/benches/dict_compare.rs +++ b/vortex-array/benches/dict_compare.rs @@ -40,13 +40,13 @@ const LENGTH_AND_UNIQUE_VALUES: &[(usize, usize)] = &[ (10_000, 128), (10_000, 512), (10_000, 2048), - (100_000, 2), - (100_000, 4), - (100_000, 8), - (100_000, 32), - (100_000, 128), - (100_000, 512), - (100_000, 2048), + (50_000, 2), + (50_000, 4), + (50_000, 8), + (50_000, 32), + (50_000, 128), + (50_000, 512), + (50_000, 2048), ]; #[divan::bench(args = LENGTH_AND_UNIQUE_VALUES)] diff --git a/vortex-array/benches/dict_compress.rs b/vortex-array/benches/dict_compress.rs index fa53dae8c43..e8b1b0ec8c5 100644 --- a/vortex-array/benches/dict_compress.rs +++ b/vortex-array/benches/dict_compress.rs @@ -31,16 +31,17 @@ fn main() { const BENCH_ARGS: &[(usize, usize)] = &[ // length, unique_values + // Lengths sized to keep CodSpeed simulation under 1ms per benchmark. (1_000, 2), (1_000, 4), (1_000, 8), (1_000, 32), (1_000, 512), - (10_000, 2), - (10_000, 4), - (10_000, 8), - (10_000, 32), - (10_000, 512), + (2_000, 2), + (2_000, 4), + (2_000, 8), + (2_000, 32), + (2_000, 512), ]; static SESSION: LazyLock = LazyLock::new(array_session); diff --git a/vortex-array/benches/dict_mask.rs b/vortex-array/benches/dict_mask.rs index 04cf18cfd03..35e1654fabc 100644 --- a/vortex-array/benches/dict_mask.rs +++ b/vortex-array/benches/dict_mask.rs @@ -47,7 +47,8 @@ fn filter_mask(len: usize, fraction_masked: f64, rng: &mut StdRng) -> Mask { fn bench_dict_mask(bencher: Bencher, (fraction_valid, fraction_masked): (f64, f64)) { let mut rng = StdRng::seed_from_u64(0); - let len = 65_535; + // Sized to keep CodSpeed simulation under 1ms per benchmark. + let len = 4_096; let codes = PrimitiveArray::from_iter((0..len).map(|_| { if rng.random_bool(fraction_valid) { 1u64 diff --git a/vortex-array/benches/dict_unreferenced_mask.rs b/vortex-array/benches/dict_unreferenced_mask.rs index ec8653d6688..ee06b7d44a4 100644 --- a/vortex-array/benches/dict_unreferenced_mask.rs +++ b/vortex-array/benches/dict_unreferenced_mask.rs @@ -63,7 +63,8 @@ fn bench_many_codes_few_values(bencher: Bencher, num_values: i32) { fn bench_many_nulls(bencher: Bencher, fraction_valid: f64) { let mut rng = StdRng::seed_from_u64(0); - let num_codes = 65_536; + // Sized to keep CodSpeed simulation under 1ms per benchmark. + let num_codes = 32_768; let num_values = 1024i32; // Create values array diff --git a/vortex-array/benches/expr/case_when_bench.rs b/vortex-array/benches/expr/case_when_bench.rs index c7b0e96686d..fa5af4bc0ef 100644 --- a/vortex-array/benches/expr/case_when_bench.rs +++ b/vortex-array/benches/expr/case_when_bench.rs @@ -7,6 +7,7 @@ use std::sync::LazyLock; use divan::Bencher; +use mimalloc::MiMalloc; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; @@ -26,6 +27,9 @@ use vortex_array::expr::root; use vortex_buffer::Buffer; use vortex_session::VortexSession; +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + static SESSION: LazyLock = LazyLock::new(array_session); fn main() { @@ -58,7 +62,7 @@ fn make_fragmented_array(size: usize) -> ArrayRef { } /// Benchmark a simple binary CASE WHEN with varying array sizes. -#[divan::bench(args = [1000, 10000, 100000])] +#[divan::bench(args = [1000, 10000, 40000])] fn case_when_simple(bencher: Bencher, size: usize) { let array = make_struct_array(size); @@ -160,7 +164,7 @@ fn case_when_nary_equality_lookup(bencher: Bencher, size: usize) { } /// Benchmark CASE WHEN without ELSE clause (result is nullable). -#[divan::bench(args = [1000, 10000, 100000])] +#[divan::bench(args = [1000, 10000, 40000])] fn case_when_without_else(bencher: Bencher, size: usize) { let array = make_struct_array(size); @@ -180,7 +184,7 @@ fn case_when_without_else(bencher: Bencher, size: usize) { } /// Benchmark CASE WHEN where all conditions are true. -#[divan::bench(args = [1000, 10000, 100000])] +#[divan::bench(args = [1000, 10000, 40000])] fn case_when_all_true(bencher: Bencher, size: usize) { let array = make_struct_array(size); @@ -237,7 +241,7 @@ fn case_when_nary_early_dominant(bencher: Bencher, size: usize) { } /// Benchmark CASE WHEN where all conditions are false. -#[divan::bench(args = [1000, 10000, 100000])] +#[divan::bench(args = [1000, 10000, 40000])] fn case_when_all_false(bencher: Bencher, size: usize) { let array = make_struct_array(size); @@ -262,7 +266,7 @@ fn case_when_all_false(bencher: Bencher, size: usize) { /// Benchmark CASE WHEN cycling through 3 branches per row (triggers merge_row_by_row). /// Run length = 1; exercises branch 0, branch 1, and the else fallback at every 3rd row. -#[divan::bench(args = [100, 1000])] +#[divan::bench(args = [100, 400])] fn case_when_fragmented(bencher: Bencher, size: usize) { let array = make_fragmented_array(size); diff --git a/vortex-array/benches/expr/large_struct_pack.rs b/vortex-array/benches/expr/large_struct_pack.rs index 31471629cd9..40a5fa0a387 100644 --- a/vortex-array/benches/expr/large_struct_pack.rs +++ b/vortex-array/benches/expr/large_struct_pack.rs @@ -4,6 +4,7 @@ #![expect(clippy::unwrap_used)] use divan::Bencher; +use mimalloc::MiMalloc; use vortex_array::dtype::DType; use vortex_array::dtype::FieldName; use vortex_array::dtype::Nullability; @@ -13,11 +14,14 @@ use vortex_array::expr::get_item; use vortex_array::expr::pack; use vortex_array::expr::root; +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + fn main() { divan::main(); } -#[divan::bench(args = [100, 500, 1000, 2000])] +#[divan::bench(args = [100, 200, 350, 500])] fn pack_return_dtype(bencher: Bencher, num_fields: usize) { // struct with many columns let field_names: Vec = (0..num_fields) diff --git a/vortex-array/benches/expr/optimize_bench.rs b/vortex-array/benches/expr/optimize_bench.rs index 16a3f8765ba..d86e115c643 100644 --- a/vortex-array/benches/expr/optimize_bench.rs +++ b/vortex-array/benches/expr/optimize_bench.rs @@ -5,6 +5,7 @@ #![expect(clippy::cast_possible_truncation)] use divan::Bencher; +use mimalloc::MiMalloc; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; @@ -16,6 +17,9 @@ use vortex_array::expr::lit; use vortex_array::expr::or; use vortex_array::expr::root; +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + fn main() { divan::main(); } diff --git a/vortex-array/benches/expr/optimize_predicate.rs b/vortex-array/benches/expr/optimize_predicate.rs new file mode 100644 index 00000000000..731c7c2b06f --- /dev/null +++ b/vortex-array/benches/expr/optimize_predicate.rs @@ -0,0 +1,213 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks `Expression::optimize_recursive` on lookup-style pushdown predicates: +//! an id membership test (either `list_contains` or a balanced OR of equalities) conjoined +//! with timestamp range bounds and kind filters. + +#![expect(clippy::unwrap_used)] + +use std::fmt::Display; +use std::fmt::Formatter; +use std::fmt::Result as FmtResult; +use std::sync::Arc; + +use divan::Bencher; +use divan::black_box; +use mimalloc::MiMalloc; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::StructFields; +use vortex_array::expr::Expression; +use vortex_array::expr::and; +use vortex_array::expr::and_collect; +use vortex_array::expr::col; +use vortex_array::expr::eq; +use vortex_array::expr::gt_eq; +use vortex_array::expr::list_contains; +use vortex_array::expr::lit; +use vortex_array::expr::lt; +use vortex_array::expr::not_eq; +use vortex_array::expr::or_collect; +use vortex_array::extension::datetime::TimeUnit; +use vortex_array::extension::datetime::Timestamp; +use vortex_array::scalar::Scalar; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +fn main() { + divan::main(); +} + +const FIELD_ID: &str = "id"; +const FIELD_KIND: &str = "kind"; +const FIELD_TIMESTAMP: &str = "timestamp"; + +const TIMESTAMP_BASE_SECS: i64 = 1_700_000_000; +const MICROS_PER_SECOND: i64 = 1_000_000; +const KIND_INCLUDED: u8 = 1; +const KIND_EXCLUDED: u8 = 2; + +#[derive(Debug, Clone, Copy)] +enum IdPredicateShape { + InList, + BalancedOr, +} + +impl Display for IdPredicateShape { + fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { + match self { + IdPredicateShape::InList => write!(f, "in_list"), + IdPredicateShape::BalancedOr => write!(f, "balanced_or"), + } + } +} + +#[derive(Debug, Clone, Copy)] +struct PredicateCase { + id_count: usize, + shape: IdPredicateShape, +} + +impl Display for PredicateCase { + fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { + write!(f, "ids={}, shape={}", self.id_count, self.shape) + } +} + +const PREDICATE_CASES: &[PredicateCase] = &[ + PredicateCase { + id_count: 1, + shape: IdPredicateShape::InList, + }, + PredicateCase { + id_count: 1, + shape: IdPredicateShape::BalancedOr, + }, + PredicateCase { + id_count: 16, + shape: IdPredicateShape::InList, + }, + PredicateCase { + id_count: 16, + shape: IdPredicateShape::BalancedOr, + }, + PredicateCase { + id_count: 64, + shape: IdPredicateShape::InList, + }, + PredicateCase { + id_count: 64, + shape: IdPredicateShape::BalancedOr, + }, + PredicateCase { + id_count: 128, + shape: IdPredicateShape::InList, + }, + PredicateCase { + id_count: 128, + shape: IdPredicateShape::BalancedOr, + }, + PredicateCase { + id_count: 192, + shape: IdPredicateShape::InList, + }, + PredicateCase { + id_count: 192, + shape: IdPredicateShape::BalancedOr, + }, + PredicateCase { + id_count: 256, + shape: IdPredicateShape::InList, + }, + PredicateCase { + id_count: 256, + shape: IdPredicateShape::BalancedOr, + }, +]; + +fn timestamp_dtype() -> DType { + DType::Extension(Timestamp::new(TimeUnit::Microseconds, Nullability::NonNullable).erased()) +} + +fn scope() -> DType { + DType::Struct( + StructFields::new( + [FIELD_ID, FIELD_KIND, FIELD_TIMESTAMP].into(), + vec![ + DType::Utf8(Nullability::NonNullable), + DType::Primitive(PType::U8, Nullability::NonNullable), + timestamp_dtype(), + ], + ), + Nullability::NonNullable, + ) +} + +fn id_strings(id_count: usize) -> Vec { + (0..id_count) + .map(|idx| format!("00000000-0000-0000-0000-{idx:012x}")) + .collect() +} + +fn timestamp_lit(micros: i64) -> Expression { + let storage = Scalar::primitive(micros, Nullability::NonNullable); + lit(Scalar::extension_ref( + Timestamp::new(TimeUnit::Microseconds, Nullability::NonNullable).erased(), + storage, + )) +} + +fn id_in_list_filter(ids: &[String]) -> Expression { + let elements = ids + .iter() + .map(|id| Scalar::utf8(id.as_str(), Nullability::Nullable)) + .collect::>(); + let list = Scalar::list( + Arc::new(DType::Utf8(Nullability::Nullable)), + elements, + Nullability::Nullable, + ); + list_contains(lit(list), col(FIELD_ID)) +} + +fn id_balanced_or_filter(ids: &[String]) -> Expression { + or_collect(ids.iter().map(|id| eq(col(FIELD_ID), lit(id.as_str())))).unwrap() +} + +fn id_filter(ids: &[String], shape: IdPredicateShape) -> Expression { + match shape { + IdPredicateShape::InList => id_in_list_filter(ids), + IdPredicateShape::BalancedOr => id_balanced_or_filter(ids), + } +} + +fn lookup_predicate(predicate_case: PredicateCase) -> Expression { + let ids = id_strings(predicate_case.id_count); + let min_timestamp_us = TIMESTAMP_BASE_SECS * MICROS_PER_SECOND; + let max_timestamp_us = + (TIMESTAMP_BASE_SECS + 1) * MICROS_PER_SECOND + predicate_case.id_count as i64; + + let id_filter = and( + id_filter(&ids, predicate_case.shape), + not_eq(col(FIELD_KIND), lit(KIND_EXCLUDED)), + ); + + and_collect([ + id_filter, + gt_eq(col(FIELD_TIMESTAMP), timestamp_lit(min_timestamp_us)), + lt(col(FIELD_TIMESTAMP), timestamp_lit(max_timestamp_us)), + eq(col(FIELD_KIND), lit(KIND_INCLUDED)), + ]) + .unwrap() +} + +#[divan::bench(args = PREDICATE_CASES)] +fn optimize_lookup_predicate(bencher: Bencher, predicate_case: &PredicateCase) { + let scope = scope(); + let predicate = lookup_predicate(*predicate_case); + + bencher.bench(|| black_box(predicate.optimize_recursive(&scope))); +} diff --git a/vortex-array/benches/filter_fixed_width.rs b/vortex-array/benches/filter_fixed_width.rs index 01334108848..eddeb367092 100644 --- a/vortex-array/benches/filter_fixed_width.rs +++ b/vortex-array/benches/filter_fixed_width.rs @@ -15,6 +15,7 @@ use std::sync::LazyLock; use divan::Bencher; +use mimalloc::MiMalloc; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; @@ -28,6 +29,9 @@ use vortex_buffer::BitBuffer; use vortex_mask::Mask; use vortex_session::VortexSession; +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + fn main() { LazyLock::force(&SESSION); divan::main(); @@ -36,7 +40,8 @@ fn main() { static SESSION: LazyLock = LazyLock::new(array_session); // Keep each case small: the sweep has 24 cases and the targeted sections add only seven more. -const LEN: usize = 16_384; +// Sized to keep CodSpeed simulation under 1ms per benchmark. +const LEN: usize = 4_096; const DENSITIES: &[f64] = &[0.01, 0.5, 0.8, 0.95]; const CACHED_DENSITIES: &[f64] = &[0.01, 0.1]; diff --git a/vortex-array/benches/filter_list.rs b/vortex-array/benches/filter_list.rs new file mode 100644 index 00000000000..bbe38839728 --- /dev/null +++ b/vortex-array/benches/filter_list.rs @@ -0,0 +1,108 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks filtering lists with different lengths and outer-mask locality. + +#![expect(clippy::cast_possible_truncation, clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::RecursiveCanonical; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(array_session); + +const LIST_LENGTH: usize = 512; +const N_ELEMENTS_SHORT: usize = 32; +const N_ELEMENTS_LONG: usize = 2_048; +const MASK_LENGTH_CLUSTERED: usize = 8; + +#[derive(Clone, Copy, Debug)] +enum MaskSetup { + AllSelected, + Prefix, + Repeated, + Clustered, + Sparse, + EdgeSpanning, +} + +const MASK_SETUPS: &[MaskSetup] = &[ + MaskSetup::AllSelected, + MaskSetup::Prefix, + MaskSetup::Repeated, + MaskSetup::Clustered, + MaskSetup::Sparse, + MaskSetup::EdgeSpanning, +]; + +fn primitive_list(width: usize) -> ArrayRef { + let element_count = LIST_LENGTH * width; + let elements = PrimitiveArray::from_iter(0..element_count as u32).into_array(); + let offsets = + Buffer::from_iter((0..=LIST_LENGTH).map(|index| (index * width) as u32)).into_array(); + ListArray::try_new(elements, offsets, Validity::NonNullable) + .unwrap() + .into_array() +} + +fn selection_mask(setup: MaskSetup) -> Mask { + match setup { + MaskSetup::AllSelected => Mask::new_true(LIST_LENGTH), + MaskSetup::Prefix => Mask::from_slices(LIST_LENGTH, vec![(0, MASK_LENGTH_CLUSTERED)]), + MaskSetup::Repeated => Mask::from_indices(LIST_LENGTH, (0..LIST_LENGTH).step_by(5)), + MaskSetup::Clustered => { + let start = (LIST_LENGTH - MASK_LENGTH_CLUSTERED) / 2; + Mask::from_slices(LIST_LENGTH, vec![(start, start + MASK_LENGTH_CLUSTERED)]) + } + MaskSetup::Sparse => Mask::from_indices(LIST_LENGTH, [LIST_LENGTH / 2]), + MaskSetup::EdgeSpanning => Mask::from_indices(LIST_LENGTH, [0, LIST_LENGTH - 1]), + } +} + +fn run(bencher: Bencher, array: ArrayRef, mask: Mask) { + bencher + .with_inputs(|| (array.clone(), mask.clone(), SESSION.create_execution_ctx())) + .bench_values(|(array, mask, mut ctx)| { + divan::black_box( + array + .filter(mask) + .unwrap() + .execute::(&mut ctx) + .unwrap(), + ); + }); +} + +#[divan::bench(args = MASK_SETUPS)] +fn filter_list_primitive_short(bencher: Bencher, setup: MaskSetup) { + run( + bencher, + primitive_list(N_ELEMENTS_SHORT), + selection_mask(setup), + ); +} + +#[divan::bench(args = MASK_SETUPS)] +fn filter_list_primitive_wide(bencher: Bencher, setup: MaskSetup) { + run( + bencher, + primitive_list(N_ELEMENTS_LONG), + selection_mask(setup), + ); +} diff --git a/vortex-array/benches/interleave.rs b/vortex-array/benches/interleave.rs index d92c033c2a4..a04267c0d92 100644 --- a/vortex-array/benches/interleave.rs +++ b/vortex-array/benches/interleave.rs @@ -1,8 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Benchmarks the Vortex [`Interleave`](vortex_array::arrays::Interleave) boolean execute path on a -//! focused set of configurations: +//! Benchmarks the Vortex [`Interleave`](vortex_array::arrays::Interleave) Boolean and primitive +//! execute paths on a focused set of configurations: //! //! - `round_robin`, 2 children: a merge — `array_index = i % N`, `row_index = i / N`. //! - `random`, 2 children: fully random `(array_index, row_index)` per output row. @@ -26,7 +26,9 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::InterleaveArray; +use vortex_array::arrays::PrimitiveArray; use vortex_buffer::Buffer; fn main() { @@ -130,6 +132,45 @@ fn vortex_inputs(combo: Combo) -> (Vec, Buffer, Buffer) { (values, array_indices, row_indices) } +fn primitive_inputs( + combo: Combo, + mixed_constants: bool, +) -> (Vec, Buffer, Buffer) { + let mut rng = StdRng::seed_from_u64(0); + let values = (0..combo.branches) + .map(|branch| { + if mixed_constants && branch % 2 == 0 { + ConstantArray::new(branch as u64, ARRAY_SIZE).into_array() + } else if combo.nullable { + PrimitiveArray::from_option_iter( + (0..ARRAY_SIZE) + .map(|row| ((row + branch) % 8 != 0).then_some((row ^ branch) as u64)), + ) + .into_array() + } else { + PrimitiveArray::from_iter((0..ARRAY_SIZE).map(|row| (row ^ branch) as u64)) + .into_array() + } + }) + .collect(); + + let branch = Uniform::new(0u32, u32::try_from(combo.branches).unwrap()).unwrap(); + let row = Uniform::new(0u32, u32::try_from(ARRAY_SIZE).unwrap()).unwrap(); + let array_indices = (0..ARRAY_SIZE) + .map(|i| match combo.pattern { + Pattern::Random => rng.sample(branch), + Pattern::RoundRobin => u32::try_from(i % combo.branches).unwrap(), + }) + .collect(); + let row_indices = (0..ARRAY_SIZE) + .map(|i| match combo.pattern { + Pattern::Random => rng.sample(row), + Pattern::RoundRobin => u32::try_from((i / combo.branches) % ARRAY_SIZE).unwrap(), + }) + .collect(); + (values, array_indices, row_indices) +} + #[divan::bench(args = combos())] fn vortex(bencher: Bencher, combo: Combo) { let (values, array_indices, row_indices) = vortex_inputs(combo); @@ -149,3 +190,32 @@ fn vortex(bencher: Bencher, combo: Combo) { }) .bench_refs(|(array, ctx)| array.clone().execute::(ctx)); } + +fn bench_primitive(bencher: Bencher, combo: Combo, mixed_constants: bool) { + let (values, array_indices, row_indices) = primitive_inputs(combo, mixed_constants); + let session = array_session(); + bencher + .with_inputs(|| { + ( + InterleaveArray::try_new( + values.clone(), + array_indices.clone().into_array(), + row_indices.clone().into_array(), + ) + .unwrap() + .into_array(), + session.create_execution_ctx(), + ) + }) + .bench_refs(|(array, ctx)| array.clone().execute::(ctx)); +} + +#[divan::bench(args = combos())] +fn primitive(bencher: Bencher, combo: Combo) { + bench_primitive(bencher, combo, false); +} + +#[divan::bench(args = combos())] +fn primitive_mixed_constants(bencher: Bencher, combo: Combo) { + bench_primitive(bencher, combo, true); +} diff --git a/vortex-array/benches/like.rs b/vortex-array/benches/like.rs index 68219724717..f618e226fd6 100644 --- a/vortex-array/benches/like.rs +++ b/vortex-array/benches/like.rs @@ -14,7 +14,6 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::VarBinViewArray; -use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::scalar_fn::fns::like::Like; use vortex_array::scalar_fn::fns::like::LikeOptions; @@ -22,7 +21,8 @@ fn main() { divan::main(); } -const ARRAY_SIZE: usize = 2_048; +// Sized to keep CodSpeed simulation under 1ms per benchmark. +const ARRAY_SIZE: usize = 1_024; /// Random lowercase strings of 4..=24 bytes, some with a `hello` infix. fn strings() -> ArrayRef { @@ -47,15 +47,13 @@ fn bench_like(bencher: Bencher, pattern: &str, options: LikeOptions) { bencher .with_inputs(|| { ( - Like.try_new_array( - ARRAY_SIZE, + Like::try_new( + array.clone(), + ConstantArray::new(pattern, ARRAY_SIZE).into_array(), options, - [ - array.clone(), - ConstantArray::new(pattern, ARRAY_SIZE).into_array(), - ], ) - .unwrap(), + .unwrap() + .into_array(), session.create_execution_ctx(), ) }) @@ -97,12 +95,9 @@ fn like_per_row_patterns(bencher: Bencher) { bencher .with_inputs(|| { ( - Like.try_new_array( - ARRAY_SIZE, - LikeOptions::default(), - [array.clone(), patterns.clone()], - ) - .unwrap(), + Like::try_new(array.clone(), patterns.clone(), LikeOptions::default()) + .unwrap() + .into_array(), session.create_execution_ctx(), ) }) diff --git a/vortex-array/benches/list_length.rs b/vortex-array/benches/list_length.rs index 6148b54f8f9..7312d7b84ef 100644 --- a/vortex-array/benches/list_length.rs +++ b/vortex-array/benches/list_length.rs @@ -40,7 +40,8 @@ const BASE_LIST_SIZE: usize = 8; const SMALL: usize = 100; const MEDIUM: usize = 10_000; -const LARGE: usize = 1_000_000; +// Sized to keep the CodSpeed simulation under 1ms per benchmark. +const LARGE: usize = 20_000; /// A uniformly-random partition of `num_lists * LIST_SIZE` elements into `num_lists` lists, /// plus a validity mask with ~1/8 of lists null at random positions. diff --git a/vortex-array/benches/list_sum.rs b/vortex-array/benches/list_sum.rs index 6092b73a807..4a2caf65c09 100644 --- a/vortex-array/benches/list_sum.rs +++ b/vortex-array/benches/list_sum.rs @@ -37,9 +37,10 @@ static SESSION: LazyLock = LazyLock::new(vortex_array::array_sess const BASE_LIST_SIZE: usize = 8; +// Sized to keep the CodSpeed simulation under 1ms per benchmark. const SMALL: usize = 100; -const MEDIUM: usize = 10_000; -const LARGE: usize = 1_000_000; +const MEDIUM: usize = 250; +const LARGE: usize = 500; /// A uniformly-random partition of `num_lists * BASE_LIST_SIZE` elements into `num_lists` /// lists, plus a validity mask with ~1/8 of lists null at random positions. diff --git a/vortex-array/benches/listview_builder_extend.rs b/vortex-array/benches/listview_builder_extend.rs index 99f35e9fe82..821e49b37c2 100644 --- a/vortex-array/benches/listview_builder_extend.rs +++ b/vortex-array/benches/listview_builder_extend.rs @@ -25,18 +25,17 @@ fn main() { divan::main(); } +// Sized to keep the CodSpeed simulation under 1ms per benchmark. const ZCTL_ARGS: &[(usize, usize)] = &[ // num_lists, list_size + (250, 32), (1_000, 8), - (1_000, 64), - (10_000, 8), ]; const NON_ZCTL_ARGS: &[(usize, usize)] = &[ // num_lists, list_size (1_000, 8), (1_000, 32), - (10_000, 8), ]; fn make_listview( @@ -60,17 +59,21 @@ fn make_listview( #[divan::bench(args = ZCTL_ARGS)] fn extend_from_array_zctl(bencher: Bencher, (num_lists, list_size): (usize, usize)) { - let source = make_listview(num_lists, list_size, list_size, false); - debug_assert!(source.is_zero_copy_to_list()); + // `ListViewArray::new` never detects zero-copy eligibility, so mark the flag explicitly to + // exercise the ZCTL extend path; debug builds validate the invariants. + let source = unsafe { + make_listview(num_lists, list_size, list_size, false).with_zero_copy_to_list(true) + }; let source = source.into_array(); bencher.with_inputs(|| &source).bench_refs(|source| { let mut ctx = array_session().create_execution_ctx(); - let mut builder = ListViewBuilder::::with_capacity( + let mut builder = ListViewBuilder::::with_capacity_in( Arc::new(DType::Primitive(I32, NonNullable)), NonNullable, num_lists * list_size, num_lists, + ctx.allocator(), ); source.append_to_builder(&mut builder, &mut ctx).unwrap(); divan::black_box(builder.finish_into_listview()) @@ -89,11 +92,12 @@ fn extend_from_array_non_zctl_overlapping( bencher.with_inputs(|| &source).bench_refs(|source| { let mut ctx = array_session().create_execution_ctx(); - let mut builder = ListViewBuilder::::with_capacity( + let mut builder = ListViewBuilder::::with_capacity_in( Arc::new(DType::Primitive(I32, NonNullable)), Nullable, num_lists * list_size, num_lists, + ctx.allocator(), ); source.append_to_builder(&mut builder, &mut ctx).unwrap(); divan::black_box(builder.finish_into_listview()) diff --git a/vortex-array/benches/listview_rebuild.rs b/vortex-array/benches/listview_rebuild.rs index 2e61b969244..b1c1f7676ce 100644 --- a/vortex-array/benches/listview_rebuild.rs +++ b/vortex-array/benches/listview_rebuild.rs @@ -19,7 +19,7 @@ use vortex_array::arrays::ListViewArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; -use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::arrays::listview::ListViewRebuildMode; use vortex_array::dtype::FieldNames; use vortex_array::validity::Validity; @@ -201,7 +201,7 @@ fn i32_large(bencher: Bencher) { #[divan::bench] fn varbinview_large(bencher: Bencher) { - let lv = make_varbinview_lv(5, 1_024, 1_024); + let lv = make_varbinview_lv(4, 512, 512); bencher .with_inputs(|| (&lv, SESSION.create_execution_ctx())) .bench_refs(|(lv, ctx)| { diff --git a/vortex-array/benches/patches_lookup.rs b/vortex-array/benches/patches_lookup.rs index cbf753e668f..262e3144495 100644 --- a/vortex-array/benches/patches_lookup.rs +++ b/vortex-array/benches/patches_lookup.rs @@ -19,7 +19,8 @@ fn main() { const ARRAY_LEN: usize = 1_000_000; const NUM_PATCHES: usize = 100; -const NUM_QUERIES: usize = 1_000; +// Sized to keep CodSpeed simulation under 1ms per benchmark. +const NUM_QUERIES: usize = 100; const PATCH_LOW: usize = 100_000; const PATCH_HIGH: usize = 110_000; diff --git a/vortex-array/benches/piecewise_sequence_take_primitive.rs b/vortex-array/benches/piecewise_sequence_take_primitive.rs new file mode 100644 index 00000000000..5933058caa6 --- /dev/null +++ b/vortex-array/benches/piecewise_sequence_take_primitive.rs @@ -0,0 +1,164 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Primitive take benchmarks for [`PiecewiseSequenceArray`] length representations. +//! +//! The `reify_constant_lengths_then_take` case simulates removing the constant-length +//! specialization: it starts with constant lengths, materializes them to a dense primitive array, +//! and then executes the primitive take through the non-constant-length path. + +#![allow(clippy::cast_possible_truncation)] +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use divan::counter::BytesCount; +use rand::SeedableRng; +use rand::prelude::*; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PiecewiseSequenceArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_session::VortexSession; + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(array_session); + +const SOURCE_LEN: usize = 512 * 1024; +// Sized to keep CodSpeed simulation under 1ms at run length 1. +const OUTPUT_LEN: usize = 4 * 1024; +const RUN_LENGTHS: &[usize] = &[1, 4, 16, 64, 256]; + +#[divan::bench(args = RUN_LENGTHS)] +fn optimized_constant_lengths(bencher: Bencher, run_length: usize) { + let values = values(); + let indices = piecewise_indices( + run_length, + lengths_constant(run_length, run_count(run_length)), + ); + + bencher + .counter(BytesCount::of_many::(OUTPUT_LEN)) + .with_inputs(|| (&values, indices.clone(), SESSION.create_execution_ctx())) + .bench_refs(|(values, indices, ctx)| { + divan::black_box( + values + .clone() + .take(indices.clone()) + .unwrap() + .execute::(ctx) + .unwrap(), + ); + }); +} + +#[divan::bench(args = RUN_LENGTHS)] +fn prebuilt_dense_lengths(bencher: Bencher, run_length: usize) { + let values = values(); + let indices = piecewise_indices(run_length, lengths_dense(run_length, run_count(run_length))); + + bencher + .counter(BytesCount::of_many::(OUTPUT_LEN)) + .with_inputs(|| (&values, indices.clone(), SESSION.create_execution_ctx())) + .bench_refs(|(values, indices, ctx)| { + divan::black_box( + values + .clone() + .take(indices.clone()) + .unwrap() + .execute::(ctx) + .unwrap(), + ); + }); +} + +#[divan::bench(args = RUN_LENGTHS)] +fn reify_constant_lengths_then_take(bencher: Bencher, run_length: usize) { + let values = values(); + let run_count = run_count(run_length); + let starts = starts(run_count, run_length).into_array(); + let lengths = lengths_constant(run_length, run_count); + let multipliers = ConstantArray::new(1u64, run_count).into_array(); + + bencher + .counter(BytesCount::of_many::(OUTPUT_LEN)) + .with_inputs(|| { + ( + &values, + starts.clone(), + lengths.clone(), + multipliers.clone(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(values, starts, lengths, multipliers, ctx)| { + let lengths = lengths + .clone() + .execute::(ctx) + .unwrap() + .into_array(); + + // SAFETY: starts, reified lengths, and unit multipliers have the same run count, are + // non-null unsigned integer arrays, and expand to OUTPUT_LEN by construction. + let indices = unsafe { + PiecewiseSequenceArray::new_unchecked( + starts.clone(), + lengths, + multipliers.clone(), + OUTPUT_LEN, + ) + } + .into_array(); + + divan::black_box( + values + .clone() + .take(indices) + .unwrap() + .execute::(ctx) + .unwrap(), + ); + }); +} + +fn values() -> ArrayRef { + PrimitiveArray::from_iter((0..SOURCE_LEN).map(|idx| idx as i64)).into_array() +} + +fn run_count(run_length: usize) -> usize { + assert_eq!(OUTPUT_LEN % run_length, 0); + OUTPUT_LEN / run_length +} + +fn piecewise_indices(run_length: usize, lengths: ArrayRef) -> ArrayRef { + let run_count = run_count(run_length); + let starts = starts(run_count, run_length).into_array(); + let multipliers = ConstantArray::new(1u64, run_count).into_array(); + + // SAFETY: starts, lengths, and unit multipliers have the same run count, are non-null unsigned + // integer arrays, and expand to OUTPUT_LEN by construction. + unsafe { PiecewiseSequenceArray::new_unchecked(starts, lengths, multipliers, OUTPUT_LEN) } + .into_array() +} + +fn lengths_constant(run_length: usize, run_count: usize) -> ArrayRef { + ConstantArray::new(run_length as u64, run_count).into_array() +} + +fn lengths_dense(run_length: usize, run_count: usize) -> ArrayRef { + PrimitiveArray::from_iter((0..run_count).map(|_| run_length as u64)).into_array() +} + +fn starts(run_count: usize, run_length: usize) -> PrimitiveArray { + let mut rng = StdRng::seed_from_u64(0x51ced); + let max_start = SOURCE_LEN - run_length; + PrimitiveArray::from_iter((0..run_count).map(|_| rng.random_range(0..=max_start) as u64)) +} diff --git a/vortex-array/benches/row_fn_output.rs b/vortex-array/benches/row_fn_output.rs new file mode 100644 index 00000000000..db3d9ebb2fe --- /dev/null +++ b/vortex-array/benches/row_fn_output.rs @@ -0,0 +1,393 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect( + clippy::cast_possible_truncation, + reason = "benchmark fixtures reduce values below 1024 before narrowing" +)] + +use std::marker::PhantomData; +use std::sync::LazyLock; + +use divan::Bencher; +use divan::counter::ItemsCount; +use mimalloc::MiMalloc; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::VecExecutionArgs; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::InputElement; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; +use vortex_array::scalar_fn::unstable::row::execute_rows; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(array_session); + +const ROWS: usize = 1 << 14; + +/// The shape every row function is measured on: one per-row operand against another. +const INPUT_SHAPES: &[InputShape] = &[InputShape::PerRowPerRow]; + +/// Constant operands are measured on `i64` alone, by the `*_constant` benchmarks below. +/// +/// The row executor decodes and orients a constant operand the same way whatever the row kernel +/// does with the value, so running both constant orientations under every element type and every +/// output mode measures one path repeatedly, three walltime legs at a time. One infallible pair and +/// one deferred pair keep it covered: those two differ in how a constant row reaches the kernel, +/// and that is the difference worth watching. +const CONSTANT_SHAPES: &[InputShape] = &[InputShape::PerRowConstant, InputShape::ConstantPerRow]; + +const VALIDITY_PATTERNS: &[ValidityPattern] = &[ + ValidityPattern::AllValid, + ValidityPattern::OneNullInEight, + ValidityPattern::NineNullsInTen, +]; + +#[derive(Clone, Copy, Debug)] +enum InputShape { + PerRowPerRow, + PerRowConstant, + ConstantPerRow, +} + +#[derive(Clone, Copy, Debug)] +enum ValidityPattern { + AllValid, + OneNullInEight, + NineNullsInTen, +} + +impl ValidityPattern { + fn is_valid(self, index: usize) -> bool { + match self { + Self::AllValid => true, + Self::OneNullInEight => !index.is_multiple_of(8), + Self::NineNullsInTen => index.is_multiple_of(10), + } + } +} + +trait BenchPrimitive: NativePType { + fn per_row(offset: usize) -> ArrayRef; + + fn constant() -> ArrayRef; +} + +impl BenchPrimitive for i32 { + fn per_row(offset: usize) -> ArrayRef { + PrimitiveArray::from_iter((0..ROWS).map(|index| ((index + offset) % 1024) as i32)) + .into_array() + } + + fn constant() -> ArrayRef { + ConstantArray::new(512_i32, ROWS).into_array() + } +} + +impl BenchPrimitive for i64 { + fn per_row(offset: usize) -> ArrayRef { + PrimitiveArray::from_iter((0..ROWS).map(|index| ((index + offset) % 1024) as i64)) + .into_array() + } + + fn constant() -> ArrayRef { + ConstantArray::new(512_i64, ROWS).into_array() + } +} + +/// Primitive input that forces partially valid batches through the filtered fallback. +/// +/// `DENSE_SAFE = false` selects valid-only execution. This type inherits the conservative +/// null-tolerant decode implementation, so direct valid-row execution declines the original +/// nullable input. A partially valid input must therefore use the executor's filtered fallback. +struct FilteredI64; + +// SAFETY: the view is a native slice, and its reported length is the slice length. +unsafe impl InputElement for FilteredI64 { + type Column = Buffer; + type Constant = i64; + type View<'a> = &'a [i64]; + type Elem<'a> = i64; + + const DENSE_SAFE: bool = false; + const DECODE_INFALLIBLE: bool = true; + + fn validate(dtype: &DType) -> VortexResult<()> { + ::validate(dtype) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + ::decode(array, ctx) + } + + fn decode_constant(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + ::decode_constant(array, ctx) + } + + fn get(column: &Self::Column, index: usize) -> Self::Elem<'_> { + column[index] + } + + fn get_constant(constant: &Self::Constant) -> Self::Elem<'_> { + *constant + } + + fn view(column: &Self::Column) -> Self::View<'_> { + column.as_slice() + } + + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> Self::Elem<'a> { + view[index] + } + + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> Self::Elem<'a> { + // SAFETY: forwarded from this method's contract. + unsafe { *view.get_unchecked(index) } + } +} + +#[derive(Clone)] +struct InfallibleBool(PhantomData); + +#[derive(Clone)] +struct DeferredBool(PhantomData); + +#[derive(Clone)] +struct DeferredI64; + +#[derive(Clone)] +struct FilteredOwnedI64; + +#[derive(Clone)] +struct FilteredSinkI64; + +impl RowFn for InfallibleBool { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_fn_output.infallible_bool"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_lt(rhs)) + } +} + +impl RowFn for DeferredBool { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const INFALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_fn_output.deferred_bool"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_deferred_bool::<(T, T), bool, false>( + |(lhs, rhs)| (lhs.is_lt(rhs), lhs.is_lt(T::default())), + |negative| { + vortex_ensure!( + !negative, + "deferred-bool benchmark inputs must be nonnegative" + ); + Ok(()) + }, + ) + } +} + +impl RowFn for DeferredI64 { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const INFALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_fn_output.deferred_i64"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_deferred::<(i64, i64), i64, bool>( + |(lhs, rhs)| lhs.overflowing_add(rhs), + |overflowed| { + vortex_ensure!( + !overflowed, + "deferred-i64 benchmark inputs must not overflow" + ); + Ok(()) + }, + ) + } +} + +impl RowFn for FilteredOwnedI64 { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_fn_output.filtered_owned_i64"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(FilteredI64,), i64>(|(value,)| value.wrapping_add(1)) + } +} + +impl RowFn for FilteredSinkI64 { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_fn_output.filtered_sink_i64"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(FilteredI64,), UninitElementSink, _>((), |(value,), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, value.wrapping_add(1)) } + }) + } +} + +#[vortex_bench_support::cpu_features] +#[divan::bench(types = [i32, i64], args = INPUT_SHAPES)] +fn infallible_bool(bencher: Bencher, &shape: &InputShape) { + let function = InfallibleBool::(PhantomData); + bench_row_fn(bencher, &function, make_args::(shape)); +} + +#[vortex_bench_support::cpu_features] +#[divan::bench(args = CONSTANT_SHAPES)] +fn infallible_bool_constant(bencher: Bencher, &shape: &InputShape) { + let function = InfallibleBool::(PhantomData); + bench_row_fn(bencher, &function, make_args::(shape)); +} + +#[vortex_bench_support::cpu_features] +#[divan::bench(types = [i32, i64], args = INPUT_SHAPES)] +fn deferred_bool(bencher: Bencher, &shape: &InputShape) { + let function = DeferredBool::(PhantomData); + bench_row_fn(bencher, &function, make_args::(shape)); +} + +#[vortex_bench_support::cpu_features] +#[divan::bench(args = CONSTANT_SHAPES)] +fn deferred_bool_constant(bencher: Bencher, &shape: &InputShape) { + let function = DeferredBool::(PhantomData); + bench_row_fn(bencher, &function, make_args::(shape)); +} + +#[vortex_bench_support::cpu_features] +#[divan::bench(args = INPUT_SHAPES)] +fn deferred_i64(bencher: Bencher, &shape: &InputShape) { + bench_row_fn(bencher, &DeferredI64, make_args::(shape)); +} + +#[vortex_bench_support::cpu_features] +#[divan::bench(args = VALIDITY_PATTERNS)] +fn filtered_owned_i64(bencher: Bencher, &validity: &ValidityPattern) { + bench_row_fn(bencher, &FilteredOwnedI64, make_filtered_args(validity)); +} + +#[vortex_bench_support::cpu_features] +#[divan::bench(args = VALIDITY_PATTERNS)] +fn filtered_sink_i64(bencher: Bencher, &validity: &ValidityPattern) { + bench_row_fn(bencher, &FilteredSinkI64, make_filtered_args(validity)); +} + +fn make_args(shape: InputShape) -> VecExecutionArgs { + let args = match shape { + InputShape::PerRowPerRow => vec![T::per_row(0), T::per_row(1)], + InputShape::PerRowConstant => vec![T::per_row(0), T::constant()], + InputShape::ConstantPerRow => vec![T::constant(), T::per_row(1)], + }; + + VecExecutionArgs::new(args, ROWS) +} + +fn make_filtered_args(validity: ValidityPattern) -> VecExecutionArgs { + let values = PrimitiveArray::new( + (0..ROWS).map(|index| index as i64).collect::>(), + Validity::from_iter((0..ROWS).map(|index| validity.is_valid(index))), + ) + .into_array(); + + VecExecutionArgs::new(vec![values], ROWS) +} + +fn bench_row_fn>( + bencher: Bencher, + function: &F, + args: VecExecutionArgs, +) { + bencher + .counter(ItemsCount::new(ROWS)) + .with_inputs(|| (&args, SESSION.create_execution_ctx())) + .bench_refs(|(args, ctx)| { + execute_rows(function, &EmptyOptions, *args, ctx) + .vortex_expect("row execution should succeed in benchmark") + }); +} diff --git a/vortex-array/benches/scalar_at_struct.rs b/vortex-array/benches/scalar_at_struct.rs index 8dca44e408a..2564120f12e 100644 --- a/vortex-array/benches/scalar_at_struct.rs +++ b/vortex-array/benches/scalar_at_struct.rs @@ -26,7 +26,8 @@ fn main() { } const ARRAY_SIZE: usize = 100_000; -const NUM_ACCESSES: usize = 1000; +// Sized to keep the CodSpeed simulation under 1ms per benchmark. +const NUM_ACCESSES: usize = 100; static SESSION: LazyLock = LazyLock::new(array_session); diff --git a/vortex-array/benches/scalar_subtract.rs b/vortex-array/benches/scalar_subtract.rs index ef004d2e270..50c0d903f20 100644 --- a/vortex-array/benches/scalar_subtract.rs +++ b/vortex-array/benches/scalar_subtract.rs @@ -29,16 +29,18 @@ fn main() { static SESSION: LazyLock = LazyLock::new(array_session); +#[vortex_bench_support::cpu_features] #[divan::bench] fn scalar_subtract(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); let range = Uniform::new(0i64, 100_000_000).unwrap(); - let data1 = (0..100_000) + // Chunks sized to keep the CodSpeed simulation under 1ms. + let data1 = (0..4_000) .map(|_| rng.sample(range)) .collect::>() .into_array(); - let data2 = (0..100_000) + let data2 = (0..4_000) .map(|_| rng.sample(range)) .collect::>() .into_array(); diff --git a/vortex-array/benches/search_sorted.rs b/vortex-array/benches/search_sorted.rs index 2b5f7fccbac..370290bf1c9 100644 --- a/vortex-array/benches/search_sorted.rs +++ b/vortex-array/benches/search_sorted.rs @@ -15,27 +15,49 @@ fn main() { divan::main(); } +/// One search over 65536 elements is ~16 comparisons, small enough that the measurement is mostly +/// harness overhead and the two implementations below become indistinguishable. Search a batch per +/// iteration instead; the targets differ, so the calls cannot be folded together. +const SEARCH_TARGETS: usize = 64; + #[divan::bench] fn binary_search_std(bencher: Bencher) { - let (sorted_array, target) = fixture(); + let (sorted_array, targets) = fixture(); bencher - .with_inputs(|| (&sorted_array, &target)) - .bench_refs(|(array, target)| array.binary_search(target)); + .with_inputs(|| (&sorted_array, &targets)) + .bench_refs(|(array, targets)| { + let mut found = 0; + for target in targets.iter() { + found += array.binary_search(target).unwrap_or_else(|idx| idx); + } + found + }); } #[divan::bench] fn binary_search_vortex(bencher: Bencher) { - let (sorted_array, target) = fixture(); + let (sorted_array, targets) = fixture(); bencher - .with_inputs(|| (&sorted_array, &target)) - .bench_refs(|(array, target)| array.search_sorted(target, SearchSortedSide::Left).unwrap()); + .with_inputs(|| (&sorted_array, &targets)) + .bench_refs(|(array, targets)| { + let mut found = 0; + for target in targets.iter() { + found += array + .search_sorted(target, SearchSortedSide::Left) + .unwrap() + .to_index(); + } + found + }); } -fn fixture() -> (Vec, i32) { +fn fixture() -> (Vec, Vec) { let mut rng = StdRng::seed_from_u64(0); let range = Uniform::new(0, 65_536).unwrap(); let mut data: Vec = (0..65_536).map(|_| rng.sample(range)).collect(); data.sort(); - (data, rng.sample(range)) + let targets = (0..SEARCH_TARGETS).map(|_| rng.sample(range)).collect(); + + (data, targets) } diff --git a/vortex-array/benches/take_chunked.rs b/vortex-array/benches/take_chunked.rs new file mode 100644 index 00000000000..3350030533e --- /dev/null +++ b/vortex-array/benches/take_chunked.rs @@ -0,0 +1,356 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks [`ChunkedArray`] take routing across ordering, chunk count, duplicate density, +//! nullability, and value-shape configurations. + +#![expect(clippy::unwrap_used)] + +use std::fmt::Display; +use std::fmt::Formatter; + +use divan::Bencher; +use divan::counter::ItemsCount; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::RecursiveCanonical; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::dtype::FieldNames; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; + +fn main() { + divan::main(); +} + +const SOURCE_LEN: usize = 1 << 20; +const NESTED_WIDTH: usize = 8; + +#[derive(Clone, Copy)] +enum Pattern { + Sorted, + Shuffled, + Repeated, + Duplicate90, + Duplicate99, + SameChunk, + AlternatingChunks, + GroupedChunks, +} + +impl Display for Pattern { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + Self::Sorted => "sorted", + Self::Shuffled => "shuffled", + Self::Repeated => "repeated", + Self::Duplicate90 => "duplicate90", + Self::Duplicate99 => "duplicate99", + Self::SameChunk => "same_chunk", + Self::AlternatingChunks => "alternating_chunks", + Self::GroupedChunks => "grouped_chunks", + }) + } +} + +#[derive(Clone, Copy)] +enum ValueKind { + Primitive, + Struct8, + FixedSizeList8, +} + +impl Display for ValueKind { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + Self::Primitive => "primitive", + Self::Struct8 => "struct8", + Self::FixedSizeList8 => "fsl8", + }) + } +} + +#[derive(Clone, Copy)] +enum IndexValidity { + NonNullable, + AllValid, + TenPercentNull, +} + +impl Display for IndexValidity { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + Self::NonNullable => "nonnull", + Self::AllValid => "nullable_all_valid", + Self::TenPercentNull => "nullable_10pct_null", + }) + } +} + +#[derive(Clone, Copy)] +struct Case { + group: &'static str, + pattern: Pattern, + value_kind: ValueKind, + validity: IndexValidity, + chunks: usize, + indices: usize, +} + +impl Case { + const fn new(group: &'static str, pattern: Pattern, chunks: usize, indices: usize) -> Self { + Self { + group, + pattern, + value_kind: ValueKind::Primitive, + validity: IndexValidity::NonNullable, + chunks, + indices, + } + } + + const fn with_value_kind(mut self, value_kind: ValueKind) -> Self { + self.value_kind = value_kind; + self + } + + const fn with_validity(mut self, validity: IndexValidity) -> Self { + self.validity = validity; + self + } +} + +impl Display for Case { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{}/{}/{}/{}/chunks={}/indices={}", + self.group, self.pattern, self.value_kind, self.validity, self.chunks, self.indices + ) + } +} + +// Chunk and index counts are sized to keep CodSpeed simulation under 1ms per case. +fn cases() -> Vec { + let mut cases = Vec::new(); + + for pattern in [Pattern::Sorted, Pattern::Shuffled] { + for chunks in [4, 16] { + for indices in [250, 1_000] { + cases.push(Case::new("core", pattern, chunks, indices)); + } + } + } + + for chunks in [256, 2_048] { + for indices in [1, 16, 64] { + cases.push(Case::new("small_m", Pattern::Shuffled, chunks, indices)); + } + } + + for pattern in [ + Pattern::Repeated, + Pattern::Duplicate90, + Pattern::Duplicate99, + ] { + cases.push(Case::new("duplicates", pattern, 16, 1_000)); + } + + for pattern in [Pattern::Shuffled, Pattern::Duplicate99] { + for value_kind in [ + ValueKind::Primitive, + ValueKind::Struct8, + ValueKind::FixedSizeList8, + ] { + cases.push(Case::new("value_shape", pattern, 16, 100).with_value_kind(value_kind)); + } + } + + for chunks in [4, 16] { + for validity in [IndexValidity::AllValid, IndexValidity::TenPercentNull] { + cases.push( + Case::new("nullable", Pattern::Sorted, chunks, 1_000).with_validity(validity), + ); + } + } + + for pattern in [ + Pattern::SameChunk, + Pattern::AlternatingChunks, + Pattern::GroupedChunks, + ] { + cases.push(Case::new("routing", pattern, 32, 1_000)); + } + + cases +} + +fn primitive_values(start: usize, end: usize) -> ArrayRef { + PrimitiveArray::from_iter((start..end).map(|value| u64::try_from(value).unwrap())).into_array() +} + +fn chunk_values(value_kind: ValueKind, start: usize, end: usize) -> ArrayRef { + match value_kind { + ValueKind::Primitive => primitive_values(start, end), + ValueKind::Struct8 => { + let fields: Vec<_> = (0..NESTED_WIDTH) + .map(|field| { + let field_offset = u64::try_from(field * SOURCE_LEN).unwrap(); + PrimitiveArray::from_iter( + (start..end).map(|value| u64::try_from(value).unwrap() + field_offset), + ) + .into_array() + }) + .collect(); + StructArray::try_new( + FieldNames::from(["f0", "f1", "f2", "f3", "f4", "f5", "f6", "f7"]), + fields, + end - start, + Validity::NonNullable, + ) + .unwrap() + .into_array() + } + ValueKind::FixedSizeList8 => { + let element_start = start * NESTED_WIDTH; + let element_end = end * NESTED_WIDTH; + FixedSizeListArray::new( + primitive_values(element_start, element_end), + u32::try_from(NESTED_WIDTH).unwrap(), + Validity::NonNullable, + end - start, + ) + .into_array() + } + } +} + +fn chunked_values(case: Case) -> ArrayRef { + let chunk_len = SOURCE_LEN / case.chunks; + let chunks = (0..case.chunks) + .map(|chunk_idx| { + let start = chunk_idx * chunk_len; + let end = start + chunk_len; + chunk_values(case.value_kind, start, end) + }) + .collect::>(); + let dtype = chunks[0].dtype().clone(); + ChunkedArray::try_new(chunks, dtype).unwrap().into_array() +} + +fn advance_random(state: &mut u64) -> u64 { + *state = state + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + *state +} + +fn index_values(case: Case) -> Vec { + let mut state = 0x4d59_5df4_d0f3_3173u64; + let mut indices = match case.pattern { + Pattern::Sorted | Pattern::Shuffled => (0..case.indices) + .map(|_| advance_random(&mut state) % u64::try_from(SOURCE_LEN).unwrap()) + .collect(), + Pattern::Repeated => vec![u64::try_from(SOURCE_LEN / 2).unwrap(); case.indices], + Pattern::Duplicate90 | Pattern::Duplicate99 => { + let divisor = match case.pattern { + Pattern::Duplicate90 => 10, + Pattern::Duplicate99 => 100, + _ => unreachable!(), + }; + let unique_count = (case.indices / divisor).max(1); + (0..case.indices) + .map(|_| { + let slot = usize::try_from( + advance_random(&mut state) % u64::try_from(unique_count).unwrap(), + ) + .unwrap(); + u64::try_from(slot * SOURCE_LEN / unique_count).unwrap() + }) + .collect() + } + Pattern::SameChunk => { + let chunk_len = SOURCE_LEN / case.chunks; + let chunk_start = (case.chunks / 2) * chunk_len; + (0..case.indices) + .map(|_| { + let local = usize::try_from( + advance_random(&mut state) % u64::try_from(chunk_len).unwrap(), + ) + .unwrap(); + u64::try_from(chunk_start + local).unwrap() + }) + .collect() + } + Pattern::AlternatingChunks => { + let chunk_len = SOURCE_LEN / case.chunks; + (0..case.indices) + .map(|position| { + let chunk_idx = if position.is_multiple_of(2) { + 0 + } else { + case.chunks - 1 + }; + let local = position / 2 % chunk_len; + u64::try_from(chunk_idx * chunk_len + local).unwrap() + }) + .collect() + } + Pattern::GroupedChunks => { + let chunk_len = SOURCE_LEN / case.chunks; + (0..case.indices) + .map(|position| { + let chunk_idx = position * case.chunks / case.indices; + let local = usize::try_from( + advance_random(&mut state) % u64::try_from(chunk_len).unwrap(), + ) + .unwrap(); + u64::try_from(chunk_idx * chunk_len + local).unwrap() + }) + .collect() + } + }; + + if matches!(case.pattern, Pattern::Sorted) { + indices.sort_unstable(); + } + indices +} + +fn take_indices(case: Case) -> ArrayRef { + let indices = index_values(case); + match case.validity { + IndexValidity::NonNullable => PrimitiveArray::from_iter(indices).into_array(), + IndexValidity::AllValid => { + PrimitiveArray::new(Buffer::from(indices), Validity::AllValid).into_array() + } + IndexValidity::TenPercentNull => PrimitiveArray::new( + Buffer::from(indices), + Validity::from_iter((0..case.indices).map(|index| !index.is_multiple_of(10))), + ) + .into_array(), + } +} + +#[divan::bench(args = cases())] +fn take(bencher: Bencher, case: Case) { + let values = chunked_values(case); + let indices = take_indices(case); + let session = array_session(); + + bencher + .counter(ItemsCount::new(case.indices)) + .with_inputs(|| (indices.clone(), session.create_execution_ctx())) + .bench_refs(|(indices, ctx)| { + values + .take(indices.clone()) + .unwrap() + .execute::(ctx) + .unwrap() + }); +} diff --git a/vortex-array/benches/take_fsl.rs b/vortex-array/benches/take_fsl.rs index 5dc491c28c5..7a89f0eb47c 100644 --- a/vortex-array/benches/take_fsl.rs +++ b/vortex-array/benches/take_fsl.rs @@ -6,6 +6,7 @@ //! Parameterized over: //! - Number of indices to take //! - Fixed size list length (elements per list) +//! - Element byte width #![expect(clippy::cast_possible_truncation)] #![expect(clippy::unwrap_used)] @@ -13,16 +14,30 @@ use std::sync::LazyLock; use divan::Bencher; +use divan::counter::BytesCount; +use num_traits::FromPrimitive; use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::array_session; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::PiecewiseSequenceArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use vortex_array::dtype::IntegerPType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::half::f16; +use vortex_array::match_smallest_offset_type; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; use vortex_session::VortexSession; fn main() { @@ -36,21 +51,51 @@ static SESSION: LazyLock = LazyLock::new(array_session); const NUM_LISTS: usize = 500; /// Number of indices to take. -const NUM_INDICES: &[usize] = &[100, 1_000]; +/// +/// Together with [`LIST_SIZES`] these are sized so the widest type at the largest list size stays +/// under 1ms per iteration under codspeed simulation. +const NUM_INDICES: &[usize] = &[10, 100]; -/// Fixed size list lengths (elements per list). -const LIST_SIZES: &[usize] = &[16, 64, 256, 1024, 4096]; +/// Fixed size list lengths (elements per list). See [`NUM_INDICES`]. +const LIST_SIZES: &[usize] = &[16, 64, 128, 256]; + +/// F16 list lengths for isolating the per-index, piecewise, and manual range-copy strategies. +const F16_STRATEGY_LIST_SIZES: &[usize] = &[1, 2, 4, 8, 16, 64, 128, 256, 512, 1024, 2048]; + +/// F16 strategy benchmarks keep a smaller take width so the forced slow strategies stay cheap. +const F16_STRATEGY_NUM_INDICES: &[usize] = &[10]; /// Creates a FixedSizeListArray with the given list size and number of lists. -fn create_fsl(list_size: usize, num_lists: usize) -> FixedSizeListArray { +fn create_fsl(list_size: usize, num_lists: usize) -> FixedSizeListArray +where + T: NativePType + FromPrimitive, +{ + create_fsl_with_validity::(list_size, num_lists, Validity::NonNullable) +} + +fn create_fsl_with_validity( + list_size: usize, + num_lists: usize, + validity: Validity, +) -> FixedSizeListArray +where + T: NativePType + FromPrimitive, +{ + let total_elements = list_size * num_lists; + let elements: Buffer = (0..total_elements) + .map(|idx| T::from_u16((idx % 251) as u16).unwrap()) + .collect(); + FixedSizeListArray::new(elements.into_array(), list_size as u32, validity, num_lists) +} + +fn create_i64_fsl_with_validity( + list_size: usize, + num_lists: usize, + validity: Validity, +) -> FixedSizeListArray { let total_elements = list_size * num_lists; let elements: Buffer = (0..total_elements as i64).collect(); - FixedSizeListArray::new( - elements.into_array(), - list_size as u32, - Validity::NonNullable, - num_lists, - ) + FixedSizeListArray::new(elements.into_array(), list_size as u32, validity, num_lists) } /// Creates random indices for taking from the array. @@ -61,39 +106,157 @@ fn create_random_indices(num_indices: usize, max_index: usize) -> Buffer { .collect() } -#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] -fn take_fsl_random(bencher: Bencher, num_indices: usize) { - let fsl = create_fsl(LIST_SIZE, NUM_LISTS); - let indices = create_random_indices(num_indices, NUM_LISTS); - let indices_array = indices.into_array(); +/// Number of chunks used for the chunked-elements and chunked-of-FSL benchmarks. +const NUM_CHUNKS: usize = 8; + +/// Take widths for the chunked benchmarks. +/// +/// Together with [`CHUNKED_LIST_SIZES`] these are sized so that the slowest current take path +/// (per-element expansion through the generic chunked take) stays under 1ms per iteration under +/// codspeed simulation, which reports roughly an order of magnitude above local wall time. +const CHUNKED_NUM_INDICES: &[usize] = &[64]; + +/// Fixed size list lengths for the chunked benchmarks. See [`CHUNKED_NUM_INDICES`]. +const CHUNKED_LIST_SIZES: &[usize] = &[8, 16, 32]; + +/// Creates a `FixedSizeListArray` whose elements child is a `ChunkedArray` with `NUM_CHUNKS` +/// chunks whose boundaries do not line up with list boundaries. +fn create_fsl_with_chunked_elements(list_size: usize, num_lists: usize) -> FixedSizeListArray { + let total_elements = list_size * num_lists; + let chunk_len = total_elements.div_ceil(NUM_CHUNKS); + let mut chunks = Vec::new(); + let mut start = 0usize; + while start < total_elements { + let end = (start + chunk_len).min(total_elements); + let chunk: Buffer = (start as i64..end as i64).collect(); + chunks.push(chunk.into_array()); + start = end; + } + let dtype = chunks[0].dtype().clone(); + let elements = ChunkedArray::try_new(chunks, dtype).unwrap().into_array(); + FixedSizeListArray::new(elements, list_size as u32, Validity::NonNullable, num_lists) +} + +/// Creates a `ChunkedArray` of `NUM_CHUNKS` `FixedSizeListArray` chunks. +fn create_chunked_fsl(list_size: usize, num_lists: usize) -> ArrayRef { + let lists_per_chunk = num_lists.div_ceil(NUM_CHUNKS); + let mut chunks = Vec::new(); + let mut start = 0usize; + while start < num_lists { + let end = (start + lists_per_chunk).min(num_lists); + let elements: Buffer = + ((start * list_size) as i64..(end * list_size) as i64).collect(); + chunks.push( + FixedSizeListArray::new( + elements.into_array(), + list_size as u32, + Validity::NonNullable, + end - start, + ) + .into_array(), + ); + start = end; + } + let dtype = chunks[0].dtype().clone(); + ChunkedArray::try_new(chunks, dtype).unwrap().into_array() +} + +fn bench_take_array( + bencher: Bencher, + num_indices: usize, + array: ArrayRef, + sorted: bool, +) { + let mut indices = create_random_indices(num_indices, NUM_LISTS) + .as_ref() + .to_vec(); + if sorted { + indices.sort_unstable(); + } + let indices_array = Buffer::from(indices).into_array(); bencher - .with_inputs(|| (&fsl, &indices_array, SESSION.create_execution_ctx())) + .counter(BytesCount::of_many::(num_indices * LIST_SIZE)) + .with_inputs(|| (&array, &indices_array, SESSION.create_execution_ctx())) .bench_refs(|(array, indices, execution_ctx)| { array - .clone() - .take(indices.clone()) + .take((*indices).clone()) .unwrap() .execute::(execution_ctx) .unwrap() }); } +#[divan::bench(args = CHUNKED_NUM_INDICES, consts = CHUNKED_LIST_SIZES)] +fn take_fsl_chunked_elements_random(bencher: Bencher, num_indices: usize) { + let fsl = create_fsl_with_chunked_elements(LIST_SIZE, NUM_LISTS); + bench_take_array::(bencher, num_indices, fsl.into_array(), false); +} + +#[divan::bench(args = CHUNKED_NUM_INDICES, consts = CHUNKED_LIST_SIZES)] +fn take_fsl_chunked_elements_sorted(bencher: Bencher, num_indices: usize) { + let fsl = create_fsl_with_chunked_elements(LIST_SIZE, NUM_LISTS); + bench_take_array::(bencher, num_indices, fsl.into_array(), true); +} + +#[divan::bench(args = CHUNKED_NUM_INDICES, consts = CHUNKED_LIST_SIZES)] +fn take_chunked_fsl_random(bencher: Bencher, num_indices: usize) { + let array = create_chunked_fsl(LIST_SIZE, NUM_LISTS); + bench_take_array::(bencher, num_indices, array, false); +} + +#[divan::bench(args = CHUNKED_NUM_INDICES, consts = CHUNKED_LIST_SIZES)] +fn take_chunked_fsl_sorted(bencher: Bencher, num_indices: usize) { + let array = create_chunked_fsl(LIST_SIZE, NUM_LISTS); + bench_take_array::(bencher, num_indices, array, true); +} + #[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] -fn take_fsl_nullable_random(bencher: Bencher, num_indices: usize) { - let total_elements = LIST_SIZE * NUM_LISTS; - let elements: Buffer = (0..total_elements as i64).collect(); +fn take_fsl_random(bencher: Bencher, num_indices: usize) { + let fsl = create_i64_fsl_with_validity(LIST_SIZE, NUM_LISTS, Validity::NonNullable); + bench_take_fsl_random::(bencher, num_indices, fsl); +} - // Create validity with ~10% nulls - let mut rng = StdRng::seed_from_u64(123); - let validity = Validity::from_iter((0..NUM_LISTS).map(|_| rng.random_ratio(9, 10))); +#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] +fn take_fsl_f16_random(bencher: Bencher, num_indices: usize) { + take_fsl_random_typed::(bencher, num_indices); +} - let fsl = FixedSizeListArray::new(elements.into_array(), LIST_SIZE as u32, validity, NUM_LISTS); +#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] +fn take_fsl_u8_random(bencher: Bencher, num_indices: usize) { + take_fsl_random_typed::(bencher, num_indices); +} +#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] +fn take_fsl_u32_random(bencher: Bencher, num_indices: usize) { + take_fsl_random_typed::(bencher, num_indices); +} + +#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] +fn take_fsl_u64_random(bencher: Bencher, num_indices: usize) { + take_fsl_random_typed::(bencher, num_indices); +} + +fn take_fsl_random_typed(bencher: Bencher, num_indices: usize) +where + T: NativePType + FromPrimitive, +{ + let fsl = create_fsl::(LIST_SIZE, NUM_LISTS); + bench_take_fsl_random::(bencher, num_indices, fsl); +} + +fn bench_take_fsl_random( + bencher: Bencher, + num_indices: usize, + fsl: FixedSizeListArray, +) where + T: NativePType, +{ let indices = create_random_indices(num_indices, NUM_LISTS); let indices_array = indices.into_array(); bencher + .counter(BytesCount::of_many::(num_indices * LIST_SIZE)) .with_inputs(|| (&fsl, &indices_array, SESSION.create_execution_ctx())) .bench_refs(|(array, indices, execution_ctx)| { array @@ -104,3 +267,186 @@ fn take_fsl_nullable_random(bencher: Bencher, num_indice .unwrap() }); } + +#[divan::bench(args = F16_STRATEGY_NUM_INDICES, consts = F16_STRATEGY_LIST_SIZES)] +fn take_fsl_f16_force_per_index(bencher: Bencher, num_indices: usize) { + let fsl = create_fsl::(LIST_SIZE, NUM_LISTS); + let indices = create_random_indices(num_indices, NUM_LISTS); + + bencher + .counter(BytesCount::of_many::(num_indices * LIST_SIZE)) + .with_inputs(|| (&fsl, &indices, SESSION.create_execution_ctx())) + .bench_refs(|(array, indices, execution_ctx)| { + match_smallest_offset_type!(array.elements().len(), |E| { + take_fsl_f16_per_index_strategy::(array, indices) + }) + .into_array() + .execute::(execution_ctx) + .unwrap() + }); +} + +#[divan::bench(args = F16_STRATEGY_NUM_INDICES, consts = F16_STRATEGY_LIST_SIZES)] +fn take_fsl_f16_force_piecewise_sequence( + bencher: Bencher, + num_indices: usize, +) { + let fsl = create_fsl::(LIST_SIZE, NUM_LISTS); + let indices = create_random_indices(num_indices, NUM_LISTS); + + bencher + .counter(BytesCount::of_many::(num_indices * LIST_SIZE)) + .with_inputs(|| (&fsl, &indices, SESSION.create_execution_ctx())) + .bench_refs(|(array, indices, execution_ctx)| { + take_fsl_f16_piecewise_sequence_strategy::(array, indices) + .into_array() + .execute::(execution_ctx) + .unwrap() + }); +} + +#[divan::bench(args = F16_STRATEGY_NUM_INDICES, consts = F16_STRATEGY_LIST_SIZES)] +fn take_fsl_f16_force_manual_range_copy( + bencher: Bencher, + num_indices: usize, +) { + let fsl = create_fsl::(LIST_SIZE, NUM_LISTS); + let indices = create_random_indices(num_indices, NUM_LISTS); + + bencher + .counter(BytesCount::of_many::(num_indices * LIST_SIZE)) + .with_inputs(|| (&fsl, &indices, SESSION.create_execution_ctx())) + .bench_refs(|(array, indices, execution_ctx)| { + take_fsl_f16_manual_range_copy_strategy::(array, indices, execution_ctx) + .into_array() + .execute::(execution_ctx) + .unwrap() + }); +} + +fn take_fsl_f16_per_index_strategy( + array: &FixedSizeListArray, + indices: &Buffer, +) -> FixedSizeListArray { + let mut element_indices = BufferMut::::with_capacity(indices.len() * LIST_SIZE); + for &idx in indices.as_ref() { + let start = idx as usize * LIST_SIZE; + let end = start + LIST_SIZE; + for element_idx in start..end { + // SAFETY: capacity is exactly `indices.len() * LIST_SIZE`, and this loop appends + // exactly `LIST_SIZE` element indices for each input index. + unsafe { element_indices.push_unchecked(E::from_usize(element_idx).unwrap()) }; + } + } + + let element_indices = + PrimitiveArray::new(element_indices.freeze(), Validity::NonNullable).into_array(); + let elements = array.elements().take(element_indices).unwrap(); + + // SAFETY: `elements` was built by taking exactly `LIST_SIZE` elements per input index, so its + // length is `indices.len() * LIST_SIZE`; the output is non-nullable by construction. + unsafe { + FixedSizeListArray::new_unchecked( + elements, + LIST_SIZE as u32, + Validity::NonNullable, + indices.len(), + ) + } +} + +fn take_fsl_f16_piecewise_sequence_strategy( + array: &FixedSizeListArray, + indices: &Buffer, +) -> FixedSizeListArray { + let starts = indices + .as_ref() + .iter() + .map(|&idx| idx * LIST_SIZE as u64) + .collect::>(); + let run_count = starts.len(); + let starts = PrimitiveArray::from_iter(starts).into_array(); + let lengths = ConstantArray::new(LIST_SIZE as u64, run_count).into_array(); + let multipliers = ConstantArray::new(1u64, run_count).into_array(); + + // SAFETY: benchmark indices are generated in-bounds; lengths and multiplier 1 are + // non-nullable unsigned constants; output length is exactly `indices.len() * LIST_SIZE`. + let element_indices = unsafe { + PiecewiseSequenceArray::new_unchecked( + starts, + lengths, + multipliers, + indices.len() * LIST_SIZE, + ) + } + .into_array(); + let elements = array.elements().take(element_indices).unwrap(); + + // SAFETY: each generated run has width `LIST_SIZE`, and there is one run per input index, + // so `elements.len() == indices.len() * LIST_SIZE`. + unsafe { + FixedSizeListArray::new_unchecked( + elements, + LIST_SIZE as u32, + Validity::NonNullable, + indices.len(), + ) + } +} + +fn take_fsl_f16_manual_range_copy_strategy( + array: &FixedSizeListArray, + indices: &Buffer, + execution_ctx: &mut ExecutionCtx, +) -> FixedSizeListArray { + let elements = array + .elements() + .clone() + .execute::(execution_ctx) + .unwrap(); + let source = elements.as_slice::(); + let mut values = BufferMut::::with_capacity(indices.len() * LIST_SIZE); + + for &idx in indices.as_ref() { + let start = idx as usize * LIST_SIZE; + values.extend_from_slice(&source[start..start + LIST_SIZE]); + } + + // SAFETY: the buffer was filled with exactly `LIST_SIZE` copied values per input index, so it + // has the element length required by the constructed FSL. + unsafe { + FixedSizeListArray::new_unchecked( + PrimitiveArray::new(values.freeze(), Validity::NonNullable).into_array(), + LIST_SIZE as u32, + Validity::NonNullable, + indices.len(), + ) + } +} + +#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] +fn take_fsl_nullable_random(bencher: Bencher, num_indices: usize) { + // Create validity with ~10% nulls + let mut rng = StdRng::seed_from_u64(123); + let validity = Validity::from_iter((0..NUM_LISTS).map(|_| rng.random_ratio(9, 10))); + + let fsl = create_i64_fsl_with_validity(LIST_SIZE, NUM_LISTS, validity); + bench_take_fsl_random::(bencher, num_indices, fsl); +} + +#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] +fn take_fsl_f16_nullable_random(bencher: Bencher, num_indices: usize) { + take_fsl_nullable_random_typed::(bencher, num_indices); +} + +fn take_fsl_nullable_random_typed(bencher: Bencher, num_indices: usize) +where + T: NativePType + FromPrimitive, +{ + // Create validity with ~10% nulls + let mut rng = StdRng::seed_from_u64(123); + let validity = Validity::from_iter((0..NUM_LISTS).map(|_| rng.random_ratio(9, 10))); + + let fsl = create_fsl_with_validity::(LIST_SIZE, NUM_LISTS, validity); + bench_take_fsl_random::(bencher, num_indices, fsl); +} diff --git a/vortex-array/benches/take_patches.rs b/vortex-array/benches/take_patches.rs index d3b99e8ea1b..cae0176df89 100644 --- a/vortex-array/benches/take_patches.rs +++ b/vortex-array/benches/take_patches.rs @@ -27,6 +27,9 @@ fn main() { static SESSION: LazyLock = LazyLock::new(array_session); +/// Sized to keep CodSpeed simulation under 1ms per benchmark. +const ARRAY_LEN: usize = 4096; + const BENCH_ARGS: &[(f64, f64)] = &[ // patches_sparsity, index_multiple (0.1, 1.0), @@ -46,7 +49,7 @@ const BENCH_ARGS: &[(f64, f64)] = &[ #[divan::bench(args = BENCH_ARGS)] fn take_search(bencher: Bencher, (patches_sparsity, index_multiple): (f64, f64)) { let mut rng = StdRng::seed_from_u64(0); - let patches = fixture(65536, patches_sparsity, &mut rng); + let patches = fixture(ARRAY_LEN, patches_sparsity, &mut rng); let indices = indices( patches.array_len(), (patches.array_len() as f64 * index_multiple) as usize, @@ -67,7 +70,7 @@ fn take_search(bencher: Bencher, (patches_sparsity, index_multiple): (f64, f64)) #[divan::bench(args = BENCH_ARGS)] fn take_search_chunked(bencher: Bencher, (patches_sparsity, index_multiple): (f64, f64)) { let mut rng = StdRng::seed_from_u64(0); - let patches = fixture_with_chunk_offsets(65536, patches_sparsity, &mut rng); + let patches = fixture_with_chunk_offsets(ARRAY_LEN, patches_sparsity, &mut rng); let indices = indices( patches.array_len(), (patches.array_len() as f64 * index_multiple) as usize, @@ -88,7 +91,7 @@ fn take_search_chunked(bencher: Bencher, (patches_sparsity, index_multiple): (f6 #[divan::bench(args = BENCH_ARGS)] fn take_map(bencher: Bencher, (patches_sparsity, index_multiple): (f64, f64)) { let mut rng = StdRng::seed_from_u64(0); - let patches = fixture(65536, patches_sparsity, &mut rng); + let patches = fixture(ARRAY_LEN, patches_sparsity, &mut rng); let indices = indices( patches.array_len(), (patches.array_len() as f64 * index_multiple) as usize, @@ -125,7 +128,7 @@ fn fixture_with_chunk_offsets(len: usize, sparsity: f64, rng: &mut StdRng) -> Pa let sparse_len = patch_indices.len(); let values = Buffer::from_iter((0..sparse_len).map(|x| x as u64)).into_array(); - const PATCH_CHUNK_SIZE: usize = 1024; + const PATCH_CHUNK_SIZE: usize = 256; let chunk_offsets: Vec = (0..len) .step_by(PATCH_CHUNK_SIZE) .map(|chunk_start| { diff --git a/vortex-array/benches/take_primitive.rs b/vortex-array/benches/take_primitive.rs index 471363b99e4..a3787d55a0d 100644 --- a/vortex-array/benches/take_primitive.rs +++ b/vortex-array/benches/take_primitive.rs @@ -1,9 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Benchmarks comparing [`PVector`] take vs [`DictArray`] canonicalization. -//! -//! Both are tracked by number of indices/codes for fair comparison. +//! Benchmarks for primitive take and [`DictArray`] canonicalization. #![expect(clippy::cast_possible_truncation)] #![expect(clippy::unwrap_used)] @@ -11,6 +9,7 @@ use std::sync::LazyLock; use divan::Bencher; +use divan::counter::ItemsCount; use rand::distr::Uniform; use rand::prelude::*; use rand_distr::Zipf; @@ -27,14 +26,38 @@ fn main() { divan::main(); } -/// Number of indices to take. -const NUM_INDICES: &[usize] = &[1_000, 10_000, 100_000]; +/// Number of indices to take. The top tier is sized to keep CodSpeed simulation under 1ms. +const NUM_INDICES: &[usize] = &[1_000, 10_000, 25_000]; + +/// Large enough to measure both cache-resident and streaming dictionary decoding. +const GT_NUM_INDICES: &[usize] = &[1_000_000, 16_000_000]; /// Size of the source vector / dictionary values. const VECTOR_SIZE: &[usize] = &[16, 256, 2048, 8192]; static SESSION: LazyLock = LazyLock::new(array_session); +#[divan::bench(sample_count = 200)] +fn primitive_take_u32(bencher: Bencher) { + // Sized to keep CodSpeed simulation under 1ms. + const NUM_INDICES: usize = 50_000; + + let values = PrimitiveArray::from_iter(0u32..256).into_array(); + let indices = + PrimitiveArray::from_iter((0..NUM_INDICES).map(|index| index as u32 & 255)).into_array(); + + bencher + .counter(ItemsCount::new(NUM_INDICES)) + .with_inputs(|| (indices.clone(), SESSION.create_execution_ctx())) + .bench_values(|(indices, mut ctx)| { + values + .take(indices) + .unwrap() + .execute::(&mut ctx) + .unwrap() + }); +} + // --- DictArray canonicalization benchmarks --- #[divan::bench(args = NUM_INDICES, consts = VECTOR_SIZE, sample_count = 100_000)] @@ -70,3 +93,28 @@ fn dict_canonicalize_zipfian(bencher: Bencher, num_indi .with_inputs(|| (&dict, SESSION.create_execution_ctx())) .bench_refs(|(dict, ctx)| (*dict).clone().into_array().execute::(ctx)); } + +/// StatPopGen GT distribution: 75.21% 0, 3.85% 1, 0.45% 2, and 20.49% null. +#[vortex_bench_support::cpu_features] +#[divan::bench(args = GT_NUM_INDICES)] +fn dict_canonicalize_gt_u8(bencher: Bencher, num_indices: usize) { + let values = PrimitiveArray::from_option_iter([Some(0u8), Some(1), Some(2), None]); + let range = Uniform::new(0u16, 10_000).unwrap(); + let codes = PrimitiveArray::from_iter( + StdRng::seed_from_u64(0) + .sample_iter(range) + .take(num_indices) + .map(|sample| match sample { + 0..7521 => 0u8, + 7521..7906 => 1, + 7906..7951 => 2, + _ => 3, + }), + ); + let dict = DictArray::try_new(codes.into_array(), values.into_array()).unwrap(); + + bencher + .counter(ItemsCount::new(num_indices)) + .with_inputs(|| (&dict, SESSION.create_execution_ctx())) + .bench_refs(|(dict, ctx)| (*dict).clone().into_array().execute::(ctx)); +} diff --git a/vortex-array/benches/take_slices_to_buffer_matrix.rs b/vortex-array/benches/take_slices_to_buffer_matrix.rs new file mode 100644 index 00000000000..21cb539ce32 --- /dev/null +++ b/vortex-array/benches/take_slices_to_buffer_matrix.rs @@ -0,0 +1,311 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Microbenchmarks for primitive `take_slices_to_buffer` copy-loop variants. +//! +//! The matrix covers: +//! - append via `BufferMut::extend_from_slice`, indexed cursor copy, and advancing pointer copy +//! into spare output capacity +//! - ordinary checked slicing vs a preverification pass followed by unchecked slicing +//! - fixed-width short slices at the run counts used by the FSL take benchmarks + +#![allow(clippy::cast_possible_truncation)] +#![expect(clippy::unwrap_used)] + +use divan::Bencher; +use divan::counter::BytesCount; +use rand::SeedableRng; +use rand::prelude::*; +use rand_distr::Distribution; +use rand_distr::Normal; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; + +fn main() { + divan::main(); +} + +const SOURCE_LENS: &[usize] = &[1_000, 10_000, 100_000]; +const SLICE_COUNT: usize = 50; +const FIXED_16_SLICE_COUNTS: &[usize] = &[100, 1_000]; +const FIXED_16_SOURCE_RUNS: usize = 500; + +type TakeSlicesFn = fn(&[u16], &[usize], &[usize], usize) -> Buffer; + +#[derive(Clone)] +struct Case { + values: Vec, + starts: Vec, + lengths: Vec, + output_len: usize, +} + +#[divan::bench(args = SOURCE_LENS)] +fn extend_safe(bencher: Bencher, source_len: usize) { + let case = Case::new(source_len); + bench_case(bencher, case, take_extend_safe); +} + +#[divan::bench(args = SOURCE_LENS)] +fn cursor_copy_safe(bencher: Bencher, source_len: usize) { + let case = Case::new(source_len); + bench_case(bencher, case, take_cursor_copy_safe); +} + +#[divan::bench(args = SOURCE_LENS)] +fn advancing_ptr_safe(bencher: Bencher, source_len: usize) { + let case = Case::new(source_len); + bench_case(bencher, case, take_advancing_ptr_safe); +} + +#[divan::bench(args = SOURCE_LENS)] +fn preverify_extend_unchecked(bencher: Bencher, source_len: usize) { + let case = Case::new(source_len); + bench_case(bencher, case, take_preverify_extend_unchecked); +} + +#[divan::bench(args = SOURCE_LENS)] +fn preverify_cursor_copy_unchecked(bencher: Bencher, source_len: usize) { + let case = Case::new(source_len); + bench_case(bencher, case, take_preverify_cursor_copy_unchecked); +} + +#[divan::bench(args = SOURCE_LENS)] +fn preverify_advancing_ptr_unchecked(bencher: Bencher, source_len: usize) { + let case = Case::new(source_len); + bench_case(bencher, case, take_preverify_advancing_ptr_unchecked); +} + +#[divan::bench(args = FIXED_16_SLICE_COUNTS)] +fn fixed_16_extend_safe(bencher: Bencher, slice_count: usize) { + let case = Case::fixed_width(slice_count, 16); + bench_case(bencher, case, take_extend_safe); +} + +#[divan::bench(args = FIXED_16_SLICE_COUNTS)] +fn fixed_16_cursor_copy_safe(bencher: Bencher, slice_count: usize) { + let case = Case::fixed_width(slice_count, 16); + bench_case(bencher, case, take_cursor_copy_safe); +} + +#[divan::bench(args = FIXED_16_SLICE_COUNTS)] +fn fixed_16_advancing_ptr_safe(bencher: Bencher, slice_count: usize) { + let case = Case::fixed_width(slice_count, 16); + bench_case(bencher, case, take_advancing_ptr_safe); +} + +fn bench_case(bencher: Bencher, case: Case, f: TakeSlicesFn) { + bencher + .counter(BytesCount::of_many::(case.output_len)) + .with_inputs(|| (&case.values, &case.starts, &case.lengths, case.output_len)) + .bench_refs(|(values, starts, lengths, output_len)| { + divan::black_box(f(values, starts, lengths, *output_len)); + }); +} + +impl Case { + fn new(source_len: usize) -> Self { + let mut rng = StdRng::seed_from_u64(source_len as u64 ^ 0x5eed_51ce); + let values = (0..source_len).map(|_| rng.random::()).collect(); + let mean = source_len as f64 / 100.0; + let normal = Normal::new(mean, (mean / 4.0).max(1.0)).unwrap(); + + let mut ranges = (0..SLICE_COUNT) + .map(|_| { + let length = normal + .sample(&mut rng) + .round() + .clamp(1.0, source_len as f64) as usize; + let start = rng.random_range(0..=source_len - length); + (start, length) + }) + .collect::>(); + ranges.sort_unstable_by_key(|&(start, _)| start); + + let output_len = ranges.iter().map(|&(_, length)| length).sum(); + let (starts, lengths) = ranges.into_iter().unzip(); + Self { + values, + starts, + lengths, + output_len, + } + } + + fn fixed_width(slice_count: usize, width: usize) -> Self { + let mut rng = StdRng::seed_from_u64(slice_count as u64 ^ 0x05ee_df16); + let source_len = FIXED_16_SOURCE_RUNS * width; + let values = (0..source_len).map(|_| rng.random::()).collect(); + let starts = (0..slice_count) + .map(|_| rng.random_range(0..FIXED_16_SOURCE_RUNS) * width) + .collect(); + let lengths = vec![width; slice_count]; + Self { + values, + starts, + lengths, + output_len: slice_count * width, + } + } +} + +fn take_extend_safe( + values: &[u16], + starts: &[usize], + lengths: &[usize], + output_len: usize, +) -> Buffer { + let mut result = BufferMut::::with_capacity(output_len); + for (&start, &length) in starts.iter().zip(lengths) { + result.extend_from_slice(&values[start..start + length]); + } + result.freeze() +} + +fn take_cursor_copy_safe( + values: &[u16], + starts: &[usize], + lengths: &[usize], + output_len: usize, +) -> Buffer { + let mut result = BufferMut::::with_capacity(output_len); + let mut cursor = 0usize; + for (&start, &length) in starts.iter().zip(lengths) { + let end = cursor.checked_add(length).unwrap(); + assert!(end <= output_len); + copy_to_spare(&mut result, cursor, &values[start..start + length]); + cursor = end; + } + assert_eq!(cursor, output_len); + + // SAFETY: the loop writes exactly `output_len` values into spare capacity. + unsafe { result.set_len(output_len) }; + result.freeze() +} + +fn take_advancing_ptr_safe( + values: &[u16], + starts: &[usize], + lengths: &[usize], + output_len: usize, +) -> Buffer { + let mut result = BufferMut::::with_capacity(output_len); + let mut cursor = 0usize; + let mut dst = result.spare_capacity_mut().as_mut_ptr().cast::(); + for (&start, &length) in starts.iter().zip(lengths) { + let end = cursor.checked_add(length).unwrap(); + assert!(end <= output_len); + let source = &values[start..start + length]; + // SAFETY: `end <= output_len` proves destination capacity, and safe slicing proves source + // bounds. + unsafe { + copy_to_uninit(dst, source); + dst = dst.add(length); + } + cursor = end; + } + assert_eq!(cursor, output_len); + + // SAFETY: the loop writes exactly `output_len` values into spare capacity. + unsafe { result.set_len(output_len) }; + result.freeze() +} + +fn take_preverify_extend_unchecked( + values: &[u16], + starts: &[usize], + lengths: &[usize], + output_len: usize, +) -> Buffer { + preverify(values.len(), starts, lengths, output_len); + + let mut result = BufferMut::::with_capacity(output_len); + for (&start, &length) in starts.iter().zip(lengths) { + // SAFETY: `preverify` checked every source range. + unsafe { + result.extend_from_slice(values.get_unchecked(start..start + length)); + } + } + result.freeze() +} + +fn take_preverify_cursor_copy_unchecked( + values: &[u16], + starts: &[usize], + lengths: &[usize], + output_len: usize, +) -> Buffer { + preverify(values.len(), starts, lengths, output_len); + + let mut result = BufferMut::::with_capacity(output_len); + let mut cursor = 0usize; + for (&start, &length) in starts.iter().zip(lengths) { + // SAFETY: `preverify` checked every source range. + let source = unsafe { values.get_unchecked(start..start + length) }; + // SAFETY: `preverify` checked the summed output length. + unsafe { copy_to_spare_unchecked(&mut result, cursor, source) }; + cursor += length; + } + + // SAFETY: the loop writes exactly `output_len` values into spare capacity. + unsafe { result.set_len(output_len) }; + result.freeze() +} + +fn take_preverify_advancing_ptr_unchecked( + values: &[u16], + starts: &[usize], + lengths: &[usize], + output_len: usize, +) -> Buffer { + preverify(values.len(), starts, lengths, output_len); + + let mut result = BufferMut::::with_capacity(output_len); + let mut dst = result.spare_capacity_mut().as_mut_ptr().cast::(); + for (&start, &length) in starts.iter().zip(lengths) { + // SAFETY: `preverify` checked every source range and the summed output length. + unsafe { + let source = values.get_unchecked(start..start + length); + copy_to_uninit(dst, source); + dst = dst.add(length); + } + } + + // SAFETY: `preverify` proves the loop writes exactly `output_len` values into spare capacity. + unsafe { result.set_len(output_len) }; + result.freeze() +} + +fn preverify(source_len: usize, starts: &[usize], lengths: &[usize], output_len: usize) { + assert_eq!(starts.len(), lengths.len()); + let mut cursor = 0usize; + for (&start, &length) in starts.iter().zip(lengths) { + let end = start.checked_add(length).unwrap(); + assert!(end <= source_len); + cursor = cursor.checked_add(length).unwrap(); + assert!(cursor <= output_len); + } + assert_eq!(cursor, output_len); +} + +fn copy_to_spare(result: &mut BufferMut, cursor: usize, source: &[u16]) { + let dst = &mut result.spare_capacity_mut()[cursor..][..source.len()]; + // SAFETY: `dst` has exactly `source.len()` spare slots and does not overlap with source. + unsafe { copy_to_uninit(dst.as_mut_ptr().cast(), source) }; +} + +unsafe fn copy_to_spare_unchecked(result: &mut BufferMut, cursor: usize, source: &[u16]) { + // SAFETY: callers ensure `cursor..cursor + source.len()` is within spare capacity. + let dst = unsafe { + result + .spare_capacity_mut() + .get_unchecked_mut(cursor..cursor + source.len()) + }; + // SAFETY: `dst` has exactly `source.len()` spare slots and does not overlap with source. + unsafe { copy_to_uninit(dst.as_mut_ptr().cast(), source) }; +} + +unsafe fn copy_to_uninit(dst: *mut u16, source: &[u16]) { + // SAFETY: callers ensure `dst` points to `source.len()` writable uninitialized u16 slots. + unsafe { std::ptr::copy_nonoverlapping(source.as_ptr(), dst, source.len()) }; +} diff --git a/vortex-array/benches/take_varbin.rs b/vortex-array/benches/take_varbin.rs new file mode 100644 index 00000000000..77676dd50de --- /dev/null +++ b/vortex-array/benches/take_varbin.rs @@ -0,0 +1,55 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use vortex_array::IntoArray; +use vortex_array::RecursiveCanonical; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::VarBinArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_buffer::Buffer; +use vortex_session::VortexSession; + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(array_session); + +const ARRAY_SIZE: usize = 20_000; +// Sized to keep CodSpeed simulation under 1ms. +const TAKE_SIZE: usize = 3_000; + +#[divan::bench] +fn take_varbin(bencher: Bencher) { + let array = VarBinArray::from_iter( + (0..ARRAY_SIZE).map(|i| Some(format!("row-{i:0>40}"))), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + + let mut rng = StdRng::seed_from_u64(0); + let indices: Buffer = (0..TAKE_SIZE) + .map(|_| rng.random_range(0..ARRAY_SIZE) as u64) + .collect(); + let indices = indices.into_array(); + + bencher + .with_inputs(|| (&array, &indices, SESSION.create_execution_ctx())) + .bench_refs(|(array, indices, ctx)| { + array + .take((*indices).clone()) + .unwrap() + .execute::(ctx) + }); +} diff --git a/vortex-array/benches/varbinview_compact.rs b/vortex-array/benches/varbinview_compact.rs index 80f7e18ec30..903bf7a76a4 100644 --- a/vortex-array/benches/varbinview_compact.rs +++ b/vortex-array/benches/varbinview_compact.rs @@ -29,10 +29,11 @@ fn main() { const ARGS: &[(usize, usize)] = &[ // (output_size, buffer_utilization_pct) - (1 << 12, 10), - (1 << 12, 90), - (1 << 14, 10), - (1 << 14, 90), + // Output sizes sized to keep CodSpeed simulation under 1ms per benchmark. + (1 << 10, 10), + (1 << 10, 90), + (1 << 11, 10), + (1 << 11, 90), ]; static SESSION: LazyLock = LazyLock::new(array_session); @@ -92,7 +93,11 @@ fn compact_sliced_impl(bencher: Bencher, (output_size, utilization_pct): (usize, /// Creates a base VarBinViewArray with mix of inlined and outlined strings. fn build_varbinview_fixture(len: usize) -> VarBinViewArray { - let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::NonNullable), len); + let mut builder = VarBinViewBuilder::with_capacity_in( + DType::Utf8(Nullability::NonNullable), + len, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); let mut rng = StdRng::seed_from_u64(42); for _ in 0..len { diff --git a/vortex-array/benches/varbinview_try_new.rs b/vortex-array/benches/varbinview_try_new.rs new file mode 100644 index 00000000000..fd87bc70a83 --- /dev/null +++ b/vortex-array/benches/varbinview_try_new.rs @@ -0,0 +1,98 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; +use std::sync::LazyLock; + +use divan::Bencher; +use mimalloc::MiMalloc; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::varbinview::BinaryView; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::validity::Validity; +use vortex_buffer::BitBuffer; +use vortex_buffer::Buffer; +use vortex_buffer::ByteBuffer; +use vortex_error::VortexExpect; +use vortex_session::VortexSession; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +fn main() { + divan::main(); +} + +const SIZES: &[usize] = &[1 << 12, 1 << 16]; + +static SESSION: LazyLock = LazyLock::new(array_session); + +fn garbage_view() -> BinaryView { + BinaryView::new_ref(13, *b"AAAA", 0xDEAD_BEEF, 0xF000_0000) +} + +fn inline_view() -> BinaryView { + BinaryView::new_inlined(b"hello world!") +} + +fn owned_nullable_views(len: usize) -> Buffer { + let garbage = garbage_view(); + let inline = inline_view(); + let views: Vec = (0..len) + .map(|i| if i % 2 == 0 { inline } else { garbage }) + .collect(); + Buffer::copy_from(views) +} + +fn owned_valid_views(len: usize) -> Buffer { + let views: Vec = vec![inline_view(); len]; + Buffer::copy_from(views) +} + +fn nullable_validity(len: usize) -> Validity { + let bits = BitBuffer::from_iter((0..len).map(|i| i % 2 == 0)); + Validity::from_bit_buffer(bits, Nullability::Nullable) +} + +#[divan::bench(args = SIZES)] +fn nullable_exclusive(bencher: Bencher, len: usize) { + let dtype = DType::Utf8(Nullability::Nullable); + let buffers: Arc<[ByteBuffer]> = Arc::new([]); + bencher + .with_inputs(|| { + ( + owned_nullable_views(len), + Arc::clone(&buffers), + dtype.clone(), + nullable_validity(len), + SESSION.create_execution_ctx(), + ) + }) + .bench_values(|(views, buffers, dtype, validity, mut ctx)| { + VarBinViewArray::try_new(views, buffers, dtype, validity, &mut ctx) + .vortex_expect("try_new must succeed") + }); +} + +#[divan::bench(args = SIZES)] +fn all_valid_exclusive(bencher: Bencher, len: usize) { + let dtype = DType::Utf8(Nullability::NonNullable); + let buffers: Arc<[ByteBuffer]> = Arc::new([]); + bencher + .with_inputs(|| { + ( + owned_valid_views(len), + Arc::clone(&buffers), + dtype.clone(), + Validity::NonNullable, + SESSION.create_execution_ctx(), + ) + }) + .bench_values(|(views, buffers, dtype, validity, mut ctx)| { + VarBinViewArray::try_new(views, buffers, dtype, validity, &mut ctx) + .vortex_expect("try_new must succeed") + }); +} diff --git a/vortex-array/benches/varbinview_zip.rs b/vortex-array/benches/varbinview_zip.rs index b2d55b1b60d..caf72abf74f 100644 --- a/vortex-array/benches/varbinview_zip.rs +++ b/vortex-array/benches/varbinview_zip.rs @@ -27,7 +27,8 @@ static SESSION: LazyLock = LazyLock::new(array_session); /// Benchmarks zip on VarBinView arrays with a highly fragmented mask (worst case for per-slice lookup paths). #[divan::bench] fn varbinview_zip_fragmented_mask(bencher: Bencher) { - let len = 65_536; + // Sized to keep CodSpeed simulation under 1ms per benchmark. + let len = 2_048; let if_true = fixture(len).into_array(); let if_false = other_fixture(len).into_array(); let mask = alternating_mask(len); @@ -52,7 +53,8 @@ fn varbinview_zip_fragmented_mask(bencher: Bencher) { /// Benchmarks zip on VarBinView arrays with blocky mask segments to contrast with the fragmented case. #[divan::bench] fn varbinview_zip_block_mask(bencher: Bencher) { - let len = 65_536; + // Sized to keep CodSpeed simulation under 1ms per benchmark. + let len = 2_048; let if_true = fixture(len).into_array(); let if_false = other_fixture(len).into_array(); let mask = block_mask(len, 128); diff --git a/vortex-array/src/aggregate_fn/accumulator_grouped.rs b/vortex-array/src/aggregate_fn/accumulator_grouped.rs index b87c04ee204..eda0acdad97 100644 --- a/vortex-array/src/aggregate_fn/accumulator_grouped.rs +++ b/vortex-array/src/aggregate_fn/accumulator_grouped.rs @@ -26,8 +26,9 @@ use crate::arrays::ChunkedArray; use crate::arrays::FixedSizeListArray; use crate::arrays::ListViewArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; -use crate::arrays::listview::ListViewArrayExt; -use crate::builders::builder_with_capacity; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use crate::arrays::listview::ListViewArraySlotsExt; +use crate::builders::builder_with_capacity_in; use crate::builtins::ArrayBuiltins; use crate::columnar::AnyColumnar; use crate::dtype::DType; @@ -272,7 +273,10 @@ impl DynGroupedAccumulator for GroupedAccumulator { } fn flush(&mut self) -> VortexResult { - let states = std::mem::take(&mut self.partials); + let mut states = std::mem::take(&mut self.partials); + if states.len() == 1 { + return Ok(states.pop().vortex_expect("checked one partial")); + } Ok(ChunkedArray::try_new(states, self.partial_dtype.clone())?.into_array()) } @@ -357,7 +361,8 @@ impl GroupedAccumulator { self.options.clone(), self.dtype.clone(), )?; - let mut states = builder_with_capacity(&self.partial_dtype, grouped.len()); + let mut states = + builder_with_capacity_in(&self.partial_dtype, grouped.len(), ctx.allocator()); let group_ranges = grouped.group_ranges(ctx)?; let group_validity = grouped.group_validity(ctx)?; @@ -417,3 +422,68 @@ fn fixed_size_list_group_ranges(groups: &FixedSizeListArray) -> GroupRanges { size: groups.list_size() as usize, } } + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use crate::ArrayRef; + use crate::IntoArray; + use crate::aggregate_fn::DynGroupedAccumulator; + use crate::aggregate_fn::GroupedAccumulator; + use crate::aggregate_fn::NumericalAggregateOpts; + use crate::aggregate_fn::fns::count::Count; + use crate::arrays::Chunked; + use crate::arrays::PrimitiveArray; + use crate::dtype::DType; + use crate::dtype::Nullability::NonNullable; + use crate::dtype::PType; + + fn accumulator() -> VortexResult> { + GroupedAccumulator::try_new( + Count, + NumericalAggregateOpts::default(), + DType::Primitive(PType::I32, NonNullable), + ) + } + + fn state(values: impl IntoIterator) -> ArrayRef { + PrimitiveArray::from_iter(values).into_array() + } + + #[test] + fn test_flush_single_partial_returns_original_array() -> VortexResult<()> { + let mut accumulator = accumulator()?; + let state = state([1, 2, 3]); + accumulator.push_result(state.clone())?; + + let flushed = accumulator.flush()?; + + assert!(ArrayRef::ptr_eq(&flushed, &state)); + Ok(()) + } + + #[test] + fn test_flush_multiple_partials_returns_chunked_array() -> VortexResult<()> { + let mut accumulator = accumulator()?; + accumulator.push_result(state([1, 2]))?; + accumulator.push_result(state([3, 4]))?; + + let flushed = accumulator.flush()?; + + assert!(flushed.is::()); + assert_eq!(flushed.len(), 4); + Ok(()) + } + + #[test] + fn test_flush_without_partials_returns_empty_chunked_array() -> VortexResult<()> { + let mut accumulator = accumulator()?; + + let flushed = accumulator.flush()?; + + assert!(flushed.is::()); + assert!(flushed.is_empty()); + Ok(()) + } +} diff --git a/vortex-array/src/aggregate_fn/combined.rs b/vortex-array/src/aggregate_fn/combined.rs index 76ad9877314..70a385add48 100644 --- a/vortex-array/src/aggregate_fn/combined.rs +++ b/vortex-array/src/aggregate_fn/combined.rs @@ -102,16 +102,6 @@ pub trait BinaryCombined: 'static + Send + Sync + Clone { ); } - /// Coerce the input type. Default: chains `right.coerce_args(left.coerce_args(input))`. - fn coerce_args( - &self, - options: &CombinedOptions, - input_dtype: &DType, - ) -> VortexResult { - let left_coerced = self.left().coerce_args(&options.0, input_dtype)?; - self.right().coerce_args(&options.1, &left_coerced) - } - /// Build the partial struct dtype that wraps the two child partials. fn partial_struct_dtype(&self, left: DType, right: DType) -> DType { DType::Struct( @@ -157,10 +147,6 @@ impl AggregateFnVTable for Combined { BinaryCombined::deserialize(&self.0, metadata, session) } - fn coerce_args(&self, options: &Self::Options, input_dtype: &DType) -> VortexResult { - BinaryCombined::coerce_args(&self.0, options, input_dtype) - } - fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { BinaryCombined::return_dtype(&self.0, input_dtype) } diff --git a/vortex-array/src/aggregate_fn/erased.rs b/vortex-array/src/aggregate_fn/erased.rs index 0cf1f7c464d..c512ea81b7e 100644 --- a/vortex-array/src/aggregate_fn/erased.rs +++ b/vortex-array/src/aggregate_fn/erased.rs @@ -80,11 +80,6 @@ impl AggregateFnRef { self.0.can_satisfy(requested) } - /// Coerce the input type for this aggregate function. - pub fn coerce_args(&self, input_dtype: &DType) -> VortexResult { - self.0.coerce_args(input_dtype) - } - /// Compute the return [`DType`] per group given the input element type. /// /// Returns `None` if the input dtype is not supported by the aggregate function. diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/fixed_size_list.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/fixed_size_list.rs index daaf9fe4cd7..7256d593029 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/fixed_size_list.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/fixed_size_list.rs @@ -8,7 +8,7 @@ use super::filter::filter_valid_rows_if_needed; use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::FixedSizeListArray; -use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; pub(super) fn check_fixed_size_list_identical( lhs: &FixedSizeListArray, diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/list.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/list.rs index 2f4e3e52d91..89d1bb7c912 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/list.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/list.rs @@ -10,7 +10,9 @@ use crate::IntoArray; use crate::arrays::ListArray; use crate::arrays::ListViewArray; use crate::arrays::list::ListArrayExt; +use crate::arrays::list::ListArraySlotsExt; use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::list_from_list_view; pub(super) fn check_list_identical( diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/map.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/map.rs new file mode 100644 index 00000000000..f1150b68213 --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/map.rs @@ -0,0 +1,27 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use super::list::check_list_identical; +use crate::ExecutionCtx; +use crate::arrays::ListView; +use crate::arrays::MapArray; +use crate::arrays::map::MapArrayExt; +use crate::arrays::map::MapArraySlotsExt; + +pub(super) fn check_map_identical( + lhs: &MapArray, + rhs: &MapArray, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if lhs.map_dtype() != rhs.map_dtype() { + return Ok(false); + } + + check_list_identical( + &lhs.entries().as_::().into_owned(), + &rhs.entries().as_::().into_owned(), + ctx, + ) +} diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs index f030cc810e2..a8392f3fa74 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs @@ -7,6 +7,7 @@ mod extension; mod filter; mod fixed_size_list; mod list; +mod map; mod primitive; mod struct_; #[cfg(test)] @@ -27,6 +28,7 @@ use self::extension::check_extension_identical; use self::filter::shared_validity_mask; use self::fixed_size_list::check_fixed_size_list_identical; use self::list::check_list_identical; +use self::map::check_map_identical; use self::primitive::check_primitive_identical; use self::struct_::check_struct_identical; use self::varbin::check_varbinview_identical; @@ -262,6 +264,7 @@ fn check_canonical_identical( } (Canonical::Struct(lhs), Canonical::Struct(rhs)) => check_struct_identical(lhs, rhs, ctx), (Canonical::List(lhs), Canonical::List(rhs)) => check_list_identical(lhs, rhs, ctx), + (Canonical::Map(lhs), Canonical::Map(rhs)) => check_map_identical(lhs, rhs, ctx), (Canonical::FixedSizeList(lhs), Canonical::FixedSizeList(rhs)) => { check_fixed_size_list_identical(lhs, rhs, ctx) } diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/list.rs b/vortex-array/src/aggregate_fn/fns/is_constant/list.rs index c33e325d463..ce684389dfd 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/list.rs +++ b/vortex-array/src/aggregate_fn/fns/is_constant/list.rs @@ -8,6 +8,7 @@ use super::is_constant; use crate::ExecutionCtx; use crate::arrays::ListViewArray; use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; /// Check if a list view array is constant by comparing each list's elements. /// diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/map.rs b/vortex-array/src/aggregate_fn/fns/is_constant/map.rs new file mode 100644 index 00000000000..63d82b463d9 --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/is_constant/map.rs @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use super::list::check_listview_constant; +use crate::ExecutionCtx; +use crate::arrays::ListView; +use crate::arrays::MapArray; +use crate::arrays::map::MapArraySlotsExt; + +pub(super) fn check_map_constant(map: &MapArray, ctx: &mut ExecutionCtx) -> VortexResult { + check_listview_constant(&map.entries().as_::().into_owned(), ctx) +} diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs index 490ede7f640..96017c3ed92 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs @@ -6,6 +6,7 @@ mod decimal; mod extension; mod fixed_size_list; mod list; +mod map; pub mod primitive; mod struct_; mod varbin; @@ -20,6 +21,7 @@ use self::decimal::check_decimal_constant; use self::extension::check_extension_constant; use self::fixed_size_list::check_fixed_size_list_constant; use self::list::check_listview_constant; +use self::map::check_map_constant; use self::primitive::check_primitive_constant; use self::struct_::check_struct_constant; use self::varbin::check_varbinview_constant; @@ -402,8 +404,12 @@ impl AggregateFnVTable for IsConstant { Canonical::Struct(s) => check_struct_constant(s, ctx)?, Canonical::Extension(e) => check_extension_constant(e, ctx)?, Canonical::List(l) => check_listview_constant(l, ctx)?, + Canonical::Map(m) => check_map_constant(m, ctx)?, Canonical::FixedSizeList(f) => check_fixed_size_list_constant(f, ctx)?, Canonical::Null(_) => true, + Canonical::Union(_) => { + todo!("TODO(connor)[Union]: implement IsConstant for Union arrays") + } Canonical::Variant(_) => { vortex_bail!("Variant arrays don't support IsConstant") } @@ -450,14 +456,58 @@ mod tests { use crate::arrays::ListArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; + use crate::builders::MapBuilder; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::FieldNames; + use crate::dtype::MapDType; use crate::dtype::Nullability; use crate::dtype::PType; use crate::expr::stats::Stat; + use crate::scalar::Scalar; use crate::validity::Validity; + type MapEntryFixture<'a> = (i32, Option<&'a str>); + type MapRowFixture<'a> = Option>>; + + fn map_array_from_rows(rows: &[MapRowFixture<'_>]) -> VortexResult { + let map_dtype = MapDType::try_new( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + false, + )?; + let dtype = DType::Map(map_dtype.clone(), Nullability::Nullable); + let mut builder = MapBuilder::::with_capacity_in( + map_dtype, + Nullability::Nullable, + rows.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + + for row in rows { + let scalar = match row { + Some(entries) => { + let entries = entries + .iter() + .map(|(key, value)| { + let key = Scalar::primitive(*key, Nullability::NonNullable); + let value = value.map_or_else( + || Scalar::null(DType::Utf8(Nullability::Nullable)), + |value| Scalar::utf8(value, Nullability::Nullable), + ); + (key, value) + }) + .collect::>(); + Scalar::try_map(dtype.clone(), entries)? + } + None => Scalar::null(dtype.clone()), + }; + builder.append_value(scalar.as_map())?; + } + + Ok(builder.finish_into_map().into_array()) + } + // Tests migrated from compute/is_constant.rs #[test] fn is_constant_min_max_no_nan() -> VortexResult<()> { @@ -681,4 +731,26 @@ mod tests { assert_eq!(is_constant(&list_array.into_array(), &mut ctx)?, expected); Ok(()) } + + #[test] + fn test_map_is_constant() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let identical = map_array_from_rows(&[ + Some(vec![(1, Some("one")), (2, None)]), + Some(vec![(1, Some("one")), (2, None)]), + ])?; + assert!(is_constant(&identical, &mut ctx)?); + + let different = map_array_from_rows(&[ + Some(vec![(1, Some("one")), (2, None)]), + Some(vec![(1, Some("one")), (3, None)]), + ])?; + assert!(!is_constant(&different, &mut ctx)?); + + let all_null = map_array_from_rows(&[None, None])?; + assert!(is_constant(&all_null, &mut ctx)?); + + Ok(()) + } } diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/primitive.rs b/vortex-array/src/aggregate_fn/fns/is_constant/primitive.rs index 510306237ad..48c7423c0a7 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/is_constant/primitive.rs @@ -20,15 +20,14 @@ pub fn compute_is_constant(values: &[T]) -> let first_value = values[0]; let first_vec = &[first_value; WIDTH]; - let mut chunks = values[1..].chunks_exact(WIDTH); - for chunk in &mut chunks { - assert_eq!(chunk.len(), WIDTH); // let the compiler know each chunk is WIDTH. + let (chunks, remainder) = values[1..].as_chunks::(); + for chunk in chunks { if first_vec != chunk { return false; } } - for value in chunks.remainder() { + for value in remainder { if !value.is_eq(first_value) { return false; } diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/struct_.rs b/vortex-array/src/aggregate_fn/fns/is_constant/struct_.rs index 9376115521f..84c829dfaca 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/struct_.rs +++ b/vortex-array/src/aggregate_fn/fns/is_constant/struct_.rs @@ -10,7 +10,7 @@ use crate::arrays::struct_::StructArrayExt; /// Check if a struct array is constant by checking each field independently. pub(super) fn check_struct_constant(s: &StructArray, ctx: &mut ExecutionCtx) -> VortexResult { - for field in s.unmasked_fields().iter() { + for field in s.iter_unmasked_fields() { if !is_constant(field, ctx)? { return Ok(false); } diff --git a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs index ada9e6eb2e5..99a7b99128a 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs @@ -246,6 +246,7 @@ impl AggregateFnVTable for IsSorted { DType::Null | DType::List(..) | DType::FixedSizeList(..) + | DType::Map(..) | DType::Struct(..) | DType::Union(..) | DType::Variant(..) @@ -263,6 +264,7 @@ impl AggregateFnVTable for IsSorted { DType::Null | DType::List(..) | DType::FixedSizeList(..) + | DType::Map(..) | DType::Struct(..) | DType::Union(..) | DType::Variant(..) diff --git a/vortex-array/src/aggregate_fn/fns/mean/mod.rs b/vortex-array/src/aggregate_fn/fns/mean/mod.rs index 20b7f15834c..a3a0f77cb45 100644 --- a/vortex-array/src/aggregate_fn/fns/mean/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mean/mod.rs @@ -7,9 +7,9 @@ use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::ExecutionCtx; +use crate::IntoArray; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::combined::BinaryCombined; @@ -19,6 +19,7 @@ use crate::aggregate_fn::combined::PairOptions; use crate::aggregate_fn::fns::count::Count; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::sum_decimal_dtype; +use crate::arrays::ConstantArray; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::DecimalDType; @@ -51,7 +52,8 @@ pub fn mean(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { /// /// Implemented as `Sum / Count` via [`BinaryCombined`]. /// -/// Booleans and primitive numeric types are cast to f64. Decimals stay decimals. +/// Booleans and primitive numeric types produce nullable `f64` results. +/// Decimals produce a nullable decimal result. #[derive(Clone, Debug)] pub struct Mean; @@ -95,9 +97,21 @@ impl BinaryCombined for Mean { vortex_bail!("grouped mean over decimals is not yet supported"); } let target = DType::Primitive(PType::F64, Nullability::Nullable); - let sum_cast = sum.cast(target.clone())?; - let count_cast = count.cast(target)?; - sum_cast.binary(count_cast, Operator::Div) + let sum = sum.cast(target.clone())?; + let count = count.cast(target.clone())?; + + let non_zero = count + .binary( + ConstantArray::new(Scalar::zero_value(&target), count.len()).into_array(), + Operator::NotEq, + )? + .fill_null(false)?; + // if count is 0, dividing by 0 below produces NaN, and we need Null. + // mask values to skip 0 so on 0 count turns into Null, dividing by + // Null is always Null + let count = count.mask(non_zero)?; + + sum.binary(count, Operator::Div) } fn finalize_scalar(&self, left_scalar: Scalar, right_scalar: Scalar) -> VortexResult { @@ -112,9 +126,8 @@ impl BinaryCombined for Mean { let sum = sum_cast.as_primitive().typed_value::(); let count = count_cast.as_primitive().typed_value::(); let value = match (sum, count) { - (None, _) | (_, None) => return Ok(Scalar::null(target)), // Sum overflowed - // A count of zero yields 0/0 = NaN, matching the array `finalize` path: nulls are - // skipped during accumulation, so an all-null input is an empty mean, not null. + // None sum means sum overflowed, 0 count means empty input + (None, _) | (_, None) | (_, Some(0.0)) => return Ok(Scalar::null(target)), (Some(s), Some(c)) => s / c, }; Ok(Scalar::primitive(value, Nullability::Nullable)) @@ -123,34 +136,6 @@ impl BinaryCombined for Mean { fn serialize(&self, _options: &CombinedOptions) -> VortexResult>> { unimplemented!("mean is not yet serializable"); } - - fn coerce_args( - &self, - _options: &PairOptions< - ::Options, - ::Options, - >, - input_dtype: &DType, - ) -> VortexResult { - // Advisory hint for query planners: where possible, cast input to the - // type we're going to compute the mean in. - Ok(coerced_input_dtype(input_dtype).unwrap_or_else(|| input_dtype.clone())) - } -} - -/// Hint for callers: what to cast the input to before accumulation. -/// -/// - Bool stays as bool — `Sum` has a native bool path and bool → f64 isn't -/// currently a direct cast in vortex. -/// - Primitive numerics → `f64` so the sum and finalize work without overflow. -/// - Decimals stay as decimals -fn coerced_input_dtype(input_dtype: &DType) -> Option { - match input_dtype { - DType::Bool(_) => Some(input_dtype.clone()), - DType::Primitive(_, n) => Some(DType::Primitive(PType::F64, *n)), - DType::Decimal(..) => Some(input_dtype.clone()), - _ => None, - } } fn mean_output_dtype(input_dtype: &DType) -> Option { @@ -217,13 +202,14 @@ mod tests { use vortex_error::VortexResult; use super::*; - use crate::IntoArray; use crate::VortexSessionExecute; + use crate::aggregate_fn::DynGroupedAccumulator; + use crate::aggregate_fn::GroupedAccumulator; use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; - use crate::arrays::ConstantArray; use crate::arrays::DecimalArray; + use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::dtype::DecimalDType; use crate::validity::Validity; @@ -315,11 +301,11 @@ mod tests { } #[test] - fn mean_all_null_returns_nan() -> VortexResult<()> { + fn mean_all_null_returns_null() -> VortexResult<()> { let array = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); let mut ctx = array_session().create_execution_ctx(); let result = mean(&array, &mut ctx)?; - assert!(result.as_primitive().as_::().is_some_and(f64::is_nan)); + assert_eq!(result.as_primitive().as_::(), None); Ok(()) } @@ -414,4 +400,65 @@ mod tests { assert_eq!(result.as_primitive().as_::(), Some(3.0)); Ok(()) } + + fn mean_nan_null() -> Vec<(Vec>, Option)> { + vec![ + (vec![Some(f64::NAN), Some(1.0), None], Some(1.0)), + (vec![Some(f64::NAN), Some(1.0), Some(3.0)], Some(2.0)), + (vec![None, None, Some(f64::NAN)], None), + (vec![None, None, None], None), + (vec![Some(1.0), Some(2.0), Some(3.0)], Some(2.0)), + ] + } + + #[test] + fn mean_combined_partials() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + for (case, (group, expected)) in mean_nan_null().into_iter().enumerate() { + let mut acc = Accumulator::try_new( + Mean::combined(), + PairOptions( + NumericalAggregateOpts::default(), + NumericalAggregateOpts::default(), + ), + DType::Primitive(PType::F64, Nullability::Nullable), + )?; + let (head, tail) = group.split_at(2); + let head = PrimitiveArray::from_option_iter(head.iter().copied()).into_array(); + let tail = PrimitiveArray::from_option_iter(tail.iter().copied()).into_array(); + acc.accumulate(&head, &mut ctx)?; + acc.accumulate(&tail, &mut ctx)?; + let result = acc.finish()?; + assert_eq!(result.as_primitive().as_::(), expected, "case {case}"); + } + Ok(()) + } + + #[test] + fn mean_grouped_finalize() -> VortexResult<()> { + let cases = mean_nan_null(); + let elements = PrimitiveArray::from_option_iter( + cases.iter().flat_map(|(group, _)| group.iter().copied()), + ) + .into_array(); + let groups = FixedSizeListArray::try_new(elements, 3, Validity::NonNullable, cases.len())?; + + let mut acc = GroupedAccumulator::try_new( + Mean::combined(), + PairOptions( + NumericalAggregateOpts::default(), + NumericalAggregateOpts::default(), + ), + DType::Primitive(PType::F64, Nullability::Nullable), + )?; + let mut ctx = array_session().create_execution_ctx(); + acc.accumulate_list(&groups.into_array(), &mut ctx)?; + let result = acc.finish()?; + + for (case, (_, expected)) in cases.into_iter().enumerate() { + let actual = result.execute_scalar(case, &mut ctx)?; + assert_eq!(actual.as_primitive().as_::(), expected, "case {case}"); + } + Ok(()) + } } diff --git a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs index 957582f9100..1c23e89221d 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs @@ -417,8 +417,10 @@ impl AggregateFnVTable for MinMax { Canonical::Decimal(d) => accumulate_decimal(partial, d, ctx), Canonical::Extension(e) => accumulate_extension(partial, e, ctx), Canonical::Null(_) => Ok(()), - Canonical::Struct(_) + Canonical::Union(_) + | Canonical::Struct(_) | Canonical::List(_) + | Canonical::Map(_) | Canonical::FixedSizeList(_) | Canonical::Variant(_) => { vortex_bail!("Unsupported canonical type for min_max: {}", batch.dtype()) diff --git a/vortex-array/src/aggregate_fn/fns/mod.rs b/vortex-array/src/aggregate_fn/fns/mod.rs index da81e67ecb9..c4e8d1b72ca 100644 --- a/vortex-array/src/aggregate_fn/fns/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mod.rs @@ -20,4 +20,5 @@ pub mod min_max; pub mod nan_count; pub mod null_count; pub mod sum; +pub mod sum_v2; pub mod uncompressed_size_in_bytes; diff --git a/vortex-array/src/aggregate_fn/fns/sum/bool.rs b/vortex-array/src/aggregate_fn/fns/sum/bool.rs index b3993840586..f15356caaa1 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/bool.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/bool.rs @@ -13,7 +13,7 @@ use crate::ExecutionCtx; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; -pub(super) fn accumulate_bool( +pub(crate) fn accumulate_bool( inner: &mut SumState, b: &BoolArray, ctx: &mut ExecutionCtx, diff --git a/vortex-array/src/aggregate_fn/fns/sum/constant.rs b/vortex-array/src/aggregate_fn/fns/sum/constant.rs index 0f366620e5c..ced92af1d68 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/constant.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/constant.rs @@ -15,7 +15,7 @@ use crate::scalar::Scalar; /// /// Returns `Ok(None)` if the scalar is null (no contribution to the sum). /// Returns a null scalar on overflow (saturation). -pub(super) fn multiply_constant( +pub(crate) fn multiply_constant( scalar: &Scalar, len: usize, return_dtype: &DType, diff --git a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs index 872e5769a01..197b8f10b04 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs @@ -23,7 +23,7 @@ use crate::scalar::DecimalValue; /// Accumulate a decimal array into the sum state. /// Returns Ok(true) if saturated (overflow), Ok(false) if not. -pub(super) fn accumulate_decimal( +pub(crate) fn accumulate_decimal( inner: &mut SumState, d: &DecimalArray, ctx: &mut ExecutionCtx, diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index efe0825d4d6..c44feb44823 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -170,7 +170,7 @@ mod tests { use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; - use crate::builders::builder_with_capacity; + use crate::builders::builder_with_capacity_in; use crate::dtype::DType; use crate::dtype::Nullability::NonNullable; use crate::dtype::Nullability::Nullable; @@ -202,7 +202,11 @@ mod tests { let sum_dtype = Sum .partial_dtype(&NumericalAggregateOpts::default(), elem_dtype) .expect("sum partial dtype"); - let mut builder = builder_with_capacity(&sum_dtype, ranges.len()); + let mut builder = builder_with_capacity_in( + &sum_dtype, + ranges.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); for (i, &(offset, size)) in ranges.iter().enumerate() { if group_valid[i] { let slice = elements.slice(offset..offset + size)?; diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 010371fee4f..76eaba22c15 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -15,10 +15,13 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use self::bool::accumulate_bool; -use self::constant::multiply_constant; -use self::decimal::accumulate_decimal; -use self::primitive::accumulate_primitive; +pub(crate) use self::bool::accumulate_bool; +pub(crate) use self::constant::multiply_constant; +pub(crate) use self::decimal::accumulate_decimal; +pub(crate) use self::primitive::accumulate_primitive; +pub(crate) use self::primitive::sum_float_all; +pub(crate) use self::primitive::sum_signed_all; +pub(crate) use self::primitive::sum_unsigned_all; use crate::ArrayRef; use crate::Canonical; use crate::Columnar; @@ -351,7 +354,7 @@ pub enum SumState { }, } -fn make_zero_state(return_dtype: &DType) -> SumState { +pub(crate) fn make_zero_state(return_dtype: &DType) -> SumState { match return_dtype { DType::Primitive(ptype, _) => match ptype { PType::U8 | PType::U16 | PType::U32 | PType::U64 => SumState::Unsigned(0), @@ -367,6 +370,7 @@ fn make_zero_state(return_dtype: &DType) -> SumState { } /// Checked add for u64, returning true if overflow occurred. +#[allow(clippy::inline_always)] #[inline(always)] fn checked_add_u64(acc: &mut u64, val: u64) -> bool { match acc.checked_add(val) { @@ -379,6 +383,7 @@ fn checked_add_u64(acc: &mut u64, val: u64) -> bool { } /// Checked add for i64, returning true if overflow occurred. +#[allow(clippy::inline_always)] #[inline(always)] fn checked_add_i64(acc: &mut i64, val: i64) -> bool { match acc.checked_add(val) { @@ -476,10 +481,14 @@ mod tests { .checked_add(&rhs.as_primitive()) .map(Scalar::from) .unwrap_or_else(|| Scalar::null(sum_dtype.as_nullable())), - DType::Decimal(..) => lhs + // Add widens the result precision, so restate the sum in the accumulator's own + // decimal type, treating a value that no longer fits as an overflow. + DType::Decimal(decimal_dtype, _) => lhs .as_decimal() - .checked_binary_numeric(&rhs.as_decimal(), NumericOperator::Add) - .map(Scalar::from) + .checked_binary_numeric(&rhs.as_decimal(), NumericOperator::Add)? + .and_then(|scalar| scalar.as_decimal().decimal_value()) + .filter(|value| value.fits_in_precision(*decimal_dtype)) + .map(|value| Scalar::decimal(value, *decimal_dtype, Nullable)) .unwrap_or_else(|| Scalar::null(sum_dtype.as_nullable())), _ => unreachable!("Sum will always be a decimal or a primitive dtype"), }) diff --git a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs index 87d8da4b143..930a353f6e1 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs @@ -21,7 +21,7 @@ use crate::match_each_native_ptype; /// than 64 bits cannot overflow the 64-bit accumulator: `2^16 * (2^32 - 1) < 2^64`. const SUM_CHUNK: usize = 1 << 16; -pub(super) fn accumulate_primitive( +pub(crate) fn accumulate_primitive( inner: &mut SumState, p: &PrimitiveArray, ctx: &mut ExecutionCtx, @@ -66,7 +66,7 @@ fn accumulate_primitive_all( /// Sum the values of a float slice into an `f64` accumulator. When `skip_nans` is set, NaN values /// are skipped to match the scalar `sum` semantics; otherwise any NaN poisons the accumulator to /// NaN. Floats cannot overflow the accumulator, so this never reports saturation. -pub(super) fn sum_float_all(acc: &mut f64, slice: &[T], skip_nans: bool) { +pub(crate) fn sum_float_all(acc: &mut f64, slice: &[T], skip_nans: bool) { if skip_nans { for &v in slice { if !v.is_nan() { @@ -84,7 +84,7 @@ pub(super) fn sum_float_all(acc: &mut f64, slice: &[T], skip_nan /// chunks of [`SUM_CHUNK`] with a single checked add per chunk, which lets the inner loop vectorize /// to packed widening adds. `u64` input keeps a per-element checked add since a chunk of `u64`s /// could itself overflow. Returns `true` on overflow. -pub(super) fn sum_unsigned_all(acc: &mut u64, slice: &[T]) -> bool +pub(crate) fn sum_unsigned_all(acc: &mut u64, slice: &[T]) -> bool where T: NativePType + AsPrimitive, { @@ -106,7 +106,7 @@ where } /// Signed counterpart of [`sum_unsigned_all`]. -pub(super) fn sum_signed_all(acc: &mut i64, slice: &[T]) -> bool +pub(crate) fn sum_signed_all(acc: &mut i64, slice: &[T]) -> bool where T: NativePType + AsPrimitive, { diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/grouped.rs new file mode 100644 index 00000000000..5c5e1d7fa6a --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/grouped.rs @@ -0,0 +1,192 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::BitBuffer; +use vortex_buffer::BitBufferMut; +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use super::SumV2; +use super::sum_v2_partial_fields; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::aggregate_fn::AggregateFnRef; +use crate::aggregate_fn::GroupRanges; +use crate::aggregate_fn::GroupedArray; +use crate::aggregate_fn::fns::sum::sum_float_all; +use crate::aggregate_fn::fns::sum::sum_signed_all; +use crate::aggregate_fn::fns::sum::sum_unsigned_all; +use crate::aggregate_fn::kernels::DynGroupedAggregateKernel; +use crate::arrays::BoolArray; +use crate::arrays::Primitive; +use crate::arrays::PrimitiveArray; +use crate::arrays::StructArray; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::match_each_native_ptype; +use crate::validity::Validity; + +/// Encoding-specific grouped [`SumV2`] kernel for primitive element arrays. +#[derive(Debug)] +pub(crate) struct PrimitiveGroupedSumV2EncodingKernel; + +impl DynGroupedAggregateKernel for PrimitiveGroupedSumV2EncodingKernel { + fn grouped_aggregate( + &self, + aggregate_fn: &AggregateFnRef, + groups: &GroupedArray, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some(options) = aggregate_fn.as_opt::() else { + return Ok(None); + }; + try_grouped_sum(groups, ctx, options.skip_nans) + } +} + +/// Grouped [`SumV2`] implementation for canonical primitive elements. +/// +/// Reuses the scalar primitive-sum reductions ([`sum_unsigned_all`]/[`sum_signed_all`]/ +/// [`sum_float_all`]) so the per-group semantics match scalar `sum_v2` exactly (overflow saturates +/// to a null sum, NaNs are skipped). The element validity mask is materialized once and sliced +/// per group, rather than the per-group accumulator setup of the generic fallback path. +pub(super) fn try_grouped_sum( + groups: &GroupedArray, + ctx: &mut ExecutionCtx, + skip_nans: bool, +) -> VortexResult> { + if !groups.elements().is::() { + return Ok(None); + } + let elements = groups.elements().clone().downcast::(); + let group_ranges = groups.group_ranges(ctx)?; + let group_validity = groups.group_validity(ctx)?; + + Ok(Some(grouped_sum( + &elements, + &group_ranges, + &group_validity, + ctx, + skip_nans, + )?)) +} + +/// Sum each group described by `group_ranges` (element `(offset, size)` pairs), one sum per group. +fn grouped_sum( + elements: &PrimitiveArray, + group_ranges: &GroupRanges, + group_validity: &Mask, + ctx: &mut ExecutionCtx, + skip_nans: bool, +) -> VortexResult { + let elem_mask = elements + .as_ref() + .validity()? + .execute_mask(elements.as_ref().len(), ctx)?; + let all_valid = elem_mask.all_true(); + + let (sums, is_overflow, is_empty) = match_each_native_ptype!(elements.ptype(), + unsigned: |T| { + let values = elements.as_slice::(); + collect_sums::( + values, group_ranges, group_validity, &elem_mask, all_valid, sum_unsigned_all) + }, + signed: |T| { + let values = elements.as_slice::(); + collect_sums::( + values, group_ranges, group_validity, &elem_mask, all_valid, sum_signed_all) + }, + floating: |T| { + let values = elements.as_slice::(); + collect_sums::( + values, group_ranges, group_validity, &elem_mask, all_valid, + |acc, slice| { sum_float_all(acc, slice, skip_nans); false }) + } + ); + + let partial_fields = sum_v2_partial_fields(sums.dtype().clone()); + + // SAFETY: all three children have one value per group and match `partial_fields`; the struct + // validity is derived from the same group count. + Ok(unsafe { + StructArray::new_unchecked( + vec![ + sums.into_array(), + BoolArray::new(is_overflow, Validity::NonNullable).into_array(), + BoolArray::new(is_empty, Validity::NonNullable).into_array(), + ], + partial_fields, + group_validity.len(), + Validity::from_mask(group_validity.clone(), Nullability::Nullable), + ) + } + .into_array()) +} + +/// Reduce each group's element slice into a non-null sum, overflow bitmap, and empty bitmap. +fn collect_sums( + values: &[T], + group_ranges: &GroupRanges, + group_validity: &Mask, + elem_mask: &Mask, + all_valid: bool, + sum_run: impl Fn(&mut A, &[T]) -> bool, +) -> (PrimitiveArray, BitBuffer, BitBuffer) { + let group_count = group_ranges.len(); + let mut is_overflow = BitBufferMut::new_unset(group_count); + let mut is_empty = BitBufferMut::new_unset(group_count); + let sums = group_ranges.iter().enumerate().map(|(i, (offset, size))| { + if !group_validity.value(i) { + return A::default(); + } + let mut acc = A::default(); + let (overflow, any_valid) = if all_valid { + (sum_run(&mut acc, &values[offset..offset + size]), size > 0) + } else { + sum_masked_group(&mut acc, values, offset, size, elem_mask, &sum_run) + }; + if overflow { + // SAFETY: `i` comes from enumerating `group_ranges`, and the bitmap has one bit per + // group. + unsafe { is_overflow.set_unchecked(i) }; + } + if !any_valid { + // SAFETY: `i` comes from enumerating `group_ranges`, and the bitmap has one bit per + // group. + unsafe { is_empty.set_unchecked(i) }; + } + acc + }); + let sums = PrimitiveArray::from_iter(sums); + (sums, is_overflow.freeze(), is_empty.freeze()) +} + +/// Sum valid runs in one group, returning `(overflow, any_valid)`. +fn sum_masked_group( + acc: &mut A, + values: &[T], + offset: usize, + size: usize, + elem_mask: &Mask, + sum_run: &impl Fn(&mut A, &[T]) -> bool, +) -> (bool, bool) { + match elem_mask { + Mask::AllTrue(_) => (sum_run(acc, &values[offset..offset + size]), size > 0), + Mask::AllFalse(_) => (false, false), + Mask::Values(mask_values) => { + let validity = mask_values + .bit_buffer() + .as_view() + .slice(offset..offset + size); + let mut any_valid = false; + for (start, end) in validity.set_slices() { + any_valid = true; + if sum_run(acc, &values[offset + start..offset + end]) { + return (true, true); + } + } + (false, any_valid) + } + } +} diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs new file mode 100644 index 00000000000..e1287448a8f --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs @@ -0,0 +1,422 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +mod grouped; + +pub(crate) use grouped::PrimitiveGroupedSumV2EncodingKernel; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ArrayRef; +use crate::ArrayView; +use crate::Canonical; +use crate::Columnar; +use crate::ExecutionCtx; +use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateFnId; +use crate::aggregate_fn::AggregateFnVTable; +use crate::aggregate_fn::DynAccumulator; +use crate::aggregate_fn::NumericalAggregateOpts; +use crate::aggregate_fn::fns::sum::Sum; +use crate::aggregate_fn::fns::sum::SumState; +use crate::aggregate_fn::fns::sum::accumulate_bool; +use crate::aggregate_fn::fns::sum::accumulate_decimal; +use crate::aggregate_fn::fns::sum::accumulate_primitive; +use crate::aggregate_fn::fns::sum::make_zero_state; +use crate::aggregate_fn::fns::sum::multiply_constant; +use crate::arrays::Struct; +use crate::arrays::struct_::StructArrayExt; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::FieldName; +use crate::dtype::FieldNames; +use crate::dtype::Nullability; +use crate::dtype::StructFields; +use crate::expr::stats::Precision; +use crate::expr::stats::Stat; +use crate::expr::stats::StatsProviderExt; +use crate::scalar::DecimalValue; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::operators::Operator; +use crate::validity::Validity; + +const SUM_FIELD: &str = "sum"; +const IS_OVERFLOW_FIELD: &str = "is_overflow"; +const IS_EMPTY_FIELD: &str = "is_empty"; + +/// Return the SQL-style sum of an array. +/// +/// Unlike [`sum`](crate::aggregate_fn::fns::sum::sum), this returns null when the array has no +/// valid values. +pub fn sum_v2(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let mut acc = Accumulator::try_new( + SumV2, + NumericalAggregateOpts::default(), + array.dtype().clone(), + )?; + acc.accumulate(array, ctx)?; + acc.finish() +} + +/// Sum an array, returning null when it has no valid values or if the sum overflows. +/// +/// This aggregate intentionally has a distinct ID and partial representation from the legacy +/// [`Sum`]. Keeping `vortex.sum` unchanged preserves the scalar partials stored by older Vortex +/// files, while `SumV2` can use an explicit `{ sum, is_overflow, is_empty }` state. +/// +/// NaN handling for float inputs is controlled by [`NumericalAggregateOpts`]. With `skip_nans` +/// (the default), NaN values contribute nothing but still make the input non-empty. Otherwise, +/// any NaN value poisons the sum to NaN. +#[derive(Clone, Copy, Debug)] +pub struct SumV2; + +impl AggregateFnVTable for SumV2 { + type Options = NumericalAggregateOpts; + type Partial = SumV2Partial; + + fn id(&self) -> AggregateFnId { + static ID: CachedId = CachedId::new("vortex.sum_v2"); + *ID + } + + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + Ok(Some(options.serialize())) + } + + fn deserialize( + &self, + metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + NumericalAggregateOpts::deserialize(metadata) + } + + fn return_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { + Sum.return_dtype(options, input_dtype) + } + + fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { + self.return_dtype(options, input_dtype) + .map(sum_v2_partial_dtype) + } + + fn empty_partial( + &self, + options: &Self::Options, + input_dtype: &DType, + ) -> VortexResult { + let return_dtype = self + .return_dtype(options, input_dtype) + .ok_or_else(|| vortex_err!("Unsupported sum_v2 dtype: {}", input_dtype))?; + let sum = make_zero_state(&return_dtype); + Ok(SumV2Partial { + return_dtype, + sum, + is_overflow: false, + is_empty: true, + skip_nans: options.skip_nans, + }) + } + + fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { + let (other_sum, other_is_overflow, other_is_empty) = decode_partial_scalar(other)?; + validate_sum_field_dtype(&other_sum, &partial.return_dtype)?; + + if partial.is_overflow { + return Ok(()); + } + if other_is_overflow { + partial.is_overflow = true; + partial.is_empty = false; + return Ok(()); + } + if other_is_empty { + return Ok(()); + } + + partial.is_overflow = checked_add_sum_state(&mut partial.sum, &other_sum)?; + partial.is_empty = false; + Ok(()) + } + + fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + Ok(Scalar::struct_( + sum_v2_partial_dtype(partial.return_dtype.clone()), + vec![ + sum_state_scalar(partial, Nullability::NonNullable), + Scalar::bool(partial.is_overflow, Nullability::NonNullable), + Scalar::bool(partial.is_empty, Nullability::NonNullable), + ], + )) + } + + fn reset(&self, partial: &mut Self::Partial) { + partial.sum = make_zero_state(&partial.return_dtype); + partial.is_overflow = false; + partial.is_empty = true; + } + + fn is_saturated(&self, partial: &Self::Partial) -> bool { + partial.is_overflow || matches!(&partial.sum, SumState::Float(value) if value.is_nan()) + } + + fn try_accumulate( + &self, + partial: &mut Self::Partial, + batch: &ArrayRef, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + if partial.skip_nans || !matches!(&partial.sum, SumState::Float(_)) { + return Ok(false); + } + + match batch.statistics().get_as::(Stat::NaNCount) { + Precision::Exact(0) => Ok(false), + Precision::Exact(_) => { + let SumState::Float(sum) = &mut partial.sum else { + unreachable!("checked float sum state") + }; + *sum = f64::NAN; + partial.is_empty = false; + Ok(true) + } + _ => Ok(false), + } + } + + fn accumulate( + &self, + partial: &mut Self::Partial, + batch: &Columnar, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if partial.is_overflow { + return Ok(()); + } + + if let Columnar::Constant(constant) = batch { + if !constant.scalar().is_null() && !constant.is_empty() { + partial.is_empty = false; + } + if partial.skip_nans + && constant + .scalar() + .as_primitive_opt() + .is_some_and(|primitive| primitive.is_nan()) + { + return Ok(()); + } + if let Some(product) = + multiply_constant(constant.scalar(), constant.len(), &partial.return_dtype)? + { + if product.is_null() { + partial.is_overflow = true; + partial.is_empty = false; + } else { + partial.is_overflow = checked_add_sum_state(&mut partial.sum, &product)?; + partial.is_empty = false; + } + } + return Ok(()); + } + + let any_valid = partial.is_empty && has_valid_value(batch, ctx)?; + let result = match batch { + Columnar::Canonical(canonical) => match canonical { + Canonical::Primitive(array) => { + accumulate_primitive(&mut partial.sum, array, ctx, partial.skip_nans) + } + Canonical::Bool(array) => accumulate_bool(&mut partial.sum, array, ctx), + Canonical::Decimal(array) => accumulate_decimal(&mut partial.sum, array, ctx), + _ => vortex_bail!("Unsupported canonical type for sum_v2: {}", batch.dtype()), + }, + Columnar::Constant(_) => unreachable!(), + }; + + if any_valid { + partial.is_empty = false; + } + if result? { + partial.is_overflow = true; + partial.is_empty = false; + } + Ok(()) + } + + fn finalize(&self, partials: ArrayRef) -> VortexResult { + if let Some(partials) = partials.as_opt::() { + return finalize_struct(partials); + } + + let sum = partials.get_item(SUM_FIELD)?; + let is_invalid = partials + .get_item(IS_OVERFLOW_FIELD)? + .binary(partials.get_item(IS_EMPTY_FIELD)?, Operator::Or)? + .fill_null(true)?; + sum.mask(is_invalid.not()?) + } + + fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { + if partial.is_overflow || partial.is_empty { + return Ok(Scalar::null(partial.return_dtype.as_nullable())); + } + Ok(sum_state_scalar(partial, Nullability::Nullable)) + } +} + +fn finalize_struct(partials: ArrayView<'_, Struct>) -> VortexResult { + let sum = partials.unmasked_field_by_name(SUM_FIELD)?.clone(); + let is_overflow = partials.unmasked_field_by_name(IS_OVERFLOW_FIELD)?.clone(); + let is_empty = partials.unmasked_field_by_name(IS_EMPTY_FIELD)?.clone(); + + let is_invalid = is_overflow + .binary(is_empty, Operator::Or)? + .fill_null(true)?; + let mut is_valid = is_invalid.not()?; + match partials.struct_validity() { + Validity::NonNullable | Validity::AllValid => {} + validity => { + is_valid = is_valid.binary(validity.to_array(partials.len()), Operator::And)?; + } + } + + sum.mask(is_valid) +} + +/// In-memory state for SumV2 accumulation. +pub struct SumV2Partial { + return_dtype: DType, + sum: SumState, + is_overflow: bool, + is_empty: bool, + skip_nans: bool, +} + +fn has_valid_value(batch: &Columnar, ctx: &mut ExecutionCtx) -> VortexResult { + let (validity, len) = match batch { + Columnar::Canonical(Canonical::Primitive(array)) => { + (array.as_ref().validity()?, array.as_ref().len()) + } + Columnar::Canonical(Canonical::Bool(array)) => { + (array.as_ref().validity()?, array.as_ref().len()) + } + Columnar::Canonical(Canonical::Decimal(array)) => { + (array.as_ref().validity()?, array.as_ref().len()) + } + Columnar::Canonical(_) => return Ok(false), + Columnar::Constant(constant) => { + return Ok(!constant.is_empty() && !constant.scalar().is_null()); + } + }; + Ok(validity.execute_mask(len, ctx)?.true_count() > 0) +} + +fn decode_partial_scalar(scalar: Scalar) -> VortexResult<(Scalar, bool, bool)> { + vortex_ensure!(!scalar.is_null(), "SumV2 partial must not be null"); + + let Some(fields) = scalar.as_struct_opt() else { + vortex_bail!("SumV2 partial must be a struct, got {}", scalar.dtype()); + }; + let sum = fields + .field(SUM_FIELD) + .ok_or_else(|| vortex_err!("SumV2 partial is missing the sum field"))?; + let is_overflow = bool::try_from( + &fields + .field(IS_OVERFLOW_FIELD) + .ok_or_else(|| vortex_err!("SumV2 partial is missing the is_overflow field"))?, + )?; + let is_empty = bool::try_from( + &fields + .field(IS_EMPTY_FIELD) + .ok_or_else(|| vortex_err!("SumV2 partial is missing the is_empty field"))?, + )?; + + Ok((sum, is_overflow, is_empty)) +} + +fn validate_sum_field_dtype(sum: &Scalar, return_dtype: &DType) -> VortexResult<()> { + vortex_ensure!( + sum.dtype().nullability() == Nullability::NonNullable + && sum.dtype().eq_ignore_nullability(return_dtype), + "SumV2 partial value has dtype {}, expected {}", + sum.dtype(), + return_dtype.as_nonnullable(), + ); + Ok(()) +} + +fn checked_add_sum_state(state: &mut SumState, other: &Scalar) -> VortexResult { + Ok(match state { + SumState::Unsigned(sum) => checked_add_u64(sum, u64::try_from(other)?), + SumState::Signed(sum) => checked_add_i64(sum, i64::try_from(other)?), + SumState::Float(sum) => { + *sum += f64::try_from(other)?; + false + } + SumState::Decimal { value, dtype } => { + let other = DecimalValue::try_from(other)?; + match value.checked_add(&other) { + Some(result) if result.fits_in_precision(*dtype) => { + *value = result; + false + } + Some(_) | None => true, + } + } + }) +} + +fn sum_v2_partial_dtype(sum_dtype: DType) -> DType { + DType::Struct(sum_v2_partial_fields(sum_dtype), Nullability::Nullable) +} + +fn sum_v2_partial_fields(sum_dtype: DType) -> StructFields { + StructFields::new( + FieldNames::from_iter([ + FieldName::from(SUM_FIELD), + FieldName::from(IS_OVERFLOW_FIELD), + FieldName::from(IS_EMPTY_FIELD), + ]), + vec![ + sum_dtype.as_nonnullable(), + DType::Bool(Nullability::NonNullable), + DType::Bool(Nullability::NonNullable), + ], + ) +} + +fn sum_state_scalar(partial: &SumV2Partial, nullability: Nullability) -> Scalar { + match &partial.sum { + SumState::Unsigned(value) => Scalar::primitive(*value, nullability), + SumState::Signed(value) => Scalar::primitive(*value, nullability), + SumState::Float(value) => Scalar::primitive(*value, nullability), + SumState::Decimal { value, dtype } => Scalar::decimal(*value, *dtype, nullability), + } +} + +fn checked_add_u64(sum: &mut u64, value: u64) -> bool { + match sum.checked_add(value) { + Some(result) => { + *sum = result; + false + } + None => true, + } +} + +fn checked_add_i64(sum: &mut i64, value: i64) -> bool { + match sum.checked_add(value) { + Some(result) => { + *sum = result; + false + } + None => true, + } +} + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs new file mode 100644 index 00000000000..65e54a5bb6a --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs @@ -0,0 +1,384 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use prost::Message; +use rstest::rstest; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_proto::expr as pb; + +use super::SumV2; +use super::sum_v2; +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateFnRef; +use crate::aggregate_fn::AggregateFnVTable; +use crate::aggregate_fn::AggregateFnVTableExt; +use crate::aggregate_fn::DynAccumulator; +use crate::aggregate_fn::DynGroupedAccumulator; +use crate::aggregate_fn::GroupedAccumulator; +use crate::aggregate_fn::NumericalAggregateOpts; +use crate::aggregate_fn::fns::sum::Sum; +use crate::array_session; +use crate::arrays::BoolArray; +use crate::arrays::ChunkedArray; +use crate::arrays::ConstantArray; +use crate::arrays::DecimalArray; +use crate::arrays::ListViewArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::StructArray; +use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::FieldNames; +use crate::dtype::Nullability; +use crate::dtype::Nullability::Nullable; +use crate::dtype::PType; +use crate::expr::stats::Precision; +use crate::expr::stats::Stat; +use crate::scalar::Scalar; +use crate::scalar::ScalarValue; +use crate::validity::Validity; + +fn sum_with_options(array: &ArrayRef, options: NumericalAggregateOpts) -> VortexResult { + let mut accumulator = Accumulator::try_new(SumV2, options, array.dtype().clone())?; + accumulator.accumulate(array, &mut array_session().create_execution_ctx())?; + accumulator.finish() +} + +#[test] +fn distinct_id_and_struct_partial() -> VortexResult<()> { + let options = NumericalAggregateOpts::default(); + let aggregate = SumV2.bind(options); + assert_eq!(aggregate.id().as_ref(), "vortex.sum_v2"); + + let input_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let partial_dtype = SumV2 + .partial_dtype(&options, &input_dtype) + .expect("supported sum_v2 dtype"); + assert_eq!(partial_dtype.nullability(), Nullable); + let fields = partial_dtype.as_struct_fields(); + assert_eq!(fields.names().as_ref(), &["sum", "is_overflow", "is_empty"]); + assert_eq!( + fields.field("sum"), + Some(DType::Primitive(PType::I64, Nullability::NonNullable)) + ); + + let proto = aggregate.serialize_proto()?; + let decoded = pb::AggregateFn::decode(proto.encode_to_vec().as_slice())?; + let round_tripped = AggregateFnRef::from_proto(&decoded, &array_session())?; + assert_eq!(round_tripped, aggregate); + Ok(()) +} + +#[test] +fn legacy_sum_is_unchanged() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let options = NumericalAggregateOpts::default(); + assert_eq!(Sum.bind(options).id().as_ref(), "vortex.sum"); + assert!(matches!( + Sum.partial_dtype(&options, &dtype), + Some(DType::Primitive(..)) + )); + + let mut legacy = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype.clone())?; + let mut v2 = Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype)?; + + assert_eq!( + legacy.finish()?.as_primitive().typed_value::(), + Some(0) + ); + assert!(v2.finish()?.is_null()); + Ok(()) +} + +#[rstest] +#[case::primitive(PrimitiveArray::from_option_iter([None::, None]).into_array())] +#[case::float(PrimitiveArray::from_option_iter::([None, None]).into_array())] +#[case::bool(BoolArray::from_iter([None::, None]).into_array())] +#[case::constant( + ConstantArray::new(Scalar::null(DType::Primitive(PType::U32, Nullable)), 4).into_array() +)] +#[case::decimal( + DecimalArray::new( + buffer![0i64, 0], + DecimalDType::new(10, 2), + Validity::AllInvalid, + ) + .into_array() +)] +fn all_null_is_null(#[case] array: ArrayRef) -> VortexResult<()> { + assert!(sum_v2(&array, &mut array_session().create_execution_ctx())?.is_null()); + Ok(()) +} + +#[test] +fn valid_zero_is_not_empty() -> VortexResult<()> { + let zeros = PrimitiveArray::from_iter([0i32, 0]).into_array(); + assert_eq!( + sum_v2(&zeros, &mut array_session().create_execution_ctx())? + .as_primitive() + .typed_value::(), + Some(0) + ); + + let false_values = BoolArray::from_iter([false, false]).into_array(); + assert_eq!( + sum_v2(&false_values, &mut array_session().create_execution_ctx())? + .as_primitive() + .typed_value::(), + Some(0) + ); + Ok(()) +} + +#[test] +fn all_nan_is_nonempty_zero_when_skipped() -> VortexResult<()> { + let array = + PrimitiveArray::new(buffer![f64::NAN, f64::NAN], Validity::NonNullable).into_array(); + assert_eq!( + sum_v2(&array, &mut array_session().create_execution_ctx())? + .as_primitive() + .typed_value::(), + Some(0.0) + ); + + let included = sum_with_options(&array, NumericalAggregateOpts::include_nans())?; + assert!( + included + .as_primitive() + .typed_value::() + .expect("non-null sum") + .is_nan() + ); + Ok(()) +} + +#[test] +fn ignores_ambiguous_legacy_sum_stat() -> VortexResult<()> { + let array = PrimitiveArray::from_option_iter([None::, None]).into_array(); + array + .statistics() + .set(Stat::Sum, Precision::Exact(ScalarValue::from(42i64))); + + assert!(sum_v2(&array, &mut array_session().create_execution_ctx())?.is_null()); + Ok(()) +} + +#[test] +fn overflow_is_explicit_and_absorbing() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); + let mut accumulator = Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype)?; + let batch = PrimitiveArray::new(buffer![i64::MAX, 1, 7], Validity::NonNullable).into_array(); + accumulator.accumulate(&batch, &mut array_session().create_execution_ctx())?; + assert!(accumulator.is_saturated()); + + let partial = accumulator.partial_scalar()?; + let fields = partial.as_struct(); + assert_eq!( + fields + .field("sum") + .and_then(|sum| sum.as_primitive().typed_value::()), + Some(i64::MAX) + ); + assert_eq!( + fields + .field("is_overflow") + .and_then(|flag| flag.as_bool().value()), + Some(true) + ); + assert_eq!( + fields + .field("is_empty") + .and_then(|flag| flag.as_bool().value()), + Some(false) + ); + assert!(accumulator.finish()?.is_null()); + Ok(()) +} + +#[test] +fn empty_chunk_is_a_merge_identity() -> VortexResult<()> { + let value = PrimitiveArray::from_option_iter([Some(7i32)]).into_array(); + let empty = PrimitiveArray::from_option_iter([None::]).into_array(); + let chunked = + ChunkedArray::try_new(vec![value, empty], DType::Primitive(PType::I32, Nullable))? + .into_array(); + + assert_eq!( + sum_v2(&chunked, &mut array_session().create_execution_ctx())? + .as_primitive() + .typed_value::(), + Some(7) + ); + Ok(()) +} + +#[test] +fn combine_partials_empty_is_identity() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); + let mut empty = Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype.clone())?; + let empty_partial = empty.partial_scalar()?; + + empty.combine_partials(empty_partial.clone())?; + assert!(empty.final_scalar()?.is_null()); + + let mut value = Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype.clone())?; + let batch = PrimitiveArray::from_iter([7i64]).into_array(); + value.accumulate(&batch, &mut array_session().create_execution_ctx())?; + let value_partial = value.partial_scalar()?; + + empty.combine_partials(value_partial.clone())?; + assert_eq!( + empty.final_scalar()?.as_primitive().typed_value::(), + Some(7) + ); + + let mut value_then_empty = + Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype)?; + value_then_empty.combine_partials(value_partial)?; + value_then_empty.combine_partials(empty_partial)?; + assert_eq!( + value_then_empty + .final_scalar()? + .as_primitive() + .typed_value::(), + Some(7) + ); + Ok(()) +} + +#[test] +fn combine_partials_overflow_is_absorbing() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); + let mut max = Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype.clone())?; + let max_batch = PrimitiveArray::from_iter([i64::MAX]).into_array(); + max.accumulate(&max_batch, &mut array_session().create_execution_ctx())?; + + let mut one = Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype.clone())?; + let one_batch = PrimitiveArray::from_iter([1i64]).into_array(); + one.accumulate(&one_batch, &mut array_session().create_execution_ctx())?; + + let mut combined = + Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype.clone())?; + combined.combine_partials(max.partial_scalar()?)?; + combined.combine_partials(one.partial_scalar()?)?; + assert!(combined.is_saturated()); + assert!(combined.final_scalar()?.is_null()); + + combined.combine_partials(max.partial_scalar()?)?; + let overflow_partial = combined.partial_scalar()?; + let fields = overflow_partial.as_struct(); + assert_eq!( + fields + .field("sum") + .and_then(|sum| sum.as_primitive().typed_value::()), + Some(i64::MAX) + ); + assert_eq!( + fields + .field("is_overflow") + .and_then(|flag| flag.as_bool().value()), + Some(true) + ); + assert_eq!( + fields + .field("is_empty") + .and_then(|flag| flag.as_bool().value()), + Some(false) + ); + + let mut propagated = Accumulator::try_new(SumV2, NumericalAggregateOpts::default(), dtype)?; + propagated.combine_partials(overflow_partial)?; + assert!(propagated.is_saturated()); + assert!(propagated.final_scalar()?.is_null()); + Ok(()) +} + +#[test] +fn grouped_primitive_tracks_empty_overflow_and_null_groups() -> VortexResult<()> { + let elements = + PrimitiveArray::from_option_iter([Some(5i64), None, Some(i64::MAX), Some(1)]).into_array(); + let groups = ListViewArray::try_new( + elements, + buffer![0i32, 0, 1, 2, 0].into_array(), + buffer![1i32, 0, 1, 2, 1].into_array(), + Validity::from_iter([true, true, true, true, false]), + )? + .into_array(); + let mut accumulator = GroupedAccumulator::try_new( + SumV2, + NumericalAggregateOpts::default(), + DType::Primitive(PType::I64, Nullable), + )?; + let mut ctx = array_session().create_execution_ctx(); + accumulator.accumulate_list(&groups, &mut ctx)?; + + let result = accumulator.finish()?; + let expected = + PrimitiveArray::from_option_iter([Some(5i64), None, None, None, None]).into_array(); + assert_arrays_eq!(&result, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn grouped_bool_fallback_tracks_empty_groups() -> VortexResult<()> { + let elements = BoolArray::from_iter([Some(true), Some(true), None, None]).into_array(); + let groups = ListViewArray::try_new( + elements, + buffer![0i32, 2, 2].into_array(), + buffer![2i32, 0, 2].into_array(), + Validity::NonNullable, + )? + .into_array(); + let mut accumulator = GroupedAccumulator::try_new( + SumV2, + NumericalAggregateOpts::default(), + DType::Bool(Nullable), + )?; + let mut ctx = array_session().create_execution_ctx(); + accumulator.accumulate_list(&groups, &mut ctx)?; + + let result = accumulator.finish()?; + let expected = PrimitiveArray::from_option_iter([Some(2u64), None, None]).into_array(); + assert_arrays_eq!(&result, &expected, &mut ctx); + Ok(()) +} + +#[rstest] +#[case::all_valid( + Validity::AllValid, + [Some(10i64), None, None, Some(40)] +)] +#[case::all_invalid(Validity::AllInvalid, [None, None, None, None])] +#[case::some_invalid( + Validity::from_iter([true, true, true, false]), + [Some(10i64), None, None, None] +)] +fn finalize_struct_applies_partial_and_struct_validity( + #[case] validity: Validity, + #[case] expected: [Option; 4], +) -> VortexResult<()> { + let partials = StructArray::try_new( + FieldNames::from(["sum", "is_overflow", "is_empty"]), + vec![ + PrimitiveArray::from_iter([10i64, 20, 30, 40]).into_array(), + BoolArray::from_iter([false, true, false, false]).into_array(), + BoolArray::from_iter([false, false, true, false]).into_array(), + ], + 4, + validity, + )? + .into_array(); + + let result = SumV2.finalize(partials)?; + let expected = PrimitiveArray::from_option_iter(expected).into_array(); + assert_arrays_eq!( + &result, + &expected, + &mut array_session().create_execution_ctx() + ); + Ok(()) +} diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/fixed_size_list.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/fixed_size_list.rs index 7f8d55b07c1..03512cb8345 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/fixed_size_list.rs +++ b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/fixed_size_list.rs @@ -8,7 +8,7 @@ use super::uncompressed_size_in_bytes_u64; use super::validity_uncompressed_size_in_bytes; use crate::ExecutionCtx; use crate::arrays::FixedSizeListArray; -use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; pub(super) fn fixed_size_list_uncompressed_size_in_bytes( array: &FixedSizeListArray, diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/list_view.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/list_view.rs index f75fcb92025..3021230d076 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/list_view.rs +++ b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/list_view.rs @@ -8,7 +8,7 @@ use super::uncompressed_size_in_bytes_u64; use super::validity_uncompressed_size_in_bytes; use crate::ExecutionCtx; use crate::arrays::ListViewArray; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::ListViewRebuildMode; pub(super) fn list_view_uncompressed_size_in_bytes( diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs index 1d04a074d6f..a23bcfdc2c9 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs @@ -9,6 +9,7 @@ mod list_view; mod null; mod primitive; mod struct_; +mod union; mod varbinview; use std::mem::size_of; @@ -21,6 +22,7 @@ use list_view::list_view_uncompressed_size_in_bytes; use null::null_uncompressed_size_in_bytes; use primitive::primitive_uncompressed_size_in_bytes; use struct_::struct_uncompressed_size_in_bytes; +use union::union_uncompressed_size_in_bytes; use varbinview::varbinview_uncompressed_size_in_bytes; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -43,6 +45,8 @@ use crate::aggregate_fn::EmptyOptions; use crate::array::ArrayView; use crate::arrays::Constant; use crate::arrays::ConstantArray; +use crate::arrays::ListView; +use crate::arrays::map::MapArraySlotsExt; use crate::arrays::varbinview::BinaryView; use crate::dtype::DType; use crate::dtype::DecimalType; @@ -197,8 +201,13 @@ pub(crate) fn canonical_uncompressed_size_in_bytes( Canonical::Decimal(array) => decimal_uncompressed_size_in_bytes(array, ctx), Canonical::VarBinView(array) => varbinview_uncompressed_size_in_bytes(array, ctx), Canonical::List(array) => list_view_uncompressed_size_in_bytes(array, ctx), + Canonical::Map(array) => list_view_uncompressed_size_in_bytes( + &array.entries().as_::().into_owned(), + ctx, + ), Canonical::FixedSizeList(array) => fixed_size_list_uncompressed_size_in_bytes(array, ctx), Canonical::Struct(array) => struct_uncompressed_size_in_bytes(array, ctx), + Canonical::Union(array) => union_uncompressed_size_in_bytes(array, ctx), Canonical::Extension(array) => extension_uncompressed_size_in_bytes(array, ctx), Canonical::Variant(_) => { vortex_bail!("UncompressedSizeInBytes is not supported for Variant arrays") @@ -229,11 +238,15 @@ pub(crate) fn constant_uncompressed_size_in_bytes( array.len(), array.scalar().as_binary().value().map(|value| value.len()), )?, - DType::List(..) | DType::FixedSizeList(..) | DType::Struct(..) | DType::Extension(_) => { + DType::List(..) + | DType::Map(..) + | DType::FixedSizeList(..) + | DType::Struct(..) + | DType::Union(..) + | DType::Extension(_) => { let canonical = array.array().clone().execute::(ctx)?; return canonical_uncompressed_size_in_bytes(&canonical, ctx); } - DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), DType::Variant(_) => { vortex_bail!("UncompressedSizeInBytes is not supported for Variant arrays") } @@ -286,6 +299,10 @@ fn supports_uncompressed_size_in_bytes(dtype: &DType) -> bool { DType::List(element_dtype, _) | DType::FixedSizeList(element_dtype, ..) => { supports_uncompressed_size_in_bytes(element_dtype) } + DType::Map(map_dtype, _) => { + supports_uncompressed_size_in_bytes(&map_dtype.key_dtype()) + && supports_uncompressed_size_in_bytes(&map_dtype.value_dtype()) + } DType::Struct(fields, _) => fields .fields() .all(|field| supports_uncompressed_size_in_bytes(&field)), @@ -321,6 +338,7 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; + use crate::RecursiveCanonical; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnVTable; @@ -339,14 +357,17 @@ mod tests { use crate::arrays::NullArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; + use crate::arrays::UnionArray; use crate::arrays::VarBinViewArray; use crate::arrays::VariantArray; - use crate::builders::builder_with_capacity; + use crate::arrays::listview::ListViewRebuildMode; + use crate::builders::builder_with_capacity_in; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::dtype::PType; + use crate::dtype::UnionVariants; use crate::expr::stats::Precision; use crate::expr::stats::Stat; use crate::expr::stats::StatsProvider; @@ -356,15 +377,30 @@ mod tests { use crate::scalar::ScalarValue; use crate::validity::Validity; + /// The size the array occupies once rebuilt through the canonical builders, which is the + /// layout [`UncompressedSizeInBytes`] is defined against: the builders normalize physical + /// widths that the input array is free to choose differently, picking the smallest decimal + /// value type for a precision and `u64` list-view offsets and sizes. + /// + /// Builders no longer canonicalize their children, so the finished array is only canonical at + /// the top level - recursively canonicalize it before measuring. fn materialized_uncompressed_size_in_bytes(array: &ArrayRef) -> u64 { - let mut builder = builder_with_capacity(array.dtype(), array.len()); + let mut ctx = array_session().create_execution_ctx(); + let mut builder = builder_with_capacity_in( + array.dtype(), + array.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); array - .append_to_builder( - builder.as_mut(), - &mut array_session().create_execution_ctx(), - ) + .append_to_builder(builder.as_mut(), &mut ctx) .vortex_expect("appended"); - builder.finish().nbytes() + builder + .finish() + .execute::(&mut ctx) + .vortex_expect("recursively canonicalized") + .0 + .into_array() + .nbytes() } fn aggregate(array: &ArrayRef) -> VortexResult { @@ -496,9 +532,19 @@ mod tests { let array = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable).into_array(); + // These lists are out of order and leave element 1 unreferenced, which the builder + // round-trip inside `materialized_uncompressed_size_in_bytes` now keeps. Compare against + // the exact layout instead, which is what "materialized" means here. + let mut ctx = array_session().create_execution_ctx(); + let exact = array + .clone() + .execute::(&mut ctx)? + .rebuild(ListViewRebuildMode::MakeExact, &mut ctx)? + .into_array(); + assert_eq!( aggregate(&array)?, - materialized_uncompressed_size_in_bytes(&array) + materialized_uncompressed_size_in_bytes(&exact) ); Ok(()) } @@ -536,6 +582,26 @@ mod tests { Ok(()) } + #[test] + fn union_sums_type_ids_and_sparse_children() -> VortexResult<()> { + let type_ids = PrimitiveArray::from_iter([5_u8, 9, 5]).into_array(); + let numbers = PrimitiveArray::from_iter([10_i32, 0, 30]).into_array(); + let flags = BoolArray::from_iter([false, true, false]).into_array(); + let expected = aggregate(&type_ids)? + aggregate(&numbers)? + aggregate(&flags)?; + let variants = UnionVariants::try_new( + ["number", "flag"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Bool(Nullability::NonNullable), + ], + vec![5, 9], + )?; + let array = UnionArray::try_new(type_ids, variants, vec![numbers, flags])?.into_array(); + + assert_eq!(aggregate(&array)?, expected); + Ok(()) + } + #[test] fn extension_matches_materialized_size() -> VortexResult<()> { let storage = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(); diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/union.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/union.rs new file mode 100644 index 00000000000..cf43ed4e186 --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/union.rs @@ -0,0 +1,26 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use super::uncompressed_size_in_bytes_u64; +use crate::ExecutionCtx; +use crate::arrays::UnionArray; +use crate::arrays::union::UnionArrayExt; +use crate::arrays::union::UnionArraySlotsExt; + +pub(super) fn union_uncompressed_size_in_bytes( + array: &UnionArray, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mut size = uncompressed_size_in_bytes_u64(array.type_ids(), ctx)?; + + for child in array.iter_children() { + size = size + .checked_add(uncompressed_size_in_bytes_u64(child, ctx)?) + .ok_or_else(|| vortex_err!("uncompressed size in bytes overflowed u64"))?; + } + + Ok(size) +} diff --git a/vortex-array/src/aggregate_fn/session.rs b/vortex-array/src/aggregate_fn/session.rs index 72ef78a400c..65b00d81031 100644 --- a/vortex-array/src/aggregate_fn/session.rs +++ b/vortex-array/src/aggregate_fn/session.rs @@ -4,6 +4,7 @@ use std::any::Any; use std::sync::Arc; +use vortex_session::ArcSwapMap; use vortex_session::SessionExt; use vortex_session::SessionGuard; use vortex_session::SessionVar; @@ -32,10 +33,11 @@ use crate::aggregate_fn::fns::nan_count::NanCount; use crate::aggregate_fn::fns::null_count::NullCount; use crate::aggregate_fn::fns::sum::PrimitiveGroupedSumEncodingKernel; use crate::aggregate_fn::fns::sum::Sum; +use crate::aggregate_fn::fns::sum_v2::PrimitiveGroupedSumV2EncodingKernel; +use crate::aggregate_fn::fns::sum_v2::SumV2; use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; use crate::aggregate_fn::kernels::DynAggregateKernel; use crate::aggregate_fn::kernels::DynGroupedAggregateKernel; -use crate::arc_swap_map::ArcSwapMap; use crate::array::ArrayId; use crate::array::VTable; use crate::arrays::Chunked; @@ -54,12 +56,11 @@ use crate::dtype::DType; /// [`VortexSession`](vortex_session::VortexSession). #[derive(Clone, Debug)] pub struct AggregateFnSession { - registry: ArcSwapMap, + registry: AggregateFnRegistry, - kernels: ArcSwapMap, - grouped_kernels: ArcSwapMap, - grouped_encoding_kernels: - ArcSwapMap, + kernels: AggregateKernelRegistry, + grouped_kernels: GroupedKernelRegistry, + grouped_encoding_kernels: GroupedEncodingKernelRegistry, } impl SessionVar for AggregateFnSession { @@ -75,13 +76,23 @@ impl SessionVar for AggregateFnSession { type AggregateKernelKey = (ArrayId, Option); type GroupedEncodingKernelKey = (ArrayId, AggregateFnId); +/// Registry of aggregate function plugins, keyed by aggregate function id. +type AggregateFnRegistry = ArcSwapMap; +/// Registry of aggregate kernels, keyed by encoding and optional aggregate function. +type AggregateKernelRegistry = ArcSwapMap; +/// Registry of encoding-agnostic grouped aggregate kernels, keyed by aggregate function id. +type GroupedKernelRegistry = ArcSwapMap; +/// Registry of grouped aggregate kernels, keyed by encoding and aggregate function. +type GroupedEncodingKernelRegistry = + ArcSwapMap; + impl Default for AggregateFnSession { fn default() -> Self { let this = Self { - registry: ArcSwapMap::default(), - kernels: ArcSwapMap::default(), - grouped_kernels: ArcSwapMap::default(), - grouped_encoding_kernels: ArcSwapMap::default(), + registry: AggregateFnRegistry::default(), + kernels: AggregateKernelRegistry::default(), + grouped_kernels: GroupedKernelRegistry::default(), + grouped_encoding_kernels: GroupedEncodingKernelRegistry::default(), }; // Register the built-in aggregate functions @@ -102,6 +113,7 @@ impl Default for AggregateFnSession { this.register(NanCount); this.register(NullCount); this.register(Sum); + this.register(SumV2); this.register(UncompressedSizeInBytes); // Register the built-in aggregate kernels. @@ -117,6 +129,11 @@ impl Default for AggregateFnSession { Sum.id(), &PrimitiveGroupedSumEncodingKernel, ); + this.register_grouped_encoding_kernel( + Primitive.id(), + SumV2.id(), + &PrimitiveGroupedSumV2EncodingKernel, + ); this } diff --git a/vortex-array/src/aggregate_fn/typed.rs b/vortex-array/src/aggregate_fn/typed.rs index b399477c0a5..b5e95febea9 100644 --- a/vortex-array/src/aggregate_fn/typed.rs +++ b/vortex-array/src/aggregate_fn/typed.rs @@ -39,7 +39,6 @@ pub(super) trait DynAggregateFn: 'static + Send + Sync + super::sealed::Sealed { fn id(&self) -> AggregateFnId; fn options_any(&self) -> &dyn Any; - fn coerce_args(&self, input_dtype: &DType) -> VortexResult; fn can_satisfy(&self, requested: &AggregateFnRef) -> AggregateFnSatisfaction; fn return_dtype(&self, input_dtype: &DType) -> Option; fn state_dtype(&self, input_dtype: &DType) -> Option; @@ -64,11 +63,13 @@ pub(super) struct AggregateFnInner { } impl DynAggregateFn for AggregateFnInner { + #[allow(clippy::inline_always)] #[inline(always)] fn as_any(&self) -> &dyn Any { self } + #[allow(clippy::inline_always)] #[inline(always)] fn id(&self) -> AggregateFnId { V::id(&self.vtable) @@ -78,10 +79,6 @@ impl DynAggregateFn for AggregateFnInner { &self.options } - fn coerce_args(&self, input_dtype: &DType) -> VortexResult { - V::coerce_args(&self.vtable, &self.options, input_dtype) - } - fn can_satisfy(&self, requested: &AggregateFnRef) -> AggregateFnSatisfaction { V::can_satisfy(&self.vtable, &self.options, requested) } diff --git a/vortex-array/src/aggregate_fn/vtable.rs b/vortex-array/src/aggregate_fn/vtable.rs index 3f0a4cf567c..24449bc7572 100644 --- a/vortex-array/src/aggregate_fn/vtable.rs +++ b/vortex-array/src/aggregate_fn/vtable.rs @@ -60,15 +60,6 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { vortex_bail!("Aggregate function {} is not deserializable", self.id()); } - /// Coerce the input type for this aggregate function. - /// - /// This is optionally used by Vortex users when performing type coercion over a Vortex - /// expression. The default implementation returns the input type unchanged. - fn coerce_args(&self, options: &Self::Options, input_dtype: &DType) -> VortexResult { - let _ = options; - Ok(input_dtype.clone()) - } - /// Return whether this stored aggregate can satisfy `requested`. /// /// The default implementation only treats exactly equal aggregate functions as satisfying the diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index c3dc2e0eed7..76c0b33cd18 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -34,15 +34,10 @@ use crate::array::ArrayId; use crate::array::ArrayInner; use crate::array::ArraySlots; use crate::array::DynArrayData; -use crate::arrays::Bool; use crate::arrays::Constant; use crate::arrays::DictArray; use crate::arrays::FilterArray; -use crate::arrays::Null; -use crate::arrays::Primitive; use crate::arrays::SliceArray; -use crate::arrays::VarBin; -use crate::arrays::VarBinView; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::dtype::DType; @@ -89,12 +84,14 @@ impl ArrayRef { } /// Returns a reference to the `dyn DynArrayData` inside the inner. + #[allow(clippy::inline_always)] #[inline(always)] pub(crate) fn dyn_array(&self) -> &dyn DynArrayData { &self.0.data } /// Returns a mutable reference to the inner if this is the sole owner. + #[allow(clippy::inline_always)] #[inline(always)] pub(crate) fn inner_mut(&mut self) -> Option<&mut ArrayInner> { Arc::get_mut(&mut self.0) @@ -124,6 +121,7 @@ impl ArrayRef { /// /// # Safety /// The caller must guarantee the concrete type behind `dyn DynArrayData` is `ArrayData`. + #[allow(clippy::inline_always)] #[inline(always)] pub(crate) unsafe fn downcast_inner_unchecked( self, @@ -306,6 +304,10 @@ impl ArrayRef { /// Returns whether all items in the array are valid. pub fn all_valid(&self, ctx: &mut ExecutionCtx) -> VortexResult { + if self.is_empty() { + return Ok(true); + } + match self.validity()? { Validity::NonNullable | Validity::AllValid => Ok(true), Validity::AllInvalid => Ok(false), @@ -315,6 +317,10 @@ impl ArrayRef { /// Returns whether the array is all invalid. pub fn all_invalid(&self, ctx: &mut ExecutionCtx) -> VortexResult { + if self.is_empty() { + return Ok(true); + } + match self.validity()? { Validity::NonNullable | Validity::AllValid => Ok(false), Validity::AllInvalid => Ok(true), @@ -443,15 +449,6 @@ impl ArrayRef { nbytes } - /// Returns whether this array is an arrow encoding. - pub fn is_arrow(&self) -> bool { - self.is::() - || self.is::() - || self.is::() - || self.is::() - || self.is::() - } - /// Whether the array is of a canonical encoding. pub fn is_canonical(&self) -> bool { self.is::() @@ -685,29 +682,46 @@ impl ArrayRef { // ArrayVisitor delegation methods + /// Returns an iterator over the children of the array: its non-None slots in order. + pub fn children_iter(&self) -> impl Iterator { + self.0.slots.iter().filter_map(|s| s.as_ref()) + } + /// Returns the children of the array. pub fn children(&self) -> Vec { - self.0.data.children(self) + self.children_iter().cloned().collect() } /// Returns the number of children of the array. pub fn nchildren(&self) -> usize { - self.0.data.nchildren(self) + self.children_iter().count() } /// Returns the nth child of the array without allocating a Vec. + /// + /// Returns `None` if the index is out of bounds. pub fn nth_child(&self, idx: usize) -> Option { - self.0.data.nth_child(self, idx) + self.children_iter().nth(idx).cloned() } - /// Returns the names of the children of the array. + /// Returns the names of the children of the array: the slot names of the non-None slots + /// in order. pub fn children_names(&self) -> Vec { - self.0.data.children_names(self) + self.0 + .slots + .iter() + .enumerate() + .filter(|(_, s)| s.is_some()) + .map(|(slot_idx, _)| self.slot_name(slot_idx)) + .collect() } /// Returns the array's children with their names. pub fn named_children(&self) -> Vec<(String, ArrayRef)> { - self.0.data.named_children(self) + self.children_names() + .into_iter() + .zip(self.children_iter().cloned()) + .collect() } /// Returns the data buffers of the array. @@ -780,6 +794,7 @@ impl ArrayRef { } impl IntoArray for ArrayRef { + #[allow(clippy::inline_always)] #[inline(always)] fn into_array(self) -> ArrayRef { self diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index 3eea66c6a68..8b9b2812bfa 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -53,6 +53,63 @@ use crate::hash::ArrayHash; /// heap allocation in the common case. pub type ArraySlots = SmallVec<[Option; 4]>; +/// A borrowed run of required slots, e.g. the variadic tail of a slot layout. +/// +/// Wraps a `&[Option]` whose entries are guaranteed present by encoding +/// validation, exposing them as `&ArrayRef` without per-call-site unwrapping. +#[derive(Clone, Copy, Debug)] +pub struct SlotSlice<'a> { + slots: &'a [Option], + expect: &'static str, +} + +impl<'a> SlotSlice<'a> { + /// Wrap a slice of slots that validation guarantees are all present. + /// + /// `expect` names the slot run in the panic message if a slot is unexpectedly absent. + pub fn new(slots: &'a [Option], expect: &'static str) -> Self { + Self { slots, expect } + } + + /// The number of slots in the run. + pub fn len(&self) -> usize { + self.slots.len() + } + + /// Returns `true` if the run contains no slots. + pub fn is_empty(&self) -> bool { + self.slots.is_empty() + } + + /// Returns the slot at `idx`, or `None` if out of bounds. + pub fn get(&self, idx: usize) -> Option<&'a ArrayRef> { + self.slots + .get(idx) + .map(|slot| slot.as_ref().vortex_expect(self.expect)) + } + + /// Iterate the slots in order. + pub fn iter(&self) -> impl ExactSizeIterator + use<'a> { + let expect = self.expect; + self.slots + .iter() + .map(move |slot| slot.as_ref().vortex_expect(expect)) + } + + /// Clone every slot into an owned `Vec`. + pub fn to_vec(&self) -> Vec { + self.iter().cloned().collect() + } +} + +impl std::ops::Index for SlotSlice<'_> { + type Output = ArrayRef; + + fn index(&self, idx: usize) -> &Self::Output { + self.slots[idx].as_ref().vortex_expect(self.expect) + } +} + /// The public API trait for all Vortex arrays. /// /// This trait is sealed and cannot be implemented outside of `vortex-array`. @@ -80,23 +137,6 @@ pub(crate) trait DynArrayData: 'static + private::Sealed + Send + Sync + Debug { // --- Visitor methods (formerly in ArrayVisitor) --- - /// Returns the children of the array. - fn children(&self, this: &ArrayRef) -> Vec; - - /// Returns the number of children of the array. - fn nchildren(&self, this: &ArrayRef) -> usize; - - /// Returns the nth child of the array without allocating a Vec. - /// - /// Returns `None` if the index is out of bounds. - fn nth_child(&self, this: &ArrayRef, idx: usize) -> Option; - - /// Returns the names of the children of the array. - fn children_names(&self, this: &ArrayRef) -> Vec; - - /// Returns the array's children with their names. - fn named_children(&self, this: &ArrayRef) -> Vec<(String, ArrayRef)>; - /// Returns the buffers of the array. fn buffers(&self, this: &ArrayRef) -> Vec; @@ -267,35 +307,6 @@ impl DynArrayData for ArrayData { Ok(()) } - fn children(&self, this: &ArrayRef) -> Vec { - let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; - (0..V::nchildren(view)).map(|i| V::child(view, i)).collect() - } - - fn nchildren(&self, this: &ArrayRef) -> usize { - let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; - V::nchildren(view) - } - - fn nth_child(&self, this: &ArrayRef, idx: usize) -> Option { - let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; - (idx < V::nchildren(view)).then(|| V::child(view, idx)) - } - - fn children_names(&self, this: &ArrayRef) -> Vec { - let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; - (0..V::nchildren(view)) - .map(|i| V::child_name(view, i)) - .collect() - } - - fn named_children(&self, this: &ArrayRef) -> Vec<(String, ArrayRef)> { - let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; - (0..V::nchildren(view)) - .map(|i| (V::child_name(view, i), V::child(view, i))) - .collect() - } - fn buffers(&self, this: &ArrayRef) -> Vec { let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; (0..V::nbuffers(view)) diff --git a/vortex-array/src/array/plugin.rs b/vortex-array/src/array/plugin.rs index 66845eb9a0a..5b4fca39854 100644 --- a/vortex-array/src/array/plugin.rs +++ b/vortex-array/src/array/plugin.rs @@ -6,7 +6,9 @@ use std::fmt::Debug; use std::fmt::Formatter; use std::sync::Arc; +use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_session::VortexSession; use crate::ArrayRef; @@ -21,38 +23,127 @@ use crate::serde::ArrayChildren; /// Reference-counted array plugin. pub type ArrayPluginRef = Arc; -/// Registry trait for ID-based deserialization of arrays. +/// The wire representation produced by an in-memory array's serializer. /// -/// Plugins are registered in the session by their [`ArrayId`]. When a serialized array is -/// encountered, the session resolves the ID to the plugin and calls [`deserialize`] to reconstruct -/// the value as an [`ArrayRef`]. +/// A serializer may reuse the in-memory array's buffers and children with [`Self::from_array`], +/// or return different parts when an older wire representation requires a lossless structural +/// downgrade. +#[derive(Clone, Debug)] +pub struct ArraySerialization { + /// The concrete array ID to write on the wire. + pub serialized_id: ArrayId, + /// Encoding-specific metadata written into the array node. + pub metadata: Vec, + /// Top-level buffers written for this array node. + pub buffers: Vec, + /// Child arrays to serialize recursively. + pub children: Vec, +} + +impl ArraySerialization { + /// Create a wire representation from an ID, metadata, buffers, and children. + pub fn new( + serialized_id: ArrayId, + metadata: Vec, + buffers: Vec, + children: Vec, + ) -> Self { + Self { + serialized_id, + metadata, + buffers, + children, + } + } + + /// Reuse an in-memory array's buffers and children with the supplied serialized metadata. + pub fn from_array(serialized_id: ArrayId, array: &ArrayRef, metadata: Vec) -> Self { + Self::new(serialized_id, metadata, array.buffers(), array.children()) + } +} + +/// The borrowed wire components passed to an array deserializer. +pub struct ArrayDeserialization<'a> { + /// The exact array ID found on the wire. + pub serialized_id: ArrayId, + /// The logical dtype supplied by the containing format. + pub dtype: &'a DType, + /// The logical array length supplied by the containing format. + pub len: usize, + /// Encoding-specific metadata from the array node. + pub metadata: &'a [u8], + /// Top-level buffers referenced by the array node. + pub buffers: &'a [BufferHandle], + /// Lazily decoded child arrays referenced by the array node. + pub children: &'a dyn ArrayChildren, +} + +impl<'a> ArrayDeserialization<'a> { + /// Create borrowed deserialization input from a wire ID and its serialized components. + pub fn new( + serialized_id: ArrayId, + dtype: &'a DType, + len: usize, + metadata: &'a [u8], + buffers: &'a [BufferHandle], + children: &'a dyn ArrayChildren, + ) -> Self { + Self { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } + } +} + +/// Registry trait for serializing and deserializing an in-memory array representation. /// -/// [`deserialize`]: ArrayPlugin::deserialize +/// A plugin has one [`id`](Self::id) for the in-memory representation and one or more +/// [`serialized_ids`](Self::serialized_ids) for wire representations. Its serializer chooses the +/// wire representation, and the serialization context validates that the chosen ID is permitted +/// before it is written. +/// +/// Every serialized ID is also registered for deserialization. A current plugin may therefore +/// deserialize several historical IDs into the same in-memory representation. A reader that +/// predates a newer ID has no registration for it and reports it as unknown instead of silently +/// interpreting an unsupported representation. pub trait ArrayPlugin: 'static + Send + Sync { - /// Returns the ID for this array encoding. - /// - /// During serde, this is the key the registry uses to find - /// this plugin instance and call the appropriate method on it. + /// Returns the ID of the in-memory array representation handled by this plugin. fn id(&self) -> ArrayId; - /// Serialize the array metadata. + /// Returns the serialized array IDs understood by this plugin, ordered oldest to newest. /// - /// This function will only be called for arrays where the encoding ID matches that of this - /// plugin. - fn serialize(&self, array: &ArrayRef, session: &VortexSession) - -> VortexResult>>; + /// The default uses the in-memory ID as the sole wire ID. Override this for an in-memory array + /// that has multiple serialized variants. IDs retained only for reading may also be included; + /// the single serializer need not select them. + fn serialized_ids(&self) -> Vec { + vec![self.id()] + } - /// Deserialize an array from serialized components. + /// Serialize `array` to its wire representation. /// - /// The returned array doesn't necessary have to match this plugin's encoding ID. This is - /// useful for implementing back-compat logic and deserializing arrays into the new version. + /// This function is called only for arrays whose in-memory encoding matches [`id`](Self::id). + /// The returned ID must be declared by [`serialized_ids`](Self::serialized_ids). Return + /// `Ok(None)` when the array cannot be serialized. + fn serialize( + &self, + array: &ArrayRef, + session: &VortexSession, + ) -> VortexResult>; + + /// Deserialize one recognized wire representation into the current in-memory array. + /// + /// `serialized_id` identifies the exact representation encountered on disk. The returned + /// array does not necessarily have to use this plugin's in-memory ID; this supports legacy + /// representations that are normalized into another current in-memory array. Implementations + /// must validate the contract of that exact ID rather than accepting every form understood by + /// the current in-memory representation under an older ID. fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, session: &VortexSession, ) -> VortexResult; @@ -80,26 +171,36 @@ impl ArrayPlugin for V { &self, array: &ArrayRef, session: &VortexSession, - ) -> VortexResult>> { - assert_eq!( + ) -> VortexResult> { + vortex_ensure!( + self.id() == array.encoding_id(), + "array plugin {} cannot serialize in-memory array {}", self.id(), array.encoding_id(), - "Invoked for incorrect array ID" ); - V::serialize(array.as_::(), session) + Ok(V::serialize(array.as_::(), session)? + .map(|metadata| ArraySerialization::from_array(self.id(), array, metadata))) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, session: &VortexSession, ) -> VortexResult { + vortex_ensure!( + self.id() == parts.serialized_id, + "array plugin {} does not recognize serialized ID {}", + self.id(), + parts.serialized_id, + ); Ok(Array::::try_from_parts(V::deserialize( - self, dtype, len, metadata, buffers, children, session, + self, + parts.dtype, + parts.len, + parts.metadata, + parts.buffers, + parts.children, + session, )?)? .into_array()) } diff --git a/vortex-array/src/array/typed.rs b/vortex-array/src/array/typed.rs index 4769073b82c..a9def0318bf 100644 --- a/vortex-array/src/array/typed.rs +++ b/vortex-array/src/array/typed.rs @@ -348,6 +348,7 @@ impl Array { } /// Returns the internal [`ArrayRef`]. + #[allow(clippy::inline_always)] #[inline(always)] pub fn as_array(&self) -> &ArrayRef { &self.inner @@ -361,6 +362,7 @@ impl Array { } /// Downcast the inner `ArrayRef` to `&ArrayData`. + #[allow(clippy::inline_always)] #[inline(always)] fn downcast_inner(&self) -> &ArrayData { let any = self.inner.dyn_array().as_any(); @@ -508,6 +510,7 @@ impl AsRef for Array { } impl IntoArray for Array { + #[allow(clippy::inline_always)] #[inline(always)] fn into_array(self) -> ArrayRef { self.inner diff --git a/vortex-array/src/array/vtable/mod.rs b/vortex-array/src/array/vtable/mod.rs index 344aac64a88..3078a6ae9f5 100644 --- a/vortex-array/src/array/vtable/mod.rs +++ b/vortex-array/src/array/vtable/mod.rs @@ -115,45 +115,6 @@ pub trait VTable: 'static + Clone + Sized + Send + Sync + Debug { buffers: &[BufferHandle], ) -> VortexResult>; - /// Returns the number of children in the array. - /// - /// The default counts non-None slots. - fn nchildren(array: ArrayView<'_, Self>) -> usize { - array.slots().iter().filter(|s| s.is_some()).count() - } - - /// Returns the child at the given index. - /// - /// The default returns the `idx`-th non-None slot. - /// - /// # Panics - /// Panics if `idx >= nchildren(array)`. - fn child(array: ArrayView<'_, Self>, idx: usize) -> ArrayRef { - array - .slots() - .iter() - .filter_map(|s| s.clone()) - .nth(idx) - .vortex_expect("child index out of bounds") - } - - /// Returns the name of the child at the given index. - /// - /// The default returns the slot name of the `idx`-th non-None slot. - /// - /// # Panics - /// Panics if `idx >= nchildren(array)`. - fn child_name(array: ArrayView<'_, Self>, idx: usize) -> String { - array - .slots() - .iter() - .enumerate() - .filter(|(_, s)| s.is_some()) - .nth(idx) - .map(|(slot_idx, _)| Self::slot_name(array, slot_idx)) - .vortex_expect("child_name index out of bounds") - } - /// Serialize encoding metadata into a byte buffer for IPC or file storage. /// /// Return `None` if the array cannot be serialized by this encoding. Buffers and children are diff --git a/vortex-array/src/arrays/arbitrary.rs b/vortex-array/src/arrays/arbitrary.rs index 5d5d9f60bf9..ec5fad69135 100644 --- a/vortex-array/src/arrays/arbitrary.rs +++ b/vortex-array/src/arrays/arbitrary.rs @@ -28,10 +28,13 @@ use crate::builders::ArrayBuilder; use crate::builders::DecimalBuilder; use crate::builders::FixedSizeListBuilder; use crate::builders::ListViewBuilder; +use crate::builders::MapBuilder; use crate::dtype::DType; use crate::dtype::IntegerPType; +use crate::dtype::MapDType; use crate::dtype::NativePType; use crate::dtype::Nullability; +use crate::dtype::OffsetBuilderPType; use crate::dtype::PType; use crate::match_each_decimal_value_type; use crate::scalar::Scalar; @@ -141,7 +144,11 @@ fn random_array_chunk( d @ DType::Decimal(decimal, n) => { let elem_len = chunk_len.unwrap_or(u.int_in_range(0..=20)?); match_each_decimal_value_type!(DecimalType::smallest_decimal_value_type(decimal), |D| { - let mut builder = DecimalBuilder::new::(*decimal, *n); + let mut builder = DecimalBuilder::new_in::( + *decimal, + *n, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); for _i in 0..elem_len { let random_decimal = random_scalar(u, d)?; builder.append_scalar(&random_decimal).vortex_expect( @@ -157,6 +164,9 @@ fn random_array_chunk( DType::FixedSizeList(elem_dtype, list_size, null) => { random_fixed_size_list(u, elem_dtype, *list_size, *null, chunk_len) } + DType::Map(map_dtype, nullability) => { + random_map(u, map_dtype.clone(), *nullability, chunk_len) + } DType::Struct(sdt, n) => { let first_array = sdt .fields() @@ -197,6 +207,46 @@ fn random_array_chunk( } } +fn random_map( + u: &mut Unstructured, + map_dtype: MapDType, + nullability: Nullability, + chunk_len: Option, +) -> Result { + let array_length = chunk_len.unwrap_or(u.int_in_range(0..=20)?); + let key_dtype = map_dtype.key_dtype(); + let value_dtype = map_dtype.value_dtype(); + let dtype = DType::Map(map_dtype.clone(), nullability); + let mut builder = MapBuilder::::with_capacity_in( + map_dtype, + nullability, + array_length, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + + for _ in 0..array_length { + if nullability == Nullability::Nullable && u.arbitrary::()? { + builder.append_null(); + } else { + let entry_count = u.int_in_range(0..=20)?; + let entries = (0..entry_count) + .map(|_| { + let key = random_scalar(u, &key_dtype)?; + let value = random_scalar(u, &value_dtype)?; + Ok((key, value)) + }) + .collect::>>()?; + let scalar = Scalar::try_map(dtype.clone(), entries) + .vortex_expect("generated map scalar should be valid"); + builder + .append_scalar(&scalar) + .vortex_expect("generated map scalar should append"); + } + } + + Ok(builder.finish_into_map().into_array()) +} + /// Creates a random fixed-size list array. /// /// If the `chunk_len` is specified, the length of the array will be equal to the chunk length. @@ -209,8 +259,13 @@ fn random_fixed_size_list( ) -> Result { let array_length = chunk_len.unwrap_or(u.int_in_range(0..=20)?); - let mut builder = - FixedSizeListBuilder::with_capacity(Arc::clone(elem_dtype), list_size, null, array_length); + let mut builder = FixedSizeListBuilder::with_capacity_in( + Arc::clone(elem_dtype), + list_size, + null, + array_length, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); for _ in 0..array_length { if null == Nullability::Nullable && u.arbitrary::()? { @@ -238,17 +293,11 @@ fn random_list( // Worst-case total elements: each list can have up to 20 elements. let max_total_elements = array_length as u64 * 20; - match u.int_in_range(0..=5)? { - 0 if i16::max_value_as_u64() >= max_total_elements => { - random_list_with_offset_type::(u, elem_dtype, null, array_length) - } - 1 if i32::max_value_as_u64() >= max_total_elements => { + match u.int_in_range(0..=3)? { + 0 if i32::max_value_as_u64() >= max_total_elements => { random_list_with_offset_type::(u, elem_dtype, null, array_length) } - 3 if u16::max_value_as_u64() >= max_total_elements => { - random_list_with_offset_type::(u, elem_dtype, null, array_length) - } - 4 if u32::max_value_as_u64() >= max_total_elements => { + 1 if u32::max_value_as_u64() >= max_total_elements => { random_list_with_offset_type::(u, elem_dtype, null, array_length) } // i64 and u64 always fit; also the fallback for when narrower types don't. @@ -262,15 +311,20 @@ fn random_list( } } -/// Creates a random list array with the given [`IntegerPType`] for the internal offsets child. -fn random_list_with_offset_type( +/// Creates a random list array with the given [`OffsetBuilderPType`] for the internal offsets child. +fn random_list_with_offset_type( u: &mut Unstructured, elem_dtype: &Arc, null: Nullability, array_length: usize, ) -> Result { - let mut builder = - ListViewBuilder::::with_capacity(Arc::clone(elem_dtype), null, array_length, 10); + let mut builder = ListViewBuilder::::with_capacity_in( + Arc::clone(elem_dtype), + null, + array_length, + 10, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); for _ in 0..array_length { if null == Nullability::Nullable && u.arbitrary::()? { diff --git a/vortex-array/src/arrays/bool/array.rs b/vortex-array/src/arrays/bool/array.rs index 2dd08c55832..5b1083e7f88 100644 --- a/vortex-array/src/arrays/bool/array.rs +++ b/vortex-array/src/arrays/bool/array.rs @@ -23,16 +23,19 @@ use crate::array::ArrayParts; use crate::array::TypedArrayRef; use crate::array::child_to_validity; use crate::array::validity_to_child; +use crate::array_slots; use crate::arrays::Bool; use crate::arrays::BoolArray; use crate::buffer::BufferHandle; use crate::dtype::DType; use crate::validity::Validity; -/// The validity bitmap indicating which elements are non-null. -pub(super) const VALIDITY_SLOT: usize = 0; -pub(super) const NUM_SLOTS: usize = 1; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"]; +#[array_slots(Bool)] +pub struct BoolSlots { + /// The validity bitmap indicating which elements are non-null. + #[slot(0)] + pub validity: Option, +} /// Inner data for a boolean array that stores true/false values in a compact bit-packed format. /// @@ -94,7 +97,7 @@ pub trait BoolArrayExt: TypedArrayRef { fn validity(&self) -> Validity { child_to_validity( - self.as_ref().slots()[VALIDITY_SLOT].as_ref(), + self.as_ref().slots()[BoolSlots::VALIDITY].as_ref(), self.nullability(), ) } @@ -110,7 +113,7 @@ pub trait BoolArrayExt: TypedArrayRef { } fn maybe_execute_mask(&self, ctx: &mut ExecutionCtx) -> VortexResult> { - let all_valid = match &self.validity() { + let all_valid = match &BoolArrayExt::validity(self) { Validity::NonNullable | Validity::AllValid => true, Validity::AllInvalid => false, Validity::Array(a) => a.statistics().compute_min::(ctx).unwrap_or(false), @@ -125,8 +128,7 @@ pub trait BoolArrayExt: TypedArrayRef { } fn to_mask_fill_null_false(&self, ctx: &mut ExecutionCtx) -> Mask { - let validity_mask = self - .validity() + let validity_mask = BoolArrayExt::validity(self) .execute_mask(self.as_ref().len(), ctx) .vortex_expect("Failed to compute validity mask"); let buffer = match validity_mask { @@ -367,9 +369,12 @@ mod tests { use std::iter::once; use std::iter::repeat_n; + use vortex_buffer::Alignment; use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; + use vortex_buffer::ByteBuffer; use vortex_buffer::buffer; + use vortex_error::VortexResult; use crate::IntoArray; use crate::VortexSessionExecute; @@ -378,6 +383,7 @@ mod tests { use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; use crate::assert_arrays_eq; + use crate::buffer::BufferHandle; use crate::patches::Patches; use crate::validity::Validity; @@ -471,6 +477,25 @@ mod tests { assert_arrays_eq!(sliced, BoolArray::from_iter([true; 8]), &mut ctx); } + #[test] + fn slice_aligned_host_handle_at_unaligned_byte() -> VortexResult<()> { + let bits: ByteBuffer = buffer![0b1010_1100_u8, 0b0110_1001, 0]; + let bits = bits.aligned(Alignment::of::()); + let array = + BoolArray::new_handle(BufferHandle::new_host(bits), 0, 16, Validity::NonNullable) + .into_array(); + + let sliced = array.slice(9..15)?; + + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!( + sliced, + BoolArray::from_iter([false, false, true, false, true, true]), + &mut ctx + ); + Ok(()) + } + #[test] fn patch_bools_owned() { let mut ctx = array_session().create_execution_ctx(); diff --git a/vortex-array/src/arrays/bool/compute/filter.rs b/vortex-array/src/arrays/bool/compute/filter.rs index 98f52ed11ee..106bd9650dd 100644 --- a/vortex-array/src/arrays/bool/compute/filter.rs +++ b/vortex-array/src/arrays/bool/compute/filter.rs @@ -129,6 +129,7 @@ fn filter_pext_fallback(src: &BitBuffer, mask_buf: &BitBuffer, true_count: usize /// Extracted so the same logic is shared between the software and hardware paths. /// Uses raw pointer writes instead of Vec::push to eliminate bounds checks /// in the hot loop — we know the exact output size from true_count. +#[allow(clippy::inline_always)] #[inline(always)] #[allow(clippy::cast_possible_truncation)] fn filter_inner( @@ -227,6 +228,7 @@ fn filter_inner( /// lookup table per mask byte. Each byte PEXT is a single table lookup with no /// data dependencies between bytes, making this faster than the parallel-prefix /// approach (~12ns vs ~18ns per word). +#[allow(clippy::inline_always)] #[inline(always)] pub fn pext_fallback(src: u64, mask: u64) -> u64 { pext_byte_lut(src, mask) @@ -268,6 +270,7 @@ static BYTE_PEXT_LUT: &[u8; 256 * 256] = &{ }; /// Byte-level PEXT using precomputed lookup table. +#[allow(clippy::inline_always)] #[inline(always)] fn pext_byte_lut(src: u64, mask: u64) -> u64 { let src_bytes = src.to_le_bytes(); diff --git a/vortex-array/src/arrays/bool/compute/slice.rs b/vortex-array/src/arrays/bool/compute/slice.rs index 0b38d2a6117..f88eb418da5 100644 --- a/vortex-array/src/arrays/bool/compute/slice.rs +++ b/vortex-array/src/arrays/bool/compute/slice.rs @@ -10,18 +10,23 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::Bool; use crate::arrays::BoolArray; -use crate::arrays::bool::BoolArrayExt; use crate::arrays::slice::SliceReduce; +use crate::buffer::BufferHandle; impl SliceReduce for Bool { fn slice(array: ArrayView<'_, Bool>, range: Range) -> VortexResult> { - let bit_buffer = array.to_bit_buffer().slice(range.clone()); + let (byte_start, meta) = array.meta.slice(range.clone()); + let byte_end = byte_start + meta.byte_len(); + + let bits = if let Some(host) = array.bits.as_host_opt() { + BufferHandle::new_host(host.slice_unaligned(byte_start..byte_end)) + } else { + array.bits.slice(byte_start..byte_end) + }; let validity = array.validity()?.slice(range)?; - // Safety: - // range is verified in the callers and is the same for both bits and validity. - let array = unsafe { BoolArray::new_unchecked(bit_buffer, validity).into_array() }; + let array = BoolArray::try_new_from_handle(bits, meta.offset(), meta.len(), validity)?; - Ok(Some(array)) + Ok(Some(array.into_array())) } } diff --git a/vortex-array/src/arrays/bool/compute/take.rs b/vortex-array/src/arrays/bool/compute/take.rs index 8469c7e25c2..529450f6639 100644 --- a/vortex-array/src/arrays/bool/compute/take.rs +++ b/vortex-array/src/arrays/bool/compute/take.rs @@ -4,23 +4,32 @@ use itertools::Itertools as _; use num_traits::AsPrimitive; use vortex_buffer::BitBuffer; +use vortex_buffer::BitBufferMut; use vortex_buffer::BitBufferView; use vortex_buffer::get_bit; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_mask::Mask; use crate::ArrayRef; +use crate::Columnar; use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::Bool; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; +use crate::arrays::PiecewiseSequence; use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; use crate::arrays::dict::TakeExecute; +use crate::arrays::piecewise_sequence::constant_unsigned_usize; +use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::builtins::ArrayBuiltins; +use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_integer_ptype; +use crate::match_each_unsigned_integer_ptype; use crate::scalar::Scalar; impl TakeExecute for Bool { @@ -29,6 +38,12 @@ impl TakeExecute for Bool { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)? + { + return Ok(Some(taken)); + } + let indices_nulls_zeroed = match indices.validity()?.execute_mask(indices.len(), ctx)? { Mask::AllTrue(_) => indices.clone(), Mask::AllFalse(_) => { @@ -55,6 +70,49 @@ impl TakeExecute for Bool { } } +fn take_contiguous_ranges( + array: ArrayView<'_, Bool>, + indices: ArrayView<'_, PiecewiseSequence>, + indices_ref: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { + return Ok(None); + }; + let source = array.to_bit_buffer(); + let output_len = indices_ref.len(); + let buffer = match lengths { + Columnar::Constant(lengths) => { + let length = constant_unsigned_usize(&lengths); + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + take_bit_slices_constant_length( + &source, + starts.as_slice::(), + length, + output_len, + )? + }) + } + Columnar::Canonical(lengths) => { + let lengths = lengths.into_primitive(); + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + take_bit_slices( + &source, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + )? + }) + }) + } + }; + + Ok(Some( + BoolArray::new(buffer, array.validity()?.take(indices_ref)?).into_array(), + )) +} + fn take_valid_indices>(bools: BitBufferView<'_>, indices: &[I]) -> BitBuffer { // For boolean arrays that roughly fit into a single page (at least, on Linux), it's worth // the overhead to convert to a Vec. @@ -82,6 +140,58 @@ fn take_bool_impl>(bools: BitBufferView<'_>, indices: &[I] }) } +fn take_bit_slices_constant_length( + source: &BitBuffer, + starts: &[S], + length: usize, + output_len: usize, +) -> VortexResult +where + S: UnsignedPType, +{ + let computed_len = starts + .len() + .checked_mul(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); + + let mut values = BitBufferMut::with_capacity(output_len); + for start in starts { + let start = start.as_(); + values.append_buffer(&source.slice(start..).slice(..length)); + } + + Ok(values.freeze()) +} + +fn take_bit_slices( + source: &BitBuffer, + starts: &[S], + lengths: &[L], + output_len: usize, +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, +{ + let mut values = BitBufferMut::with_capacity(output_len); + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = start.as_(); + let length = length.as_(); + values.append_buffer(&source.slice(start..).slice(..length)); + } + + vortex_ensure!( + values.len() == output_len, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + values.len() + ); + Ok(values.freeze()) +} + #[cfg(test)] mod test { use rstest::rstest; diff --git a/vortex-array/src/arrays/bool/mod.rs b/vortex-array/src/arrays/bool/mod.rs index ea51e6b5f29..96856c66a75 100644 --- a/vortex-array/src/arrays/bool/mod.rs +++ b/vortex-array/src/arrays/bool/mod.rs @@ -5,8 +5,10 @@ mod array; mod patch; pub use array::BoolArrayExt; +pub use array::BoolArraySlotsExt; pub use array::BoolData; pub use array::BoolDataParts; +pub use array::BoolSlots; pub(crate) mod compute; diff --git a/vortex-array/src/arrays/bool/vtable/mod.rs b/vortex-array/src/arrays/bool/vtable/mod.rs index 8a9cd961da8..393414924b3 100644 --- a/vortex-array/src/arrays/bool/vtable/mod.rs +++ b/vortex-array/src/arrays/bool/vtable/mod.rs @@ -21,7 +21,7 @@ use crate::array::ArrayView; use crate::array::VTable; use crate::array::child_to_validity; use crate::arrays::bool::BoolData; -use crate::arrays::bool::array::SLOT_NAMES; +use crate::arrays::bool::array::BoolSlots; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::BoolBuilder; @@ -146,7 +146,7 @@ impl VTable for Bool { data.bits.len() * 8 ); - let validity = child_to_validity(slots[0].as_ref(), *nullability); + let validity = child_to_validity(slots[BoolSlots::VALIDITY].as_ref(), *nullability); if let Some(validity_len) = validity.maybe_len() { vortex_ensure!( validity_len == len, @@ -189,7 +189,7 @@ impl VTable for Bool { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + BoolSlots::NAMES[idx].to_string() } fn append_to_builder( diff --git a/vortex-array/src/arrays/chunked/array.rs b/vortex-array/src/arrays/chunked/array.rs index 3963d6ba950..807d526c17b 100644 --- a/vortex-array/src/arrays/chunked/array.rs +++ b/vortex-array/src/arrays/chunked/array.rs @@ -10,7 +10,6 @@ use std::fmt::Display; use std::fmt::Formatter; use futures::stream; -use smallvec::SmallVec; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -24,6 +23,7 @@ use crate::IntoArray; use crate::array::Array; use crate::array::ArrayParts; use crate::array::TypedArrayRef; +use crate::array_slots; use crate::arrays::Chunked; use crate::arrays::PrimitiveArray; use crate::dtype::DType; @@ -35,8 +35,16 @@ use crate::stream::ArrayStream; use crate::stream::ArrayStreamAdapter; use crate::validity::Validity; -pub(super) const CHUNK_OFFSETS_SLOT: usize = 0; -pub(super) const CHUNKS_OFFSET: usize = 1; +/// Slot layout of a [`Chunked`] array: `[chunk_offsets, chunks...]`. +#[array_slots(Chunked)] +pub struct ChunkedSlots { + /// The non-nullable `u64` array of cumulative chunk offsets. + #[slot(0)] + pub chunk_offsets: ArrayRef, + /// The chunk arrays, each sharing the outer dtype. + #[slot(1..)] + pub chunks: Vec, +} #[derive(Clone, Debug)] pub struct ChunkedData { @@ -53,24 +61,27 @@ impl Display for ChunkedData { pub trait ChunkedArrayExt: TypedArrayRef { fn chunk_offsets_array(&self) -> &ArrayRef { - self.as_ref().slots()[CHUNK_OFFSETS_SLOT] + self.as_ref().slots()[ChunkedSlots::CHUNK_OFFSETS] .as_ref() .vortex_expect("validated chunk offsets slot") } fn nchunks(&self) -> usize { - self.as_ref().slots().len().saturating_sub(CHUNKS_OFFSET) + self.as_ref() + .slots() + .len() + .saturating_sub(ChunkedSlots::CHUNKS_OFFSET) } fn chunk(&self, idx: usize) -> &ArrayRef { - self.as_ref().slots()[CHUNKS_OFFSET + idx] + self.as_ref().slots()[ChunkedSlots::CHUNKS_OFFSET + idx] .as_ref() .vortex_expect("validated chunk slot") } fn iter_chunks<'a>(&'a self) -> Box + 'a> { Box::new( - self.as_ref().slots()[CHUNKS_OFFSET..] + self.as_ref().slots()[ChunkedSlots::CHUNKS_OFFSET..] .iter() .map(|slot| slot.as_ref().vortex_expect("validated chunk slot")), ) @@ -84,7 +95,8 @@ pub trait ChunkedArrayExt: TypedArrayRef { Box::new(self.iter_chunks().filter(|chunk| !chunk.is_empty())) } - fn chunk_offsets(&self) -> &[usize] { + /// Returns the cached chunk boundary offsets. + fn chunk_offset_values(&self) -> &[usize] { &self.chunk_offsets } @@ -93,12 +105,12 @@ pub trait ChunkedArrayExt: TypedArrayRef { index <= self.as_ref().len(), "Index out of bounds of the array" ); - let chunk_offsets = self.chunk_offsets(); - let index_chunk = chunk_offsets + let chunk_offset_values = self.chunk_offset_values(); + let index_chunk = chunk_offset_values .search_sorted(&index, SearchSortedSide::Right)? .to_ends_index(self.nchunks() + 1) .saturating_sub(1); - let chunk_start = chunk_offsets[index_chunk]; + let chunk_start = chunk_offset_values[index_chunk]; let index_in_chunk = index - chunk_start; Ok((index_chunk, index_in_chunk)) } @@ -123,35 +135,18 @@ impl ChunkedData { pub(super) fn new(chunk_offsets: Vec) -> Self { Self { chunk_offsets, - next_builder_slot: CHUNKS_OFFSET, - } - } - - pub(super) fn compute_chunk_offsets(chunks: &[ArrayRef]) -> Vec { - let mut chunk_offsets = Vec::with_capacity(chunks.len() + 1); - chunk_offsets.push(0); - let mut curr_offset = 0; - for chunk in chunks { - curr_offset += chunk.len(); - chunk_offsets.push(curr_offset); + next_builder_slot: ChunkedSlots::CHUNKS_OFFSET, } - chunk_offsets } - pub(super) fn make_slots(chunk_offsets: &[usize], chunks: &[ArrayRef]) -> ArraySlots { + pub(super) fn make_chunk_offsets_array(chunk_offsets: &[usize]) -> ArrayRef { let mut chunk_offsets_buf = BufferMut::::with_capacity(chunk_offsets.len()); for &offset in chunk_offsets { let offset = u64::try_from(offset) .vortex_expect("chunk offset must fit in u64 for serialization"); unsafe { chunk_offsets_buf.push_unchecked(offset) } } - - let mut slots = SmallVec::with_capacity(1 + chunks.len()); - slots.push(Some( - PrimitiveArray::new(chunk_offsets_buf.freeze(), Validity::NonNullable).into_array(), - )); - slots.extend(chunks.iter().map(|c| Some(c.clone()))); - slots + PrimitiveArray::new(chunk_offsets_buf.freeze(), Validity::NonNullable).into_array() } /// Validates the components that would be used to create a `ChunkedArray`. @@ -169,6 +164,32 @@ impl ChunkedData { } impl Array { + fn parts_from_chunks( + chunks: impl IntoIterator, + dtype: DType, + ) -> VortexResult> { + let chunks = chunks.into_iter(); + let (lower, _) = chunks.size_hint(); + let mut slots = ArraySlots::with_capacity(ChunkedSlots::CHUNKS_OFFSET + lower); + slots.push(None); + let mut chunk_offsets = Vec::with_capacity(lower + 1); + chunk_offsets.push(0); + let mut len = 0usize; + + for chunk in chunks { + if VALIDATE && chunk.dtype() != &dtype { + vortex_bail!(MismatchedTypes: &dtype, chunk.dtype()); + } + len += chunk.len(); + chunk_offsets.push(len); + slots.push(Some(chunk)); + } + + slots[ChunkedSlots::CHUNK_OFFSETS] = + Some(ChunkedData::make_chunk_offsets_array(&chunk_offsets)); + Ok(ArrayParts::new(Chunked, dtype, len, ChunkedData::new(chunk_offsets)).with_slots(slots)) + } + pub(super) fn with_next_builder_slot(mut self, next_builder_slot: usize) -> Self { if let Some(data) = self.data_mut() { data.next_builder_slot = next_builder_slot; @@ -190,10 +211,8 @@ impl Array { } /// Constructs a new `ChunkedArray`. - pub fn try_new(chunks: Vec, dtype: DType) -> VortexResult { - ChunkedData::validate(&chunks, &dtype)?; - // SAFETY just validated on previous line. - Ok(unsafe { Self::new_unchecked(chunks, dtype) }) + pub fn try_new(chunks: impl IntoIterator, dtype: DType) -> VortexResult { + Array::try_from_parts(Self::parts_from_chunks::(chunks, dtype)?) } pub fn rechunk( @@ -253,15 +272,10 @@ impl Array { /// # Safety /// /// All chunks must have exactly the same [`DType`] as the provided `dtype`. - pub unsafe fn new_unchecked(chunks: Vec, dtype: DType) -> Self { - let len = chunks.iter().map(|chunk| chunk.len()).sum(); - let chunk_offsets = ChunkedData::compute_chunk_offsets(&chunks); - unsafe { - Array::from_parts_unchecked( - ArrayParts::new(Chunked, dtype, len, ChunkedData::new(chunk_offsets.clone())) - .with_slots(ChunkedData::make_slots(&chunk_offsets, &chunks)), - ) - } + pub unsafe fn new_unchecked(chunks: impl IntoIterator, dtype: DType) -> Self { + let parts = Self::parts_from_chunks::(chunks, dtype) + .vortex_expect("unchecked chunked construction cannot fail"); + unsafe { Array::from_parts_unchecked(parts) } } } diff --git a/vortex-array/src/arrays/chunked/compute/filter.rs b/vortex-array/src/arrays/chunked/compute/filter.rs index 5f50ce18387..1c8fafd3203 100644 --- a/vortex-array/src/arrays/chunked/compute/filter.rs +++ b/vortex-array/src/arrays/chunked/compute/filter.rs @@ -71,7 +71,7 @@ fn filter_slices( let chunk_filters = chunk_filters(array, slices)?; // Now, apply the chunk filter to every slice. - for (chunk, chunk_filter) in array.iter_chunks().zip(chunk_filters.into_iter()) { + for (chunk, chunk_filter) in array.iter_chunks().zip(chunk_filters) { match chunk_filter { // All => preserve the entire chunk unfiltered. ChunkFilter::All => result.push(chunk.clone()), @@ -91,7 +91,7 @@ pub(crate) fn chunk_filters( array: ArrayView<'_, Chunked>, slices: impl Iterator, ) -> VortexResult> { - let chunk_offsets = array.chunk_offsets(); + let chunk_offsets = array.chunk_offset_values(); let mut chunk_filters = vec![ChunkFilter::None; array.nchunks()]; @@ -153,7 +153,7 @@ fn filter_indices( let mut current_chunk_id = 0; let mut chunk_indices = BufferMut::with_capacity(array.nchunks()); - let chunk_offsets = array.chunk_offsets(); + let chunk_offsets = array.chunk_offset_values(); for set_index in indices { let (chunk_id, index) = find_chunk_idx(set_index, chunk_offsets)?; @@ -273,7 +273,7 @@ mod test { DType::Primitive(PType::I32, Nullability::NonNullable), ).unwrap())] #[case(ChunkedArray::try_new( - (0..10).map(|i| buffer![i as i64, i as i64 + 10, i as i64 + 20].into_array()).collect(), + (0..10).map(|i| buffer![i as i64, i as i64 + 10, i as i64 + 20].into_array()), DType::Primitive(PType::I64, Nullability::NonNullable), ).unwrap())] fn test_filter_chunked_conformance(#[case] chunked: ChunkedArray) { diff --git a/vortex-array/src/arrays/chunked/compute/mask.rs b/vortex-array/src/arrays/chunked/compute/mask.rs index 01170a232bb..f05988de05c 100644 --- a/vortex-array/src/arrays/chunked/compute/mask.rs +++ b/vortex-array/src/arrays/chunked/compute/mask.rs @@ -10,8 +10,6 @@ use crate::array::ArrayView; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; use crate::arrays::chunked::ChunkedArrayExt; -use crate::arrays::scalar_fn::ScalarFnFactoryExt; -use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::fns::mask::Mask as MaskExpr; use crate::scalar_fn::fns::mask::MaskKernel; @@ -21,7 +19,7 @@ impl MaskKernel for Chunked { mask: &ArrayRef, _ctx: &mut ExecutionCtx, ) -> VortexResult> { - let chunk_offsets = array.chunk_offsets(); + let chunk_offsets = array.chunk_offset_values(); let new_chunks: Vec = array .iter_chunks() .enumerate() @@ -29,7 +27,7 @@ impl MaskKernel for Chunked { let start = chunk_offsets[i]; let end = chunk_offsets[i + 1]; let chunk_mask = mask.slice(start..end)?; - MaskExpr.try_new_array(chunk.len(), EmptyOptions, [chunk.clone(), chunk_mask]) + MaskExpr::try_new(chunk.clone(), chunk_mask).map(IntoArray::into_array) }) .collect::>()?; @@ -78,7 +76,7 @@ mod test { DType::Primitive(PType::U8, Nullability::NonNullable), ).unwrap())] #[case(ChunkedArray::try_new( - (0..20).map(|i| buffer![i as f32, i as f32 + 0.5].into_array()).collect(), + (0..20).map(|i| buffer![i as f32, i as f32 + 0.5].into_array()), DType::Primitive(PType::F32, Nullability::NonNullable), ).unwrap())] fn test_mask_chunked_conformance(#[case] chunked: ChunkedArray) { diff --git a/vortex-array/src/arrays/chunked/compute/mod.rs b/vortex-array/src/arrays/chunked/compute/mod.rs index f7c9a4e9d61..66d64595814 100644 --- a/vortex-array/src/arrays/chunked/compute/mod.rs +++ b/vortex-array/src/arrays/chunked/compute/mod.rs @@ -21,10 +21,12 @@ mod tests { use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::ChunkedArray; + use crate::arrays::DecimalArray; use crate::arrays::PrimitiveArray; use crate::compute::conformance::binary_numeric::test_binary_numeric_array; use crate::compute::conformance::consistency::test_array_consistency; use crate::dtype::DType; + use crate::dtype::DecimalDType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -46,7 +48,7 @@ mod tests { ).unwrap())] // Many chunks #[case::many_small_chunks(ChunkedArray::try_new( - (0..10).map(|i| buffer![i as i64, i as i64 + 10, i as i64 + 20].into_array()).collect(), + (0..10).map(|i| buffer![i as i64, i as i64 + 10, i as i64 + 20].into_array()), DType::Primitive(PType::I64, Nullability::NonNullable), ).unwrap())] // Edge cases @@ -132,7 +134,7 @@ mod tests { DType::Primitive(PType::I32, Nullability::NonNullable), ).unwrap())] #[case::chunked_many_small_chunks(ChunkedArray::try_new( - (0..10).map(|i| buffer![i * 10, i * 10 + 1].into_array()).collect(), + (0..10).map(|i| buffer![i * 10, i * 10 + 1].into_array()), DType::Primitive(PType::I32, Nullability::NonNullable), ).unwrap())] #[case::chunked_nullable(ChunkedArray::try_new( @@ -158,6 +160,13 @@ mod tests { ], DType::Primitive(PType::I32, Nullability::NonNullable), ).unwrap())] + #[case::chunked_decimal(ChunkedArray::try_new( + vec![ + DecimalArray::from_iter::([100, 250], DecimalDType::new(10, 2)).into_array(), + DecimalArray::from_iter::([300, 400, 500], DecimalDType::new(10, 2)).into_array(), + ], + DType::Decimal(DecimalDType::new(10, 2), Nullability::NonNullable), + ).unwrap())] fn test_chunked_binary_numeric(#[case] array: ChunkedArray) { test_binary_numeric_array( &array.into_array(), diff --git a/vortex-array/src/arrays/chunked/compute/take.rs b/vortex-array/src/arrays/chunked/compute/take.rs index 1b867981236..62bcda40c14 100644 --- a/vortex-array/src/arrays/chunked/compute/take.rs +++ b/vortex-array/src/arrays/chunked/compute/take.rs @@ -1,62 +1,110 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use itertools::Itertools; +use num_traits::AsPrimitive; +use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_mask::Mask; use crate::ArrayRef; use crate::Canonical; +use crate::Columnar; use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; +use crate::arrays::ConstantArray; +use crate::arrays::FixedSizeListArray; +use crate::arrays::PiecewiseSequence; +use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::dict::TakeExecute; +use crate::arrays::piecewise_sequence::constant_unsigned_usize; +use crate::arrays::piecewise_sequence::maybe_contiguous_slices; +use crate::arrays::primitive::PrimitiveArrayExt; +use crate::builders::ArrayBuilder; +use crate::builders::builder_with_capacity_in; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::PType; use crate::executor::ExecutionCtx; +use crate::match_each_unsigned_integer_ptype; use crate::validity::Validity; -// TODO(joe): this is pretty unoptimized but better than before. We want canonical using a builder -// we also want to return a chunked array ideally. -fn take_chunked( +/// Flattens per-chunk take/filter results into a single array. Flat dtypes append directly into +/// a canonical builder; nested dtypes collect the chunks and canonicalize them as a chunked +/// array, which reuses the chunks' children zero-copy. +enum ChunkFlattener { + Builder(Box), + Chunks(Vec), +} + +impl ChunkFlattener { + fn new(dtype: &DType, capacity: usize, allocator: &BufferAllocatorRef) -> Self { + if dtype.is_nested() { + Self::Chunks(Vec::new()) + } else { + Self::Builder(builder_with_capacity_in(dtype, capacity, allocator)) + } + } + + fn push(&mut self, chunk: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()> { + match self { + Self::Builder(builder) => chunk.append_to_builder(builder.as_mut(), ctx), + Self::Chunks(chunks) => { + chunks.push(chunk); + Ok(()) + } + } + } + + fn finish(self, dtype: &DType, ctx: &mut ExecutionCtx) -> VortexResult { + match self { + Self::Builder(mut builder) => Ok(builder.finish()), + // SAFETY: every chunk is a filter or take of a chunk of a chunked array with `dtype`, + // which leaves the dtype unchanged. + Self::Chunks(chunks) => unsafe { ChunkedArray::new_unchecked(chunks, dtype.clone()) } + .into_array() + .execute::(ctx) + .map(IntoArray::into_array), + } + } +} + +fn take_chunked_via_sort( array: ArrayView<'_, Chunked>, - indices: &ArrayRef, + indices: &PrimitiveArray, + indices_mask: Mask, ctx: &mut ExecutionCtx, ) -> VortexResult { - let indices = indices - .cast(DType::Primitive(PType::U64, indices.dtype().nullability()))? - .execute::(ctx)?; - - let indices_mask = indices - .as_ref() - .validity()? - .execute_mask(indices.as_ref().len(), ctx)?; let indices_values = indices.as_slice::(); let n = indices_values.len(); - - // 1. Sort (value, orig_pos) pairs so indices for the same chunk are contiguous. - // Skip null indices — their final_take slots stay 0 and are masked null by validity. let mut pairs: Vec<(u64, usize)> = indices_values .iter() .enumerate() - .filter(|&(i, _)| indices_mask.value(i)) - .map(|(i, &v)| (v, i)) + .filter(|&(position, _)| indices_mask.value(position)) + .map(|(position, &index)| (index, position)) .collect(); pairs.sort_unstable(); - // 2. Fused pass: walk sorted pairs against chunk boundaries. - // - Dedup inline → build per-chunk filter masks - // - Scatter final_take[orig_pos] = dedup_idx for every pair - let chunk_offsets = array.chunk_offsets(); - let nchunks = array.nchunks(); - let mut chunks = Vec::with_capacity(nchunks); - let mut final_take = BufferMut::::with_capacity(n); - final_take.push_n(0u64, n); + if let Some(&(index, _)) = pairs.last() { + let index = usize::try_from(index)?; + if index >= array.len() { + vortex_bail!(OutOfBounds: index, 0, array.len()); + } + } + let chunk_offsets = array.chunk_offset_values(); + let nchunks = array.nchunks(); + let mut flattener = ChunkFlattener::new(array.dtype(), pairs.len(), ctx.allocator()); + let mut final_take = BufferMut::::zeroed(n); let mut cursor = 0usize; let mut dedup_idx = 0u64; @@ -65,80 +113,942 @@ fn take_chunked( let chunk_end = chunk_offsets[chunk_idx + 1]; let chunk_len = chunk_end - chunk_start; let chunk_end_u64 = u64::try_from(chunk_end)?; - let range_end = cursor + pairs[cursor..].partition_point(|&(v, _)| v < chunk_end_u64); let chunk_pairs = &pairs[cursor..range_end]; if !chunk_pairs.is_empty() { - let mut local_indices: Vec = Vec::new(); - for (i, &(val, orig_pos)) in chunk_pairs.iter().enumerate() { - if cursor + i > 0 && val != pairs[cursor + i - 1].0 { + let mut local_indices = Vec::new(); + for (i, &(value, original_position)) in chunk_pairs.iter().enumerate() { + if cursor + i > 0 && value != pairs[cursor + i - 1].0 { dedup_idx += 1; } - let local = usize::try_from(val)? - chunk_start; - if local_indices.last() != Some(&local) { - local_indices.push(local); + let local_index = usize::try_from(value)? - chunk_start; + if local_indices.last() != Some(&local_index) { + local_indices.push(local_index); } - final_take[orig_pos] = dedup_idx; + final_take[original_position] = dedup_idx; } - let filter_mask = Mask::from_indices(chunk_len, local_indices); - chunks.push(array.chunk(chunk_idx).filter(filter_mask)?); + flattener.push( + array + .chunk(chunk_idx) + .filter(Mask::from_indices(chunk_len, local_indices))?, + ctx, + )?; } cursor = range_end; } - // SAFETY: every chunk came from a filter on a chunk with the same base dtype, - // unioned with the index nullability. - let flat = unsafe { ChunkedArray::new_unchecked(chunks, array.dtype().clone()) } - .into_array() - // TODO(joe): can we relax this. - .execute::(ctx)? - .into_array(); + let flat = flattener.finish(array.dtype(), ctx)?; + let take_validity = Validity::from_mask(indices_mask, indices.dtype().nullability()); + flat.take(PrimitiveArray::new(final_take.freeze(), take_validity).into_array()) +} - // 4. Single take to restore original order and expand duplicates. - // Carry the original index validity so null indices produce null outputs. - let take_validity = Validity::from_mask( - indices - .as_ref() - .validity()? - .execute_mask(indices.as_ref().len(), ctx)?, - indices.dtype().nullability(), - ); +fn valid_indices_are_monotonic(indices: &[u64], indices_mask: &Mask) -> bool { + let mut previous = None; + for (position, &index) in indices.iter().enumerate() { + if !indices_mask.value(position) { + continue; + } + if previous.is_some_and(|previous| index < previous) { + return false; + } + previous = Some(index); + } + true +} + +// TODO(joe): we want to return a chunked array ideally. +fn take_chunked( + array: ArrayView<'_, Chunked>, + indices: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let indices = indices + .cast(DType::Primitive(PType::U64, indices.dtype().nullability()))? + .execute::(ctx)?; + + let indices_mask = indices + .as_ref() + .validity()? + .execute_mask(indices.as_ref().len(), ctx)?; + let indices_values = indices.as_slice::(); + let n = indices_values.len(); + let chunk_offsets = array.chunk_offset_values(); + let nchunks = array.nchunks(); + + // Strictly increasing non-nullable indices select every row at most once and in order, so the + // per-chunk gathers concatenate straight into the result with no dedup or reorder take. + if !indices.as_ref().dtype().is_nullable() && indices_values.is_sorted_by(|a, b| a < b) { + return take_chunked_sorted(array, indices_values); + } + + // For a small unsorted fixed-size-list take over a small number of chunks, sorting has lower + // fixed overhead and lets the nested elements filter in source order. Larger, monotonic, and + // non-nested takes use bucketing. + if matches!(array.dtype(), DType::FixedSizeList(..)) + && n <= 64 + && nchunks <= 16 + && !valid_indices_are_monotonic(indices_values, &indices_mask) + { + return take_chunked_via_sort(array, &indices, indices_mask, ctx); + } + + // Route each valid index into its source chunk. Within each bucket, preserve the request order + // so the taken chunks can be assembled and then restored to the original cross-chunk order. + let mut buckets = vec![Vec::<(u64, usize)>::new(); nchunks]; + let mut monotonic = true; + let mut last_index = None; + let mut sorted_chunk_idx = 0; + + for (original_position, &index) in indices_values.iter().enumerate() { + if !indices_mask.value(original_position) { + continue; + } + + let index = usize::try_from(index)?; + if index >= array.len() { + vortex_bail!(OutOfBounds: index, 0, array.len()); + } + + let still_sorted = last_index.is_none_or(|last_index| index >= last_index); + let chunk_idx = if monotonic && still_sorted { + while chunk_offsets[sorted_chunk_idx + 1] <= index { + sorted_chunk_idx += 1; + } + sorted_chunk_idx + } else { + monotonic = false; + chunk_offsets.partition_point(|&offset| offset <= index) - 1 + }; + last_index = Some(index); + + let local_index = u64::try_from(index - chunk_offsets[chunk_idx])?; + buckets[chunk_idx].push((local_index, original_position)); + } + + let mut flattener = + ChunkFlattener::new(array.dtype(), indices_mask.true_count(), ctx.allocator()); + let mut final_take = + (!monotonic || indices.dtype().is_nullable()).then(|| BufferMut::::zeroed(n)); + let mut grouped_position = 0u64; + + for (chunk_idx, bucket) in buckets.into_iter().enumerate() { + if bucket.is_empty() { + continue; + } + + let mut local_indices = BufferMut::::with_capacity(bucket.len()); + for (local_index, original_position) in bucket { + local_indices.push(local_index); + if let Some(final_take) = &mut final_take { + final_take[original_position] = grouped_position; + } + grouped_position += 1; + } + + let local_indices = + PrimitiveArray::new(local_indices.freeze(), Validity::NonNullable).into_array(); + flattener.push(array.chunk(chunk_idx).take(local_indices)?, ctx)?; + } + + // TODO(joe): can we relax this. + let flat = flattener.finish(array.dtype(), ctx)?; + + // Non-nullable monotonic indices are already in the same order as the assembled chunks, so no + // final reorder is needed. + let Some(final_take) = final_take else { + return Ok(flat); + }; + + // A single take restores original order and expands duplicates. Carrying the original index + // validity makes null indices produce null outputs. + let take_validity = Validity::from_mask(indices_mask, indices.dtype().nullability()); flat.take(PrimitiveArray::new(final_take.freeze(), take_validity).into_array()) } +/// Take for strictly increasing, non-nullable indices: every row is selected at most once and in +/// order, so the per-chunk gathers concatenate directly into the result with no bucketing, dedup, +/// or reorder take. +fn take_chunked_sorted( + array: ArrayView<'_, Chunked>, + indices_values: &[u64], +) -> VortexResult { + let chunk_offsets = array.chunk_offset_values(); + let mut chunks = Vec::new(); + let mut cursor = 0usize; + + // Skip straight to the chunk holding the first index instead of walking every leading chunk; + // `<=` steps past empty chunks sharing the same offset. + let first_chunk = match indices_values.first() { + Some(&first) => { + let first = usize::try_from(first)?; + chunk_offsets + .partition_point(|&offset| offset <= first) + .saturating_sub(1) + } + None => 0, + }; + + for chunk_idx in first_chunk..array.nchunks() { + if cursor == indices_values.len() { + break; + } + let chunk_start = chunk_offsets[chunk_idx]; + let chunk_end = chunk_offsets[chunk_idx + 1]; + let chunk_end_u64 = u64::try_from(chunk_end)?; + + let range_end = cursor + indices_values[cursor..].partition_point(|&v| v < chunk_end_u64); + if range_end > cursor { + let chunk_start_u64 = u64::try_from(chunk_start)?; + let mut local_indices = BufferMut::::with_capacity(range_end - cursor); + for &val in &indices_values[cursor..range_end] { + local_indices.push(val - chunk_start_u64); + } + let local_indices = + PrimitiveArray::new(local_indices.freeze(), Validity::NonNullable).into_array(); + chunks.push(array.chunk(chunk_idx).take(local_indices)?); + } + cursor = range_end; + } + if cursor != indices_values.len() { + vortex_bail!( + OutOfBounds: usize::try_from(indices_values[cursor])?, 0, array.as_ref().len() + ); + } + + if chunks.len() == 1 { + return Ok(chunks.swap_remove(0)); + } + // SAFETY: every chunk is a take of a chunk of `array` with non-nullable indices, which + // leaves the dtype unchanged. + Ok(unsafe { ChunkedArray::new_unchecked(chunks, array.dtype().clone()) }.into_array()) +} + impl TakeExecute for Chunked { fn take( array: ArrayView<'_, Chunked>, indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { + if array.nchunks() == 1 { + return array.chunk(0).take(indices.clone()).map(Some); + } + + if let Some(taken) = take_chunked_fsl(array, indices, ctx)? { + return Ok(Some(taken)); + } + + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_piecewise_chunked(array, piecewise_indices, ctx)? + { + return Ok(Some(taken)); + } + take_chunked(array, indices, ctx).map(Some) } } +/// Rewrites take over a chunked fixed-size list array as take over a single +/// [`FixedSizeListArray`]. The swizzle is structural: FSL-encoded chunks execute to themselves, +/// so their elements chain zero-copy, and the FSL take then gathers per-chunk element runs +/// instead of taking each chunk. +fn take_chunked_fsl( + array: ArrayView<'_, Chunked>, + indices: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + if !matches!(array.dtype(), DType::FixedSizeList(..)) { + return Ok(None); + } + + let fsl = array.as_ref().clone().execute::(ctx)?; + fsl.into_array().take(indices.clone()).map(Some) +} + +/// A per-chunk gather plan: chunk-local sub-piece runs to take from one chunk, in output order. +#[derive(Default)] +struct ChunkGather { + starts: Vec, + lengths: Vec, + total: usize, +} + +/// Take for [`PiecewiseSequence`] indices with unit multipliers. +/// +/// Each piece is a contiguous index run, so instead of expanding one index per element, pieces +/// are split at chunk boundaries and gathered from each chunk with a chunk-local +/// `PiecewiseSequenceArray`. Sub-pieces that visit chunks in non-decreasing order concatenate +/// directly into the result; otherwise the gathered chunks are flattened once and a second +/// piecewise take restores output order. +fn take_piecewise_chunked( + array: ArrayView<'_, Chunked>, + indices: ArrayView<'_, PiecewiseSequence>, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { + return Ok(None); + }; + + let output_len = indices.as_ref().len(); + let starts: Vec = match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + starts + .as_slice::() + .iter() + .map(|&start| start.as_()) + .collect() + }); + let lengths: Vec = match lengths { + Columnar::Constant(lengths) => vec![constant_unsigned_usize(&lengths); starts.len()], + Columnar::Canonical(lengths) => { + let lengths = lengths.into_primitive(); + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + lengths + .as_slice::() + .iter() + .map(|&length| length.as_()) + .collect() + }) + } + }; + + let array_len = array.as_ref().len(); + let chunk_offsets = array.chunk_offset_values(); + let nchunks = array.nchunks(); + + let mut plans: Vec = Vec::new(); + plans.resize_with(nchunks, ChunkGather::default); + // Sub-pieces in output order: (chunk index, offset within that chunk's gathered output, run + // length). + let mut sub_pieces: Vec<(usize, usize, usize)> = Vec::with_capacity(starts.len()); + let mut monotonic = true; + let mut prev_chunk = 0usize; + let mut total_len = 0usize; + + for (&start, &length) in starts.iter().zip_eq(&lengths) { + if length == 0 { + continue; + } + let end = start + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + if end > array_len { + vortex_bail!(OutOfBounds: end - 1, 0, array_len); + } + total_len = total_len + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + + // Locate the chunk containing `start`; `<=` skips empty chunks sharing the same offset. + let mut chunk_idx = chunk_offsets.partition_point(|&offset| offset <= start) - 1; + let mut cursor = start; + let mut remaining = length; + while remaining > 0 { + while chunk_offsets[chunk_idx + 1] <= cursor { + chunk_idx += 1; + } + let run = remaining.min(chunk_offsets[chunk_idx + 1] - cursor); + monotonic &= chunk_idx >= prev_chunk; + prev_chunk = chunk_idx; + + let plan = &mut plans[chunk_idx]; + // Consecutive sub-pieces in the same chunk stay adjacent in its gathered output, so + // they merge into one reorder run. + match sub_pieces.last_mut() { + Some((last_chunk, _, last_run)) if *last_chunk == chunk_idx => *last_run += run, + _ => sub_pieces.push((chunk_idx, plan.total, run)), + } + // A run that continues exactly where the chunk's previous run ended extends it, so + // contiguous spans gather as one run. + let local_start = (cursor - chunk_offsets[chunk_idx]) as u64; + match (plan.starts.last(), plan.lengths.last_mut()) { + (Some(&prev_start), Some(prev_len)) if prev_start + *prev_len == local_start => { + *prev_len += run as u64; + } + _ => { + plan.starts.push(local_start); + plan.lengths.push(run as u64); + } + } + plan.total += run; + + cursor += run; + remaining -= run; + } + } + + vortex_ensure!( + total_len == output_len, + "PiecewiseSequenceArray expanded length {total_len} does not match declared length {output_len}" + ); + + // Chunks visited in order: the per-chunk gathers already concatenate into the result. + if monotonic { + let mut gathered = Vec::new(); + for (chunk_idx, plan) in plans.into_iter().enumerate() { + if plan.starts.is_empty() { + continue; + } + gathered.push(gather_chunk(array.chunk(chunk_idx), plan)?); + } + let result = if gathered.len() == 1 { + gathered.swap_remove(0) + } else { + // SAFETY: every gathered chunk is a take of a chunk with dtype `array.dtype()`, and + // the non-nullable piecewise indices leave the dtype unchanged. + unsafe { ChunkedArray::new_unchecked(gathered, array.dtype().clone()) }.into_array() + }; + return Ok(Some(result)); + } + + // Out-of-order sub-pieces: flatten the per-chunk gathers, then take the sub-piece runs from + // the flattened result in output order. + let mut bases = vec![0usize; nchunks]; + let mut running = 0usize; + let mut flattener = ChunkFlattener::new(array.dtype(), output_len, ctx.allocator()); + for (chunk_idx, plan) in plans.into_iter().enumerate() { + if plan.starts.is_empty() { + continue; + } + bases[chunk_idx] = running; + running += plan.total; + flattener.push(gather_chunk(array.chunk(chunk_idx), plan)?, ctx)?; + } + let flat = flattener.finish(array.dtype(), ctx)?; + + let mut reorder_starts = Vec::with_capacity(sub_pieces.len()); + let mut reorder_lengths = Vec::with_capacity(sub_pieces.len()); + for &(chunk_idx, offset, run) in &sub_pieces { + reorder_starts.push((bases[chunk_idx] + offset) as u64); + reorder_lengths.push(run as u64); + } + flat.take(contiguous_runs(reorder_starts, reorder_lengths, output_len)) + .map(Some) +} + +/// Gathers a chunk's planned runs: a single run is a zero-copy slice, multiple runs become a +/// chunk-local piecewise take. +fn gather_chunk(chunk: &ArrayRef, plan: ChunkGather) -> VortexResult { + if let [start] = plan.starts.as_slice() { + let start = usize::try_from(*start)?; + return chunk.slice(start..start + plan.total); + } + chunk.take(contiguous_runs(plan.starts, plan.lengths, plan.total)) +} + +/// Builds a `PiecewiseSequenceArray` of contiguous (unit multiplier) runs whose lengths sum to +/// `total`. +fn contiguous_runs(starts: Vec, lengths: Vec, total: usize) -> ArrayRef { + let count = starts.len(); + debug_assert_eq!(count, lengths.len()); + let starts = PrimitiveArray::new(Buffer::from(starts), Validity::NonNullable).into_array(); + let lengths = match lengths.first() { + Some(&first) if lengths.iter().all(|&length| length == first) => { + ConstantArray::new(first, count).into_array() + } + _ => PrimitiveArray::new(Buffer::from(lengths), Validity::NonNullable).into_array(), + }; + let multipliers = ConstantArray::new(1u64, count).into_array(); + // SAFETY: starts, lengths, and multipliers are non-nullable u64 arrays of equal length, and + // `total` is the sum of the lengths. + unsafe { PiecewiseSequenceArray::new_unchecked(starts, lengths, multipliers, total) } + .into_array() +} + #[cfg(test)] -mod test { +mod tests { + use vortex_buffer::Buffer; use vortex_buffer::bitbuffer; use vortex_buffer::buffer; use vortex_error::VortexResult; + use crate::ArrayRef; + use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; + use crate::arrays::ConstantArray; + use crate::arrays::FixedSizeListArray; + use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::assert_arrays_eq; use crate::compute::conformance::take::test_take_conformance; + use crate::dtype::DType; use crate::dtype::FieldNames; use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::scalar::Scalar; use crate::validity::Validity; + fn chunked_i32() -> VortexResult { + ChunkedArray::try_new( + vec![ + buffer![0i32, 1, 2, 3, 4].into_array(), + buffer![5i32, 6, 7, 8, 9].into_array(), + buffer![10i32, 11, 12, 13, 14].into_array(), + ], + DType::Primitive(PType::I32, Nullability::NonNullable), + ) + } + + fn contiguous_pieces(starts: &[u64], lengths: &[u64]) -> VortexResult { + let len = usize::try_from(lengths.iter().sum::())?; + Ok(PiecewiseSequenceArray::try_new( + starts.iter().copied().collect::>().into_array(), + lengths + .iter() + .copied() + .collect::>() + .into_array(), + ConstantArray::new(1u64, starts.len()).into_array(), + len, + )? + .into_array()) + } + + #[test] + fn test_take_piecewise_monotonic_spanning_chunks() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = chunked_i32()?; + + // The second piece crosses the first chunk boundary, the third spans the last two chunks. + let indices = contiguous_pieces(&[1, 4, 9], &[3, 4, 6])?; + let result = arr.take(indices)?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6, 7, 9, 10, 11, 12, 13, 14]), + &mut ctx + ); + Ok(()) + } + + #[test] + fn test_take_piecewise_interleaved() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = chunked_i32()?; + + // Pieces visit chunks out of order, forcing the reorder take. + let indices = contiguous_pieces(&[12, 2, 7, 0], &[3, 2, 3, 1])?; + let result = arr.take(indices)?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([12i32, 13, 14, 2, 3, 7, 8, 9, 0]), + &mut ctx + ); + Ok(()) + } + + #[test] + fn test_take_piecewise_whole_array() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = chunked_i32()?; + + let indices = contiguous_pieces(&[0], &[15])?; + let result = arr.take(indices)?; + + assert_arrays_eq!(result, PrimitiveArray::from_iter(0i32..15), &mut ctx); + Ok(()) + } + + #[test] + fn test_take_piecewise_across_empty_chunk() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = ChunkedArray::try_new( + vec![ + buffer![0i32, 1, 2, 3, 4].into_array(), + PrimitiveArray::empty::(Nullability::NonNullable).into_array(), + buffer![5i32, 6, 7, 8, 9].into_array(), + ], + DType::Primitive(PType::I32, Nullability::NonNullable), + )?; + + let indices = contiguous_pieces(&[3], &[4])?; + let result = arr.take(indices)?; + + assert_arrays_eq!(result, PrimitiveArray::from_iter([3i32, 4, 5, 6]), &mut ctx); + Ok(()) + } + + #[test] + fn test_take_piecewise_adjacent_pieces_merge() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = chunked_i32()?; + + // The first two pieces are contiguous within one chunk, so they merge into a single run + // that gathers as a zero-copy slice. + let indices = contiguous_pieces(&[0, 2, 7], &[2, 2, 2])?; + let result = arr.take(indices)?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([0i32, 1, 2, 3, 7, 8]), + &mut ctx + ); + Ok(()) + } + + #[test] + fn test_take_piecewise_same_chunk_out_of_order() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = chunked_i32()?; + + // Non-contiguous pieces in the same chunk must stay separate runs in piece order. + let indices = contiguous_pieces(&[2, 0], &[2, 2])?; + let result = arr.take(indices)?; + + assert_arrays_eq!(result, PrimitiveArray::from_iter([2i32, 3, 0, 1]), &mut ctx); + Ok(()) + } + + #[test] + fn test_take_piecewise_zero_length_pieces() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = chunked_i32()?; + + let indices = contiguous_pieces(&[5, 0, 2], &[0, 4, 0])?; + let result = arr.take(indices)?; + + assert_arrays_eq!(result, PrimitiveArray::from_iter([0i32, 1, 2, 3]), &mut ctx); + Ok(()) + } + + #[test] + fn test_take_piecewise_out_of_bounds() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = chunked_i32()?; + + let indices = contiguous_pieces(&[12], &[5])?; + let result = arr + .take(indices) + .and_then(|taken| taken.execute::(&mut ctx)); + + assert!(result.is_err()); + Ok(()) + } + + #[test] + fn test_take_piecewise_non_unit_multiplier() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = chunked_i32()?; + + // Multiplier 2 falls back to the generic path but must stay correct. + let indices = PiecewiseSequenceArray::try_new( + buffer![0u64, 1].into_array(), + buffer![5u64, 3].into_array(), + buffer![2u64, 4].into_array(), + 8, + )? + .into_array(); + let result = arr.take(indices)?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([0i32, 2, 4, 6, 8, 1, 5, 9]), + &mut ctx + ); + Ok(()) + } + + #[test] + fn test_take_piecewise_nullable_values() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = ChunkedArray::try_new( + vec![ + PrimitiveArray::from_option_iter([Some(0i32), None, Some(2)]).into_array(), + PrimitiveArray::from_option_iter([None, Some(4i32), Some(5)]).into_array(), + ], + DType::Primitive(PType::I32, Nullability::Nullable), + )?; + + let indices = contiguous_pieces(&[4, 1], &[2, 3])?; + let result = arr.take(indices)?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_option_iter([Some(4i32), Some(5), None, Some(2), None]), + &mut ctx + ); + Ok(()) + } + + #[test] + fn test_take_fsl_over_chunked_elements() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // Chunk boundaries at 8 and 13 do not line up with the list size of 3. + let elements = ChunkedArray::try_new( + vec![ + PrimitiveArray::from_iter(0i32..8).into_array(), + PrimitiveArray::from_iter(8i32..13).into_array(), + PrimitiveArray::from_iter(13i32..18).into_array(), + ], + DType::Primitive(PType::I32, Nullability::NonNullable), + )? + .into_array(); + let fsl = FixedSizeListArray::try_new(elements, 3, Validity::NonNullable, 6)?; + + let indices = buffer![4u64, 0, 2, 5, 1, 4].into_array(); + let result = fsl.take(indices.clone())?; + let expected = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..18).into_array(), + 3, + Validity::NonNullable, + 6, + )? + .take(indices)?; + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_take_fsl_over_chunked_elements_conformance() -> VortexResult<()> { + let elements = ChunkedArray::try_new( + vec![ + PrimitiveArray::from_iter(0i32..8).into_array(), + PrimitiveArray::from_iter(8i32..13).into_array(), + PrimitiveArray::from_iter(13i32..18).into_array(), + ], + DType::Primitive(PType::I32, Nullability::NonNullable), + )? + .into_array(); + let fsl = FixedSizeListArray::try_new(elements, 3, Validity::NonNullable, 6)?; + test_take_conformance( + &fsl.into_array(), + &mut array_session().create_execution_ctx(), + ); + Ok(()) + } + + #[test] + fn test_take_chunked_fsl() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let c0 = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..6).into_array(), + 2, + Validity::NonNullable, + 3, + )?; + let c1 = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(6i32..10).into_array(), + 2, + Validity::NonNullable, + 2, + )?; + let dtype = c0.dtype().clone(); + let arr = ChunkedArray::try_new(vec![c0.into_array(), c1.into_array()], dtype)?; + + let indices = buffer![4u64, 0, 3, 3, 1].into_array(); + let result = arr.take(indices.clone())?; + let expected = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..10).into_array(), + 2, + Validity::NonNullable, + 5, + )? + .take(indices)?; + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_take_chunked_fsl_nullable() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let c0 = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..6).into_array(), + 2, + Validity::Array(bitbuffer![1 0 1].into_array()), + 3, + )?; + let c1 = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(6i32..10).into_array(), + 2, + Validity::AllValid, + 2, + )?; + let dtype = c0.dtype().clone(); + let arr = ChunkedArray::try_new(vec![c0.into_array(), c1.into_array()], dtype)?; + + let indices = + PrimitiveArray::from_option_iter([Some(4u64), None, Some(0), Some(1)]).into_array(); + let result = arr.take(indices.clone())?; + let expected = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..10).into_array(), + 2, + Validity::Array(bitbuffer![1 0 1 1 1].into_array()), + 5, + )? + .take(indices)?; + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_take_chunked_fsl_non_fsl_chunk() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let c0 = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..6).into_array(), + 2, + Validity::NonNullable, + 3, + )?; + let c1 = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(6i32..10).into_array(), + 2, + Validity::NonNullable, + 2, + )?; + let dtype = c0.dtype().clone(); + // Wrap one chunk in a nested chunked array so it is not FSL-encoded; the dtype-based + // swizzle must still handle it. + let nested = ChunkedArray::try_new(vec![c1.into_array()], dtype.clone())?; + let arr = ChunkedArray::try_new(vec![c0.into_array(), nested.into_array()], dtype)?; + + let indices = buffer![4u64, 0, 3, 1].into_array(); + let result = arr.take(indices.clone())?; + let expected = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..10).into_array(), + 2, + Validity::NonNullable, + 5, + )? + .take(indices)?; + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_take_chunked_fsl_constant_chunk() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let c0 = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..6).into_array(), + 2, + Validity::AllValid, + 3, + )?; + let dtype = c0.dtype().clone(); + // A constant chunk is only logically FSL; the swizzle canonicalizes it in place. + let c1 = ConstantArray::new(Scalar::null(dtype.clone()), 2).into_array(); + let arr = ChunkedArray::try_new(vec![c0.into_array(), c1], dtype)?; + + let indices = buffer![4u64, 0, 3, 1].into_array(); + let result = arr.take(indices.clone())?; + let expected = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..10).into_array(), + 2, + Validity::Array(bitbuffer![1 1 1 0 0].into_array()), + 5, + )? + .take(indices)?; + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_take_piecewise_chunked_fsl() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let c0 = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..6).into_array(), + 2, + Validity::NonNullable, + 3, + )?; + let c1 = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(6i32..10).into_array(), + 2, + Validity::NonNullable, + 2, + )?; + let dtype = c0.dtype().clone(); + let arr = ChunkedArray::try_new(vec![c0.into_array(), c1.into_array()], dtype)?; + let reference = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..10).into_array(), + 2, + Validity::NonNullable, + 5, + )?; + + // A monotonic run crossing the chunk boundary. + let indices = contiguous_pieces(&[1], &[3])?; + assert_arrays_eq!( + arr.take(indices.clone())?, + reference.take(indices)?, + &mut ctx + ); + + // Out-of-order pieces forcing the reorder take. + let indices = contiguous_pieces(&[3, 0], &[2, 2])?; + assert_arrays_eq!( + arr.take(indices.clone())?, + reference.take(indices)?, + &mut ctx + ); + Ok(()) + } + + #[test] + fn test_take_chunked_struct_nested_flatten() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let s0 = StructArray::try_new( + ["a"].into(), + vec![buffer![0i32, 1, 2].into_array()], + 3, + Validity::NonNullable, + )?; + let s1 = StructArray::try_new( + ["a"].into(), + vec![buffer![3i32, 4].into_array()], + 2, + Validity::NonNullable, + )?; + let dtype = s0.dtype().clone(); + let arr = ChunkedArray::try_new(vec![s0.into_array(), s1.into_array()], dtype)?; + + let result = arr.take(buffer![4u64, 0, 2, 4].into_array())?; + let expected = StructArray::try_new( + ["a"].into(), + vec![buffer![4i32, 0, 2, 4].into_array()], + 4, + Validity::NonNullable, + )?; + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_take_chunked_fsl_conformance() -> VortexResult<()> { + let c0 = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..6).into_array(), + 2, + Validity::NonNullable, + 3, + )?; + let c1 = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(6i32..10).into_array(), + 2, + Validity::NonNullable, + 2, + )?; + let dtype = c0.dtype().clone(); + let arr = ChunkedArray::try_new(vec![c0.into_array(), c1.into_array()], dtype)?; + test_take_conformance( + &arr.into_array(), + &mut array_session().create_execution_ctx(), + ); + Ok(()) + } + #[test] fn test_take() { let mut ctx = array_session().create_execution_ctx(); @@ -264,6 +1174,65 @@ mod test { Ok(()) } + #[test] + fn test_take_shuffled_duplicates_with_empty_chunks() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let empty = PrimitiveArray::empty::(Nullability::NonNullable).into_array(); + let arr = ChunkedArray::try_new( + vec![ + empty.clone(), + buffer![0i32, 1].into_array(), + empty.clone(), + buffer![2i32, 3].into_array(), + empty, + ], + PrimitiveArray::empty::(Nullability::NonNullable) + .dtype() + .clone(), + )?; + + let result = arr.take(buffer![3u64, 0, 2, 0, 3, 1, 2].into_array())?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([3i32, 0, 2, 0, 3, 1, 2]), + &mut ctx + ); + Ok(()) + } + + #[test] + fn test_take_small_shuffled_fixed_size_lists() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let first = FixedSizeListArray::new( + buffer![0i32, 1, 2, 3].into_array(), + 2, + Validity::NonNullable, + 2, + ) + .into_array(); + let second = FixedSizeListArray::new( + buffer![4i32, 5, 6, 7].into_array(), + 2, + Validity::NonNullable, + 2, + ) + .into_array(); + let dtype = first.dtype().clone(); + let array = ChunkedArray::try_new(vec![first, second], dtype)?; + + let result = array.take(buffer![3u64, 0, 2, 1, 3].into_array())?; + let expected = FixedSizeListArray::new( + buffer![6i32, 7, 0, 1, 4, 5, 2, 3, 6, 7].into_array(), + 2, + Validity::NonNullable, + 5, + ); + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + #[test] fn test_take_shuffled_large() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); @@ -289,10 +1258,7 @@ mod test { indices.swap(i, j); } - let indices_arr = PrimitiveArray::new( - vortex_buffer::Buffer::from(indices.clone()), - Validity::NonNullable, - ); + let indices_arr = PrimitiveArray::new(Buffer::from(indices.clone()), Validity::NonNullable); let result = arr.take(indices_arr.into_array())?; // Verify every element. @@ -340,6 +1306,54 @@ mod test { Ok(()) } + #[test] + fn test_take_sorted_indices() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = chunked_i32()?; + + // Strictly increasing indices spanning chunks hit the sorted fast path. + let indices = buffer![1u64, 4, 5, 9, 10, 14].into_array(); + let result = arr.take(indices)?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([1i32, 4, 5, 9, 10, 14]), + &mut ctx + ); + Ok(()) + } + + #[test] + fn test_take_out_of_bounds() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = chunked_i32()?; + + // Sorted indices hit the fast path, unsorted the generic path; both must fail. + for indices in [buffer![0u64, 100], buffer![100u64, 0]] { + let result = arr + .take(indices.into_array()) + .and_then(|taken| taken.execute::(&mut ctx)); + assert!(result.is_err(), "expected out-of-bounds error"); + } + Ok(()) + } + + #[test] + fn test_take_all_null_indices() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = chunked_i32()?; + + let indices = PrimitiveArray::from_option_iter([None::, None]).into_array(); + let result = arr.take(indices)?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_option_iter([None::, None]), + &mut ctx + ); + Ok(()) + } + #[test] fn test_take_chunked_conformance() { let a = buffer![1i32, 2, 3].into_array(); diff --git a/vortex-array/src/arrays/chunked/mod.rs b/vortex-array/src/arrays/chunked/mod.rs index 4bd6cb43b7e..3bd361d66a1 100644 --- a/vortex-array/src/arrays/chunked/mod.rs +++ b/vortex-array/src/arrays/chunked/mod.rs @@ -4,6 +4,8 @@ mod array; pub use array::ChunkedArrayExt; pub use array::ChunkedData; +pub use array::ChunkedSlots; +pub use array::ChunkedSlotsView; pub use vtable::ChunkedArray; pub(crate) mod compute; diff --git a/vortex-array/src/arrays/chunked/tests.rs b/vortex-array/src/arrays/chunked/tests.rs index a32d03e207b..e0755108e9c 100644 --- a/vortex-array/src/arrays/chunked/tests.rs +++ b/vortex-array/src/arrays/chunked/tests.rs @@ -24,7 +24,7 @@ use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::dict_test::gen_dict_primitive_chunks; use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; -use crate::builders::builder_with_capacity; +use crate::builders::builder_with_capacity_in; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -54,7 +54,8 @@ fn builder_kernel_path_canonicalizes_primitive_chunks() { let dtype = array.dtype().clone(); let len = array.len(); - let builder = builder_with_capacity(&dtype, len); + let builder = + builder_with_capacity_in(&dtype, len, vortex_buffer::BufferAllocatorRef::static_ref()); // Clone the array into the builder path — the test also holds `array` so refcount > 1 on // entry, which previously caused `take_slot_unchecked` to silently keep slots populated. let mut builder = execute_into_builder(array.clone(), builder, &mut ctx).unwrap(); @@ -93,7 +94,8 @@ fn builder_kernel_nested_chunked_of_chunked() { let dtype = outer.dtype().clone(); let len = outer.len(); - let builder = builder_with_capacity(&dtype, len); + let builder = + builder_with_capacity_in(&dtype, len, vortex_buffer::BufferAllocatorRef::static_ref()); let mut builder = execute_into_builder(outer, builder, &mut ctx).unwrap(); let output = builder.finish(); @@ -120,13 +122,15 @@ fn builder_kernel_path_repeated_shared_chunked_dict_execution() { .into_array(); let first = { - let builder = builder_with_capacity(&dtype, len); + let builder = + builder_with_capacity_in(&dtype, len, vortex_buffer::BufferAllocatorRef::static_ref()); let mut builder = execute_into_builder(array.clone(), builder, &mut ctx).unwrap(); builder.finish() }; let second = { - let builder = builder_with_capacity(&dtype, len); + let builder = + builder_with_capacity_in(&dtype, len, vortex_buffer::BufferAllocatorRef::static_ref()); let mut builder = execute_into_builder(array, builder, &mut ctx).unwrap(); builder.finish() }; @@ -177,7 +181,11 @@ fn execute_path_nested_chunked_dict_of_dict_into_canonical() { let keep_alive = outer.clone(); let expected = { - let mut builder = builder_with_capacity(outer.dtype(), outer.len()); + let mut builder = builder_with_capacity_in( + outer.dtype(), + outer.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); inner_1 .append_to_builder(builder.as_mut(), &mut ctx) .unwrap(); @@ -213,7 +221,7 @@ fn with_slot_rewrites_chunk_and_offsets() { let array = array.as_::(); assert_eq!(array.nchunks(), 3); - assert_eq!(array.chunk_offsets(), [0, 3, 6, 9]); + assert_eq!(array.chunk_offset_values(), [0, 3, 6, 9]); assert_arrays_eq!( array.chunk(0).clone(), PrimitiveArray::from_iter([1u64, 2, 3]), @@ -376,10 +384,12 @@ pub fn pack_nested_structs() -> VortexResult<()> { )? .into_array(); let canonical_struct = chunked.execute::(&mut ctx)?; - let canonical_varbin = canonical_struct.unmasked_fields()[0] + let canonical_varbin = canonical_struct + .unmasked_field(0) .clone() .execute::(&mut ctx)?; - let original_varbin = struct_array.unmasked_fields()[0] + let original_varbin = struct_array + .unmasked_field(0) .clone() .execute::(&mut ctx)?; assert_arrays_eq!(original_varbin, canonical_varbin, &mut ctx); diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index 12c3bbcd9e4..a521fd0e6cc 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools as _; -use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -21,16 +20,15 @@ use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::VariantArray; use crate::arrays::chunked::ChunkedArrayExt; -use crate::arrays::fixed_size_list::FixedSizeListArrayExt; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::ListViewRebuildMode; -use crate::arrays::variant::VariantArrayExt; +use crate::arrays::variant::VariantArraySlotsExt; use crate::builders::builder_with_capacity_in; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; -use crate::memory::HostAllocatorExt; use crate::validity::Validity; pub(super) fn _canonicalize( @@ -71,7 +69,7 @@ pub(super) fn _canonicalize( } DType::Variant(_) => Canonical::Variant(pack_variant_chunks(owned_chunks, ctx)?), _ => { - let mut builder = builder_with_capacity_in(ctx.allocator(), array.dtype(), array.len()); + let mut builder = builder_with_capacity_in(array.dtype(), array.len(), ctx.allocator()); array.array().append_to_builder(builder.as_mut(), ctx)?; builder.finish_into_canonical(ctx) } @@ -102,11 +100,13 @@ fn pack_variant_chunks( .try_collect()?; let outer_dtype = variant_chunks[0].dtype().clone(); - let core_chunks = variant_chunks - .iter() - .map(|chunk| chunk.core_storage().clone()) - .collect(); - let core_storage = ChunkedArray::try_new(core_chunks, outer_dtype)?.into_array(); + let core_storage = ChunkedArray::try_new( + variant_chunks + .iter() + .map(|chunk| chunk.core_storage().clone()), + outer_dtype, + )? + .into_array(); let shredded = match variant_chunks[0].shredded() { None => { @@ -178,10 +178,10 @@ fn swizzle_list_chunks( // We (somewhat arbitrarily) choose `u64` for our offsets and sizes here. These can always be // narrowed later by the compressor. let allocator = ctx.allocator(); - let mut offsets = allocator.allocate_typed::(len)?; - let mut sizes = allocator.allocate_typed::(len)?; - let offsets_out: &mut [u64] = offsets.as_mut_slice_typed::()?; - let sizes_slice_out: &mut [u64] = sizes.as_mut_slice_typed::()?; + let mut offsets = allocator.zeroed::(len); + let mut sizes = allocator.zeroed::(len); + let offsets_out = offsets.as_mut_slice(); + let sizes_slice_out = sizes.as_mut_slice(); let mut next_list = 0usize; for chunk in chunks { @@ -227,16 +227,8 @@ fn swizzle_list_chunks( unsafe { ChunkedArray::new_unchecked(list_elements_chunks, elem_dtype.clone()) } .into_array(); - let offsets = PrimitiveArray::new( - Buffer::::from_byte_buffer(offsets.freeze()), - Validity::NonNullable, - ) - .into_array(); - let sizes = PrimitiveArray::new( - Buffer::::from_byte_buffer(sizes.freeze()), - Validity::NonNullable, - ) - .into_array(); + let offsets = PrimitiveArray::new(offsets.freeze(), Validity::NonNullable).into_array(); + let sizes = PrimitiveArray::new(sizes.freeze(), Validity::NonNullable).into_array(); // SAFETY: // - `offsets` and `sizes` are non-nullable u64 arrays of the same length @@ -284,11 +276,17 @@ fn swizzle_fixed_size_list_chunks( #[cfg(test)] mod tests { + use std::alloc::Layout; + use std::ptr::NonNull; use std::sync::Arc; use std::sync::LazyLock; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; + use allocator_api2::alloc::AllocError; + use allocator_api2::alloc::Allocator; + use allocator_api2::alloc::Global; + use vortex_buffer::BufferAllocatorRef; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -309,17 +307,14 @@ mod tests { use crate::arrays::VarBinViewArray; use crate::arrays::VariantArray; use crate::arrays::struct_::StructArrayExt; - use crate::arrays::variant::VariantArrayExt; + use crate::arrays::variant::VariantArraySlotsExt; use crate::assert_arrays_eq; use crate::dtype::DType::List; use crate::dtype::DType::Primitive; use crate::dtype::DType::Variant as VariantDType; use crate::dtype::Nullability::NonNullable; use crate::dtype::PType::I32; - use crate::memory::DefaultHostAllocator; - use crate::memory::HostAllocator; use crate::memory::MemorySessionExt; - use crate::memory::WritableHostBuffer; use crate::scalar::Scalar; use crate::validity::Validity; @@ -331,14 +326,16 @@ mod tests { allocations: Arc, } - impl HostAllocator for CountingAllocator { - fn allocate( - &self, - len: usize, - alignment: vortex_buffer::Alignment, - ) -> VortexResult { + // SAFETY: this forwards memory operations to Global and only counts allocations. + unsafe impl Allocator for CountingAllocator { + fn allocate(&self, layout: Layout) -> Result, AllocError> { self.allocations.fetch_add(1, Ordering::Relaxed); - DefaultHostAllocator.allocate(len, alignment) + Global.allocate(layout) + } + + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + // SAFETY: ptr and layout came from Global. + unsafe { Global.deallocate(ptr, layout) } } } @@ -347,12 +344,13 @@ mod tests { } fn variant_core(values: impl IntoIterator) -> VortexResult { - let chunks = values - .into_iter() - .map(|value| ConstantArray::new(variant_scalar(value), 1).into_array()) - .collect(); - - Ok(ChunkedArray::try_new(chunks, VariantDType(NonNullable))?.into_array()) + Ok(ChunkedArray::try_new( + values + .into_iter() + .map(|value| ConstantArray::new(variant_scalar(value), 1).into_array()), + VariantDType(NonNullable), + )? + .into_array()) } fn variant_chunk(values: impl IntoIterator) -> VortexResult { @@ -660,9 +658,10 @@ mod tests { #[test] fn list_canonicalize_uses_memory_session_allocator() { let allocations = Arc::new(AtomicUsize::new(0)); - let session = crate::array_session().with_allocator(Arc::new(CountingAllocator { - allocations: Arc::clone(&allocations), - })); + let session = + crate::array_session().with_allocator(BufferAllocatorRef::new(CountingAllocator { + allocations: Arc::clone(&allocations), + })); let mut ctx = session.create_execution_ctx(); let l1 = ListArray::try_new( diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index 2eabd883742..59e89f031e1 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -32,8 +32,7 @@ use crate::array::with_empty_buffers; use crate::arrays::PrimitiveArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::chunked::ChunkedData; -use crate::arrays::chunked::array::CHUNK_OFFSETS_SLOT; -use crate::arrays::chunked::array::CHUNKS_OFFSET; +use crate::arrays::chunked::array::ChunkedSlots; use crate::arrays::chunked::compute::rules::PARENT_RULES; use crate::arrays::chunked::vtable::canonical::_canonicalize; use crate::buffer::BufferHandle; @@ -88,7 +87,7 @@ impl VTable for Chunked { !slots.is_empty(), "ChunkedArray must have at least a chunk offsets slot" ); - let chunk_offsets = slots[CHUNK_OFFSETS_SLOT] + let chunk_offsets = slots[ChunkedSlots::CHUNK_OFFSETS] .as_ref() .vortex_expect("validated chunk offsets slot"); vortex_ensure!( @@ -103,10 +102,10 @@ impl VTable for Chunked { data.chunk_offsets.len() ); vortex_ensure!( - data.chunk_offsets.len() == slots.len() - CHUNKS_OFFSET + 1, + data.chunk_offsets.len() == slots.len() - ChunkedSlots::CHUNKS_OFFSET + 1, "ChunkedArray chunk offsets length {} does not match {} chunks", data.chunk_offsets.len(), - slots.len() - CHUNKS_OFFSET + slots.len() - ChunkedSlots::CHUNKS_OFFSET ); vortex_ensure!( data.chunk_offsets @@ -125,7 +124,7 @@ impl VTable for Chunked { .tuple_windows() .enumerate() { - let chunk = slots[CHUNKS_OFFSET + idx] + let chunk = slots[ChunkedSlots::CHUNKS_OFFSET + idx] .as_ref() .vortex_expect("validated chunk slot"); vortex_ensure!( @@ -193,7 +192,7 @@ impl VTable for Chunked { let nchunks = children.len() - 1; let chunk_offsets = children.get( - CHUNK_OFFSETS_SLOT, + ChunkedSlots::CHUNK_OFFSETS, &DType::Primitive(PType::U64, Nullability::NonNullable), nchunks + 1, )?; @@ -219,7 +218,11 @@ impl VTable for Chunked { .enumerate() { let chunk_len = end - start; - slots.push(Some(children.get(idx + CHUNKS_OFFSET, dtype, chunk_len)?)); + slots.push(Some(children.get( + idx + ChunkedSlots::CHUNKS_OFFSET, + dtype, + chunk_len, + )?)); } Ok(ArrayParts::new( @@ -244,13 +247,19 @@ impl VTable for Chunked { fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { match idx { - CHUNK_OFFSETS_SLOT => "chunk_offsets".to_string(), - n => format!("chunks[{}]", n - CHUNKS_OFFSET), + ChunkedSlots::CHUNK_OFFSETS => "chunk_offsets".to_string(), + n => format!("chunks[{}]", n - ChunkedSlots::CHUNKS_OFFSET), } } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { match array.dtype() { + DType::Union(..) => { + todo!( + "TODO(connor)[Union]: canonicalize chunked Union arrays by packing type IDs and \ + every sparse child along identical chunk boundaries" + ) + } // Struct, List, FixedSizeList, and Variant need child swizzling that the builder path // cannot express. DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) | DType::Variant(..) => { @@ -259,7 +268,7 @@ impl VTable for Chunked { } // For all other types, use the builder path via AppendChild. _ => { - let slot_idx = array.next_builder_slot.max(CHUNKS_OFFSET); + let slot_idx = array.next_builder_slot.max(ChunkedSlots::CHUNKS_OFFSET); if slot_idx < array.slots().len() { Ok(ExecutionResult::append_child( array.with_next_builder_slot(slot_idx + 1), diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index e160bcc16a9..e6a656afa51 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -20,13 +21,15 @@ use crate::arrays::DecimalArray; use crate::arrays::ExtensionArray; use crate::arrays::FixedSizeListArray; use crate::arrays::ListViewArray; +use crate::arrays::MapArray; use crate::arrays::NullArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; +use crate::arrays::UnionArray; use crate::arrays::VarBinViewArray; use crate::arrays::VariantArray; use crate::arrays::varbinview::BinaryView; -use crate::builders::builder_with_capacity; +use crate::builders::builder_with_capacity_in; use crate::dtype::DType; use crate::dtype::DecimalType; use crate::dtype::Nullability; @@ -125,7 +128,21 @@ pub(crate) fn constant_canonicalize( array.len(), )) } - DType::List(..) => Canonical::List(constant_canonical_list_array(scalar, array.len())), + DType::List(..) => Canonical::List(constant_canonical_list_array( + scalar, + array.len(), + ctx.allocator(), + )), + DType::Map(map_dtype, nullability) => { + let entries_scalar = Scalar::try_new( + DType::List(Arc::new(map_dtype.entries_dtype()), *nullability), + scalar.value().cloned(), + )?; + Canonical::Map(MapArray::try_new( + map_dtype.clone(), + constant_canonical_list_array(&entries_scalar, array.len(), ctx.allocator()), + )?) + } DType::FixedSizeList(element_dtype, list_size, _) => { let value = scalar.as_list(); @@ -135,6 +152,7 @@ pub(crate) fn constant_canonicalize( *list_size, value.dtype().nullability(), array.len(), + ctx.allocator(), )) } DType::Struct(struct_dtype, _) => { @@ -164,7 +182,7 @@ pub(crate) fn constant_canonicalize( StructArray::new_unchecked(fields, struct_dtype.clone(), array.len(), validity) }) } - DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), + DType::Union(..) => Canonical::Union(UnionArray::constant(scalar, array.len())?), DType::Variant(_) => Canonical::Variant(VariantArray::try_new( array.array().clone().into_array(), None, @@ -236,18 +254,23 @@ fn constant_canonical_byte_view( /// /// We basically just project the list scalar value into list view components. If the caller wants /// a fully decompressed and non-overlapping array, they can rebuild the array. -fn constant_canonical_list_array(scalar: &Scalar, len: usize) -> ListViewArray { +fn constant_canonical_list_array( + scalar: &Scalar, + len: usize, + allocator: &BufferAllocatorRef, +) -> ListViewArray { let list = scalar.as_list(); // Since "canonicalize" only applies to the top level array, we can simply have 1 scalar in our // child `elements` and have all list views point to that scalar. let elements = if let Some(elements) = list.elements() { // Extract the list elements out of the scalar into a new array. - let mut builder = builder_with_capacity( + let mut builder = builder_with_capacity_in( list.dtype() .as_list_element_opt() .vortex_expect("list scalar somehow did not have a list DType"), list.len(), + allocator, ); for scalar in &elements { builder @@ -290,13 +313,15 @@ fn constant_canonical_fixed_size_list_array( list_size: u32, list_nullability: Nullability, len: usize, + allocator: &BufferAllocatorRef, ) -> FixedSizeListArray { match values { None => { // Even though the scalar is null, we still have to allocate the correct amount of space // for the given `DType`. let elements_len = list_size as usize * len; - let mut element_builder = builder_with_capacity(element_dtype, elements_len); + let mut element_builder = + builder_with_capacity_in(element_dtype, elements_len, allocator); element_builder.append_defaults(elements_len); let elements = element_builder.finish(); @@ -307,7 +332,8 @@ fn constant_canonical_fixed_size_list_array( } } Some(values) => { - let mut elements_builder = builder_with_capacity(element_dtype, len * values.len()); + let mut elements_builder = + builder_with_capacity_in(element_dtype, len * values.len(), allocator); for _ in 0..len { for v in &values { @@ -350,7 +376,8 @@ mod tests { use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; - use crate::arrays::listview::ListViewArrayExt; + use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; + use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::ListViewRebuildMode; use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; diff --git a/vortex-array/src/arrays/constant/vtable/mod.rs b/vortex-array/src/arrays/constant/vtable/mod.rs index a2057077335..c0796a0e7d0 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -5,7 +5,7 @@ use std::fmt::Debug; use std::hash::Hash; use std::hash::Hasher; -use vortex_buffer::ByteBufferMut; +use itertools::Itertools; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -26,6 +26,7 @@ use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::VTable; use crate::array::unsupported_buffer_replacement; +use crate::arrays::ExtensionArray; use crate::arrays::constant::ConstantData; use crate::arrays::constant::compute::rules::PARENT_RULES; use crate::arrays::constant::vtable::canonical::constant_canonicalize; @@ -33,14 +34,21 @@ use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::BoolBuilder; use crate::builders::DecimalBuilder; +use crate::builders::FixedSizeListBuilder; +use crate::builders::ListViewBuilder; use crate::builders::NullBuilder; use crate::builders::PrimitiveBuilder; use crate::builders::VarBinViewBuilder; +use crate::builders::builder_with_capacity_in; use crate::canonical::Canonical; use crate::dtype::DType; +use crate::dtype::OffsetBuilderPType; use crate::match_each_decimal_value; +use crate::match_each_listview_builder; use crate::match_each_native_ptype; +use crate::match_each_varbin_builder; use crate::scalar::DecimalValue; +use crate::scalar::ListScalar; use crate::scalar::Scalar; use crate::scalar::ScalarValue; use crate::serde::ArrayChildren; @@ -98,7 +106,7 @@ impl VTable for Constant { fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { match idx { 0 => BufferHandle::new_host( - ScalarValue::to_proto_bytes::(array.scalar.value()).freeze(), + ScalarValue::to_proto_bytes::>(array.scalar.value()).into(), ), _ => vortex_panic!("ConstantArray buffer index {idx} out of bounds"), } @@ -222,38 +230,166 @@ impl VTable for Constant { }); } DType::Utf8(_) => { - append_value_or_nulls::(builder, scalar.is_null(), n, |b| { - let typed = scalar.as_utf8(); - let value = typed - .value() - .vortex_expect("non-null utf8 scalar must have a value"); - b.append_n_values(value.as_bytes(), n); - }); + if let Some(result) = match_each_varbin_builder!(builder, |builder| { + builder.append_scalar_repeated(scalar, n) + }) { + result?; + } else { + append_value_or_nulls::(builder, scalar.is_null(), n, |b| { + let value = scalar + .as_utf8() + .value() + .vortex_expect("non-null utf8 scalar must have a value"); + b.append_n_values(value.as_bytes(), n); + }); + } } DType::Binary(_) => { - append_value_or_nulls::(builder, scalar.is_null(), n, |b| { - let typed = scalar.as_binary(); - let value = typed - .value() - .vortex_expect("non-null binary scalar must have a value"); - b.append_n_values(value, n); - }); + if let Some(result) = match_each_varbin_builder!(builder, |builder| { + builder.append_scalar_repeated(scalar, n) + }) { + result?; + } else { + append_value_or_nulls::(builder, scalar.is_null(), n, |b| { + let value = scalar + .as_binary() + .value() + .vortex_expect("non-null binary scalar must have a value"); + b.append_n_values(value, n); + }); + } + } + DType::List(..) => append_constant_list_run(array, n, builder, ctx)?, + DType::Extension(ext_dtype) => { + // An extension array is its storage wearing a dtype, so a run of identical values + // is a constant storage array, which stays constant-encoded in the builder. + // Canonicalizing instead would materialize the storage: see the note in + // `constant_canonicalize` about `ExtensionConstantRule`. + let storage = ConstantArray::new(scalar.as_extension().to_storage_scalar(), n); + ExtensionArray::new(ext_dtype.clone(), storage.into_array()) + .into_array() + .append_to_builder(builder, ctx)? } - // TODO: add fast paths for DType::Struct, DType::List, DType::FixedSizeList, DType::Extension. - _ => { - let canonical = array - .array() - .clone() - .execute::(ctx)? - .into_array(); - canonical.append_to_builder(builder, ctx)?; + DType::FixedSizeList(..) => { + append_constant_fixed_size_list_run(array, n, builder, ctx)? } + // The remaining dtypes canonicalize cheaply: a constant struct canonicalizes to + // constant fields, and a constant map to views sharing one copy of the entries, so + // appending the canonical array preserves the run's economy. + // TODO: add a fast path for DType::Union once it has a builder. + _ => append_via_canonical(array, builder, ctx)?, } Ok(()) } } +/// Appends the constant list `array` as one run sharing a single copy of its elements. +/// +/// The list's elements materialize once, and +/// [`ListViewBuilder::append_array_as_repeated_list`] points the run's `n` views at that one +/// copy. Only a list-view builder has a layout that can share elements; any other builder for a +/// list dtype - a [`ListBuilder`](crate::builders::ListBuilder), whose offsets can only describe +/// contiguous lists - appends the canonical run instead. +fn append_constant_list_run( + array: ArrayView<'_, Constant>, + n: usize, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let scalar = array.scalar(); + match match_each_listview_builder!(builder, |b| append_repeated_list_run( + b, + scalar.as_list(), + n, + ctx + )) { + Some(result) => result, + None => append_via_canonical(array, builder, ctx), + } +} + +/// Appends the list `scalar` to a [`ListViewBuilder`] `n` times, storing its elements once. +fn append_repeated_list_run( + builder: &mut ListViewBuilder, + scalar: ListScalar, + n: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + if n == 0 { + return Ok(()); + } + + let Some(elements) = scalar.elements() else { + // A null run stores no elements at all. + builder.append_nulls(n); + return Ok(()); + }; + + let mut elements_builder = + builder_with_capacity_in(scalar.element_dtype(), elements.len(), ctx.allocator()); + for element in &elements { + elements_builder.append_scalar(element)?; + } + + builder.append_array_as_repeated_list(&elements_builder.finish(), n, ctx) +} + +/// Appends the constant fixed-size-list `array` as its list's elements tiled `n` times. +/// +/// The list's elements materialize into a tile once - a single [`ConstantArray`] when they are +/// all the same scalar, so that the tiling costs nothing - and +/// [`FixedSizeListBuilder::append_array_as_repeated_list`] shares that one tile across the run. +fn append_constant_fixed_size_list_run( + array: ArrayView<'_, Constant>, + n: usize, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let Some(builder) = builder.as_any_mut().downcast_mut::() else { + return append_via_canonical(array, builder, ctx); + }; + + if n == 0 { + return Ok(()); + } + + let scalar = array.scalar().as_list(); + let Some(elements) = scalar.elements() else { + // A null run stores no elements of its own, only the placeholders the builder writes. + builder.append_nulls(n); + return Ok(()); + }; + + let tile = match elements.iter().all_equal_value() { + Ok(uniform) => ConstantArray::new(uniform.clone(), elements.len()).into_array(), + Err(_) => { + let mut tile_builder = + builder_with_capacity_in(builder.element_dtype(), elements.len(), ctx.allocator()); + for element in &elements { + tile_builder.append_scalar(element)?; + } + tile_builder.finish() + } + }; + + builder.append_array_as_repeated_list(&tile, n, ctx) +} + +/// Appends `array` by canonicalizing it first, for the dtypes with no fast path of their own. +fn append_via_canonical( + array: ArrayView<'_, Constant>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let canonical = array + .array() + .clone() + .execute::(ctx)? + .into_array(); + canonical.append_to_builder(builder, ctx) +} + /// Downcasts `builder` to `B`, then either appends `n` nulls or calls `fill` with the typed /// builder depending on `is_null`. /// @@ -278,19 +414,36 @@ fn append_value_or_nulls( #[cfg(test)] mod tests { + use std::sync::Arc; + use rstest::rstest; use vortex_error::VortexResult; use crate::IntoArray; use crate::VortexSessionExecute; + use crate::arrays::Chunked; + use crate::arrays::Constant; use crate::arrays::ConstantArray; + use crate::arrays::Extension; + use crate::arrays::FixedSizeList; + use crate::arrays::ListView; + use crate::arrays::Struct; + use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::constant::vtable::canonical::constant_canonicalize; + use crate::arrays::extension::ExtensionArrayExt; + use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; + use crate::arrays::listview::ListViewArraySlotsExt; + use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; - use crate::builders::builder_with_capacity; + use crate::builders::ArrayBuilder; + use crate::builders::ListBuilder; + use crate::builders::builder_with_capacity_in; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; + use crate::extension::datetime::Date; + use crate::extension::datetime::TimeUnit; use crate::scalar::Scalar; /// Appends `array` into a fresh builder and asserts the result matches `constant_canonicalize`. @@ -298,7 +451,11 @@ mod tests { let mut ctx = crate::array_session().create_execution_ctx(); let expected = constant_canonicalize(array.as_view(), &mut ctx)?.into_array(); - let mut builder = builder_with_capacity(array.dtype(), array.len()); + let mut builder = builder_with_capacity_in( + array.dtype(), + array.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); array .into_array() .append_to_builder(builder.as_mut(), &mut ctx)?; @@ -397,6 +554,88 @@ mod tests { )) } + #[rstest] + #[case::non_empty(vec![Scalar::from(1i32), Scalar::from(2i32)], 4)] + #[case::empty(vec![], 3)] + #[case::n_zero(vec![Scalar::from(1i32)], 0)] + fn test_list_constant_append( + #[case] elements: Vec, + #[case] n: usize, + ) -> VortexResult<()> { + let scalar = Scalar::list( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + elements, + Nullability::NonNullable, + ); + assert_append_matches_canonical(ConstantArray::new(scalar, n)) + } + + #[test] + fn test_null_list_constant_append() -> VortexResult<()> { + let dtype = DType::List( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + Nullability::Nullable, + ); + assert_append_matches_canonical(ConstantArray::new(Scalar::null(dtype), 3)) + } + + /// A run of identical lists appended into a list-view builder shares one copy of its elements + /// across the whole run. + #[test] + fn test_list_constant_append_keeps_one_copy_of_the_elements() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let scalar = Scalar::list( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + vec![Scalar::from(1i32), Scalar::from(2i32), Scalar::from(3i32)], + Nullability::NonNullable, + ); + let array = ConstantArray::new(scalar, 1_000); + + let mut builder = builder_with_capacity_in( + array.dtype(), + array.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + array + .into_array() + .append_to_builder(builder.as_mut(), &mut ctx)?; + let result = builder.finish(); + + assert_eq!( + result.as_::().elements().len(), + 3, + "the run's elements should be stored once, not once per row", + ); + Ok(()) + } + + /// A `ListBuilder`'s offsets can only describe contiguous lists, so a constant run cannot + /// share its elements there and takes the canonical path instead. + #[test] + fn test_list_constant_append_into_list_builder() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let element_dtype: Arc = + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)); + let scalar = Scalar::list( + Arc::clone(&element_dtype), + vec![Scalar::from(1i32), Scalar::from(2i32)], + Nullability::NonNullable, + ); + let array = ConstantArray::new(scalar, 4).into_array(); + + let mut builder = ListBuilder::::with_capacity_in( + element_dtype, + Nullability::NonNullable, + 0, + 0, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + array.append_to_builder(&mut builder, &mut ctx)?; + + assert_arrays_eq!(&builder.finish(), &array, &mut ctx); + Ok(()) + } + #[test] fn test_struct_constant_append() -> VortexResult<()> { let fields = StructFields::new( @@ -425,4 +664,131 @@ mod tests { let dtype = DType::Struct(fields, Nullability::Nullable); assert_append_matches_canonical(ConstantArray::new(Scalar::null(dtype), 4)) } + + /// A run of identical structs should leave each field constant-encoded rather than materialize + /// a value per row per field. + #[test] + fn test_struct_constant_append_keeps_fields_constant() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let fields = StructFields::new( + ["x", "y"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::NonNullable), + ], + ); + let scalar = Scalar::struct_( + DType::Struct(fields, Nullability::NonNullable), + [ + Scalar::primitive(42i32, Nullability::NonNullable), + Scalar::utf8("hi", Nullability::NonNullable), + ], + ); + let array = ConstantArray::new(scalar, 1_000); + + let mut builder = builder_with_capacity_in( + array.dtype(), + array.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + array + .into_array() + .append_to_builder(builder.as_mut(), &mut ctx)?; + let result = builder.finish(); + + let struct_array = result.as_::(); + for field in 0..2 { + assert!( + struct_array.unmasked_field(field).is::(), + "field {field} should have stayed constant-encoded", + ); + } + Ok(()) + } + + #[rstest] + #[case::non_uniform(vec![Scalar::from(1i32), Scalar::from(2i32)])] + #[case::uniform(vec![Scalar::from(7i32), Scalar::from(7i32)])] + fn test_fixed_size_list_constant_append(#[case] elements: Vec) -> VortexResult<()> { + let scalar = Scalar::fixed_size_list( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + elements, + Nullability::NonNullable, + ); + assert_append_matches_canonical(ConstantArray::new(scalar, 4)) + } + + #[test] + fn test_null_fixed_size_list_constant_append() -> VortexResult<()> { + let dtype = DType::FixedSizeList( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + 2, + Nullability::Nullable, + ); + assert_append_matches_canonical(ConstantArray::new(Scalar::null(dtype), 3)) + } + + /// A fixed-size list whose elements are all the same scalar tiles a constant array, so the + /// tile's chunks stay constant-encoded rather than materializing a value per row. + #[test] + fn test_uniform_fixed_size_list_constant_append_keeps_elements_constant() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let scalar = Scalar::fixed_size_list( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + vec![Scalar::from(7i32), Scalar::from(7i32)], + Nullability::NonNullable, + ); + let array = ConstantArray::new(scalar, 1_000); + + let mut builder = builder_with_capacity_in( + array.dtype(), + array.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + array + .into_array() + .append_to_builder(builder.as_mut(), &mut ctx)?; + let result = builder.finish(); + + let elements = result.as_::().elements().clone(); + assert!( + elements + .as_::() + .iter_chunks() + .all(|chunk| chunk.is::()), + "a uniform tile should have stayed constant-encoded", + ); + Ok(()) + } + + #[test] + fn test_extension_constant_append() -> VortexResult<()> { + let scalar = Scalar::extension::(TimeUnit::Days, Scalar::from(Some(42i32))); + assert_append_matches_canonical(ConstantArray::new(scalar, 5)) + } + + /// An extension array is its storage wearing a dtype, so a run of identical values should leave + /// the storage constant-encoded. + #[test] + fn test_extension_constant_append_keeps_storage_constant() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let scalar = Scalar::extension::(TimeUnit::Days, Scalar::from(Some(42i32))); + let array = ConstantArray::new(scalar, 1_000); + + let mut builder = builder_with_capacity_in( + array.dtype(), + array.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + array + .into_array() + .append_to_builder(builder.as_mut(), &mut ctx)?; + let result = builder.finish(); + + assert!( + result.as_::().storage_array().is::(), + "the storage should have stayed constant-encoded", + ); + Ok(()) + } } diff --git a/vortex-array/src/arrays/constant/vtable/operations.rs b/vortex-array/src/arrays/constant/vtable/operations.rs index fdd2dd67345..e3568a9c39f 100644 --- a/vortex-array/src/arrays/constant/vtable/operations.rs +++ b/vortex-array/src/arrays/constant/vtable/operations.rs @@ -18,3 +18,47 @@ impl OperationsVTable for Constant { Ok(array.scalar.clone()) } } + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::arrays::ConstantArray; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::dtype::UnionVariants; + use crate::scalar::Scalar; + + #[test] + fn scalar_at_preserves_union_scalar() -> VortexResult<()> { + let variants = UnionVariants::try_new( + ["int", "string"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::Nullable), + DType::Utf8(Nullability::NonNullable), + ], + vec![5, 9], + )?; + + let scalar = Scalar::union( + variants.clone(), + 5, + Scalar::primitive(42_i32, Nullability::Nullable), + Nullability::Nullable, + )?; + let array = ConstantArray::new(scalar.clone(), 3).into_array(); + let mut ctx = crate::array_session().create_execution_ctx(); + + assert_eq!(array.execute_scalar(1, &mut ctx)?, scalar); + + let null = Scalar::null(DType::Union(variants, Nullability::Nullable)); + let array = ConstantArray::new(null.clone(), 3).into_array(); + + assert_eq!(array.execute_scalar(1, &mut ctx)?, null); + + Ok(()) + } +} diff --git a/vortex-array/src/arrays/decimal/array.rs b/vortex-array/src/arrays/decimal/array.rs index 1f4d6c0c277..f59ca7774a9 100644 --- a/vortex-array/src/arrays/decimal/array.rs +++ b/vortex-array/src/arrays/decimal/array.rs @@ -25,6 +25,7 @@ use crate::array::ArrayParts; use crate::array::TypedArrayRef; use crate::array::child_to_validity; use crate::array::validity_to_child; +use crate::array_slots; use crate::arrays::Decimal; use crate::arrays::DecimalArray; use crate::arrays::PrimitiveArray; @@ -42,10 +43,12 @@ use crate::match_each_unsigned_integer_ptype; use crate::patches::Patches; use crate::validity::Validity; -/// The validity bitmap indicating which elements are non-null. -pub(super) const VALIDITY_SLOT: usize = 0; -pub(super) const NUM_SLOTS: usize = 1; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"]; +#[array_slots(Decimal)] +pub struct DecimalSlots { + /// The validity bitmap indicating which elements are non-null. + #[slot(0)] + pub validity: Option, +} /// A decimal array that stores fixed-precision decimal numbers with configurable scale. /// @@ -144,12 +147,12 @@ pub trait DecimalArrayExt: TypedArrayRef { } fn validity_child(&self) -> Option<&ArrayRef> { - self.as_ref().slots()[VALIDITY_SLOT].as_ref() + self.as_ref().slots()[DecimalSlots::VALIDITY].as_ref() } fn validity(&self) -> Validity { child_to_validity( - self.as_ref().slots()[VALIDITY_SLOT].as_ref(), + self.as_ref().slots()[DecimalSlots::VALIDITY].as_ref(), self.nullability(), ) } @@ -606,7 +609,7 @@ where ) } - for (idx, value) in patch_indices.iter().zip_eq(patch_values.into_iter()) { + for (idx, value) in patch_indices.iter().zip_eq(patch_values) { buffer[idx.as_() - patch_indices_offset] = ::from(value).vortex_expect( "values of a given DecimalDType are representable in all compatible NativeDecimalType", ); diff --git a/vortex-array/src/arrays/decimal/compute/cast.rs b/vortex-array/src/arrays/decimal/compute/cast.rs index 20cac1d5106..8770a9bc8f4 100644 --- a/vortex-array/src/arrays/decimal/compute/cast.rs +++ b/vortex-array/src/arrays/decimal/compute/cast.rs @@ -216,41 +216,48 @@ where { let values = array.buffer::(); let values = values.as_slice(); - let cast_plan = DecimalCastPlan::::new(from_decimal_dtype, to_decimal_dtype); + let cast_result = match DecimalCastPlan::::new(from_decimal_dtype, to_decimal_dtype) { + DecimalCastPlan::SameScale { min, max } => { + cast_decimal_buffer(values, valid_values, |value| { + let value = ::from(value)?; + (value >= min && value <= max).then_some(value) + }) + } + cast_plan => cast_decimal_buffer(values, valid_values, |value| cast_plan.cast(value)), + }; + let buffer = cast_result.map_err(|idx| { + decimal_cast_error::(values[idx], from_decimal_dtype, to_decimal_dtype) + })?; - let buffer = match valid_values { + Ok(DecimalArray::new(buffer, to_decimal_dtype, validity).into_array()) +} + +fn cast_decimal_buffer( + values: &[F], + valid_values: &Mask, + cast: impl Fn(F) -> Option, +) -> Result, usize> +where + F: NativeDecimalType, + T: NativeDecimalType, +{ + let mut buffer = BufferMut::::with_capacity(values.len()); + match valid_values { Mask::AllTrue(_) => { - let mut buffer = BufferMut::::with_capacity(values.len()); - values - .try_map_into(&mut buffer.spare_capacity_mut()[..values.len()], |value| { - cast_plan.cast(value) - }) - .map_err(|idx| { - decimal_cast_error::(values[idx], from_decimal_dtype, to_decimal_dtype) - })?; - // SAFETY: try_map_into initializes every lane before returning Ok. - unsafe { buffer.set_len(values.len()) }; - buffer.freeze() + values.try_map_into(&mut buffer.spare_capacity_mut()[..values.len()], &cast)?; } - Mask::AllFalse(_) => BufferMut::::zeroed(values.len()).freeze(), + Mask::AllFalse(_) => return Ok(BufferMut::::zeroed(values.len()).freeze()), Mask::Values(mask) => { - let mut buffer = BufferMut::::with_capacity(values.len()); - values - .try_map_masked_into( - mask.bit_buffer(), - &mut buffer.spare_capacity_mut()[..values.len()], - |value| cast_plan.cast(value), - ) - .map_err(|idx| { - decimal_cast_error::(values[idx], from_decimal_dtype, to_decimal_dtype) - })?; - // SAFETY: try_map_masked_into initializes every lane before returning Ok. - unsafe { buffer.set_len(values.len()) }; - buffer.freeze() + values.try_map_masked_into( + mask.bit_buffer(), + &mut buffer.spare_capacity_mut()[..values.len()], + &cast, + )?; } - }; - - Ok(DecimalArray::new(buffer, to_decimal_dtype, validity).into_array()) + } + // SAFETY: the selected map kernel initialized every lane before returning Ok. + unsafe { buffer.set_len(values.len()) }; + Ok(buffer.freeze()) } #[cold] diff --git a/vortex-array/src/arrays/decimal/compute/fixed_width.rs b/vortex-array/src/arrays/decimal/compute/fixed_width.rs new file mode 100644 index 00000000000..8ec0987e0f8 --- /dev/null +++ b/vortex-array/src/arrays/decimal/compute/fixed_width.rs @@ -0,0 +1,36 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::ByteBuffer; +use vortex_error::VortexResult; + +use crate::array::ArrayView; +use crate::arrays::Decimal; +use crate::arrays::DecimalArray; +use crate::arrays::fixed_width::FixedWidthArray; +use crate::buffer::BufferHandle; +use crate::validity::Validity; + +impl FixedWidthArray for Decimal { + fn byte_width(array: ArrayView<'_, Self>) -> usize { + array.values_type().byte_width() + } + + fn values(array: ArrayView<'_, Self>) -> ByteBuffer { + array.buffer_handle().to_host_sync() + } + + fn with_values( + array: ArrayView<'_, Self>, + values: ByteBuffer, + _len: usize, + validity: Validity, + ) -> VortexResult { + DecimalArray::try_new_handle( + BufferHandle::new_host(values), + array.values_type(), + array.decimal_dtype(), + validity, + ) + } +} diff --git a/vortex-array/src/arrays/decimal/compute/mod.rs b/vortex-array/src/arrays/decimal/compute/mod.rs index b2f21fa8398..a6e1cfb2848 100644 --- a/vortex-array/src/arrays/decimal/compute/mod.rs +++ b/vortex-array/src/arrays/decimal/compute/mod.rs @@ -4,9 +4,9 @@ mod between; mod cast; mod fill_null; +mod fixed_width; mod mask; pub mod rules; -mod take; #[cfg(test)] mod tests { @@ -17,8 +17,10 @@ mod tests { use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::DecimalArray; + use crate::compute::conformance::binary_numeric::test_binary_numeric_array; use crate::compute::conformance::consistency::test_array_consistency; use crate::dtype::DecimalDType; + use crate::dtype::NativeDecimalType; use crate::validity::Validity; #[rstest] @@ -60,4 +62,60 @@ mod tests { &mut array_session().create_execution_ctx(), ); } + + #[rstest] + #[case::decimal_array(DecimalArray::new( + buffer![100i128, 200i128, 300i128, 400i128, 500i128], + DecimalDType::new(19, 2), + Validity::NonNullable, + ))] + #[case::decimal_nullable(DecimalArray::new( + buffer![1000i128, 2000i128, 3000i128, 4000i128, 5000i128], + DecimalDType::new(19, 3), + Validity::from_iter([true, false, true, true, false]), + ))] + #[case::decimal_small_precision(DecimalArray::new( + buffer![10i128, 20i128, 30i128], + DecimalDType::new(5, 1), + Validity::NonNullable, + ))] + #[case::decimal_narrow_storage(DecimalArray::new( + buffer![10i8, 20i8, 30i8], + DecimalDType::new(5, 1), + Validity::NonNullable, + ))] + #[case::decimal_widened_carry(DecimalArray::new( + buffer![ + ::MAX_BY_PRECISION[38], + ::MIN_BY_PRECISION[38], + ], + DecimalDType::new(38, 0), + Validity::NonNullable, + ))] + #[case::decimal_single(DecimalArray::new( + buffer![42i128], + DecimalDType::new(10, 0), + Validity::NonNullable, + ))] + #[case::decimal_large_scale(DecimalArray::new( + buffer![123456789012345i128, 987654321098765i128], + DecimalDType::new(20, 10), + Validity::NonNullable, + ))] + #[case::decimal_negative(DecimalArray::new( + buffer![-100i128, -200i128, 300i128, -400i128], + DecimalDType::new(10, 2), + Validity::NonNullable, + ))] + #[case::decimal_negative_scale(DecimalArray::new( + buffer![5i64, -3i64, 600i64], + DecimalDType::new(10, -2), + Validity::NonNullable, + ))] + fn test_decimal_binary_numeric(#[case] array: DecimalArray) { + test_binary_numeric_array( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); + } } diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs deleted file mode 100644 index ce786f8e4ab..00000000000 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ /dev/null @@ -1,139 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use vortex_buffer::Buffer; -use vortex_error::VortexResult; - -use crate::ArrayRef; -use crate::IntoArray; -use crate::array::ArrayView; -use crate::arrays::Decimal; -use crate::arrays::DecimalArray; -use crate::arrays::PrimitiveArray; -use crate::arrays::dict::TakeExecute; -use crate::dtype::IntegerPType; -use crate::dtype::NativeDecimalType; -use crate::executor::ExecutionCtx; -use crate::match_each_decimal_value_type; -use crate::match_each_integer_ptype; - -impl TakeExecute for Decimal { - fn take( - array: ArrayView<'_, Decimal>, - indices: &ArrayRef, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - let indices = indices.clone().execute::(ctx)?; - let validity = array.validity()?.take(&indices.clone().into_array())?; - - // TODO(joe): if the true count of take indices validity is low, only take array values with - // valid indices. - let decimal = match_each_decimal_value_type!(array.values_type(), |D| { - match_each_integer_ptype!(indices.ptype(), |I| { - let buffer = - take_to_buffer::(indices.as_slice::(), array.buffer::().as_slice()); - // SAFETY: Take operation preserves decimal dtype and creates valid buffer. - // Validity is computed correctly from the parent array and indices. - unsafe { DecimalArray::new_unchecked(buffer, array.decimal_dtype(), validity) } - }) - }); - - Ok(Some(decimal.into_array())) - } -} - -fn take_to_buffer(indices: &[I], values: &[T]) -> Buffer { - indices.iter().map(|idx| values[idx.as_()]).collect() -} - -#[cfg(test)] -mod tests { - use rstest::rstest; - use vortex_buffer::Buffer; - use vortex_buffer::buffer; - - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::array_session; - use crate::arrays::DecimalArray; - use crate::arrays::PrimitiveArray; - use crate::assert_arrays_eq; - use crate::compute::conformance::take::test_take_conformance; - use crate::dtype::DecimalDType; - use crate::validity::Validity; - - #[test] - fn test_take() { - let mut ctx = array_session().create_execution_ctx(); - let ddtype = DecimalDType::new(19, 1); - let array = DecimalArray::new( - buffer![10i128, 11i128, 12i128, 13i128], - ddtype, - Validity::NonNullable, - ); - - let indices = buffer![0, 2, 3].into_array(); - let taken = array.take(indices).unwrap(); - - let expected = DecimalArray::from_iter([10i128, 12, 13], ddtype); - assert_arrays_eq!(expected, taken, &mut ctx); - } - - #[test] - fn test_take_null_indices() { - let mut ctx = array_session().create_execution_ctx(); - let ddtype = DecimalDType::new(19, 1); - let array = DecimalArray::new( - buffer![i128::MAX, 11i128, 12i128, 13i128], - ddtype, - Validity::NonNullable, - ); - - let indices = PrimitiveArray::from_option_iter([None, Some(2), Some(3)]).into_array(); - let taken = array.take(indices).unwrap(); - - let expected = DecimalArray::from_option_iter([None, Some(12i128), Some(13)], ddtype); - assert_arrays_eq!(expected, taken, &mut ctx); - } - - #[rstest] - #[case(DecimalArray::new( - buffer![100i128, 200i128, 300i128, 400i128, 500i128], - DecimalDType::new(19, 2), - Validity::NonNullable, - ))] - #[case(DecimalArray::new( - buffer![10i64, 20i64, 30i64, 40i64, 50i64], - DecimalDType::new(10, 1), - Validity::NonNullable, - ))] - #[case(DecimalArray::new( - buffer![1i32, 2i32, 3i32, 4i32, 5i32], - DecimalDType::new(5, 0), - Validity::NonNullable, - ))] - #[case(DecimalArray::new( - buffer![1000i128, 2000i128, 3000i128, 4000i128, 5000i128], - DecimalDType::new(19, 3), - Validity::from_iter([true, false, true, true, false]), - ))] - #[case(DecimalArray::new( - buffer![42i128], - DecimalDType::new(19, 0), - Validity::NonNullable, - ))] - #[case({ - let values: Vec = (0..100).map(|i| i * 1000).collect(); - DecimalArray::new( - Buffer::from_iter(values), - DecimalDType::new(19, 4), - Validity::NonNullable, - ) - })] - fn test_take_decimal_conformance(#[case] array: DecimalArray) { - test_take_conformance( - &array.into_array(), - &mut array_session().create_execution_ctx(), - ); - } -} diff --git a/vortex-array/src/arrays/decimal/mod.rs b/vortex-array/src/arrays/decimal/mod.rs index d8af8529e87..a806957e62d 100644 --- a/vortex-array/src/arrays/decimal/mod.rs +++ b/vortex-array/src/arrays/decimal/mod.rs @@ -3,8 +3,10 @@ mod array; pub use array::DecimalArrayExt; +pub use array::DecimalArraySlotsExt; pub use array::DecimalData; pub use array::DecimalDataParts; +pub use array::DecimalSlots; pub use vtable::DecimalArray; pub(crate) mod compute; diff --git a/vortex-array/src/arrays/decimal/utils.rs b/vortex-array/src/arrays/decimal/utils.rs index a1eafde0fe3..3c7304e6fda 100644 --- a/vortex-array/src/arrays/decimal/utils.rs +++ b/vortex-array/src/arrays/decimal/utils.rs @@ -3,11 +3,90 @@ use itertools::Itertools; use itertools::MinMaxResult; +use vortex_buffer::Buffer; use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_mask::Mask; use crate::arrays::DecimalArray; +use crate::arrays::decimal::DecimalArrayExt; use crate::dtype::DecimalType; +use crate::dtype::NativeDecimalType; use crate::dtype::i256; +use crate::match_each_decimal_value_type; + +/// Return the array's unscaled values widened to `W`, which must be at least as wide as the +/// array's storage type. +pub(crate) fn widened_buffer(array: &DecimalArray) -> Buffer { + if array.values_type() == W::DECIMAL_TYPE { + return array.buffer::(); + } + match_each_decimal_value_type!(array.values_type(), |T| { + array + .buffer::() + .iter() + .map(|v| W::from(*v).vortex_expect("widening decimal cast must succeed")) + .collect() + }) +} + +/// Return the array's unscaled values converted to exactly `W`, whatever the array's +/// storage type. Zero-copy when the array is already stored at `W`. +/// +/// Widening is lossless. Narrowing fails for any *valid* value that does not fit `W`. +/// Null slots may hold arbitrary bytes and never fail; their contents in the returned +/// buffer are likewise arbitrary and must not be read. +pub fn converted_buffer( + array: &DecimalArray, + validity: &Mask, +) -> VortexResult> { + // Widening can never fail, so it needs no validation pass. + if array.values_type() <= W::DECIMAL_TYPE { + return Ok(widened_buffer(array)); + } + match_each_decimal_value_type!(array.values_type(), |T| { + let src = array.buffer::(); + match validity { + Mask::AllTrue(_) => { + // Keeping the overflow scan branchless and vectorizable. Only on overflow + // do we rescan for diagnostics. + let any_overflow = src.iter().fold(false, |acc, v| acc | W::from(*v).is_none()); + if any_overflow { + let (i, v) = src + .iter() + .enumerate() + .find(|&(_, v)| W::from(*v).is_none()) + .vortex_expect("overflow scan found an overflowing value"); + vortex_bail!( + "decimal value {v} at index {i} does not fit {}", + W::DECIMAL_TYPE + ); + } + } + Mask::AllFalse(_) => return Ok(Buffer::zeroed(src.len())), + Mask::Values(values) => { + let any_overflow = src.iter().fold(false, |acc, v| acc | W::from(*v).is_none()); + if any_overflow { + for (i, v) in src.iter().enumerate() { + if values.value(i) && W::from(*v).is_none() { + vortex_bail!( + "decimal value {v} at index {i} does not fit {}", + W::DECIMAL_TYPE + ); + } + } + } + } + } + // The convert pass is infallible: every valid value fits, and out-of-range null-slot + // garbage is mapped to the default value. Callers must still ignore null slots. + Ok(src + .iter() + .map(|v| W::from(*v).unwrap_or_default()) + .collect()) + }) +} macro_rules! try_downcast { ($array:expr, from: $src:ty, to: $($dst:ty),*) => {{ diff --git a/vortex-array/src/arrays/decimal/vtable/mod.rs b/vortex-array/src/arrays/decimal/vtable/mod.rs index 1ed839b813e..5f2e087e983 100644 --- a/vortex-array/src/arrays/decimal/vtable/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/mod.rs @@ -8,7 +8,6 @@ use vortex_buffer::Alignment; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; -use vortex_error::vortex_panic; use vortex_session::VortexSession; use crate::ArrayParts; @@ -19,6 +18,7 @@ use crate::array::Array; use crate::array::ArrayView; use crate::array::VTable; use crate::arrays::decimal::DecimalData; +use crate::arrays::fixed_width::vtable as fixed_width; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::DecimalBuilder; @@ -27,7 +27,6 @@ use crate::dtype::DecimalType; use crate::dtype::NativeDecimalType; use crate::match_each_decimal_value_type; use crate::serde::ArrayChildren; -use crate::validity::Validity; mod kernel; mod operations; mod validity; @@ -38,7 +37,7 @@ use vortex_session::registry::CachedId; use crate::EqMode; use crate::array::ArrayId; -use crate::arrays::decimal::array::SLOT_NAMES; +use crate::arrays::decimal::array::DecimalSlots; use crate::arrays::decimal::compute::rules::RULES; use crate::hash::ArrayEq; use crate::hash::ArrayHash; @@ -85,17 +84,11 @@ impl VTable for Decimal { } fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { - match idx { - 0 => array.values.clone(), - _ => vortex_panic!("DecimalArray buffer index {idx} out of bounds"), - } + fixed_width::buffer("DecimalArray", &array.values, idx) } fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { - match idx { - 0 => Some("values".to_string()), - _ => None, - } + fixed_width::buffer_name(idx) } fn with_buffers( @@ -103,13 +96,8 @@ impl VTable for Decimal { array: ArrayView<'_, Self>, buffers: &[BufferHandle], ) -> VortexResult> { - vortex_ensure!( - buffers.len() == 1, - "Expected 1 buffer, got {}", - buffers.len() - ); let mut data = array.data().clone(); - data.values = buffers[0].clone(); + data.values = fixed_width::single_buffer(buffers)?; Ok( ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) .with_slots(array.slots().iter().cloned().collect()), @@ -145,7 +133,8 @@ impl VTable for Decimal { data.len(), len ); - let validity = crate::array::child_to_validity(slots[0].as_ref(), *nullability); + let validity = + crate::array::child_to_validity(slots[DecimalSlots::VALIDITY].as_ref(), *nullability); if let Some(validity_len) = validity.maybe_len() { vortex_ensure!( validity_len == len, @@ -164,25 +153,14 @@ impl VTable for Decimal { dtype: &DType, len: usize, metadata: &[u8], - buffers: &[BufferHandle], children: &dyn ArrayChildren, _session: &VortexSession, ) -> VortexResult> { let metadata = DecimalMetadata::decode(metadata)?; - if buffers.len() != 1 { - vortex_bail!("Expected 1 buffer, got {}", buffers.len()); - } - let values = buffers[0].clone(); - - let validity = if children.is_empty() { - Validity::from(dtype.nullability()) - } else if children.len() == 1 { - let validity = children.get(0, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - vortex_bail!("Expected 0 or 1 child, got {}", children.len()); - }; + let values = fixed_width::single_buffer(buffers)?; + + let validity = fixed_width::deserialize_validity(dtype.nullability(), len, children)?; let Some(decimal_dtype) = dtype.as_decimal_opt() else { vortex_bail!("Expected Decimal dtype, got {:?}", dtype) @@ -202,7 +180,7 @@ impl VTable for Decimal { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + DecimalSlots::NAMES[idx].to_string() } fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { diff --git a/vortex-array/src/arrays/dict/array.rs b/vortex-array/src/arrays/dict/array.rs index cefc5a381fd..beed9088e6f 100644 --- a/vortex-array/src/arrays/dict/array.rs +++ b/vortex-array/src/arrays/dict/array.rs @@ -45,8 +45,10 @@ pub struct DictMetadata { #[array_slots(Dict)] pub struct DictSlots { /// The codes array mapping each element to a dictionary entry. + #[slot(0)] pub codes: ArrayRef, /// The dictionary values array containing the unique values. + #[slot(1)] pub values: ArrayRef, } @@ -294,8 +296,10 @@ mod test { use crate::arrays::ChunkedArray; use crate::arrays::DictArray; use crate::arrays::PrimitiveArray; + use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; - use crate::builders::builder_with_capacity; + use crate::builders::VarBinBuilder; + use crate::builders::builder_with_capacity_in; use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::Nullability::NonNullable; @@ -440,6 +444,25 @@ mod test { .into_array() } + #[test] + fn test_dict_utf8_append_to_varbin_builder() -> VortexResult<()> { + let values = VarBinViewArray::from_iter_str(["zero", "one", "two"]); + let dict = DictArray::try_new(buffer![2u8, 0, 2, 1].into_array(), values.into_array())?; + let expected = VarBinViewArray::from_iter_str(["two", "zero", "two", "one"]); + let mut ctx = array_session().create_execution_ctx(); + let mut builder = VarBinBuilder::::with_capacity_in( + dict.dtype().clone(), + dict.len(), + ctx.allocator(), + ); + + dict.into_array() + .append_to_builder(&mut builder, &mut ctx)?; + + assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx); + Ok(()) + } + #[test] fn test_dict_array_from_primitive_chunks() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); @@ -447,9 +470,10 @@ mod test { let chunk_count = 2; let array = make_dict_primitive_chunks::(len, 2, chunk_count); - let mut builder = builder_with_capacity( + let mut builder = builder_with_capacity_in( &DType::Primitive(PType::U64, NonNullable), len * chunk_count, + ctx.allocator(), ); array.append_to_builder( builder.as_mut(), diff --git a/vortex-array/src/arrays/dict/compute/like.rs b/vortex-array/src/arrays/dict/compute/like.rs index a4ca7c1f76e..da51e1cc0b3 100644 --- a/vortex-array/src/arrays/dict/compute/like.rs +++ b/vortex-array/src/arrays/dict/compute/like.rs @@ -11,7 +11,6 @@ use crate::array::ArrayView; use crate::arrays::ConstantArray; use crate::arrays::dict::DictArrayExt; use crate::arrays::dict::DictArraySlotsExt; -use crate::arrays::scalar_fn::ScalarFnFactoryExt; use crate::optimizer::ArrayOptimizer; use crate::scalar_fn::fns::like::Like; use crate::scalar_fn::fns::like::LikeOptions; @@ -30,8 +29,8 @@ impl LikeReduce for Dict { if let Some(pattern) = pattern.as_constant() { let pattern = ConstantArray::new(pattern, array.values().len()).into_array(); - let values = Like - .try_new_array(pattern.len(), options, [array.values().clone(), pattern])? + let values = Like::try_new(array.values().clone(), pattern, options)? + .into_array() .optimize()?; // SAFETY: LIKE preserves the len of the values, so codes are still pointing at @@ -62,7 +61,6 @@ mod tests { use crate::arrays::DictArray; use crate::arrays::VarBinArray; use crate::arrays::dict::compute::like::ConstantArray; - use crate::arrays::scalar_fn::ScalarFnFactoryExt; use crate::assert_arrays_eq; use crate::optimizer::ArrayOptimizer; use crate::scalar_fn::fns::like::Like; @@ -78,8 +76,8 @@ mod tests { .into_array(); let pattern = ConstantArray::new("hello%", 4).into_array(); - let result = Like - .try_new_array(4, LikeOptions::default(), [dict, pattern])? + let result = Like::try_new(dict, pattern, LikeOptions::default())? + .into_array() .optimize()?; assert_arrays_eq!( diff --git a/vortex-array/src/arrays/dict/compute/mask.rs b/vortex-array/src/arrays/dict/compute/mask.rs index 05bb3cf091f..b129fdd19c5 100644 --- a/vortex-array/src/arrays/dict/compute/mask.rs +++ b/vortex-array/src/arrays/dict/compute/mask.rs @@ -9,18 +9,12 @@ use crate::array::ArrayView; use crate::arrays::Dict; use crate::arrays::DictArray; use crate::arrays::dict::DictArraySlotsExt; -use crate::arrays::scalar_fn::ScalarFnFactoryExt; -use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::fns::mask::Mask as MaskExpr; use crate::scalar_fn::fns::mask::MaskReduce; impl MaskReduce for Dict { fn mask(array: ArrayView<'_, Dict>, mask: &ArrayRef) -> VortexResult> { - let masked_codes = MaskExpr.try_new_array( - array.codes().len(), - EmptyOptions, - [array.codes().clone(), mask.clone()], - )?; + let masked_codes = MaskExpr::try_new(array.codes().clone(), mask.clone())?.into_array(); // SAFETY: masking codes doesn't change dict invariants Ok(Some(unsafe { DictArray::new_unchecked(masked_codes, array.values().clone()).into_array() diff --git a/vortex-array/src/arrays/dict/compute/mod.rs b/vortex-array/src/arrays/dict/compute/mod.rs index 5c013b2d321..8bea82e6132 100644 --- a/vortex-array/src/arrays/dict/compute/mod.rs +++ b/vortex-array/src/arrays/dict/compute/mod.rs @@ -32,11 +32,13 @@ impl TakeExecute for Dict { _ctx: &mut ExecutionCtx, ) -> VortexResult> { let codes = array.codes().take(indices.clone())?; - // SAFETY: selecting codes doesn't change the invariants of DictArray - // Preserve all_values_referenced since taking codes doesn't affect which values are referenced - Ok(Some(unsafe { - DictArray::new_unchecked(codes, array.values().clone()).into_array() - })) + + // SAFETY: Selection preserves the integer code type and non-null code bounds, so every code + // indexes the unchanged values. `new_unchecked` resets `all_values_referenced` because + // selection can remove a value's last reference. + let taken = unsafe { DictArray::new_unchecked(codes, array.values().clone()) }; + + Ok(Some(taken.into_array())) } } @@ -44,11 +46,12 @@ impl FilterReduce for Dict { fn filter(array: ArrayView<'_, Dict>, mask: &Mask) -> VortexResult> { let codes = array.codes().filter(mask.clone())?; - // SAFETY: filtering codes doesn't change invariants - // Preserve all_values_referenced since filtering codes doesn't affect which values are referenced - Ok(Some(unsafe { - DictArray::new_unchecked(codes, array.values().clone()).into_array() - })) + // SAFETY: Selection preserves the integer code type and non-null code bounds, so every code + // indexes the unchanged values. `new_unchecked` resets `all_values_referenced` because + // selection can remove a value's last reference. + let filtered = unsafe { DictArray::new_unchecked(codes, array.values().clone()) }; + + Ok(Some(filtered.into_array())) } } diff --git a/vortex-array/src/arrays/dict/compute/rules.rs b/vortex-array/src/arrays/dict/compute/rules.rs index a15f2641eab..d43c6deeb75 100644 --- a/vortex-array/src/arrays/dict/compute/rules.rs +++ b/vortex-array/src/arrays/dict/compute/rules.rs @@ -106,101 +106,104 @@ impl ArrayParentReduceRule for DictionaryScalarFnValuesPushDownRule { parent: ArrayView<'_, ScalarFn>, child_idx: usize, ) -> VortexResult> { - // Check that the scalar function can actually be pushed down. - let sig = parent.scalar_fn().signature(); + let scalar_fn = parent.scalar_fn(); + let signature = scalar_fn.signature(); - // Don't push down pack expressions since we might want to unpack them in exporters - // later. - if parent.scalar_fn().is::() { + // Preserve pack expressions so exporters can unpack them later. + if scalar_fn.is::() { return Ok(None); } - // Don't push down cast operations — CastReduceAdaptor handles these eagerly. - // If it declined (returned None), we must fall through to the canonical path - // rather than creating a lazy cast inside the dictionary values. - if parent.scalar_fn().is::() { + // CastReduceAdaptor handles casts eagerly. If it declines the rewrite, leave the cast on + // the dictionary instead of creating a lazy cast over its values. + if scalar_fn.is::() { return Ok(None); } - // If the dictionary has less codes than values don't push down this might - // happen if the dictionary is sliced. + // A sliced dictionary can have more values than code rows. Do not increase the work in + // that case. if array.values().len() > array.codes().len() { return Ok(None); } - // If the scalar function is fallible, we cannot push it down since it may fail over a - // value that isn't referenced by any code. - if !array.all_values_referenced && sig.is_fallible() { + // A fallible function could fail on an unreferenced value that row-wise evaluation would + // never visit. + if !array.has_all_values_referenced() && !signature.is_infallible() { tracing::trace!( "Not pushing down fallible scalar function {} over dictionary with sparse codes {}", - parent.scalar_fn(), + scalar_fn, Dict.id(), ); return Ok(None); } - // Check that all siblings are constant - // TODO(ngates): we can also support other dictionaries if the values are the same! - if !parent + // TODO(ngates): Support dictionary siblings when their values match. + let other_children_are_constant = parent .iter_children() .enumerate() - .all(|(idx, c)| idx == child_idx || c.is::()) - { + .all(|(idx, child)| idx == child_idx || child.is::()); + if !other_children_are_constant { return Ok(None); } - // If the scalar function is null-sensitive, then we cannot push it down to values if - // we have any nulls in the codes. - if array.codes().dtype().is_nullable() + // Before this rewrite, a null code supplies null for this argument while the constant + // arguments retain their values. After the rewrite, the null code masks the function's + // result. Those are equivalent only for a strict function. + let codes_have_nulls = array.codes().dtype().is_nullable() && !matches!( array.codes().validity()?, Validity::NonNullable | Validity::AllValid - ) - && sig.is_null_sensitive() - { + ); + if codes_have_nulls && !signature.is_strict() { tracing::trace!( - "Not pushing down null-sensitive scalar function {} over dictionary with null codes {}", - parent.scalar_fn(), + "Not pushing down non-strict scalar function {} over dictionary with null codes {}", + scalar_fn, Dict.id(), ); return Ok(None); } - // Now we push the parent scalar function into the dictionary values. let values_len = array.values().len(); - let mut new_children = Vec::with_capacity(parent.nchildren()); + let mut value_children = Vec::with_capacity(parent.nchildren()); for (idx, child) in parent.iter_children().enumerate() { if idx == child_idx { - new_children.push(array.values().clone()); + value_children.push(array.values().clone()); } else { let scalar = child.as_::().scalar().clone(); - new_children.push(ConstantArray::new(scalar, values_len).into_array()); + value_children.push(ConstantArray::new(scalar, values_len).into_array()); } } - let new_values = ScalarFnArray::try_new(parent.scalar_fn().clone(), new_children)? + let transformed_values = ScalarFnArray::try_new(scalar_fn.clone(), value_children)? .into_array() .optimize()?; - // We can only push down null-sensitive functions when we have all-valid codes. - // In these cases, we cannot have the codes influence the nullability of the output DType. - // Therefore, we cast the codes to be non-nullable and then cast the dictionary output - // back to nullable if needed. - if sig.is_null_sensitive() && array.codes().dtype().is_nullable() { - let new_codes = array.codes().cast(array.codes().dtype().as_nonnullable())?; - let new_dict = unsafe { - DictArray::new_unchecked(new_codes, new_values) + // A non-strict function reaches this point only when the codes are all valid, but their + // dtype may still be nullable. Remove that declared nullability while rebuilding the + // dictionary, then cast its output to the function's declared dtype. + if !signature.is_strict() && array.codes().dtype().is_nullable() { + let non_nullable_codes = array.codes().cast(array.codes().dtype().as_nonnullable())?; + + // SAFETY: The validity guard proves that the codes contain no nulls. Removing their + // declared nullability preserves every code, and `transformed_values` has one entry + // for each original dictionary value. + let transformed_dict = unsafe { + DictArray::new_unchecked(non_nullable_codes, transformed_values) .set_all_values_referenced(array.has_all_values_referenced()) } .into_array(); - return Ok(Some(new_dict.cast(parent.dtype().clone())?)); + + return Ok(Some(transformed_dict.cast(parent.dtype().clone())?)); } - Ok(Some(unsafe { - DictArray::new_unchecked(array.codes().clone(), new_values) + // SAFETY: The codes are unchanged and `transformed_values` has one entry for each original + // dictionary value, so code bounds and `all_values_referenced` remain unchanged. + let transformed_dict = unsafe { + DictArray::new_unchecked(array.codes().clone(), transformed_values) .set_all_values_referenced(array.has_all_values_referenced()) - .into_array() - })) + }; + + Ok(Some(transformed_dict.into_array())) } } @@ -271,18 +274,21 @@ mod tests { use crate::arrays::BoolArray; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; + use crate::arrays::ConstantArray; use crate::arrays::Dict; use crate::arrays::DictArray; use crate::arrays::PrimitiveArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::dict::DictArrayExt; use crate::arrays::dict::DictArraySlotsExt; - use crate::arrays::scalar_fn::ScalarFnFactoryExt; use crate::assert_arrays_eq; + use crate::dtype::Nullability; use crate::executor::VortexSessionExecute; use crate::optimizer::ArrayOptimizer; - use crate::scalar_fn::EmptyOptions; + use crate::scalar::Scalar; + use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::not::Not; + use crate::scalar_fn::fns::operators::Operator; #[test] #[allow(clippy::disallowed_methods)] @@ -345,13 +351,45 @@ mod tests { } .into_array(); - let result = Not - .try_new_array(dict.len(), EmptyOptions, [dict])? - .optimize()?; + let result = Not::try_new(dict)?.into_array().optimize()?; let result = result.as_::(); assert!(result.has_all_values_referenced()); Ok(()) } + + #[test] + fn kleene_and_dict_values_pushdown_counterexample() -> VortexResult<()> { + let codes = PrimitiveArray::from_option_iter([Some(0u8), Some(1), None]).into_array(); + let values = BoolArray::from_iter([true, false]).into_array(); + let dict = DictArray::try_new(codes, values)?.into_array(); + let const_false = + ConstantArray::new(Scalar::bool(false, Nullability::NonNullable), 3).into_array(); + let expr = Binary::try_new(dict, const_false, Operator::And)?.into_array(); + + let mut ctx = array_session().create_execution_ctx(); + // Kleene AND: null AND false == false, so all three rows must be valid `false`. + let expected = expr.clone().execute::(&mut ctx)?.into_array(); + let optimized = expr.optimize()?; + assert_arrays_eq!(optimized, expected, &mut ctx); + Ok(()) + } + + #[test] + fn kleene_or_dict_values_pushdown_counterexample() -> VortexResult<()> { + let codes = PrimitiveArray::from_option_iter([Some(0u8), Some(1), None]).into_array(); + let values = BoolArray::from_iter([true, false]).into_array(); + let dict = DictArray::try_new(codes, values)?.into_array(); + let const_true = + ConstantArray::new(Scalar::bool(true, Nullability::NonNullable), 3).into_array(); + let expr = Binary::try_new(dict, const_true, Operator::Or)?.into_array(); + + let mut ctx = array_session().create_execution_ctx(); + // Kleene OR: null OR true == true, so all three rows must be valid `true`. + let expected = expr.clone().execute::(&mut ctx)?.into_array(); + let optimized = expr.optimize()?; + assert_arrays_eq!(optimized, expected, &mut ctx); + Ok(()) + } } diff --git a/vortex-array/src/arrays/dict/execute.rs b/vortex-array/src/arrays/dict/execute.rs index 47a9a051666..098ce932e04 100644 --- a/vortex-array/src/arrays/dict/execute.rs +++ b/vortex-array/src/arrays/dict/execute.rs @@ -6,10 +6,10 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; +use crate::ArrayView; use crate::Canonical; use crate::CanonicalView; use crate::ExecutionCtx; -use crate::IntoArray; use crate::arrays::Bool; use crate::arrays::BoolArray; use crate::arrays::Decimal; @@ -20,17 +20,22 @@ use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; use crate::arrays::ListView; use crate::arrays::ListViewArray; +use crate::arrays::Map; +use crate::arrays::MapArray; +use crate::arrays::Null; use crate::arrays::NullArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::Struct; use crate::arrays::StructArray; +use crate::arrays::Union; +use crate::arrays::UnionArray; use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; use crate::arrays::VariantArray; use crate::arrays::dict::TakeExecute; use crate::arrays::dict::TakeReduce; -use crate::arrays::variant::VariantArrayExt; +use crate::arrays::variant::VariantArraySlotsExt; /// Take from a canonical array using indices (codes), returning a new canonical array. /// @@ -38,28 +43,29 @@ use crate::arrays::variant::VariantArrayExt; /// by looking up each code in the values array. pub(crate) fn take_canonical( values: CanonicalView, - codes: &PrimitiveArray, + codes: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> VortexResult { - let values = Canonical::from(values); Ok(match values { - Canonical::Null(a) => Canonical::Null(take_null(&a, codes)), - Canonical::Bool(a) => Canonical::Bool(take_bool(&a, codes, ctx)?), - Canonical::Primitive(a) => Canonical::Primitive(take_primitive(&a, codes, ctx)), - Canonical::Decimal(a) => Canonical::Decimal(take_decimal(&a, codes, ctx)), - Canonical::VarBinView(a) => Canonical::VarBinView(take_varbinview(&a, codes, ctx)), - Canonical::List(a) => Canonical::List(take_listview(&a, codes, ctx)), - Canonical::FixedSizeList(a) => { - Canonical::FixedSizeList(take_fixed_size_list(&a, codes, ctx)) + CanonicalView::Null(a) => Canonical::Null(take_null(a, codes)), + CanonicalView::Bool(a) => Canonical::Bool(take_bool(a, codes, ctx)?), + CanonicalView::Primitive(a) => Canonical::Primitive(take_primitive(a, codes, ctx)), + CanonicalView::Decimal(a) => Canonical::Decimal(take_decimal(a, codes, ctx)), + CanonicalView::VarBinView(a) => Canonical::VarBinView(take_varbinview(a, codes, ctx)), + CanonicalView::List(a) => Canonical::List(take_listview(a, codes, ctx)), + CanonicalView::Map(a) => Canonical::Map(take_map(a, codes, ctx)), + CanonicalView::FixedSizeList(a) => { + Canonical::FixedSizeList(take_fixed_size_list(a, codes, ctx)) } - Canonical::Struct(a) => Canonical::Struct(take_struct(&a, codes)), - Canonical::Extension(a) => Canonical::Extension(take_extension(&a, codes, ctx)), - Canonical::Variant(a) => { - let indices = codes.clone().into_array(); + CanonicalView::Struct(a) => Canonical::Struct(take_struct(a, codes)), + CanonicalView::Union(a) => Canonical::Union(take_union(a, codes)), + CanonicalView::Extension(a) => Canonical::Extension(take_extension(a, codes, ctx)), + CanonicalView::Variant(a) => { + let indices = codes.array().clone(); let taken_core_storage = a.core_storage().take(indices.clone())?; let taken_shredded = a .shredded() - .map(|shredded| shredded.take(indices.clone())) + .map(|shredded| shredded.take(indices)) .transpose()?; Canonical::Variant(VariantArray::try_new(taken_core_storage, taken_shredded)?) } @@ -67,32 +73,29 @@ pub(crate) fn take_canonical( } /// Take for NullArray is trivial - just create a new NullArray with the new length. -fn take_null(_array: &NullArray, codes: &PrimitiveArray) -> NullArray { +fn take_null(_array: ArrayView<'_, Null>, codes: ArrayView<'_, Primitive>) -> NullArray { NullArray::new(codes.len()) } -// TODO(joe): use dict_bool_take fn take_bool( - array: &BoolArray, - codes: &PrimitiveArray, + array: ArrayView<'_, Bool>, + codes: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> VortexResult { - let codes_ref = codes.clone().into_array(); - let array = array.as_view(); - Ok(::take(array, &codes_ref, ctx)? + let codes_ref = codes.array(); + Ok(::take(array, codes_ref, ctx)? .vortex_expect("take bool should not return None") .as_::() .into_owned()) } fn take_primitive( - array: &PrimitiveArray, - codes: &PrimitiveArray, + array: ArrayView<'_, Primitive>, + codes: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> PrimitiveArray { - let codes_ref = codes.clone().into_array(); - let array = array.as_view(); - ::take(array, &codes_ref, ctx) + let codes_ref = codes.array(); + ::take(array, codes_ref, ctx) .vortex_expect("take primitive array") .vortex_expect("take primitive should not return None") .as_::() @@ -100,13 +103,12 @@ fn take_primitive( } fn take_decimal( - array: &DecimalArray, - codes: &PrimitiveArray, + array: ArrayView<'_, Decimal>, + codes: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> DecimalArray { - let codes_ref = codes.clone().into_array(); - let array = array.as_view(); - ::take(array, &codes_ref, ctx) + let codes_ref = codes.array(); + ::take(array, codes_ref, ctx) .vortex_expect("take decimal array") .vortex_expect("take decimal should not return None") .as_::() @@ -114,13 +116,12 @@ fn take_decimal( } fn take_varbinview( - array: &VarBinViewArray, - codes: &PrimitiveArray, + array: ArrayView<'_, VarBinView>, + codes: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> VarBinViewArray { - let codes_ref = codes.clone().into_array(); - let array = array.as_view(); - ::take(array, &codes_ref, ctx) + let codes_ref = codes.array(); + ::take(array, codes_ref, ctx) .vortex_expect("take varbinview array") .vortex_expect("take varbinview should not return None") .as_::() @@ -128,51 +129,68 @@ fn take_varbinview( } fn take_listview( - array: &ListViewArray, - codes: &PrimitiveArray, + array: ArrayView<'_, ListView>, + codes: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> ListViewArray { - let codes_ref = codes.clone().into_array(); - let array = array.as_view(); - ::take(array, &codes_ref, ctx) + let codes_ref = codes.array(); + ::take(array, codes_ref, ctx) .vortex_expect("take listview execute") .vortex_expect("ListView TakeExecute should not return None") .as_::() .into_owned() } +fn take_map( + array: ArrayView<'_, Map>, + codes: ArrayView<'_, Primitive>, + ctx: &mut ExecutionCtx, +) -> MapArray { + ::take(array, codes.array(), ctx) + .vortex_expect("take map execute") + .vortex_expect("Map TakeExecute should not return None") + .as_::() + .into_owned() +} + fn take_fixed_size_list( - array: &FixedSizeListArray, - codes: &PrimitiveArray, + array: ArrayView<'_, FixedSizeList>, + codes: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> FixedSizeListArray { - let codes_ref = codes.clone().into_array(); - let array = array.as_view(); - ::take(array, &codes_ref, ctx) + let codes_ref = codes.array(); + ::take(array, codes_ref, ctx) .vortex_expect("take fixed size list array") .vortex_expect("take fixed size list should not return None") .as_::() .into_owned() } -fn take_struct(array: &StructArray, codes: &PrimitiveArray) -> StructArray { - let codes_ref = codes.clone().into_array(); - let array = array.as_view(); - ::take(array, &codes_ref) +fn take_struct(array: ArrayView<'_, Struct>, codes: ArrayView<'_, Primitive>) -> StructArray { + let codes_ref = codes.array(); + ::take(array, codes_ref) .vortex_expect("take struct array") .vortex_expect("take struct should not return None") .as_::() .into_owned() } +fn take_union(array: ArrayView<'_, Union>, codes: ArrayView<'_, Primitive>) -> UnionArray { + let codes_ref = codes.array(); + ::take(array, codes_ref) + .vortex_expect("take union array") + .vortex_expect("take union should not return None") + .as_::() + .into_owned() +} + fn take_extension( - array: &ExtensionArray, - codes: &PrimitiveArray, + array: ArrayView<'_, Extension>, + codes: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> ExtensionArray { - let codes_ref = codes.clone().into_array(); - let array = array.as_view(); - ::take(array, &codes_ref, ctx) + let codes_ref = codes.array(); + ::take(array, codes_ref, ctx) .vortex_expect("take extension storage") .vortex_expect("take extension should not return None") .as_::() diff --git a/vortex-array/src/arrays/dict/take.rs b/vortex-array/src/arrays/dict/take.rs index 34274e7a77e..38f07d6b0a3 100644 --- a/vortex-array/src/arrays/dict/take.rs +++ b/vortex-array/src/arrays/dict/take.rs @@ -53,11 +53,11 @@ pub trait TakeExecute: VTable { ) -> VortexResult>; } -/// Common preconditions for take operations that apply to all arrays. +/// Short-circuits take for the inputs that need no encoding-specific work. /// -/// Returns `Some(result)` if the precondition short-circuits the take operation, -/// or `None` if the take should proceed normally. -fn precondition(array: ArrayView<'_, V>, indices: &ArrayRef) -> Option { +/// Returns `Some(result)` when the answer is already known, or `None` when take must proceed +/// normally. +fn short_circuit(array: ArrayView<'_, V>, indices: &ArrayRef) -> Option { // Fast-path for empty indices. if indices.is_empty() { let result_dtype = array @@ -97,7 +97,7 @@ where if child_idx != 1 { return Ok(None); } - if let Some(result) = precondition::(array, parent.codes()) { + if let Some(result) = short_circuit::(array, parent.codes()) { return Ok(Some(result)); } let result = ::take(array, parent.codes())?; @@ -128,7 +128,7 @@ where if child_idx != 1 { return Ok(None); } - if let Some(result) = precondition::(array, parent.codes()) { + if let Some(result) = short_circuit::(array, parent.codes()) { return Ok(Some(result)); } let result = ::take(array, parent.codes(), ctx)?; diff --git a/vortex-array/src/arrays/dict/vtable/mod.rs b/vortex-array/src/arrays/dict/vtable/mod.rs index d9ced9eeffa..2eca82b5368 100644 --- a/vortex-array/src/arrays/dict/vtable/mod.rs +++ b/vortex-array/src/arrays/dict/vtable/mod.rs @@ -3,6 +3,7 @@ use std::hash::Hasher; +use num_traits::AsPrimitive; use prost::Message; use smallvec::smallvec; use vortex_error::VortexResult; @@ -10,6 +11,8 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_error::vortex_panic; +use vortex_mask::AllOr; +use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -24,6 +27,7 @@ use crate::ArrayEq; use crate::ArrayHash; use crate::ArrayRef; use crate::Canonical; +use crate::CanonicalView; use crate::EqMode; use crate::IntoArray; use crate::array::Array; @@ -34,17 +38,23 @@ use crate::array::VTable; use crate::array::with_empty_buffers; use crate::arrays::ConstantArray; use crate::arrays::Primitive; +use crate::arrays::VarBinView; use crate::arrays::dict::DictArrayExt; use crate::arrays::dict::DictArraySlotsExt; use crate::arrays::dict::compute::rules::PARENT_RULES; use crate::arrays::dict::execute::take_canonical; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; +use crate::builders::VarBinBuilder; +use crate::builders::VarBinViewBuilder; use crate::dtype::DType; use crate::dtype::Nullability; +use crate::dtype::OffsetBuilderPType; use crate::dtype::PType; use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; +use crate::match_each_integer_ptype; +use crate::match_each_varbin_builder; use crate::require_child; use crate::scalar::Scalar; use crate::serde::ArrayChildren; @@ -205,7 +215,7 @@ impl VTable for Dict { Ok(ExecutionResult::done(take_canonical( values.as_::(), - &codes.downcast::(), + codes.as_::(), ctx, )?)) } @@ -222,8 +232,21 @@ impl VTable for Dict { ) && !codes.validity()?.definitely_all_null() { - let codes = codes.into_owned(); - let canonical = take_canonical(values, &codes, ctx)?.into_array(); + if let CanonicalView::VarBinView(values) = values + && let Some(result) = match_each_varbin_builder!(builder, |builder| { + let validity = array.validity()?.execute_mask(array.len(), ctx)?; + append_dict_to_varbin(codes, values, validity, builder) + }) + { + return result; + } + if let CanonicalView::VarBinView(values) = values + && let Some(builder) = builder.as_any_mut().downcast_mut::() + { + let validity = array.validity()?.execute_mask(array.len(), ctx)?; + return append_dict_to_varbinview(codes, values, validity, builder); + } + let canonical = take_canonical(values, codes, ctx)?.into_array(); canonical.append_to_builder(builder, ctx)?; return Ok(()); } @@ -245,3 +268,186 @@ impl VTable for Dict { PARENT_RULES.evaluate(array, parent, child_idx) } } + +/// Gathers the dictionary values straight into `builder` as views. +/// +/// The canonical route first takes the values to full logical length — allocating an intermediate +/// views buffer proportional to the row count — and then `append_varbinview_array` walks all +/// those views a second time to rebase their buffer indices onto the builder's. Gathering through +/// the codes directly instead adopts the dictionary's buffers once (deduplicated across chunks +/// sharing the dictionary, when the builder deduplicates) and writes each row's rebased view in a +/// single pass, with no byte copy at all. +fn append_dict_to_varbinview( + codes: ArrayView<'_, Primitive>, + values: ArrayView<'_, VarBinView>, + validity: Mask, + builder: &mut VarBinViewBuilder, +) -> VortexResult<()> { + let views = values.views(); + let buffers = values + .data_buffers() + .iter() + .map(|buffer| buffer.as_host().clone()) + .collect::>(); + + match_each_integer_ptype!(codes.ptype(), |C| { + let codes = codes.as_slice::(); + builder.append_views_gathered(buffers, views, &validity, |row| { + AsPrimitive::::as_(codes[row]) + }); + }); + Ok(()) +} + +/// Gathers the dictionary values straight into `builder`. +/// +/// The canonical route first takes the values to full logical length, which allocates and then +/// re-reads a views buffer proportional to the row count. The dictionary is usually far smaller +/// than the column, so resolving each code against it in place skips that intermediate entirely +/// and leaves one `memcpy` per row as the only work. +fn append_dict_to_varbin( + codes: ArrayView<'_, Primitive>, + values: ArrayView<'_, VarBinView>, + validity: Mask, + builder: &mut VarBinBuilder, +) -> VortexResult<()> +where + usize: AsPrimitive, +{ + let len = codes.as_ref().len(); + + // Resolve the dictionary's storage once so that looking up a code is an O(1) read. + let views = values.views(); + let buffers = values + .data_buffers() + .iter() + .map(|buffer| buffer.as_host().as_slice()) + .collect::>(); + + match_each_integer_ptype!(codes.ptype(), |C| { + let codes = codes.as_slice::(); + let view = |row: usize| &views[AsPrimitive::::as_(codes[row])]; + + // Both passes below resolve a row through its code, so the byte total comes from the same + // walk over the valid rows that the copy will make. + let num_bytes = match validity.bit_buffer() { + AllOr::All => (0..len).map(|row| view(row).len() as usize).sum(), + AllOr::None => { + builder.push_nulls(len); + return Ok(()); + } + AllOr::Some(bits) => { + let mut total = 0; + bits.for_each_set_index(|row| total += view(row).len() as usize); + total + } + }; + + builder.append_valid_slices(num_bytes, &validity, |row| view(row).bytes(&buffers)) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::PrimitiveArray; + use crate::arrays::VarBinViewArray; + use crate::arrays::dict::DictArray; + use crate::assert_arrays_eq; + use crate::dtype::Nullability::Nullable; + + const LONG: &str = "a string that is far too long to be inlined in a view"; + + #[test] + fn append_to_builder_gathers_through_the_dictionary() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dict = DictArray::try_new( + PrimitiveArray::from_option_iter([Some(0u32), Some(2), None, Some(1), Some(0)]) + .into_array(), + VarBinViewArray::from_iter([Some(LONG), None, Some("short")], DType::Utf8(Nullable)) + .into_array(), + )?; + + let mut builder = VarBinBuilder::::new_in(DType::Utf8(Nullable), ctx.allocator()); + dict.append_to_builder(&mut builder, &mut ctx)?; + + let expected = VarBinViewArray::from_iter( + [Some(LONG), Some("short"), None, None, Some(LONG)], + DType::Utf8(Nullable), + ); + assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx); + Ok(()) + } + + /// The view-builder path gathers views through the codes without materializing the taken + /// array, so the result must still match the canonical take — including rows whose code is + /// null and rows whose dictionary value is null. + #[test] + fn append_to_view_builder_gathers_through_the_dictionary() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dict = DictArray::try_new( + PrimitiveArray::from_option_iter([Some(0u32), Some(2), None, Some(1), Some(0)]) + .into_array(), + VarBinViewArray::from_iter([Some(LONG), None, Some("short")], DType::Utf8(Nullable)) + .into_array(), + )?; + + let mut builder = VarBinViewBuilder::with_capacity_in( + DType::Utf8(Nullable), + 8, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); + builder.append_value(LONG); + dict.append_to_builder(&mut builder, &mut ctx)?; + + let expected = VarBinViewArray::from_iter( + [ + Some(LONG), + Some(LONG), + Some("short"), + None, + None, + Some(LONG), + ], + DType::Utf8(Nullable), + ); + assert_arrays_eq!(builder.finish_into_varbinview(), expected, &mut ctx); + Ok(()) + } + + /// Two chunks gathered through the same dictionary into a deduplicating builder must adopt + /// the dictionary's data buffers once. + #[test] + fn append_to_dedup_view_builder_adopts_the_dictionary_once() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let values = VarBinViewArray::from_iter([Some(LONG), Some("short")], DType::Utf8(Nullable)); + assert_eq!(values.data_buffers().len(), 1); + + let first = DictArray::try_new( + PrimitiveArray::from_option_iter([Some(0u32), Some(1)]).into_array(), + values.clone().into_array(), + )?; + let second = DictArray::try_new( + PrimitiveArray::from_option_iter([Some(1u32), Some(0)]).into_array(), + values.into_array(), + )?; + + let mut builder = VarBinViewBuilder::with_buffer_deduplication_in( + DType::Utf8(Nullable), + 8, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); + first.append_to_builder(&mut builder, &mut ctx)?; + second.append_to_builder(&mut builder, &mut ctx)?; + assert_eq!(builder.completed_block_count(), 1); + + let expected = VarBinViewArray::from_iter( + [Some(LONG), Some("short"), Some("short"), Some(LONG)], + DType::Utf8(Nullable), + ); + assert_arrays_eq!(builder.finish_into_varbinview(), expected, &mut ctx); + Ok(()) + } +} diff --git a/vortex-array/src/arrays/extension/array.rs b/vortex-array/src/arrays/extension/array.rs index 0fd57030cdd..144e2d82622 100644 --- a/vortex-array/src/arrays/extension/array.rs +++ b/vortex-array/src/arrays/extension/array.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use smallvec::smallvec; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; @@ -11,18 +10,21 @@ use crate::EmptyArrayData; use crate::array::Array; use crate::array::ArrayParts; use crate::array::TypedArrayRef; +use crate::array_slots; use crate::arrays::Extension; use crate::dtype::DType; use crate::dtype::extension::ExtDType; use crate::dtype::extension::ExtDTypeRef; use crate::dtype::extension::ExtVTable; -/// The backing storage array for this extension array. -pub(super) const STORAGE_SLOT: usize = 0; -pub(super) const NUM_SLOTS: usize = 1; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["storage"]; +#[array_slots(Extension)] +pub struct ExtensionSlots { + /// The backing storage array for this extension array. + #[slot(0)] + pub storage: ArrayRef, +} -pub trait ExtensionArrayExt: TypedArrayRef { +pub trait ExtensionArrayExt: TypedArrayRef + ExtensionArraySlotsExt { fn ext_dtype(&self) -> &ExtDTypeRef { self.as_ref() .dtype() @@ -30,10 +32,9 @@ pub trait ExtensionArrayExt: TypedArrayRef { .vortex_expect("extension array somehow did not have an extension dtype") } + /// Returns the backing storage array. fn storage_array(&self) -> &ArrayRef { - self.as_ref().slots()[STORAGE_SLOT] - .as_ref() - .vortex_expect("ExtensionArray storage slot") + self.storage() } } impl> ExtensionArrayExt for T {} @@ -59,8 +60,12 @@ impl Array { let dtype = DType::Extension(ext_dtype); let len = storage_array.len(); - let parts = ArrayParts::new(Extension, dtype, len, EmptyArrayData) - .with_slots(smallvec![Some(storage_array)]); + let parts = ArrayParts::new(Extension, dtype, len, EmptyArrayData).with_slots( + ExtensionSlots { + storage: storage_array, + } + .into_slots(), + ); Ok(unsafe { Array::from_parts_unchecked(parts) }) } diff --git a/vortex-array/src/arrays/extension/compute/mask.rs b/vortex-array/src/arrays/extension/compute/mask.rs index 817df8d7fc8..10ae5e904bc 100644 --- a/vortex-array/src/arrays/extension/compute/mask.rs +++ b/vortex-array/src/arrays/extension/compute/mask.rs @@ -9,18 +9,13 @@ use crate::array::ArrayView; use crate::arrays::Extension; use crate::arrays::ExtensionArray; use crate::arrays::extension::ExtensionArrayExt; -use crate::arrays::scalar_fn::ScalarFnFactoryExt; -use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::fns::mask::Mask as MaskExpr; use crate::scalar_fn::fns::mask::MaskReduce; impl MaskReduce for Extension { fn mask(array: ArrayView<'_, Extension>, mask: &ArrayRef) -> VortexResult> { - let masked_storage = MaskExpr.try_new_array( - array.storage_array().len(), - EmptyOptions, - [array.storage_array().clone(), mask.clone()], - )?; + let masked_storage = + MaskExpr::try_new(array.storage_array().clone(), mask.clone())?.into_array(); Ok(Some( ExtensionArray::new( array diff --git a/vortex-array/src/arrays/extension/compute/rules.rs b/vortex-array/src/arrays/extension/compute/rules.rs index 94ebf0981c3..8169fa527d7 100644 --- a/vortex-array/src/arrays/extension/compute/rules.rs +++ b/vortex-array/src/arrays/extension/compute/rules.rs @@ -11,6 +11,7 @@ use crate::arrays::ConstantArray; use crate::arrays::Extension; use crate::arrays::ExtensionArray; use crate::arrays::Filter; +use crate::arrays::dict::TakeReduceAdaptor; use crate::arrays::extension::ExtensionArrayExt; use crate::arrays::filter::FilterReduceAdaptor; use crate::arrays::slice::SliceReduceAdaptor; @@ -50,6 +51,7 @@ pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&FilterReduceAdaptor(Extension)), ParentRuleSet::lift(&MaskReduceAdaptor(Extension)), ParentRuleSet::lift(&SliceReduceAdaptor(Extension)), + ParentRuleSet::lift(&TakeReduceAdaptor(Extension)), ]); /// Push filter operations into the storage array of an extension array. @@ -97,7 +99,6 @@ mod tests { use crate::arrays::ScalarFn; use crate::arrays::extension::ExtensionArrayExt; use crate::arrays::scalar_fn::ScalarFnArrayExt; - use crate::arrays::scalar_fn::ScalarFnFactoryExt; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -227,9 +228,9 @@ mod tests { let storage = buffer![15i64, 25, 35].into_array(); let ext_array = ExtensionArray::new(ext_dtype, storage).into_array(); - let scalar_fn_array = Binary - .try_new_array(3, Operator::Lt, [constant_ext, ext_array]) - .unwrap(); + let scalar_fn_array = Binary::try_new(constant_ext, ext_array, Operator::Lt) + .unwrap() + .into_array(); let optimized = scalar_fn_array.optimize_recursive(&SESSION).unwrap(); let scalar_fn = optimized.as_opt::().unwrap(); @@ -289,9 +290,9 @@ mod tests { let const_scalar = Scalar::extension::(EmptyMetadata, Scalar::from(25i64)); let const_array = ConstantArray::new(const_scalar, 3).into_array(); - let scalar_fn_array = Binary - .try_new_array(3, Operator::Lt, [ext_array, const_array]) - .unwrap(); + let scalar_fn_array = Binary::try_new(ext_array, const_array, Operator::Lt) + .unwrap() + .into_array(); let optimized = scalar_fn_array.optimize().unwrap(); @@ -314,9 +315,9 @@ mod tests { let ext_array2 = ExtensionArray::new(ext_dtype, storage2).into_array(); // Both children are extension arrays (not constants) - let scalar_fn_array = Binary - .try_new_array(3, Operator::Lt, [ext_array1, ext_array2]) - .unwrap(); + let scalar_fn_array = Binary::try_new(ext_array1, ext_array2, Operator::Lt) + .unwrap() + .into_array(); let optimized = scalar_fn_array.optimize().unwrap(); @@ -337,9 +338,9 @@ mod tests { // Create a non-extension constant (plain primitive) let const_array = ConstantArray::new(Scalar::from(25i64), 3).into_array(); - let scalar_fn_array = Binary - .try_new_array(3, Operator::Lt, [ext_array, const_array]) - .unwrap(); + let scalar_fn_array = Binary::try_new(ext_array, const_array, Operator::Lt) + .unwrap() + .into_array(); let optimized = scalar_fn_array.optimize().unwrap(); diff --git a/vortex-array/src/arrays/extension/compute/take.rs b/vortex-array/src/arrays/extension/compute/take.rs index 8fb0da11222..d9ac1b4f600 100644 --- a/vortex-array/src/arrays/extension/compute/take.rs +++ b/vortex-array/src/arrays/extension/compute/take.rs @@ -10,8 +10,24 @@ use crate::array::ArrayView; use crate::arrays::Extension; use crate::arrays::ExtensionArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::dict::TakeReduce; use crate::arrays::extension::ExtensionArrayExt; +impl TakeReduce for Extension { + fn take(array: ArrayView<'_, Extension>, indices: &ArrayRef) -> VortexResult> { + let taken_storage = array.storage_array().take(indices.clone())?; + Ok(Some( + ExtensionArray::new( + array + .ext_dtype() + .with_nullability(taken_storage.dtype().nullability()), + taken_storage, + ) + .into_array(), + )) + } +} + impl TakeExecute for Extension { fn take( array: ArrayView<'_, Extension>, diff --git a/vortex-array/src/arrays/extension/mod.rs b/vortex-array/src/arrays/extension/mod.rs index 8dc5972b86f..f96049d2cf8 100644 --- a/vortex-array/src/arrays/extension/mod.rs +++ b/vortex-array/src/arrays/extension/mod.rs @@ -3,6 +3,9 @@ mod array; pub use array::ExtensionArrayExt; +pub use array::ExtensionArraySlotsExt; +pub use array::ExtensionSlots; +pub use array::ExtensionSlotsView; pub use vtable::ExtensionArray; pub(crate) mod compute; diff --git a/vortex-array/src/arrays/extension/vtable/mod.rs b/vortex-array/src/arrays/extension/vtable/mod.rs index ef371e46afc..2937a2b98b9 100644 --- a/vortex-array/src/arrays/extension/vtable/mod.rs +++ b/vortex-array/src/arrays/extension/vtable/mod.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use smallvec::smallvec; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -22,8 +21,7 @@ use crate::array::ArrayView; use crate::array::VTable; use crate::array::ValidityVTableFromChild; use crate::array::with_empty_buffers; -use crate::arrays::extension::array::SLOT_NAMES; -use crate::arrays::extension::array::STORAGE_SLOT; +use crate::arrays::extension::array::ExtensionSlots; use crate::arrays::extension::compute::rules::PARENT_RULES; use crate::arrays::extension::compute::rules::RULES; use crate::buffer::BufferHandle; @@ -102,7 +100,7 @@ impl VTable for Extension { len: usize, slots: &[Option], ) -> VortexResult<()> { - let storage = slots[STORAGE_SLOT] + let storage = slots[ExtensionSlots::STORAGE] .as_ref() .vortex_expect("ExtensionArray storage slot"); vortex_ensure_eq!( @@ -178,12 +176,12 @@ impl VTable for Extension { let storage = children.get(0, ext_dtype.storage_dtype(), len)?; Ok( ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData) - .with_slots(smallvec![Some(storage)]), + .with_slots(ExtensionSlots { storage }.into_slots()), ) } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + ExtensionSlots::NAMES[idx].to_string() } fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { diff --git a/vortex-array/src/arrays/filter/array.rs b/vortex-array/src/arrays/filter/array.rs index dd811ad24a2..5370983c1a2 100644 --- a/vortex-array/src/arrays/filter/array.rs +++ b/vortex-array/src/arrays/filter/array.rs @@ -4,8 +4,6 @@ use std::fmt::Display; use std::fmt::Formatter; -use smallvec::smallvec; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; use vortex_mask::Mask; @@ -13,13 +11,15 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::array::Array; use crate::array::ArrayParts; -use crate::array::TypedArrayRef; +use crate::array_slots; use crate::arrays::Filter; -/// The source array being filtered. -pub(super) const CHILD_SLOT: usize = 0; -pub(super) const NUM_SLOTS: usize = 1; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["child"]; +#[array_slots(Filter)] +pub struct FilterSlots { + /// The source array being filtered. + #[slot(0)] + pub child: ArrayRef, +} // TODO(connor): Write docs on why we have this, and what we had in the old world so that the future // does not repeat the mistakes of the past. @@ -42,15 +42,6 @@ pub struct FilterDataParts { pub mask: Mask, } -pub trait FilterArrayExt: TypedArrayRef { - fn child(&self) -> &ArrayRef { - self.as_ref().slots()[CHILD_SLOT] - .as_ref() - .vortex_expect("validated filter child slot") - } -} -impl> FilterArrayExt for T {} - impl FilterData { pub fn new(mask: Mask) -> Self { Self { mask } @@ -96,7 +87,8 @@ impl Array { let data = FilterData::new(mask); unsafe { Array::from_parts_unchecked( - ArrayParts::new(Filter, dtype, len, data).with_slots(smallvec![Some(array)]), + ArrayParts::new(Filter, dtype, len, data) + .with_slots(FilterSlots { child: array }.into_slots()), ) } } @@ -108,7 +100,8 @@ impl Array { let data = FilterData::try_new(array.len(), mask)?; Ok(unsafe { Array::from_parts_unchecked( - ArrayParts::new(Filter, dtype, len, data).with_slots(smallvec![Some(array)]), + ArrayParts::new(Filter, dtype, len, data) + .with_slots(FilterSlots { child: array }.into_slots()), ) }) } diff --git a/vortex-array/src/arrays/filter/execute/bool.rs b/vortex-array/src/arrays/filter/execute/bool.rs index 08364e3cac5..2a9f60cb1b1 100644 --- a/vortex-array/src/arrays/filter/execute/bool.rs +++ b/vortex-array/src/arrays/filter/execute/bool.rs @@ -1,20 +1,18 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::sync::Arc; - use vortex_error::VortexExpect; -use vortex_mask::MaskValues; +use vortex_mask::MaskValuesRef; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; use crate::arrays::filter::execute::bitbuffer; use crate::arrays::filter::execute::filter_validity; -pub fn filter_bool(array: &BoolArray, mask: &Arc) -> BoolArray { +pub fn filter_bool(array: &BoolArray, mask: &MaskValuesRef) -> BoolArray { let validity = array .validity() - .vortex_expect("bool validity should be derivable"); + .vortex_expect("validity is derivable for a valid BoolArray"); let filtered_validity = filter_validity(validity, mask); let bit_buffer = array.to_bit_buffer(); @@ -24,7 +22,7 @@ pub fn filter_bool(array: &BoolArray, mask: &Arc) -> BoolArray { } #[cfg(test)] -mod test { +mod tests { use itertools::Itertools; use rstest::rstest; use vortex_mask::Mask; diff --git a/vortex-array/src/arrays/filter/execute/buffer.rs b/vortex-array/src/arrays/filter/execute/buffer.rs index 73ee3026654..b8c73863f11 100644 --- a/vortex-array/src/arrays/filter/execute/buffer.rs +++ b/vortex-array/src/arrays/filter/execute/buffer.rs @@ -1,50 +1,165 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Buffer-level filter dispatch. +//! Buffer-level filter dispatch across cached, SIMD, and scalar strategies. //! -//! Provides [`filter_buffer`] which filters a [`Buffer`] by [`MaskValues`], attempting an -//! in-place filter when the buffer has exclusive ownership. +//! Selection is an ordered ladder, not a set of independent choices. The first matching row wins: +//! +//! | Priority | Condition | Out-of-place | In-place | +//! | -------- | --------------------------------------------------- | -------------- | ------------- | +//! | 1 | cached slices: one run, or average run length >= 8 | copy runs | move runs | +//! | 2 | cached indices and density <= 0.50 | gather indices | move indices | +//! | 3 | an eligible architecture kernel | SIMD compress | SIMD compress | +//! | 4 | density reaches [`byte_compress_density_threshold`] | byte compress | bitmap walk | +//! | 5 | otherwise | bitmap walk | bitmap walk | +//! +//! In-place dispatch is attempted only for a uniquely owned buffer with density >= 0.50. SIMD +//! eligibility is summarized in [`simd_compress`]. All-true, all-false, and contiguous masks are +//! handled before buffer dispatch. + +use std::mem::size_of; use vortex_buffer::Buffer; use vortex_mask::MaskValues; +use crate::arrays::filter::execute::byte_compress; +use crate::arrays::filter::execute::simd_compress; use crate::arrays::filter::execute::slice; +const CACHED_INDICES_MAX_DENSITY: f64 = 0.5; +const IN_PLACE_MIN_DENSITY: f64 = 0.5; +const MIN_SLICES_AVERAGE_RUN_LENGTH: usize = 8; + /// Filter a [`Buffer`] by [`MaskValues`], returning a new buffer. /// -/// This will attempt to filter in-place (via [`Buffer::try_into_mut`]) when the buffer has -/// exclusive ownership, avoiding an extra allocation. -pub(super) fn filter_buffer(buffer: Buffer, mask: &MaskValues) -> Buffer { - match buffer.try_into_mut() { - Ok(mut buffer_mut) => { - let new_len = slice::filter_slice_mut_by_mask_values(buffer_mut.as_mut_slice(), mask); - buffer_mut.truncate(new_len); - buffer_mut.freeze() +/// Dense uniquely owned buffers are compacted in place; other buffers allocate a new output. +pub(crate) fn filter_buffer(buffer: Buffer, mask: &MaskValues) -> Buffer { + assert_eq!(buffer.len(), mask.len()); + + let buffer = if mask.density() >= IN_PLACE_MIN_DENSITY { + match buffer.try_into_mut() { + Ok(mut buffer_mut) => { + let new_len = filter_slice_in_place(buffer_mut.as_mut_slice(), mask); + buffer_mut.truncate(new_len); + return buffer_mut.freeze(); + } + Err(buffer) => buffer, } - // Otherwise, allocate a new buffer and fill it in. - Err(buffer) => slice::filter_slice_by_mask_values(buffer.as_slice(), mask), + } else { + buffer + }; + + filter_slice(buffer.as_slice(), mask) +} + +fn filter_slice(values: &[T], mask: &MaskValues) -> Buffer { + if let Some(slices) = useful_cached_slices(mask) { + return slice::filter_slice_by_slices(values, slices, mask.true_count()); + } + + if mask.density() <= CACHED_INDICES_MAX_DENSITY + && let Some(indices) = mask.cached_indices() + { + return slice::filter_slice_by_indices(values, indices); + } + + if let Some(filtered) = simd_compress::filter_slice_by_bitmap(values, mask) { + return filtered; + } + + if mask.density() >= byte_compress_density_threshold::() { + byte_compress::filter_buffer(values, mask) + } else { + slice::filter_slice_by_bitmap(values, mask) + } +} + +fn filter_slice_in_place(values: &mut [T], mask: &MaskValues) -> usize { + if let Some(slices) = useful_cached_slices(mask) { + return slice::filter_slice_mut_by_slices(values, slices); + } + + if mask.density() <= CACHED_INDICES_MAX_DENSITY + && let Some(indices) = mask.cached_indices() + { + return slice::filter_slice_mut_by_indices(values, indices); + } + + if let Some(new_len) = simd_compress::filter_slice_mut_by_bitmap(values, mask) { + return new_len; + } + + slice::filter_slice_mut_by_bitmap(values, mask) +} + +fn useful_cached_slices(mask: &MaskValues) -> Option<&[(usize, usize)]> { + mask.cached_slices().filter(|slices| { + slices.len() == 1 || mask.true_count() / slices.len() >= MIN_SLICES_AVERAGE_RUN_LENGTH + }) +} + +fn byte_compress_density_threshold() -> f64 { + let width = size_of::(); + + // A density at or above the table entry selects byte compress after the higher-priority + // strategies have declined the mask. These crossovers are benchmarked in + // `benches/filter_fixed_width.rs`. + // + // | Target | 1 byte | 2 bytes | 4 bytes | 8 bytes | other | + // | ------- | -----: | ------: | ------: | ------: | ----: | + // | aarch64 | 0.90 | 0.90 | 0.90 | 0.75 | 0.90 | + // | other | 0.00 | 0.50 | 0.50 | 0.75 | 0.875 | + if cfg!(target_arch = "aarch64") { + return match width { + 8 => 0.75, + _ => 0.9, + }; + } + + match width { + 1 => 0.0, + 2 | 4 => 0.5, + 8 => 0.75, + _ => 0.875, } } +/// Materialize sparse indices when enough sibling arrays will reuse the same mask. +pub(crate) fn prepare_mask_for_reuse(mask: &MaskValues, consumers: usize) { + if consumers <= 1 || mask.cached_indices().is_some() || mask.cached_slices().is_some() { + return; + } + + let density_threshold = if consumers >= 3 { 0.1 } else { 0.05 }; + if mask.density() > density_threshold { + return; + } + + if super::contiguous_values_range(mask).is_some() { + return; + } + + let _ = mask.indices(); +} + #[cfg(test)] mod tests { + use vortex_buffer::BitBuffer; use vortex_buffer::BufferMut; use vortex_buffer::buffer; use vortex_mask::Mask; use super::*; - // Helper to get `MaskValues` from a `Mask`. fn mask_values(mask: &Mask) -> &MaskValues { match mask { - Mask::Values(v) => v.as_ref(), + Mask::Values(v) => v, _ => panic!("expected Mask::Values"), } } #[test] - fn test_filter_buffer_by_indices() { + fn test_filter_buffer() { let buf = buffer![10u32, 20, 30, 40, 50]; let mask = Mask::from_iter([true, false, true, false, true]); @@ -53,17 +168,8 @@ mod tests { } #[test] - fn test_filter_indices_direct() { - let buf = buffer![100u32, 200, 300, 400]; - let mask = Mask::from_iter([true, false, true, true]); - let result = filter_buffer(buf, mask_values(&mask)); - assert_eq!(result, buffer![100u32, 300, 400]); - } - - #[test] - fn test_filter_sparse() { + fn test_filter_sparse_bitmap() { let buf = Buffer::from(BufferMut::from_iter(0u32..1000)); - // Keep every third element. let mask = Mask::from_iter((0..1000).map(|i| i % 3 == 0)); let result = filter_buffer(buf, mask_values(&mask)); @@ -72,12 +178,78 @@ mod tests { } #[test] - fn test_filter_dense() { + fn test_filter_dense_in_place() { let buf = buffer![1u32, 2, 3, 4, 5, 6, 7, 8, 9, 10]; - // Dense selection (80% selected). let mask = Mask::from_iter([true, true, true, true, false, true, true, true, false, true]); let result = filter_buffer(buf, mask_values(&mask)); assert_eq!(result, buffer![1u32, 2, 3, 4, 6, 7, 8, 10]); } + + #[test] + fn test_filter_shared_buffer_by_cached_indices() { + let buf = Buffer::from(BufferMut::from_iter(0u64..16)); + let shared = buf.clone(); + let mask = Mask::from_indices(16, [1, 5, 9, 15]); + + let result = filter_buffer(buf, mask_values(&mask)); + assert_eq!(result, buffer![1u64, 5, 9, 15]); + assert_eq!(shared.len(), 16); + } + + #[test] + fn test_filter_shared_buffer_by_cached_slices() { + let buf = Buffer::from(BufferMut::from_iter(0u32..32)); + let shared = buf.clone(); + let mask = Mask::from_slices(32, vec![(3, 15), (20, 30)]); + + let result = filter_buffer(buf, mask_values(&mask)); + let expected = (3u32..15).chain(20..30).collect::>(); + assert_eq!(result.as_slice(), expected.as_slice()); + assert_eq!(shared.len(), 32); + } + + #[test] + fn test_filter_unaligned_bitmap_words() { + const LEN: usize = 151; + const OFFSET: usize = 5; + + let backing = BitBuffer::from_iter( + std::iter::repeat_n(false, OFFSET).chain((0..LEN).map(|index| index % 7 == 2)), + ); + let mask = Mask::from_buffer(BitBuffer::new_with_offset( + backing.inner().clone(), + LEN, + OFFSET, + )); + let buf = Buffer::from(BufferMut::from_iter(0u64..LEN as u64)); + + let result = filter_buffer(buf, mask_values(&mask)); + let expected = (0u64..LEN as u64) + .filter(|value| value % 7 == 2) + .collect::>(); + assert_eq!(result.as_slice(), expected.as_slice()); + } + + #[test] + fn test_prepare_sparse_mask_for_sibling_reuse() { + let two_consumers = Mask::from_iter((0..100).map(|index| index % 20 == 0)); + let values = mask_values(&two_consumers); + assert!(values.cached_indices().is_none()); + prepare_mask_for_reuse(values, 2); + assert!(values.cached_indices().is_some()); + + let three_consumers = Mask::from_iter((0..100).map(|index| index % 10 == 0)); + let values = mask_values(&three_consumers); + assert!(values.cached_indices().is_none()); + prepare_mask_for_reuse(values, 2); + assert!(values.cached_indices().is_none()); + prepare_mask_for_reuse(values, 3); + assert!(values.cached_indices().is_some()); + + let contiguous = Mask::from_iter((0..100).map(|index| (20..25).contains(&index))); + let values = mask_values(&contiguous); + prepare_mask_for_reuse(values, 3); + assert!(values.cached_indices().is_none()); + } } diff --git a/vortex-array/src/arrays/filter/execute/byte_compress.rs b/vortex-array/src/arrays/filter/execute/byte_compress.rs index 6c58ce4f33b..032b11a56b3 100644 --- a/vortex-array/src/arrays/filter/execute/byte_compress.rs +++ b/vortex-array/src/arrays/filter/execute/byte_compress.rs @@ -1,20 +1,16 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Byte-level compress for primitive filtering using a `1 << 8 = 256`-entry lookup table. +//! Byte-level compress for fixed-width filtering using a `1 << 8 = 256`-entry lookup table. //! //! For each byte of the mask (8 bits -> 8 source elements), a precomputed //! permutation table compacts the selected bytes in a single indexed copy, //! avoiding the overhead of materializing indices or slices. -use std::mem::size_of; - use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_mask::MaskValues; -const BYTE_COMPRESS_DENSITY_THRESHOLD: f64 = 0.5; - /// For each mask byte (0..256), stores the element indices to keep and the count. /// /// `BYTE_COMPRESS_LUT[mask_byte]` = `([i0, i1, ..., i7], popcount)` where @@ -45,10 +41,10 @@ static BYTE_COMPRESS_LUT: &[([u8; 8], u8); 256] = &{ /// /// Processes the mask one byte at a time (8 source elements per byte), /// using a precomputed permutation to compact selected elements. -pub(super) fn filter_buffer(buffer: Buffer, mask: &MaskValues) -> Buffer { - debug_assert_eq!(buffer.len(), mask.len()); +pub(crate) fn filter_buffer(buffer: impl AsRef<[T]>, mask: &MaskValues) -> Buffer { + let src = buffer.as_ref(); + debug_assert_eq!(src.len(), mask.len()); - let src = buffer.as_slice(); let true_count = mask.true_count(); if true_count == 0 { @@ -59,14 +55,7 @@ pub(super) fn filter_buffer(buffer: Buffer, mask: &MaskValues) -> Bu let mask_bytes = mask_buffer.inner().as_ref(); let mask_offset = mask_buffer.offset(); - // Fast path: byte-wide values benefit from avoiding index materialization more often. Wider - // values need enough selected values to justify scanning every mask byte directly. - if size_of::() == 1 || mask.density() >= BYTE_COMPRESS_DENSITY_THRESHOLD { - return filter_bitpacked(src, mask_bytes, mask_offset, true_count); - } - - // Slow path: lower-density wide values are better handled by the generic path. - super::slice::filter_slice_by_mask_values(src, mask) + filter_bitpacked(src, mask_bytes, mask_offset, true_count) } fn filter_bitpacked( diff --git a/vortex-array/src/arrays/filter/execute/decimal.rs b/vortex-array/src/arrays/filter/execute/decimal.rs deleted file mode 100644 index 5ffb4a963b2..00000000000 --- a/vortex-array/src/arrays/filter/execute/decimal.rs +++ /dev/null @@ -1,70 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::sync::Arc; - -use vortex_error::VortexExpect; -use vortex_mask::MaskValues; - -use crate::arrays::DecimalArray; -use crate::arrays::filter::execute::buffer; -use crate::arrays::filter::execute::filter_validity; -use crate::match_each_decimal_value_type; - -pub fn filter_decimal(array: &DecimalArray, mask: &Arc) -> DecimalArray { - let filtered_validity = filter_validity( - array - .validity() - .vortex_expect("decimal validity should be derivable"), - mask, - ); - - match_each_decimal_value_type!(array.values_type(), |T| { - let filtered_buffer = buffer::filter_buffer(array.buffer::(), mask.as_ref()); - - // SAFETY: We filter both the validity and the buffer with the same mask, so they must have - // the same length, and since the buffer came from an existing and valid `DecimalArray` the - // values must all be be representable by the decimal type. - unsafe { - DecimalArray::new_unchecked(filtered_buffer, array.decimal_dtype(), filtered_validity) - } - }) -} - -#[cfg(test)] -mod test { - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::array_session; - use crate::arrays::filter::execute::decimal::DecimalArray; - use crate::compute::conformance::filter::test_filter_conformance; - use crate::dtype::DecimalDType; - - #[test] - fn test_filter_decimal128_conformance() { - let decimal_dtype = DecimalDType::new(38, 2); - let values = vec![ - Some(12345i128), - Some(67890), - Some(-12345), - Some(0), - Some(99999), - ]; - let array = DecimalArray::from_option_iter(values, decimal_dtype); - test_filter_conformance( - &array.into_array(), - &mut array_session().create_execution_ctx(), - ); - } - - #[test] - fn test_filter_decimal128_with_nulls_conformance() { - let decimal_dtype = DecimalDType::new(38, 4); - let values = vec![Some(12345i128), None, Some(-12345), Some(0), None]; - let array = DecimalArray::from_option_iter(values, decimal_dtype); - test_filter_conformance( - &array.into_array(), - &mut array_session().create_execution_ctx(), - ); - } -} diff --git a/vortex-array/src/arrays/filter/execute/fixed_size_list.rs b/vortex-array/src/arrays/filter/execute/fixed_size_list.rs index 6facbe3e38d..b73fff21a2d 100644 --- a/vortex-array/src/arrays/filter/execute/fixed_size_list.rs +++ b/vortex-array/src/arrays/filter/execute/fixed_size_list.rs @@ -1,16 +1,16 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::sync::Arc; - use vortex_error::VortexExpect; use vortex_mask::Mask; use vortex_mask::MaskIter; use vortex_mask::MaskValues; +use vortex_mask::MaskValuesRef; use crate::arrays::FixedSizeListArray; use crate::arrays::filter::execute::filter_validity; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; /// Density threshold for choosing between indices and slices representation when expanding masks. /// @@ -23,12 +23,12 @@ const MASK_EXPANSION_DENSITY_THRESHOLD: f64 = 0.05; /// mask down to the child elements array. pub fn filter_fixed_size_list( array: &FixedSizeListArray, - selection_mask: &Arc, + selection_mask: &MaskValuesRef, ) -> FixedSizeListArray { let filtered_validity = filter_validity( array .validity() - .vortex_expect("fixed-size-list validity should be derivable"), + .vortex_expect("validity is derivable for a valid FixedSizeListArray"), selection_mask, ); @@ -37,13 +37,11 @@ pub fn filter_fixed_size_list( let list_size = array.list_size(); let new_elements = { - // We want to create a new mask specialized to the underlying `elements` of the array. if list_size != 0 { - // TODO(connor): If we can push down a "indices" or "slices" selection instead that - // would be much more performant. - let elements_mask = compute_mask_for_fsl_elements(selection_mask, list_size as usize); + // TODO(connor): Push down an indices or slices selection to avoid expanding the mask. + let elements_mask = + compute_mask_for_fsl_elements(selection_mask.as_ref(), list_size as usize); - // Allow the child array to filter itself. let new_elements = elements .filter(elements_mask) .vortex_expect("FixedSizeListArray elements are guaranteed to support filter"); @@ -58,22 +56,14 @@ pub fn filter_fixed_size_list( "degenerate FixedSizeListArray is invalid" ); - // NB: The safety comment for the `list_size == 0` case is here for clarity. - - // SAFETY: We have verified that when `list_size == 0` - // - `elements` has length 0 (since it came from a valid `FixedSizeListArray`) - // - `filtered_validity` has the correct length because we filter with the same - // `selection_mask` as the array itself elements.clone() } }; - // SAFETY: We have verified that - // - The case when `list_size == 0` is safe (see above) - // - The `new_elements` array is guaranteed to have a length that is a multiple of - // `list_size` - // - `filtered_validity` has the correct length because we filter with the same - // `selection_mask` as the array itself + // SAFETY: + // - A valid zero-width array has no elements, which the zero-width branch preserves. + // - Otherwise, the expanded selection retains `list_size` elements for each selected list. + // - Filtering validity with `selection_mask` gives it the required `new_len`. unsafe { FixedSizeListArray::new_unchecked(new_elements, list_size, filtered_validity, new_len) } @@ -119,7 +109,7 @@ fn compute_mask_for_fsl_elements(selection_mask: &MaskValues, list_size: usize) } #[cfg(test)] -mod test { +mod tests { use vortex_buffer::buffer; use vortex_mask::Mask; diff --git a/vortex-array/src/arrays/filter/execute/listview.rs b/vortex-array/src/arrays/filter/execute/listview.rs index 6d1521541bc..ba15ee51998 100644 --- a/vortex-array/src/arrays/filter/execute/listview.rs +++ b/vortex-array/src/arrays/filter/execute/listview.rs @@ -1,15 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::sync::Arc; - use vortex_error::VortexExpect; use vortex_mask::Mask; -use vortex_mask::MaskValues; +use vortex_mask::MaskValuesRef; use crate::arrays::ListViewArray; use crate::arrays::filter::execute::filter_validity; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; /// [`ListViewArray`] filter implementation. /// @@ -23,7 +21,7 @@ use crate::arrays::listview::ListViewArrayExt; /// /// The trade-off is that we may keep unreferenced elements in memory, but this is acceptable since /// we're optimizing for read performance and the data isn't being copied. -pub fn filter_listview(array: &ListViewArray, selection_mask: &Arc) -> ListViewArray { +pub fn filter_listview(array: &ListViewArray, selection_mask: &MaskValuesRef) -> ListViewArray { let elements = array.elements(); let offsets = array.offsets(); let sizes = array.sizes(); @@ -31,7 +29,7 @@ pub fn filter_listview(array: &ListViewArray, selection_mask: &Arc) let new_validity = filter_validity( array .validity() - .vortex_expect("listview validity should be derivable"), + .vortex_expect("validity is derivable for a valid ListViewArray"), selection_mask, ); debug_assert!( @@ -40,8 +38,7 @@ pub fn filter_listview(array: &ListViewArray, selection_mask: &Arc) .is_none_or(|len| len == selection_mask.true_count()) ); - // Simply filter the offsets and sizes arrays. - let mask_for_filter = Mask::Values(Arc::clone(selection_mask)); + let mask_for_filter = Mask::Values(MaskValuesRef::clone(selection_mask)); let new_offsets = offsets .filter(mask_for_filter.clone()) .vortex_expect("ListViewArray offsets are guaranteed to support filter"); @@ -49,15 +46,15 @@ pub fn filter_listview(array: &ListViewArray, selection_mask: &Arc) .filter(mask_for_filter) .vortex_expect("ListViewArray sizes are guaranteed to support filter"); - // SAFETY: Filter operation maintains all `ListViewArray` invariants: + // SAFETY: // - Offsets and sizes are derived from existing valid child arrays. - // - Offsets and sizes have the same length (both filtered by `selection_mask`). - // - Validity matches the filtered array's nullability. + // - Filtering offsets, sizes, and validity with `selection_mask` preserves their row alignment. + // - The retained elements buffer keeps every referenced element in bounds. unsafe { ListViewArray::new_unchecked(elements.clone(), new_offsets, new_sizes, new_validity) } } #[cfg(test)] -mod test { +mod tests { use std::sync::LazyLock; use vortex_buffer::buffer; @@ -71,6 +68,7 @@ mod test { use crate::arrays::PrimitiveArray; use crate::arrays::filter::execute::ConstantArray; use crate::arrays::listview::ListViewArrayExt; + use crate::arrays::listview::ListViewArraySlotsExt; use crate::assert_arrays_eq; use crate::compute::conformance::filter::test_filter_conformance; use crate::validity::Validity; diff --git a/vortex-array/src/arrays/filter/execute/mod.rs b/vortex-array/src/arrays/filter/execute/mod.rs index 52b0373ac49..5dc7a2663ec 100644 --- a/vortex-array/src/arrays/filter/execute/mod.rs +++ b/vortex-array/src/arrays/filter/execute/mod.rs @@ -4,13 +4,26 @@ //! Execution logic for [`super::FilterArray`]. //! //! The main entrypoint is [`execute_filter`] which filters any [`Canonical`] array. +//! Before canonical execution, the [`Filter`] implementation of +//! [`VTable::execute`](crate::array::vtable::VTable::execute) tries these cases in order, the +//! last of which is [`execute_all_null_filter_fast_path`]: +//! +//! | Condition | Result | +//! | ------------------------------- | --------------------- | +//! | mask selects nothing | empty canonical array | +//! | mask selects everything | unchanged child | +//! | mask selects one contiguous run | zero-copy slice | +//! | child is entirely null | null constant array | +//! +//! Fixed-width canonical arrays then use the strategy ladder in [`buffer`]. -use std::sync::Arc; +use std::ops::Range; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_mask::MaskValues; +use vortex_mask::MaskValuesRef; use crate::ArrayRef; use crate::Canonical; @@ -20,62 +33,91 @@ use crate::array::ArrayView; use crate::arrays::ConstantArray; use crate::arrays::ExtensionArray; use crate::arrays::Filter; +use crate::arrays::Map; +use crate::arrays::MapArray; use crate::arrays::NullArray; use crate::arrays::VariantArray; use crate::arrays::extension::ExtensionArrayExt; -use crate::arrays::filter::FilterArrayExt; -use crate::arrays::variant::VariantArrayExt; +use crate::arrays::filter::FilterReduce; +use crate::arrays::fixed_width; +use crate::arrays::variant::VariantArraySlotsExt; use crate::scalar::Scalar; use crate::validity::Validity; mod bitbuffer; mod bool; -mod buffer; +pub(crate) mod buffer; pub(crate) mod byte_compress; -mod decimal; mod fixed_size_list; mod listview; -mod primitive; +mod simd_compress; mod slice; mod struct_; -pub mod take; +mod take; +mod union; mod varbinview; -/// A helper function that lazily filters a [`Validity`] with selection mask values. -fn filter_validity(validity: Validity, mask: &Arc) -> Validity { +/// Lazily filters a [`Validity`] with a partially selective mask. +pub(crate) fn filter_validity(validity: Validity, mask: &MaskValuesRef) -> Validity { validity - .filter(&Mask::Values(Arc::clone(mask))) - .vortex_expect("Somehow unable to wrap filter around a validity array") + .filter(&Mask::Values(MaskValuesRef::clone(mask))) + .vortex_expect("filtering validity with a partially selective mask is valid") } -/// Check for some fast-path execution conditions before calling [`execute_filter`]. -pub(super) fn execute_filter_fast_paths( - array: ArrayView<'_, Filter>, - ctx: &mut ExecutionCtx, -) -> VortexResult> { - let true_count = array.mask.true_count(); +pub(super) fn contiguous_filter_range(mask: &Mask) -> Option> { + match mask { + Mask::AllTrue(len) => (*len > 0).then_some(0..*len), + Mask::AllFalse(_) => None, + Mask::Values(values) => contiguous_values_range(values.as_ref()), + } +} - // If the mask selects nothing, the output is empty. - if true_count == 0 { - return Ok(Some(Canonical::empty(array.dtype()).into_array())); +fn contiguous_values_range(mask: &MaskValues) -> Option> { + if let Some(slices) = mask.cached_slices() { + return match slices { + [(start, end)] => Some(*start..*end), + _ => None, + }; } - // If the mask selects everything, then we can just fully decompress the whole thing. - if true_count == array.mask.len() { - return Ok(Some(array.child().clone())); + if let Some(indices) = mask.cached_indices() { + let start = *indices.first()?; + let end = indices.last()?.checked_add(1)?; + return (end - start == indices.len()).then_some(start..end); } - // Also check if the array itself is completely null, in which case we only care about the total - // number of nulls, not the values. - let child_arr = array.array(); - if child_arr + let true_count = mask.true_count(); + let start = mask.bit_buffer().set_indices().next()?; + let end = start.checked_add(true_count)?; + if end > mask.len() { + return None; + } + + // Bound uncached work by probing the candidate run from the cheaper side: counting scans + // the run itself while the backward scan covers at most the tail after it. + let contiguous = if end - start <= mask.len() - end { + mask.bit_buffer().count_range(start, end) == true_count + } else { + mask.bit_buffer().last_set_index() == Some(end - 1) + }; + contiguous.then_some(start..end) +} + +/// Check for some fast-path execution conditions before calling [`execute_filter`]. +pub(super) fn execute_all_null_filter_fast_path( + array: ArrayView<'_, Filter>, + selected_count: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let child = array.array(); + if child .validity()? - .execute_mask(child_arr.len(), ctx)? + .execute_mask(child.len(), ctx)? .true_count() == 0 { return Ok(Some( - ConstantArray::new(Scalar::null(array.dtype().clone()), true_count).into_array(), + ConstantArray::new(Scalar::null(array.dtype().clone()), selected_count).into_array(), )); } @@ -83,27 +125,29 @@ pub(super) fn execute_filter_fast_paths( } /// Filter a canonical array by a mask, returning a new canonical array. -pub(super) fn execute_filter(canonical: Canonical, mask: &Arc) -> Canonical { +pub(super) fn execute_filter(canonical: Canonical, mask: &MaskValuesRef) -> Canonical { match canonical { Canonical::Null(_) => Canonical::Null(NullArray::new(mask.true_count())), Canonical::Bool(a) => Canonical::Bool(bool::filter_bool(&a, mask)), - Canonical::Primitive(a) => Canonical::Primitive(primitive::filter_primitive(&a, mask)), - Canonical::Decimal(a) => Canonical::Decimal(decimal::filter_decimal(&a, mask)), + Canonical::Primitive(a) => Canonical::Primitive(fixed_width::filter::filter(&a, mask)), + Canonical::Decimal(a) => Canonical::Decimal(fixed_width::filter::filter(&a, mask)), Canonical::VarBinView(a) => Canonical::VarBinView(varbinview::filter_varbinview(&a, mask)), Canonical::List(a) => Canonical::List(listview::filter_listview(&a, mask)), + Canonical::Map(a) => Canonical::Map(filter_map(&a, mask)), Canonical::FixedSizeList(a) => { Canonical::FixedSizeList(fixed_size_list::filter_fixed_size_list(&a, mask)) } Canonical::Struct(a) => Canonical::Struct(struct_::filter_struct(&a, mask)), + Canonical::Union(a) => Canonical::Union(union::filter_union(&a, mask)), Canonical::Extension(a) => { let filtered_storage = a .storage_array() - .filter(Mask::Values(Arc::clone(mask))) + .filter(Mask::Values(MaskValuesRef::clone(mask))) .vortex_expect("ExtensionArray storage type somehow could not be filtered"); Canonical::Extension(ExtensionArray::new(a.ext_dtype().clone(), filtered_storage)) } Canonical::Variant(a) => { - let filter_mask = Mask::Values(Arc::clone(mask)); + let filter_mask = Mask::Values(MaskValuesRef::clone(mask)); let filtered_core_storage = a .core_storage() .filter(filter_mask.clone()) @@ -120,3 +164,87 @@ pub(super) fn execute_filter(canonical: Canonical, mask: &Arc) -> Ca } } } + +fn filter_map(array: &MapArray, mask: &MaskValuesRef) -> MapArray { + let filter_mask = Mask::Values(MaskValuesRef::clone(mask)); + let filtered = ::filter(array.as_view(), &filter_mask) + .vortex_expect("MapArray somehow could not be filtered") + .vortex_expect("Map filter reduce always produces an array"); + filtered.as_::().into_owned() +} + +#[cfg(test)] +mod tests { + use vortex_buffer::BitBuffer; + use vortex_error::VortexResult; + + use super::*; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::PrimitiveArray; + + #[test] + fn uncached_contiguous_filter_executes_as_zero_copy_slice() -> VortexResult<()> { + let array = PrimitiveArray::from_iter(0i32..128); + let original = array.to_buffer::(); + let mask = Mask::from_buffer(BitBuffer::from_iter( + (0..128).map(|index| (37..91).contains(&index)), + )); + assert!(mask.values().is_some_and(|values| { + values.cached_indices().is_none() && values.cached_slices().is_none() + })); + + let filtered = array + .into_array() + .filter(mask)? + .execute::(&mut array_session().create_execution_ctx())?; + let filtered_values = filtered.to_buffer::(); + + assert_eq!(filtered_values.as_slice(), &(37..91).collect::>()); + assert_eq!(filtered_values.as_ptr(), original.as_ptr().wrapping_add(37)); + Ok(()) + } + + #[test] + fn uncached_contiguous_range_handles_sparse_and_dense_masks() { + let cases = [ + ( + Mask::from_buffer(BitBuffer::from_iter( + (0..128).map(|index| (37..41).contains(&index)), + )), + Some(37..41), + ), + ( + Mask::from_buffer(BitBuffer::from_iter( + (0..128).map(|index| (3..125).contains(&index)), + )), + Some(3..125), + ), + ( + Mask::from_buffer(BitBuffer::from_iter( + (0..128).map(|index| (100..128).contains(&index)), + )), + Some(100..128), + ), + ( + Mask::from_buffer(BitBuffer::from_iter( + (0..128).map(|index| matches!(index, 1 | 40 | 90)), + )), + None, + ), + ( + Mask::from_buffer(BitBuffer::from_iter( + (0..128).map(|index| !matches!(index, 17 | 63)), + )), + None, + ), + ]; + + for (mask, expected) in cases { + assert!(mask.values().is_some_and(|values| { + values.cached_indices().is_none() && values.cached_slices().is_none() + })); + assert_eq!(contiguous_filter_range(&mask), expected); + } + } +} diff --git a/vortex-array/src/arrays/filter/execute/primitive.rs b/vortex-array/src/arrays/filter/execute/primitive.rs deleted file mode 100644 index b7abd8efa4b..00000000000 --- a/vortex-array/src/arrays/filter/execute/primitive.rs +++ /dev/null @@ -1,116 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::sync::Arc; - -use vortex_error::VortexExpect; -use vortex_mask::MaskValues; - -use crate::arrays::PrimitiveArray; -use crate::arrays::filter::execute::buffer; -use crate::arrays::filter::execute::byte_compress; -use crate::arrays::filter::execute::filter_validity; -use crate::dtype::NativePType; -use crate::dtype::PType; -use crate::match_each_native_ptype; -use crate::validity::Validity; - -pub fn filter_primitive(array: &PrimitiveArray, mask: &Arc) -> PrimitiveArray { - let validity = array - .validity() - .vortex_expect("primitive validity should be derivable"); - let filtered_validity = filter_validity(validity, mask); - - match array.ptype() { - // Byte-compress avoids materializing indices/slices and processes 8 elements per mask byte. - PType::U8 => filter_byte_compress::(array, filtered_validity, mask), - PType::I8 => filter_byte_compress::(array, filtered_validity, mask), - PType::U16 => filter_byte_compress::(array, filtered_validity, mask), - PType::I16 => filter_byte_compress::(array, filtered_validity, mask), - PType::U32 => filter_byte_compress::(array, filtered_validity, mask), - PType::I32 => filter_byte_compress::(array, filtered_validity, mask), - _ => match_each_native_ptype!(array.ptype(), |T| { - let filtered_buffer = buffer::filter_buffer(array.to_buffer::(), mask.as_ref()); - - // SAFETY: We filter both the validity and the buffer with the same mask, so they must - // have the same length. - unsafe { PrimitiveArray::new_unchecked(filtered_buffer, filtered_validity) } - }), - } -} - -fn filter_byte_compress( - array: &PrimitiveArray, - filtered_validity: Validity, - mask: &Arc, -) -> PrimitiveArray { - let filtered_buffer = byte_compress::filter_buffer(array.to_buffer::(), mask.as_ref()); - - // SAFETY: We filter both the validity and the buffer with the same mask, so they must have the - // same length. - unsafe { PrimitiveArray::new_unchecked(filtered_buffer, filtered_validity) } -} - -#[cfg(test)] -#[expect(clippy::cast_possible_truncation)] -mod test { - use itertools::Itertools; - use rstest::rstest; - use vortex_mask::Mask; - - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::array_session; - use crate::arrays::PrimitiveArray; - use crate::compute::conformance::filter::LARGE_SIZE; - use crate::compute::conformance::filter::MEDIUM_SIZE; - use crate::compute::conformance::filter::test_filter_conformance; - - #[test] - fn filter_run_variant_mixed_test() { - let mut ctx = array_session().create_execution_ctx(); - let mask = [true, true, false, true, true, true, false, true]; - let arr = PrimitiveArray::from_iter([1u32, 24, 54, 2, 3, 2, 3, 2]); - - let filtered = arr - .filter(Mask::from_iter(mask)) - .unwrap() - .execute::(&mut ctx) - .unwrap(); - assert_eq!( - filtered.len(), - mask.iter().filter(|x| **x).collect_vec().len() - ); - - let rust_arr = arr.as_slice::(); - assert_eq!( - filtered.as_slice::().to_vec(), - mask.iter() - .enumerate() - .filter(|(_idx, b)| **b) - .map(|m| rust_arr[m.0]) - .collect_vec() - ) - } - - #[rstest] - #[case(PrimitiveArray::from_iter([-2i8, -1, 0, 1, 2]))] - #[case(PrimitiveArray::from_iter([1u16, 2, 3, 4, 5]))] - #[case(PrimitiveArray::from_iter([-2i16, -1, 0, 1, 2]))] - #[case(PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]))] - #[case(PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4), None]))] - #[case(PrimitiveArray::from_iter([42u64]))] - #[case(PrimitiveArray::from_iter(0..MEDIUM_SIZE as i32))] - #[case(PrimitiveArray::from_option_iter( - (0..MEDIUM_SIZE).map(|i| if i % 3 == 0 { None } else { Some(i as i64) })) - )] - #[case(PrimitiveArray::from_iter(0..LARGE_SIZE as u32))] - #[case(PrimitiveArray::from_iter([0.1f32, 0.2, 0.3, 0.4, 0.5]))] - #[case(PrimitiveArray::from_option_iter([Some(1.1f64), None, Some(2.2), Some(3.3), None]))] - fn test_filter_primitive_conformance(#[case] array: PrimitiveArray) { - test_filter_conformance( - &array.into_array(), - &mut array_session().create_execution_ctx(), - ); - } -} diff --git a/vortex-array/src/arrays/filter/execute/simd_compress/mod.rs b/vortex-array/src/arrays/filter/execute/simd_compress/mod.rs new file mode 100644 index 00000000000..807259cd5e7 --- /dev/null +++ b/vortex-array/src/arrays/filter/execute/simd_compress/mod.rs @@ -0,0 +1,207 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! SIMD compress kernels for fixed-width filtering. +//! +//! Architecture-specific kernels compact selected lanes a mask word at a time. Full-width stores +//! may write trailing garbage, so out-of-place outputs reserve one vector of slack and in-place +//! kernels never store beyond the source chunk already loaded. +//! +//! The table is the complete SIMD eligibility map after cached representations have declined the +//! mask. `d` is mask density; every entry also requires `len >= 64`. A dash means that no SIMD +//! kernel exists for that width. +//! +//! | Target features | 1 byte | 2 bytes | 4 bytes | 8 bytes | +//! | ----------------------- | ------------ | ------------ | ----------- | ------------------ | +//! | x86 AVX-512 VBMI2 | `d >= 0.00` | `d >= 0.15` | `d >= 0.25` | `d >= 0.30` | +//! | x86 AVX-512F (no VBMI2) | `d >= 0.15`* | `d >= 0.25`* | `d >= 0.25` | `d >= 0.30` | +//! | x86 AVX2 | `d >= 0.15` | `d >= 0.25` | `d >= 0.25` | `d >= 0.45` | +//! | aarch64 NEON | `d >= 0.15` | `d >= 0.15` | `d >= 0.30` | `0.50 <= d < 0.80` | +//! | other | — | — | — | — | +//! +//! \* The 1- and 2-byte fallbacks require AVX2. +//! +//! On x86, the widest available kernel is considered first for each width. When no entry matches, +//! control returns to [`buffer`](super::buffer) for scalar selection. + +use std::ptr; + +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_mask::MaskValues; + +#[cfg(all(target_arch = "aarch64", not(miri)))] +mod neon; +#[cfg(test)] +mod tests; +#[cfg(all(target_arch = "x86_64", not(miri)))] +mod x86; + +const MIN_LEN: usize = 64; + +const SLACK_BYTES: usize = 64; + +/// `dst == src` for in-place kernels. +type Kernel = unsafe fn(*const u8, *mut u8, &MaskValues) -> usize; + +/// Filter a slice with a SIMD compress kernel, if one applies. +/// +/// Returns `None` when the caller should use a scalar strategy. +pub(super) fn filter_slice_by_bitmap( + values: &[T], + mask: &MaskValues, +) -> Option> { + debug_assert_eq!(values.len(), mask.len()); + let kernel = select_kernel::(mask)?; + + let true_count = mask.true_count(); + let mut out = BufferMut::::with_capacity(true_count + SLACK_BYTES / size_of::()); + // SAFETY: `select_kernel` probed the kernel's target features; `values` holds `mask.len()` + // elements and the output has capacity for every selected element plus a full vector of + // slack, so each unmasked store stays in bounds. + let written = unsafe { + kernel( + values.as_ptr().cast(), + out.spare_capacity_mut().as_mut_ptr().cast(), + mask, + ) + }; + debug_assert_eq!(written, true_count); + // SAFETY: the kernel initialized the first `true_count` elements. + unsafe { out.set_len(true_count) }; + Some(out.freeze()) +} + +/// In-place variant of [`filter_slice_by_bitmap`]: compact the selected elements to the front +/// of `values` and return the new length, if a SIMD kernel applies. +pub(super) fn filter_slice_mut_by_bitmap( + values: &mut [T], + mask: &MaskValues, +) -> Option { + debug_assert_eq!(values.len(), mask.len()); + let kernel = select_kernel::(mask)?; + + let dst = values.as_mut_ptr().cast::(); + // SAFETY: `select_kernel` probed the kernel's target features; the in-place instantiation + // compacts forward (stores never pass the positions it has already read) and keeps partial + // tail chunks off the full-width store path, so all accesses stay inside `values`. + let written = unsafe { kernel(dst.cast_const(), dst, mask) }; + debug_assert_eq!(written, mask.true_count()); + Some(written) +} + +/// Choose the widest profitable kernel available for `T`. +fn select_kernel(mask: &MaskValues) -> Option { + if mask.len() < MIN_LEN { + return None; + } + + #[cfg(all(target_arch = "x86_64", not(miri)))] + { + x86::select_kernel::(mask) + } + #[cfg(all(target_arch = "aarch64", not(miri)))] + { + neon::select_kernel::(mask) + } + #[cfg(any(not(any(target_arch = "x86_64", target_arch = "aarch64")), miri))] + { + let _ = mask; + None + } +} + +/// Build byte-shuffle rows that gather selected lanes to the front of a vector. +#[cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), not(miri)))] +#[expect( + clippy::cast_possible_truncation, + reason = "byte indices are bounded by the 8- or 16-byte table width" +)] +const fn compress_lut( + lanes: usize, + elem_size: usize, +) -> [[u8; BYTES]; ROWS] { + assert!(ROWS == 1 << lanes); + assert!(lanes * elem_size <= BYTES); + + let mut lut = [[0u8; BYTES]; ROWS]; + let mut m = 0; + while m < ROWS { + let mut out_lane = 0; + let mut lane = 0; + while lane < lanes { + if m & (1 << lane) != 0 { + let mut byte = 0; + while byte < elem_size { + lut[m][out_lane * elem_size + byte] = (lane * elem_size + byte) as u8; + byte += 1; + } + out_lane += 1; + } + lane += 1; + } + m += 1; + } + lut +} + +/// Compact a partial vector without reading past the source. +/// +/// # Safety +/// +/// The pointer contract of [`filter_slice_by_bitmap`] / [`filter_slice_mut_by_bitmap`] must hold, +/// and `bits` must only select elements that are in bounds. +#[cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), not(miri)))] +#[allow(clippy::inline_always)] +#[inline(always)] +unsafe fn compress_tail( + src: *const u8, + dst: *mut u8, + mut bits: u64, + base: usize, + mut write_pos: usize, + elem_size: usize, +) -> usize { + while bits != 0 { + let index = base + bits.trailing_zeros() as usize; + // SAFETY: `index` is in bounds per the contract above and stable compaction guarantees + // `write_pos <= index`. + unsafe { bulk_copy::(src, dst, index, 1, write_pos, elem_size) }; + write_pos += 1; + bits &= bits - 1; + } + write_pos +} + +/// Copy `word_len` elements of `elem_size` bytes for a fully-set mask word. +/// +/// # Safety +/// +/// `word_len` source elements starting at `word_start` must be in bounds. Out-of-place, `dst` +/// must not overlap `src` and must have room at `write_pos`; in-place, forward compaction must +/// guarantee `write_pos <= word_start`. +#[cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), not(miri)))] +#[allow(clippy::inline_always)] +#[inline(always)] +unsafe fn bulk_copy( + src: *const u8, + dst: *mut u8, + word_start: usize, + word_len: usize, + write_pos: usize, + elem_size: usize, +) { + // SAFETY: offsets are in bounds per the contract above. + let src_bytes = unsafe { src.add(word_start * elem_size) }; + // SAFETY: see above. + let dst_bytes = unsafe { dst.add(write_pos * elem_size) }; + if IN_PLACE { + if write_pos != word_start { + // SAFETY: source and destination may overlap while compacting forward. + unsafe { ptr::copy(src_bytes, dst_bytes, word_len * elem_size) }; + } + } else { + // SAFETY: out-of-place buffers are disjoint allocations. + unsafe { ptr::copy_nonoverlapping(src_bytes, dst_bytes, word_len * elem_size) }; + } +} diff --git a/vortex-array/src/arrays/filter/execute/simd_compress/neon.rs b/vortex-array/src/arrays/filter/execute/simd_compress/neon.rs new file mode 100644 index 00000000000..5eb3376adf9 --- /dev/null +++ b/vortex-array/src/arrays/filter/execute/simd_compress/neon.rs @@ -0,0 +1,188 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! NEON `tbl` compress kernels. +//! +//! Byte-index lookup tables drive `tbl` shuffles for 1-, 2-, 4-, and 8-byte elements. + +use core::arch::aarch64::vld1_u8; +use core::arch::aarch64::vld1q_u8; +use core::arch::aarch64::vqtbl1q_u8; +use core::arch::aarch64::vst1_u8; +use core::arch::aarch64::vst1q_u8; +use core::arch::aarch64::vtbl1_u8; + +use vortex_mask::MaskValues; + +use super::super::slice::for_each_mask_word; +use super::super::slice::low_bits_mask; +use super::Kernel; +use super::bulk_copy; +use super::compress_lut; +use super::compress_tail; + +/// Return the benchmarked density range where NEON beats the scalar strategies. +fn density_band() -> Option> { + match size_of::() { + 1 | 2 => Some(0.15..f64::INFINITY), + 4 => Some(0.30..f64::INFINITY), + 8 => Some(0.50..0.80), + _ => None, + } +} + +/// Choose the kernel for this element width, if one applies to this mask. +pub(super) fn select_kernel(mask: &MaskValues) -> Option { + if !density_band::()?.contains(&mask.density()) { + return None; + } + + match size_of::() { + 1 => Some(compress_neon_8:: as Kernel), + 2 => Some(compress_neon_16:: as Kernel), + 4 => Some(compress_neon_32:: as Kernel), + 8 => Some(compress_neon_64:: as Kernel), + _ => None, + } +} + +static IDX_LUT_8: [[u8; 8]; 256] = compress_lut::<256, 8>(8, 1); +static IDX_LUT_16: [[u8; 16]; 256] = compress_lut::<256, 16>(8, 2); +static IDX_LUT_32: [[u8; 16]; 16] = compress_lut::<16, 16>(4, 4); +static IDX_LUT_64: [[u8; 16]; 4] = compress_lut::<4, 16>(2, 8); + +/// Generate a NEON `tbl` kernel for one element width. +macro_rules! neon_compress_kernel { + ( + $word_fn:ident, + $walk_fn:ident,elem_size: + $elem_size:literal,lanes: + $lanes:literal,idx_lut: + $idx_lut:ident,load: + $load:ident,tbl: + $tbl:ident,store: + $store:ident + ) => { + /// # Safety + /// + /// The pointer contract of [`filter_slice_by_bitmap`](super::filter_slice_by_bitmap) / + /// [`filter_slice_mut_by_bitmap`](super::filter_slice_mut_by_bitmap) must hold. + #[expect( + clippy::cast_possible_truncation, + reason = "deliberate submask narrowing" + )] + #[inline] + unsafe fn $word_fn( + src: *const u8, + dst: *mut u8, + word: u64, + word_start: usize, + word_len: usize, + mut write_pos: usize, + ) -> usize { + if word == 0 { + return write_pos; + } + if word == low_bits_mask(word_len) { + // SAFETY: forwarded from the caller contract. + unsafe { + bulk_copy::(src, dst, word_start, word_len, write_pos, $elem_size) + }; + return write_pos + word_len; + } + + // Empty chunks still store garbage that the next chunk overwrites; branching here + // regresses masks near the density crossover. + let mut sub = 0; + while sub + $lanes <= word_len { + let m = ((word >> sub) & low_bits_mask($lanes)) as usize; + // SAFETY: the chunk holds `$lanes` in-bounds source elements, which is exactly + // the table width. + let chunk = unsafe { $load(src.add((word_start + sub) * $elem_size)) }; + // SAFETY: every LUT row is one table wide. + let idx = unsafe { $load($idx_lut[m].as_ptr()) }; + // SAFETY: out-of-place output has vector slack. In-place, the store ends within + // the source chunk already loaded, and later stores overwrite trailing garbage. + unsafe { $store(dst.add(write_pos * $elem_size), $tbl(chunk, idx)) }; + write_pos += m.count_ones() as usize; + sub += $lanes; + } + + if sub < word_len { + let bits = (word >> sub) & low_bits_mask(word_len - sub); + // SAFETY: forwarded from the caller contract. + write_pos = unsafe { + compress_tail::( + src, + dst, + bits, + word_start + sub, + write_pos, + $elem_size, + ) + }; + } + + write_pos + } + + /// # Safety + /// + /// The pointer contract of [`filter_slice_by_bitmap`](super::filter_slice_by_bitmap) / + /// [`filter_slice_mut_by_bitmap`](super::filter_slice_mut_by_bitmap) must hold. + pub(super) unsafe fn $walk_fn( + src: *const u8, + dst: *mut u8, + mask: &MaskValues, + ) -> usize { + let mut write_pos = 0; + for_each_mask_word(mask, |word, word_start, word_len| { + // SAFETY: forwarded from the caller contract. + write_pos = unsafe { + $word_fn::(src, dst, word, word_start, word_len, write_pos) + }; + }); + write_pos + } + }; +} + +neon_compress_kernel!( + compress_word_neon_8, compress_neon_8, + elem_size: 1, + lanes: 8, + idx_lut: IDX_LUT_8, + load: vld1_u8, + tbl: vtbl1_u8, + store: vst1_u8 +); + +neon_compress_kernel!( + compress_word_neon_16, compress_neon_16, + elem_size: 2, + lanes: 8, + idx_lut: IDX_LUT_16, + load: vld1q_u8, + tbl: vqtbl1q_u8, + store: vst1q_u8 +); + +neon_compress_kernel!( + compress_word_neon_32, compress_neon_32, + elem_size: 4, + lanes: 4, + idx_lut: IDX_LUT_32, + load: vld1q_u8, + tbl: vqtbl1q_u8, + store: vst1q_u8 +); + +neon_compress_kernel!( + compress_word_neon_64, compress_neon_64, + elem_size: 8, + lanes: 2, + idx_lut: IDX_LUT_64, + load: vld1q_u8, + tbl: vqtbl1q_u8, + store: vst1q_u8 +); diff --git a/vortex-array/src/arrays/filter/execute/simd_compress/tests.rs b/vortex-array/src/arrays/filter/execute/simd_compress/tests.rs new file mode 100644 index 00000000000..e79b70a2f94 --- /dev/null +++ b/vortex-array/src/arrays/filter/execute/simd_compress/tests.rs @@ -0,0 +1,172 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![allow(clippy::cast_possible_truncation)] + +use rstest::rstest; +use vortex_buffer::BitBuffer; +use vortex_mask::Mask; +use vortex_mask::MaskValues; + +use super::super::slice; +use super::*; + +fn mask_values(mask: &Mask) -> Option<&MaskValues> { + match mask { + Mask::Values(values) => Some(values.as_ref()), + _ => None, + } +} + +fn make_mask(len: usize, offset: usize, pattern: impl Fn(usize) -> bool) -> Mask { + let backing = + BitBuffer::from_iter(std::iter::repeat_n(false, offset).chain((0..len).map(pattern))); + Mask::from_buffer(BitBuffer::new_with_offset( + backing.inner().clone(), + len, + offset, + )) +} + +type Pattern = fn(usize) -> bool; + +fn patterns() -> [Pattern; 4] { + [ + |i| i % 3 == 0, + |i| i % 16 != 0, + |i| (i / 64) % 2 == 0, + |i| i < 3 || i % 61 == 60, + ] +} + +fn check(values: &[T], mask: &Mask) { + let Some(mask) = mask_values(mask) else { + return; + }; + let expected = slice::filter_slice_by_bitmap(values, mask); + + if let Some(actual) = filter_slice_by_bitmap(values, mask) { + assert_eq!(actual.as_slice(), expected.as_slice()); + } + + let mut compacted = values.to_vec(); + if let Some(new_len) = filter_slice_mut_by_bitmap(&mut compacted, mask) { + assert_eq!(&compacted[..new_len], expected.as_slice()); + } +} + +#[rstest] +fn simd_matches_scalar(#[values(0, 5)] offset: usize, #[values(64, 151)] len: usize) { + for pattern in patterns() { + let mask = make_mask(len, offset, pattern); + + let u8_values: Vec = (0..len).map(|i| i as u8).collect(); + let u16_values: Vec = (0..len).map(|i| i as u16).collect(); + let u32_values: Vec = (0..len).map(|i| i as u32).collect(); + let u64_values: Vec = (0..len).map(|i| i as u64).collect(); + + check(&u8_values, &mask); + check(&u16_values, &mask); + check(&u32_values, &mask); + check(&u64_values, &mask); + } +} + +#[cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), not(miri)))] +#[test] +fn engages_on_supported_cpus() { + let values: Vec = (0..256).collect(); + let mask = make_mask(256, 0, |i| i % 2 == 0); + let mask = mask_values(&mask).expect("alternating mask is mixed"); + + #[cfg(target_arch = "x86_64")] + let expected = is_x86_feature_detected!("avx2"); + #[cfg(target_arch = "aarch64")] + let expected = true; + + assert_eq!(filter_slice_by_bitmap(&values, mask).is_some(), expected); +} + +#[test] +fn declines_sparse_and_short_masks() { + let values: Vec = (0..1024).collect(); + + let sparse = make_mask(1024, 0, |i| i % 128 == 0); + let sparse = mask_values(&sparse).expect("sparse mask is mixed"); + assert!(filter_slice_by_bitmap(&values, sparse).is_none()); + + let short = make_mask(32, 0, |i| i % 2 == 0); + let short = mask_values(&short).expect("alternating mask is mixed"); + assert!(filter_slice_by_bitmap(&values[..32], short).is_none()); +} + +// AVX-512 machines need direct coverage of the otherwise-unselected AVX2 tier. +#[cfg(all(target_arch = "x86_64", not(miri)))] +#[test] +fn avx2_kernels_match_scalar() { + if !is_x86_feature_detected!("avx2") { + return; + } + + fn check_kernel( + kernel_out_of_place: Kernel, + kernel_in_place: Kernel, + values: &[T], + mask: &MaskValues, + ) { + let expected = slice::filter_slice_by_bitmap(values, mask); + + let mut out = vec![T::default(); mask.true_count() + SLACK_BYTES / size_of::()]; + // SAFETY: AVX2 was detected above and the output has a vector of slack. + let written = + unsafe { kernel_out_of_place(values.as_ptr().cast(), out.as_mut_ptr().cast(), mask) }; + assert_eq!(written, mask.true_count()); + assert_eq!(&out[..written], expected.as_slice()); + + let mut compacted = values.to_vec(); + let ptr = compacted.as_mut_ptr().cast::(); + // SAFETY: AVX2 was detected above; in-place compaction stays within the slice. + let written = unsafe { kernel_in_place(ptr.cast_const(), ptr, mask) }; + assert_eq!(written, mask.true_count()); + assert_eq!(&compacted[..written], expected.as_slice()); + } + + for pattern in patterns() { + for len in [64, 151] { + for offset in [0, 5] { + let mask = make_mask(len, offset, pattern); + let Some(mask) = mask_values(&mask) else { + continue; + }; + let u8_values: Vec = (0..len).map(|i| i as u8).collect(); + let u16_values: Vec = (0..len).map(|i| i as u16).collect(); + let u32_values: Vec = (0..len as u32).collect(); + let u64_values: Vec = (0..len as u64).collect(); + check_kernel( + x86::compress_pshufb_epi8::, + x86::compress_pshufb_epi8::, + &u8_values, + mask, + ); + check_kernel( + x86::compress_pshufb_epi16::, + x86::compress_pshufb_epi16::, + &u16_values, + mask, + ); + check_kernel( + x86::compress_avx2_epi32::, + x86::compress_avx2_epi32::, + &u32_values, + mask, + ); + check_kernel( + x86::compress_avx2_epi64::, + x86::compress_avx2_epi64::, + &u64_values, + mask, + ); + } + } + } +} diff --git a/vortex-array/src/arrays/filter/execute/simd_compress/x86.rs b/vortex-array/src/arrays/filter/execute/simd_compress/x86.rs new file mode 100644 index 00000000000..85b52e5fa19 --- /dev/null +++ b/vortex-array/src/arrays/filter/execute/simd_compress/x86.rs @@ -0,0 +1,576 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! AVX-512 `vpcompress`, AVX2 `vpermd`, and 128-bit `pshufb` compress kernels. +//! +//! See the [module docs](super) for how these fit the shared dispatch. + +use std::arch::x86_64::__m128i; +use std::arch::x86_64::_mm_loadl_epi64; +use std::arch::x86_64::_mm_loadu_si128; +use std::arch::x86_64::_mm_shuffle_epi8; +use std::arch::x86_64::_mm_storel_epi64; +use std::arch::x86_64::_mm_storeu_si128; +use std::arch::x86_64::_mm256_loadu_si256; +use std::arch::x86_64::_mm256_maskload_epi32; +use std::arch::x86_64::_mm256_maskload_epi64; +use std::arch::x86_64::_mm256_maskstore_epi32; +use std::arch::x86_64::_mm256_maskstore_epi64; +use std::arch::x86_64::_mm256_permutevar8x32_epi32; +use std::arch::x86_64::_mm256_storeu_si256; +use std::arch::x86_64::_mm512_loadu_epi8; +use std::arch::x86_64::_mm512_loadu_epi16; +use std::arch::x86_64::_mm512_loadu_epi32; +use std::arch::x86_64::_mm512_loadu_epi64; +use std::arch::x86_64::_mm512_mask_storeu_epi8; +use std::arch::x86_64::_mm512_mask_storeu_epi16; +use std::arch::x86_64::_mm512_mask_storeu_epi32; +use std::arch::x86_64::_mm512_mask_storeu_epi64; +use std::arch::x86_64::_mm512_maskz_compress_epi8; +use std::arch::x86_64::_mm512_maskz_compress_epi16; +use std::arch::x86_64::_mm512_maskz_compress_epi32; +use std::arch::x86_64::_mm512_maskz_compress_epi64; +use std::arch::x86_64::_mm512_maskz_loadu_epi8; +use std::arch::x86_64::_mm512_maskz_loadu_epi16; +use std::arch::x86_64::_mm512_maskz_loadu_epi32; +use std::arch::x86_64::_mm512_maskz_loadu_epi64; +use std::arch::x86_64::_mm512_storeu_epi8; +use std::arch::x86_64::_mm512_storeu_epi16; +use std::arch::x86_64::_mm512_storeu_epi32; +use std::arch::x86_64::_mm512_storeu_epi64; + +use vortex_mask::MaskValues; + +use super::super::slice::for_each_mask_word; +use super::super::slice::low_bits_mask; +use super::Kernel; +use super::bulk_copy; +use super::compress_lut; +use super::compress_tail; + +/// Choose the widest available kernel above its benchmarked density crossover. +/// +/// Sparse masks stay on the scalar set-bit walk. See `benches/filter_fixed_width.rs` when +/// changing these thresholds. +pub(super) fn select_kernel(mask: &MaskValues) -> Option { + let (kernel, min_density) = match size_of::() { + 1 if avx512_vbmi2() => (compress_avx512_epi8:: as Kernel, 0.0), + 2 if avx512_vbmi2() => (compress_avx512_epi16:: as Kernel, 0.15), + 4 if avx512f() => (compress_avx512_epi32:: as Kernel, 0.25), + 8 if avx512f() => (compress_avx512_epi64:: as Kernel, 0.30), + // AVX-512F without VBMI2 (e.g. Skylake-X) falls through to these too. + 1 if avx2() => (compress_pshufb_epi8:: as Kernel, 0.15), + 2 if avx2() => (compress_pshufb_epi16:: as Kernel, 0.25), + 4 if avx2() => (compress_avx2_epi32:: as Kernel, 0.25), + 8 if avx2() => (compress_avx2_epi64:: as Kernel, 0.45), + _ => return None, + }; + + (mask.density() >= min_density).then_some(kernel) +} + +fn avx512f() -> bool { + is_x86_feature_detected!("avx512f") +} + +fn avx512_vbmi2() -> bool { + is_x86_feature_detected!("avx512f") + && is_x86_feature_detected!("avx512bw") + && is_x86_feature_detected!("avx512vbmi2") +} + +fn avx2() -> bool { + is_x86_feature_detected!("avx2") +} + +/// Generate an AVX-512 kernel for one element width. +macro_rules! avx512_compress_kernel { + ( + $word_fn:ident, + $walk_fn:ident,features: + $features:literal,elem: + $elem:ty,lanes: + $lanes:literal,kmask: + $kmask:ty,loadu: + $loadu:ident,maskz_loadu: + $maskz_loadu:ident,maskz_compress: + $maskz_compress:ident,storeu: + $storeu:ident,mask_storeu: + $mask_storeu:ident + ) => { + /// # Safety + /// + /// The CPU must support the enabled target features and the pointer contract of + /// [`filter_slice_by_bitmap`](super::filter_slice_by_bitmap) / + /// [`filter_slice_mut_by_bitmap`](super::filter_slice_mut_by_bitmap) must hold. + // `allow` rather than `expect`: the cast is only lossy for sub-`u64` k-masks. + #[allow(clippy::cast_possible_truncation)] + #[target_feature(enable = $features)] + #[inline] + unsafe fn $word_fn( + src: *const u8, + dst: *mut u8, + word: u64, + word_start: usize, + word_len: usize, + mut write_pos: usize, + ) -> usize { + if word == 0 { + return write_pos; + } + if word == low_bits_mask(word_len) { + // SAFETY: forwarded from the caller contract. + unsafe { + bulk_copy::( + src, + dst, + word_start, + word_len, + write_pos, + size_of::<$elem>(), + ) + }; + return write_pos + word_len; + } + + let mut sub = 0; + while sub < word_len { + let count = (word_len - sub).min($lanes); + let k = ((word >> sub) & low_bits_mask(count)) as $kmask; + if k != 0 { + let src_ptr = unsafe { src.cast::<$elem>().add(word_start + sub) }; + let chunk = if count == $lanes { + // SAFETY: `$lanes` source elements starting at + // `word_start + sub` are in bounds. + unsafe { $loadu(src_ptr) } + } else { + // SAFETY: the load mask enables only the `count` in-bounds lanes. + unsafe { $maskz_loadu(low_bits_mask(count) as $kmask, src_ptr) } + }; + let packed = $maskz_compress(k, chunk); + let selected = k.count_ones() as usize; + let dst_ptr = unsafe { dst.cast::<$elem>().add(write_pos) }; + if !IN_PLACE || count == $lanes { + // SAFETY: out-of-place output has a vector of slack past + // `true_count`; in-place, `write_pos <= word_start + sub`, so + // `write_pos + $lanes <= word_start + sub + $lanes <= len`. + unsafe { $storeu(dst_ptr, packed) }; + } else { + // SAFETY: stores exactly the `selected` selected lanes, all + // within the `true_count` output positions. + unsafe { $mask_storeu(dst_ptr, low_bits_mask(selected) as $kmask, packed) }; + } + write_pos += selected; + } + sub += $lanes; + } + write_pos + } + + /// # Safety + /// + /// The CPU must support the enabled target features and the pointer contract of + /// [`filter_slice_by_bitmap`](super::filter_slice_by_bitmap) / + /// [`filter_slice_mut_by_bitmap`](super::filter_slice_mut_by_bitmap) must hold. + #[target_feature(enable = $features)] + pub(super) unsafe fn $walk_fn( + src: *const u8, + dst: *mut u8, + mask: &MaskValues, + ) -> usize { + let mut write_pos = 0; + for_each_mask_word(mask, |word, word_start, word_len| { + // SAFETY: forwarded from the caller contract. + write_pos = unsafe { + $word_fn::(src, dst, word, word_start, word_len, write_pos) + }; + }); + write_pos + } + }; +} + +avx512_compress_kernel!( + compress_word_avx512_epi8, compress_avx512_epi8, + features: "avx512f,avx512bw,avx512vbmi2", + elem: i8, + lanes: 64, + kmask: u64, + loadu: _mm512_loadu_epi8, + maskz_loadu: _mm512_maskz_loadu_epi8, + maskz_compress: _mm512_maskz_compress_epi8, + storeu: _mm512_storeu_epi8, + mask_storeu: _mm512_mask_storeu_epi8 +); + +avx512_compress_kernel!( + compress_word_avx512_epi16, compress_avx512_epi16, + features: "avx512f,avx512bw,avx512vbmi2", + elem: i16, + lanes: 32, + kmask: u32, + loadu: _mm512_loadu_epi16, + maskz_loadu: _mm512_maskz_loadu_epi16, + maskz_compress: _mm512_maskz_compress_epi16, + storeu: _mm512_storeu_epi16, + mask_storeu: _mm512_mask_storeu_epi16 +); + +avx512_compress_kernel!( + compress_word_avx512_epi32, compress_avx512_epi32, + features: "avx512f", + elem: i32, + lanes: 16, + kmask: u16, + loadu: _mm512_loadu_epi32, + maskz_loadu: _mm512_maskz_loadu_epi32, + maskz_compress: _mm512_maskz_compress_epi32, + storeu: _mm512_storeu_epi32, + mask_storeu: _mm512_mask_storeu_epi32 +); + +avx512_compress_kernel!( + compress_word_avx512_epi64, compress_avx512_epi64, + features: "avx512f", + elem: i64, + lanes: 8, + kmask: u8, + loadu: _mm512_loadu_epi64, + maskz_loadu: _mm512_maskz_loadu_epi64, + maskz_compress: _mm512_maskz_compress_epi64, + storeu: _mm512_storeu_epi64, + mask_storeu: _mm512_mask_storeu_epi64 +); + +/// For each mask byte, `vpermd` lane indices compacting the selected 4-byte lanes to the front +/// (trailing lanes are don't-care). +static PERM_LUT_32: [[u32; 8]; 256] = { + let mut lut = [[0u32; 8]; 256]; + let mut m = 0; + while m < 256 { + let mut out_lane = 0; + let mut bit = 0; + while bit < 8 { + if m & (1 << bit) != 0 { + lut[m][out_lane] = bit as u32; + out_lane += 1; + } + bit += 1; + } + m += 1; + } + lut +}; + +/// For each mask nibble, `vpermd` lane indices compacting the selected 8-byte lanes (as pairs +/// of 4-byte lanes) to the front. +static PERM_LUT_64: [[u32; 8]; 16] = { + let mut lut = [[0u32; 8]; 16]; + let mut m = 0; + while m < 16 { + let mut out_lane = 0; + let mut bit = 0; + while bit < 4 { + if m & (1 << bit) != 0 { + lut[m][out_lane * 2] = (bit * 2) as u32; + lut[m][out_lane * 2 + 1] = (bit * 2 + 1) as u32; + out_lane += 1; + } + bit += 1; + } + m += 1; + } + lut +}; + +/// Lane-enable vectors for `vpmaskmov` loads/stores of the first `count` 4-byte lanes. +static LANE_MASK_32: [[i32; 8]; 9] = { + let mut lut = [[0i32; 8]; 9]; + let mut count = 0; + while count <= 8 { + let mut lane = 0; + while lane < count { + lut[count][lane] = -1; + lane += 1; + } + count += 1; + } + lut +}; + +/// Lane-enable vectors for `vpmaskmov` loads/stores of the first `count` 8-byte lanes. +static LANE_MASK_64: [[i64; 4]; 5] = { + let mut lut = [[0i64; 4]; 5]; + let mut count = 0; + while count <= 4 { + let mut lane = 0; + while lane < count { + lut[count][lane] = -1; + lane += 1; + } + count += 1; + } + lut +}; + +/// Generate an AVX2 permutation-LUT kernel for one element width. +macro_rules! avx2_compress_kernel { + ( + $word_fn:ident, + $walk_fn:ident,elem: + $elem:ty,lanes: + $lanes:literal,perm_lut: + $perm_lut:ident,lane_masks: + $lane_masks:ident,maskload: + $maskload:ident,maskstore: + $maskstore:ident + ) => { + /// # Safety + /// + /// The CPU must support AVX2 and the pointer contract of + /// [`filter_slice_by_bitmap`](super::filter_slice_by_bitmap) / + /// [`filter_slice_mut_by_bitmap`](super::filter_slice_mut_by_bitmap) must hold. + #[expect( + clippy::cast_possible_truncation, + reason = "deliberate submask narrowing" + )] + #[target_feature(enable = "avx2")] + #[inline] + unsafe fn $word_fn( + src: *const u8, + dst: *mut u8, + word: u64, + word_start: usize, + word_len: usize, + mut write_pos: usize, + ) -> usize { + if word == 0 { + return write_pos; + } + if word == low_bits_mask(word_len) { + // SAFETY: forwarded from the caller contract. + unsafe { + bulk_copy::( + src, + dst, + word_start, + word_len, + write_pos, + size_of::<$elem>(), + ) + }; + return write_pos + word_len; + } + + let mut sub = 0; + while sub < word_len { + let count = (word_len - sub).min($lanes); + let m = ((word >> sub) & low_bits_mask(count)) as usize; + if m != 0 { + let src_ptr = unsafe { src.cast::<$elem>().add(word_start + sub) }; + let chunk = if count == $lanes { + // SAFETY: `$lanes` source elements starting at + // `word_start + sub` are in bounds. + unsafe { _mm256_loadu_si256(src_ptr.cast()) } + } else { + // SAFETY: the lane mask enables only the `count` in-bounds lanes. + unsafe { + $maskload( + src_ptr, + _mm256_loadu_si256($lane_masks[count].as_ptr().cast()), + ) + } + }; + // SAFETY: LUT rows are 32 bytes. + let perm = unsafe { _mm256_loadu_si256($perm_lut[m].as_ptr().cast()) }; + let packed = _mm256_permutevar8x32_epi32(chunk, perm); + let selected = m.count_ones() as usize; + let dst_ptr = unsafe { dst.cast::<$elem>().add(write_pos) }; + if !IN_PLACE || count == $lanes { + // SAFETY: out-of-place output has a vector of slack past + // `true_count`; in-place, `write_pos <= word_start + sub`, so + // `write_pos + $lanes <= word_start + sub + $lanes <= len`. + unsafe { _mm256_storeu_si256(dst_ptr.cast(), packed) }; + } else { + // SAFETY: stores exactly the `selected` selected lanes, all + // within the `true_count` output positions. + unsafe { + $maskstore( + dst_ptr, + _mm256_loadu_si256($lane_masks[selected].as_ptr().cast()), + packed, + ) + }; + } + write_pos += selected; + } + sub += $lanes; + } + write_pos + } + + /// # Safety + /// + /// The CPU must support AVX2 and the pointer contract of + /// [`filter_slice_by_bitmap`](super::filter_slice_by_bitmap) / + /// [`filter_slice_mut_by_bitmap`](super::filter_slice_mut_by_bitmap) must hold. + #[target_feature(enable = "avx2")] + pub(super) unsafe fn $walk_fn( + src: *const u8, + dst: *mut u8, + mask: &MaskValues, + ) -> usize { + let mut write_pos = 0; + for_each_mask_word(mask, |word, word_start, word_len| { + // SAFETY: forwarded from the caller contract. + write_pos = unsafe { + $word_fn::(src, dst, word, word_start, word_len, write_pos) + }; + }); + write_pos + } + }; +} + +avx2_compress_kernel!( + compress_word_avx2_epi32, compress_avx2_epi32, + elem: i32, + lanes: 8, + perm_lut: PERM_LUT_32, + lane_masks: LANE_MASK_32, + maskload: _mm256_maskload_epi32, + maskstore: _mm256_maskstore_epi32 +); + +avx2_compress_kernel!( + compress_word_avx2_epi64, compress_avx2_epi64, + elem: i64, + lanes: 4, + perm_lut: PERM_LUT_64, + lane_masks: LANE_MASK_64, + maskload: _mm256_maskload_epi64, + maskstore: _mm256_maskstore_epi64 +); + +/// Byte-index rows for `pshufb`, which always indexes a full 16-byte register even though only +/// the low 8 (1-byte elements) or all 16 (2-byte elements) bytes hold lanes. +static SHUF_LUT_8: [[u8; 16]; 256] = compress_lut::<256, 16>(8, 1); +static SHUF_LUT_16: [[u8; 16]; 256] = compress_lut::<256, 16>(8, 2); + +/// Generate an AVX2 `pshufb` kernel for 1- or 2-byte elements. +macro_rules! pshufb_compress_kernel { + ( + $word_fn:ident, + $walk_fn:ident,elem_size: + $elem_size:literal,idx_lut: + $idx_lut:ident,load: + $load:ident,store: + $store:ident + ) => { + /// # Safety + /// + /// The CPU must support AVX2 and the pointer contract of + /// [`filter_slice_by_bitmap`](super::filter_slice_by_bitmap) / + /// [`filter_slice_mut_by_bitmap`](super::filter_slice_mut_by_bitmap) must hold. + #[expect( + clippy::cast_possible_truncation, + reason = "deliberate submask narrowing" + )] + #[target_feature(enable = "avx2")] + #[inline] + unsafe fn $word_fn( + src: *const u8, + dst: *mut u8, + word: u64, + word_start: usize, + word_len: usize, + mut write_pos: usize, + ) -> usize { + if word == 0 { + return write_pos; + } + if word == low_bits_mask(word_len) { + // SAFETY: forwarded from the caller contract. + unsafe { + bulk_copy::(src, dst, word_start, word_len, write_pos, $elem_size) + }; + return write_pos + word_len; + } + + // Empty chunks still store garbage that the next chunk overwrites; branching here + // regresses masks near the density crossover. + let mut sub = 0; + while sub + 8 <= word_len { + let m = ((word >> sub) & low_bits_mask(8)) as usize; + // SAFETY: the chunk holds 8 in-bounds source elements. + let chunk = unsafe { $load(src.add((word_start + sub) * $elem_size).cast()) }; + // SAFETY: every LUT row is 16 bytes. + let idx = unsafe { _mm_loadu_si128($idx_lut[m].as_ptr().cast()) }; + // SAFETY: out-of-place output has vector slack. In-place, the store ends within + // the source chunk already loaded, and later stores overwrite trailing garbage. + unsafe { + $store( + dst.add(write_pos * $elem_size).cast::<__m128i>(), + _mm_shuffle_epi8(chunk, idx), + ) + }; + write_pos += m.count_ones() as usize; + sub += 8; + } + + if sub < word_len { + let bits = (word >> sub) & low_bits_mask(word_len - sub); + // SAFETY: forwarded from the caller contract. + write_pos = unsafe { + compress_tail::( + src, + dst, + bits, + word_start + sub, + write_pos, + $elem_size, + ) + }; + } + + write_pos + } + + /// # Safety + /// + /// The CPU must support AVX2 and the pointer contract of + /// [`filter_slice_by_bitmap`](super::filter_slice_by_bitmap) / + /// [`filter_slice_mut_by_bitmap`](super::filter_slice_mut_by_bitmap) must hold. + #[target_feature(enable = "avx2")] + pub(super) unsafe fn $walk_fn( + src: *const u8, + dst: *mut u8, + mask: &MaskValues, + ) -> usize { + let mut write_pos = 0; + for_each_mask_word(mask, |word, word_start, word_len| { + // SAFETY: forwarded from the caller contract. + write_pos = unsafe { + $word_fn::(src, dst, word, word_start, word_len, write_pos) + }; + }); + write_pos + } + }; +} + +pshufb_compress_kernel!( + compress_word_pshufb_epi8, compress_pshufb_epi8, + elem_size: 1, + idx_lut: SHUF_LUT_8, + load: _mm_loadl_epi64, + store: _mm_storel_epi64 +); + +pshufb_compress_kernel!( + compress_word_pshufb_epi16, compress_pshufb_epi16, + elem_size: 2, + idx_lut: SHUF_LUT_16, + load: _mm_loadu_si128, + store: _mm_storeu_si128 +); diff --git a/vortex-array/src/arrays/filter/execute/slice.rs b/vortex-array/src/arrays/filter/execute/slice.rs index 1528d272b28..52d1328c92e 100644 --- a/vortex-array/src/arrays/filter/execute/slice.rs +++ b/vortex-array/src/arrays/filter/execute/slice.rs @@ -3,47 +3,124 @@ //! Core slice-level filtering algorithms. //! -//! Provides both immutable and mutable (in-place) filtering of typed slices by various mask -//! representations: indices and ranges (slices). +//! Provides both immutable and mutable (in-place) filtering of typed slices by cached mask +//! representations or directly from the mask bitmap. use std::ptr; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; -use vortex_mask::MaskIter; use vortex_mask::MaskValues; -// This is modeled after the constant with the equivalent name in arrow-rs. -pub(super) const FILTER_SLICES_SELECTIVITY_THRESHOLD: f64 = 0.8; +/// Invoke `f` with each `(word, word_start, word_len)` of the mask bitmap, where `word` holds +/// the mask bits for elements `word_start..word_start + word_len` in its low `word_len` bits. +#[allow(clippy::inline_always)] +#[inline(always)] +pub(super) fn for_each_mask_word(mask: &MaskValues, mut f: impl FnMut(u64, usize, usize)) { + let bits = mask.bit_buffer(); + let unaligned = bits.unaligned_chunks(); + let lead = unaligned.lead_padding(); + let mut base = 0; + + if let Some(prefix) = unaligned.prefix() { + let len = (64 - lead).min(mask.len()); + f(prefix >> lead, base, len); + base += len; + } + + for &word in unaligned.chunks() { + f(word, base, 64); + base += 64; + } + + if let Some(suffix) = unaligned.suffix() { + let len = mask.len() - base; + f(suffix, base, len); + base += len; + } -// --------------------------------------------------------------------------- -// Immutable slice filtering -// --------------------------------------------------------------------------- + debug_assert_eq!(base, mask.len()); +} -/// Filter a slice by [`MaskValues`], dispatching to the indices or slices path based on a -/// selectivity threshold. -pub(super) fn filter_slice_by_mask_values(slice: &[T], mask: &MaskValues) -> Buffer { +/// A `u64` with the low `len` bits set. +#[inline] +pub(super) fn low_bits_mask(len: usize) -> u64 { + debug_assert!(len <= 64); + if len == 64 { + u64::MAX + } else { + (1u64 << len) - 1 + } +} + +/// Filter a slice from the mask bitmap without materializing indices or ranges. +pub(super) fn filter_slice_by_bitmap(slice: &[T], mask: &MaskValues) -> Buffer { assert_eq!( mask.len(), slice.len(), "Selection mask length must equal the buffer length" ); - match mask.threshold_iter(FILTER_SLICES_SELECTIVITY_THRESHOLD) { - MaskIter::Indices(indices) => filter_slice_by_indices(slice, indices), - MaskIter::Slices(slices) => filter_slice_by_slices(slice, slices), - } + let output_len = mask.true_count(); + let mut out = BufferMut::::with_capacity(output_len); + let src_ptr = slice.as_ptr(); + let out_ptr = out.spare_capacity_mut().as_mut_ptr().cast::(); + let mut write_pos = 0; + + for_each_mask_word(mask, |word, word_start, word_len| { + let all_selected = low_bits_mask(word_len); + debug_assert_eq!(word & !all_selected, 0); + if word == all_selected { + // SAFETY: a full mask word selects `word_len` in-bounds source values and the output + // was allocated for every selected value. + unsafe { + ptr::copy_nonoverlapping(src_ptr.add(word_start), out_ptr.add(write_pos), word_len); + } + write_pos += word_len; + } else { + let mut selected = word; + while selected != 0 { + let index = word_start + selected.trailing_zeros() as usize; + // SAFETY: set bits are limited to `word_len`, and the output was allocated for + // exactly `mask.true_count()` values. + unsafe { + out_ptr.add(write_pos).write(*src_ptr.add(index)); + } + write_pos += 1; + selected &= selected - 1; + } + } + }); + + debug_assert_eq!(write_pos, output_len); + // SAFETY: every output slot was initialized exactly once above. + unsafe { out.set_len(output_len) }; + out.freeze() } /// Filter a slice by a set of strictly increasing indices. -fn filter_slice_by_indices(slice: &[T], indices: &[usize]) -> Buffer { - Buffer::::from_trusted_len_iter(indices.iter().map(|&idx| slice[idx])) +pub(super) fn filter_slice_by_indices(slice: &[T], indices: &[usize]) -> Buffer { + let mut out = BufferMut::::with_capacity(indices.len()); + let src_ptr = slice.as_ptr(); + let out_ptr = out.spare_capacity_mut().as_mut_ptr().cast::(); + + for (write_pos, &index) in indices.iter().enumerate() { + // SAFETY: mask indices are validated when the mask is constructed and the output has one + // slot allocated for every index. + unsafe { out_ptr.add(write_pos).write(*src_ptr.add(index)) }; + } + + // SAFETY: the loop initialized every output slot. + unsafe { out.set_len(indices.len()) }; + out.freeze() } /// Filter a slice by a set of strictly increasing `(start, end)` ranges. -fn filter_slice_by_slices(slice: &[T], slices: &[(usize, usize)]) -> Buffer { - let output_len: usize = slices.iter().map(|(start, end)| end - start).sum(); - +pub(super) fn filter_slice_by_slices( + slice: &[T], + slices: &[(usize, usize)], + output_len: usize, +) -> Buffer { let mut out = BufferMut::::with_capacity(output_len); for (start, end) in slices { out.extend_from_slice(&slice[*start..*end]); @@ -52,44 +129,80 @@ fn filter_slice_by_slices(slice: &[T], slices: &[(usize, usize)]) -> Bu out.freeze() } -// --------------------------------------------------------------------------- -// Mutable (in-place) slice filtering -// --------------------------------------------------------------------------- - -/// Filter a mutable slice in-place by [`MaskValues`], returning the new valid length. -/// -/// We always use the slices path here because iterating over indices will have strictly more -/// loop iterations than slices (more branches), and the overhead of batched `ptr::copy(len)` is -/// not that high. -pub(super) fn filter_slice_mut_by_mask_values( - slice: &mut [T], - mask: &MaskValues, -) -> usize { +/// Filter a mutable slice in-place from the mask bitmap, returning the new valid length. +pub(super) fn filter_slice_mut_by_bitmap(slice: &mut [T], mask: &MaskValues) -> usize { assert_eq!( slice.len(), mask.len(), "Mask length must equal the slice length" ); - filter_slice_mut_by_slices(slice, mask.slices()) + let ptr = slice.as_mut_ptr(); + let mut write_pos = 0; + + for_each_mask_word(mask, |word, word_start, word_len| { + let all_selected = low_bits_mask(word_len); + debug_assert_eq!(word & !all_selected, 0); + if word == all_selected { + if write_pos != word_start { + // SAFETY: source and destination are in bounds and may overlap while compacting + // toward the start of the same allocation. + unsafe { ptr::copy(ptr.add(word_start), ptr.add(write_pos), word_len) }; + } + write_pos += word_len; + } else { + let mut selected = word; + while selected != 0 { + let index = word_start + selected.trailing_zeros() as usize; + if write_pos != index { + // SAFETY: set bits are limited to `word_len` and stable compaction guarantees + // `write_pos <= index`. + unsafe { ptr::copy(ptr.add(index), ptr.add(write_pos), 1) }; + } + write_pos += 1; + selected &= selected - 1; + } + } + }); + + debug_assert_eq!(write_pos, mask.true_count()); + write_pos +} + +/// Filter a mutable slice in-place by strictly increasing indices. +pub(super) fn filter_slice_mut_by_indices(slice: &mut [T], indices: &[usize]) -> usize { + let ptr = slice.as_mut_ptr(); + for (write_pos, &index) in indices.iter().enumerate() { + if write_pos != index { + // SAFETY: mask indices are in bounds and stable compaction guarantees + // `write_pos <= index`. + unsafe { ptr::copy(ptr.add(index), ptr.add(write_pos), 1) }; + } + } + indices.len() } /// Filter a mutable slice in-place by a set of `(start, end)` ranges, returning the new length. -fn filter_slice_mut_by_slices(slice: &mut [T], slices: &[(usize, usize)]) -> usize { +pub(super) fn filter_slice_mut_by_slices( + slice: &mut [T], + slices: &[(usize, usize)], +) -> usize { let mut write_pos = 0; - // For each range in the selection, copy all of the elements to the current write position. for &(start, end) in slices { let len = end - start; - // SAFETY: Slices should be within bounds. - unsafe { - ptr::copy( - slice.as_ptr().add(start), - slice.as_mut_ptr().add(write_pos), - len, - ) - }; + if write_pos != start { + // SAFETY: mask slices are in bounds and source and destination may overlap while + // compacting toward the start of the same allocation. + unsafe { + ptr::copy( + slice.as_ptr().add(start), + slice.as_mut_ptr().add(write_pos), + len, + ) + }; + } write_pos += len; } diff --git a/vortex-array/src/arrays/filter/execute/struct_.rs b/vortex-array/src/arrays/filter/execute/struct_.rs index e4a8157d2e2..eb52b9b5757 100644 --- a/vortex-array/src/arrays/filter/execute/struct_.rs +++ b/vortex-array/src/arrays/filter/execute/struct_.rs @@ -1,26 +1,28 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::sync::Arc; - use vortex_error::VortexExpect; use vortex_mask::Mask; -use vortex_mask::MaskValues; +use vortex_mask::MaskValuesRef; use crate::ArrayRef; use crate::arrays::StructArray; +use crate::arrays::filter::execute::buffer::prepare_mask_for_reuse; use crate::arrays::filter::execute::filter_validity; use crate::arrays::struct_::StructArrayExt; -pub fn filter_struct(array: &StructArray, mask: &Arc) -> StructArray { +pub fn filter_struct(array: &StructArray, mask: &MaskValuesRef) -> StructArray { + let consumers = array.struct_fields().nfields(); + prepare_mask_for_reuse(mask, consumers); + let filtered_validity = filter_validity( array .validity() - .vortex_expect("struct validity should be derivable"), + .vortex_expect("validity is derivable for a valid StructArray"), mask, ); - let mask_for_filter = Mask::Values(Arc::clone(mask)); + let mask_for_filter = Mask::Values(MaskValuesRef::clone(mask)); let fields: Vec = array .iter_unmasked_fields() .map(|field| { @@ -45,7 +47,7 @@ pub fn filter_struct(array: &StructArray, mask: &Arc) -> StructArray } #[cfg(test)] -mod test { +mod tests { use vortex_mask::Mask; use crate::IntoArray; diff --git a/vortex-array/src/arrays/filter/execute/take.rs b/vortex-array/src/arrays/filter/execute/take.rs index 6e013406d71..b8d8d2da1b6 100644 --- a/vortex-array/src/arrays/filter/execute/take.rs +++ b/vortex-array/src/arrays/filter/execute/take.rs @@ -25,7 +25,7 @@ use crate::arrays::ConstantArray; use crate::arrays::Filter; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; -use crate::arrays::filter::FilterArrayExt; +use crate::arrays::filter::FilterArraySlotsExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::executor::ExecutionCtx; diff --git a/vortex-array/src/arrays/filter/execute/take/fixed_width.rs b/vortex-array/src/arrays/filter/execute/take/fixed_width.rs index 3e508abd2c5..6518293db41 100644 --- a/vortex-array/src/arrays/filter/execute/take/fixed_width.rs +++ b/vortex-array/src/arrays/filter/execute/take/fixed_width.rs @@ -20,7 +20,7 @@ use crate::arrays::DecimalArray; use crate::arrays::Filter; use crate::arrays::PrimitiveArray; use crate::arrays::decimal::DecimalArrayExt; -use crate::arrays::filter::FilterArrayExt; +use crate::arrays::filter::FilterArraySlotsExt; use crate::dtype::IntegerPType; use crate::executor::ExecutionCtx; use crate::match_each_decimal_value_type; @@ -175,7 +175,7 @@ where take_values_by_rank(values, ranks, filtered_len, |idx| idx) } } - AllOr::None => unreachable!("empty filters are handled by take preconditions"), + AllOr::None => unreachable!("empty filters are handled by the filter short circuit"), AllOr::Some(indices) => { if let Some(indices_validity) = indices_validity { take_values_by_rank_nullable( diff --git a/vortex-array/src/arrays/filter/execute/take/rank.rs b/vortex-array/src/arrays/filter/execute/take/rank.rs index efbfa768df3..08170ab85d3 100644 --- a/vortex-array/src/arrays/filter/execute/take/rank.rs +++ b/vortex-array/src/arrays/filter/execute/take/rank.rs @@ -9,6 +9,7 @@ use vortex_error::vortex_bail; use vortex_mask::AllOr; use vortex_mask::Mask; +use super::super::contiguous_filter_range; use super::small_take_rank_lookup_len; use crate::arrays::PrimitiveArray; use crate::dtype::IntegerPType; @@ -82,7 +83,7 @@ pub(in crate::arrays::filter) fn translate_ranks( match filter.indices() { AllOr::All => translate_ranks_with(ranks, ranks_validity, filtered_len, |rank| rank), - AllOr::None => unreachable!("empty filters are handled by take preconditions"), + AllOr::None => unreachable!("empty filters are handled by the filter short circuit"), AllOr::Some(filter_indices) => { translate_ranks_with(ranks, ranks_validity, filtered_len, |rank| unsafe { *filter_indices.get_unchecked(rank) @@ -153,7 +154,5 @@ pub(in crate::arrays::filter) fn contiguous_sequential_take_range Option { - let start = filter.first()?; - let end = filter.last()?.checked_add(1)?; - (end - start == filter.true_count()).then_some(start) + contiguous_filter_range(filter).map(|range| range.start) } diff --git a/vortex-array/src/arrays/filter/execute/union.rs b/vortex-array/src/arrays/filter/execute/union.rs new file mode 100644 index 00000000000..81955d1371d --- /dev/null +++ b/vortex-array/src/arrays/filter/execute/union.rs @@ -0,0 +1,32 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexExpect; +use vortex_mask::Mask; +use vortex_mask::MaskValuesRef; + +use crate::ArrayRef; +use crate::arrays::UnionArray; +use crate::arrays::union::UnionArrayExt; +use crate::arrays::union::UnionArraySlotsExt; + +pub fn filter_union(array: &UnionArray, mask: &MaskValuesRef) -> UnionArray { + let filter_mask = Mask::Values(MaskValuesRef::clone(mask)); + + let type_ids = array + .type_ids() + .filter(filter_mask.clone()) + .vortex_expect("UnionArray type IDs are guaranteed to support filter"); + + let children: Vec = array + .iter_children() + .map(|child| { + child + .filter(filter_mask.clone()) + .vortex_expect("UnionArray children are guaranteed to support filter") + }) + .collect(); + + UnionArray::try_new(type_ids, array.variants().clone(), children) + .vortex_expect("filtered UnionArray children have consistent dtypes and lengths") +} diff --git a/vortex-array/src/arrays/filter/execute/varbinview.rs b/vortex-array/src/arrays/filter/execute/varbinview.rs index 58a348b41b0..ed490347599 100644 --- a/vortex-array/src/arrays/filter/execute/varbinview.rs +++ b/vortex-array/src/arrays/filter/execute/varbinview.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use vortex_buffer::Buffer; -use vortex_mask::MaskValues; +use vortex_mask::MaskValuesRef; use crate::arrays::VarBinViewArray; use crate::arrays::filter::execute::buffer; @@ -13,7 +13,7 @@ use crate::arrays::varbinview::BinaryView; use crate::arrays::varbinview::VarBinViewArrayExt; use crate::buffer::BufferHandle; -pub fn filter_varbinview(array: &VarBinViewArray, mask: &Arc) -> VarBinViewArray { +pub fn filter_varbinview(array: &VarBinViewArray, mask: &MaskValuesRef) -> VarBinViewArray { let filtered_validity = filter_validity(array.varbinview_validity(), mask); let views = Buffer::::from_byte_buffer(array.views_handle().as_host().clone()); @@ -32,7 +32,7 @@ pub fn filter_varbinview(array: &VarBinViewArray, mask: &Arc) -> Var } #[cfg(test)] -mod test { +mod tests { use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; diff --git a/vortex-array/src/arrays/filter/kernel.rs b/vortex-array/src/arrays/filter/kernel.rs index dbbe283138c..a62dc7910a0 100644 --- a/vortex-array/src/arrays/filter/kernel.rs +++ b/vortex-array/src/arrays/filter/kernel.rs @@ -67,11 +67,11 @@ pub trait FilterKernel: VTable { ) -> VortexResult>; } -/// Common preconditions for filter operations that apply to all arrays. +/// Short-circuits filter for the inputs that need no encoding-specific work. /// -/// Returns `Some(result)` if the precondition short-circuits the filter operation, -/// or `None` if the filter should proceed normally. -fn precondition(array: ArrayView<'_, V>, mask: &Mask) -> Option { +/// Returns `Some(result)` when the answer is already known, or `None` when the filter must proceed +/// normally. +fn short_circuit(array: ArrayView<'_, V>, mask: &Mask) -> Option { let true_count = mask.true_count(); // Fast-path for empty mask (all false). @@ -104,7 +104,7 @@ where child_idx: usize, ) -> VortexResult> { assert_eq!(child_idx, 0); - if let Some(result) = precondition::(array, parent.filter_mask()) { + if let Some(result) = short_circuit::(array, parent.filter_mask()) { return Ok(Some(result)); } ::filter(array, parent.filter_mask()) @@ -129,7 +129,7 @@ where ctx: &mut ExecutionCtx, ) -> VortexResult> { assert_eq!(child_idx, 0); - if let Some(result) = precondition::(array, parent.filter_mask()) { + if let Some(result) = short_circuit::(array, parent.filter_mask()) { return Ok(Some(result)); } ::filter(array, parent.filter_mask(), ctx) diff --git a/vortex-array/src/arrays/filter/mod.rs b/vortex-array/src/arrays/filter/mod.rs index 9510ca92f5b..fb7dc6a4141 100644 --- a/vortex-array/src/arrays/filter/mod.rs +++ b/vortex-array/src/arrays/filter/mod.rs @@ -2,12 +2,17 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod array; -pub use array::FilterArrayExt; +pub use array::FilterArraySlotsExt; pub use array::FilterData; pub use array::FilterDataParts; +pub use array::FilterSlots; +pub use array::FilterSlotsView; pub use vtable::FilterArray; mod execute; +pub(crate) use execute::buffer::filter_buffer; +pub(crate) use execute::buffer::prepare_mask_for_reuse; +pub(crate) use execute::filter_validity; mod kernel; pub use kernel::FilterExecuteAdaptor; diff --git a/vortex-array/src/arrays/filter/rules.rs b/vortex-array/src/arrays/filter/rules.rs index ffa5c64bd61..53f35c154cd 100644 --- a/vortex-array/src/arrays/filter/rules.rs +++ b/vortex-array/src/arrays/filter/rules.rs @@ -11,16 +11,24 @@ use crate::array::ArrayView; use crate::arrays::Filter; use crate::arrays::Struct; use crate::arrays::StructArray; -use crate::arrays::filter::FilterArrayExt; +use crate::arrays::filter::FilterArraySlotsExt; use crate::arrays::filter::FilterReduce; use crate::arrays::filter::FilterReduceAdaptor; +use crate::arrays::filter::execute::buffer::prepare_mask_for_reuse; +use crate::arrays::scalar_fn::ExactScalarFn; +use crate::arrays::scalar_fn::ScalarFnArrayView; use crate::arrays::struct_::StructDataParts; +use crate::builtins::ArrayBuiltins; +use crate::optimizer::rules::ArrayParentReduceRule; use crate::optimizer::rules::ArrayReduceRule; use crate::optimizer::rules::ParentRuleSet; use crate::optimizer::rules::ReduceRuleSet; +use crate::scalar_fn::fns::get_item::GetItem; -pub(super) const PARENT_RULES: ParentRuleSet = - ParentRuleSet::new(&[ParentRuleSet::lift(&FilterReduceAdaptor(Filter))]); +pub(super) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ + ParentRuleSet::lift(&FilterReduceAdaptor(Filter)), + ParentRuleSet::lift(&FilterGetItemRule), +]); pub(super) const RULES: ReduceRuleSet = ReduceRuleSet::new(&[&TrivialFilterRule, &FilterStructRule]); @@ -34,6 +42,23 @@ impl FilterReduce for Filter { } } +#[derive(Debug)] +struct FilterGetItemRule; + +impl ArrayParentReduceRule for FilterGetItemRule { + type Parent = ExactScalarFn; + + fn reduce_parent( + &self, + array: ArrayView<'_, Filter>, + parent: ScalarFnArrayView<'_, GetItem>, + _child_idx: usize, + ) -> VortexResult> { + let field = array.child().get_item(parent.options.clone())?; + Ok(Some(field.filter(array.filter_mask().clone())?)) + } +} + #[derive(Debug)] struct TrivialFilterRule; @@ -66,6 +91,11 @@ impl ArrayReduceRule for FilterStructRule { .. } = struct_array.into_owned().into_data_parts(); + if let Some(values) = mask.values() { + let consumers = fields.len(); + prepare_mask_for_reuse(values, consumers); + } + let filtered_validity = validity.filter(mask)?; let filtered_fields = fields diff --git a/vortex-array/src/arrays/filter/vtable.rs b/vortex-array/src/arrays/filter/vtable.rs index 0f54a97c574..56fc112aa16 100644 --- a/vortex-array/src/arrays/filter/vtable.rs +++ b/vortex-array/src/arrays/filter/vtable.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::hash::Hasher; -use std::sync::Arc; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -10,6 +9,7 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; use vortex_mask::Mask; +use vortex_mask::MaskValuesRef; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -28,12 +28,12 @@ use crate::array::OperationsVTable; use crate::array::VTable; use crate::array::ValidityVTable; use crate::array::with_empty_buffers; -use crate::arrays::filter::FilterArrayExt; -use crate::arrays::filter::array::CHILD_SLOT; +use crate::arrays::filter::FilterArraySlotsExt; use crate::arrays::filter::array::FilterData; -use crate::arrays::filter::array::SLOT_NAMES; +use crate::arrays::filter::array::FilterSlots; +use crate::arrays::filter::execute::contiguous_filter_range; +use crate::arrays::filter::execute::execute_all_null_filter_fast_path; use crate::arrays::filter::execute::execute_filter; -use crate::arrays::filter::execute::execute_filter_fast_paths; use crate::arrays::filter::rules::PARENT_RULES; use crate::arrays::filter::rules::RULES; use crate::buffer::BufferHandle; @@ -80,10 +80,10 @@ impl VTable for Filter { slots: &[Option], ) -> VortexResult<()> { vortex_ensure!( - slots[CHILD_SLOT].is_some(), + slots[FilterSlots::CHILD].is_some(), "FilterArray child slot must be present" ); - let child = slots[CHILD_SLOT] + let child = slots[FilterSlots::CHILD] .as_ref() .vortex_expect("validated child slot"); vortex_ensure!( @@ -128,7 +128,7 @@ impl VTable for Filter { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + FilterSlots::NAMES[idx].to_string() } fn serialize( @@ -153,15 +153,30 @@ impl VTable for Filter { } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - if let Some(canonical) = execute_filter_fast_paths(array.as_view(), ctx)? { - return Ok(ExecutionResult::done(canonical)); - } + // Match the mask once. A zero-length mask is both all true and all false, so check the + // empty output before the unfiltered one. let mask_values = match &array.mask { - Mask::Values(v) => Arc::clone(v), - _ => unreachable!("`execute_filter_fast_paths` handles AllTrue and AllFalse"), + Mask::AllFalse(_) | Mask::AllTrue(0) => { + return Ok(ExecutionResult::done( + Canonical::empty(array.dtype()).into_array(), + )); + } + Mask::AllTrue(_) => return Ok(ExecutionResult::done(array.child().clone())), + Mask::Values(values) => MaskValuesRef::clone(values), }; - let array = require_child!(array, array.child(), CHILD_SLOT => AnyCanonical); + // Filtering by one contiguous range is exactly a slice and can remain zero-copy. + if let Some(range) = contiguous_filter_range(array.filter_mask()) { + return Ok(ExecutionResult::done(array.child().slice(range)?)); + } + + if let Some(canonical) = + execute_all_null_filter_fast_path(array.as_view(), mask_values.true_count(), ctx)? + { + return Ok(ExecutionResult::done(canonical)); + } + + let array = require_child!(array, array.child(), FilterSlots::CHILD => AnyCanonical); // We rely on the optimization pass that runs prior to this execution for filter pushdown, // so now we can just execute the filter without worrying. diff --git a/vortex-array/src/arrays/fixed_size_list/array.rs b/vortex-array/src/arrays/fixed_size_list/array.rs index b6328163727..98ba45d0307 100644 --- a/vortex-array/src/arrays/fixed_size_list/array.rs +++ b/vortex-array/src/arrays/fixed_size_list/array.rs @@ -5,7 +5,6 @@ use std::fmt::Display; use std::fmt::Formatter; use std::sync::Arc; -use smallvec::smallvec; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -17,24 +16,28 @@ use crate::array::ArrayParts; use crate::array::TypedArrayRef; use crate::array::child_to_validity; use crate::array::validity_to_child; +use crate::array_slots; use crate::arrays::FixedSizeList; use crate::dtype::DType; use crate::validity::Validity; -/// The `elements` data array, where each fixed-size list scalar is a _slice_ of the `elements` -/// array, and each inner list element is a _scalar_ of the `elements` array. -/// -/// The fixed-size list scalars are contiguous (regardless of nullability for easy lookups), -/// each with equal size in memory. -pub(super) const ELEMENTS_SLOT: usize = 0; -/// The validity / null map of the array. -/// -/// Note that this null map refers to which fixed-size list scalars are null, **not** which -/// sub-elements of fixed-size list scalars are null. The `elements` array will track individual -/// value nullability. -pub(super) const VALIDITY_SLOT: usize = 1; -pub(super) const NUM_SLOTS: usize = 2; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["elements", "validity"]; +#[array_slots(FixedSizeList)] +pub struct FixedSizeListSlots { + /// The `elements` data array, where each fixed-size list scalar is a _slice_ of the `elements` + /// array, and each inner list element is a _scalar_ of the `elements` array. + /// + /// The fixed-size list scalars are contiguous (regardless of nullability for easy lookups), + /// each with equal size in memory. + #[slot(0)] + pub elements: ArrayRef, + /// The validity / null map of the array. + /// + /// Note that this null map refers to which fixed-size list scalars are null, **not** which + /// sub-elements of fixed-size list scalars are null. The `elements` array will track + /// individual value nullability. + #[slot(1)] + pub validity: Option, +} /// The canonical encoding for fixed-size list arrays. /// @@ -110,7 +113,11 @@ pub struct FixedSizeListDataParts { impl FixedSizeListData { pub(crate) fn make_slots(elements: &ArrayRef, validity: &Validity, len: usize) -> ArraySlots { - smallvec![Some(elements.clone()), validity_to_child(validity, len)] + FixedSizeListSlots { + elements: elements.clone(), + validity: validity_to_child(validity, len), + } + .into_slots() } /// Creates a new `FixedSizeListArray`. @@ -203,7 +210,7 @@ impl FixedSizeListData { } } -pub trait FixedSizeListArrayExt: TypedArrayRef { +pub trait FixedSizeListArrayExt: FixedSizeListArraySlotsExt { fn dtype_parts(&self) -> (&DType, u32, crate::dtype::Nullability) { match self.as_ref().dtype() { DType::FixedSizeList(element_dtype, list_size, nullability) => { @@ -213,12 +220,6 @@ pub trait FixedSizeListArrayExt: TypedArrayRef { } } - fn elements(&self) -> &ArrayRef { - self.as_ref().slots()[ELEMENTS_SLOT] - .as_ref() - .vortex_expect("FixedSizeListArray elements slot") - } - fn list_size(&self) -> u32 { let (_, list_size, _) = self.dtype_parts(); list_size @@ -226,7 +227,10 @@ pub trait FixedSizeListArrayExt: TypedArrayRef { fn fixed_size_list_validity(&self) -> Validity { let (_, _, nullability) = self.dtype_parts(); - child_to_validity(self.as_ref().slots()[VALIDITY_SLOT].as_ref(), nullability) + child_to_validity( + self.as_ref().slots()[FixedSizeListSlots::VALIDITY].as_ref(), + nullability, + ) } #[allow(clippy::disallowed_methods)] @@ -310,7 +314,7 @@ impl Array { pub fn into_data_parts(self) -> FixedSizeListDataParts { let dtype = self.dtype().clone(); - let elements = self.slots()[ELEMENTS_SLOT] + let elements = self.slots()[FixedSizeListSlots::ELEMENTS] .clone() .vortex_expect("FixedSizeListArray elements slot"); let validity = self.fixed_size_list_validity(); diff --git a/vortex-array/src/arrays/fixed_size_list/compute/cast.rs b/vortex-array/src/arrays/fixed_size_list/compute/cast.rs index 4ac3bd1938f..8e813209a5b 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/cast.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/cast.rs @@ -10,6 +10,7 @@ use crate::array::ArrayView; use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::scalar_fn::fns::cast::CastKernel; diff --git a/vortex-array/src/arrays/fixed_size_list/compute/mask.rs b/vortex-array/src/arrays/fixed_size_list/compute/mask.rs index 3df0413e1aa..1db564552e7 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/mask.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/mask.rs @@ -9,6 +9,7 @@ use crate::array::ArrayView; use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::scalar_fn::fns::mask::MaskReduce; use crate::validity::Validity; diff --git a/vortex-array/src/arrays/fixed_size_list/compute/slice.rs b/vortex-array/src/arrays/fixed_size_list/compute/slice.rs index 5146fd2b600..12f67334153 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/slice.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/slice.rs @@ -11,6 +11,7 @@ use crate::array::ArrayView; use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::arrays::slice::SliceReduce; impl SliceReduce for FixedSizeList { diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index 193730f62c4..48aa9ceea33 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -1,221 +1,367 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_buffer::BitBufferMut; +use itertools::Itertools; +use num_traits::AsPrimitive; +use vortex_buffer::Buffer; use vortex_buffer::BufferMut; -use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_panic; -use vortex_mask::Mask; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use crate::ArrayRef; +use crate::Columnar; use crate::IntoArray; use crate::array::ArrayView; -use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; +use crate::arrays::PiecewiseSequence; +use crate::arrays::PiecewiseSequenceArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; -use crate::arrays::bool::BoolArrayExt; use crate::arrays::dict::TakeExecute; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use crate::arrays::piecewise_sequence::constant_unsigned_usize; +use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::primitive::PrimitiveArrayExt; +use crate::builders::builder_with_capacity_in; +use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::executor::ExecutionCtx; +use crate::match_each_integer_ptype; use crate::match_each_unsigned_integer_ptype; -use crate::match_smallest_offset_type; +use crate::optimizer::ArrayOptimizer; use crate::validity::Validity; /// Take implementation for [`FixedSizeListArray`]. /// -/// Unlike `ListView`, `FixedSizeListArray` must rebuild the elements array because it requires -/// that elements start at offset 0 and be perfectly packed without gaps. We expand list indices -/// into element indices and push them down to the child elements array. +/// `FixedSizeListArray` must rebuild its elements array because selected lists need to become +/// packed from offset 0. The FSL layer translates selected list rows into ordered element runs and +/// delegates the execution strategy to the elements child via `PiecewiseSequenceArray` indices. impl TakeExecute for FixedSizeList { fn take( array: ArrayView<'_, FixedSizeList>, indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - let max_element_idx = array.elements().len(); - // Indices are non-negative; dispatch over the 4 unsigned widths (the executed array is - // reinterpreted to unsigned in `take_with_indices`). `E` is already unsigned. - match_each_unsigned_integer_ptype!(indices.dtype().as_ptype().to_unsigned(), |I| { - match_smallest_offset_type!(max_element_idx, |E| { - take_with_indices::(array, indices, ctx) - }) - }) - .map(Some) + if array.is_empty() { + return take_empty_fsl(array, indices, ctx).map(Some); + } + + take_non_empty_fsl(array, indices, ctx).map(Some) } } -/// Dispatches to the appropriate take implementation based on list size and nullability. -fn take_with_indices( +fn take_empty_fsl( array: ArrayView<'_, FixedSizeList>, indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult { - let list_size = array.list_size() as usize; + debug_assert!(array.is_empty()); - let indices_array = indices.clone().execute::(ctx)?; - // Reinterpret to unsigned so `as_slice::` (with unsigned `I`) matches; values are unchanged. - let indices_array = indices_array.reinterpret_cast(indices_array.ptype().to_unsigned()); - - // Make sure to handle degenerate case where lists have size 0 (these can take fast paths). - if list_size == 0 { - debug_assert!( - array.elements().is_empty(), - "degenerate list must have empty elements" + let new_len = indices.len(); + if new_len != 0 { + let indices_validity = indices.validity()?.execute_mask(new_len, ctx)?; + vortex_ensure!( + indices_validity.all_false(), + "cannot take valid indices from an empty FixedSizeList" ); + } - // Since there are no elements to take, we just need to take on the validity map. - let new_validity = array.validity()?.take(indices)?; - let new_len = indices_array.len(); - - Ok( - // SAFETY: list_size is 0, elements array is empty, and validity has the correct length. - unsafe { - FixedSizeListArray::new_unchecked( - array.elements().clone(), // Remember that this is an empty array. - array.list_size(), - new_validity, - new_len, - ) - } - .into_array(), + let list_size = array.list_size() as usize; + let elements_len = new_len.checked_mul(list_size).ok_or_else(|| { + vortex_err!( + "FixedSizeList take output length overflow: {new_len} lists of size {list_size}" ) + })?; + let new_elements = default_elements(array, elements_len, ctx.allocator()); + let new_validity = if new_len == 0 { + array.validity()?.take(indices)? } else { - // The result's nullability is the union of the input nullabilities. - if array.dtype().is_nullable() || indices_array.dtype().is_nullable() { - let indices_array = indices_array.as_view(); - take_nullable_fsl::(array, indices_array, ctx) - } else { - let indices_array = indices_array.as_view(); - take_non_nullable_fsl::(array, indices_array) - } + Validity::AllInvalid + }; + + // SAFETY: empty output needs no child values; otherwise the index validity mask proves every + // output row is null. Placeholder child elements have the exact length required by FSL. + unsafe { + FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len) } + .into_array() + .optimize_ctx(ctx.session()) } -/// Takes from an array when both the array and indices are non-nullable. -fn take_non_nullable_fsl( +fn take_non_empty_fsl( array: ArrayView<'_, FixedSizeList>, - indices_array: ArrayView<'_, Primitive>, + indices: &ArrayRef, + ctx: &mut ExecutionCtx, ) -> VortexResult { - let list_size = array.list_size() as usize; - let indices: &[I] = indices_array.as_slice::(); - let new_len = indices.len(); + debug_assert!(!array.is_empty()); - // Build the element indices directly without validity tracking. - let mut elements_indices = BufferMut::::with_capacity(new_len * list_size); + let DType::Primitive(ptype, _) = indices.dtype() else { + vortex_bail!("Invalid indices dtype: {}", indices.dtype()) + }; + if !ptype.is_int() { + vortex_bail!("Invalid indices dtype: {}", indices.dtype()); + } - // Build the element indices for each list. - for data_idx in indices { - let data_idx = data_idx - .to_usize() - .unwrap_or_else(|| vortex_panic!("Failed to convert index to usize: {}", data_idx)); + if array.list_size() == 0 { + return take_non_empty_degenerate_fsl(array, indices, ctx); + } - let list_start = data_idx * list_size; - let list_end = (data_idx + 1) * list_size; + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_piecewise_fsl(array, piecewise_indices, ctx)? + { + return Ok(taken); + } - // Expand the list into individual element indices. - for i in list_start..list_end { - // SAFETY: We've allocated enough space for enough indices for all `new_len` lists (that each consist of `list_size = list_end - list_start` elements), so we know we have enough capacity. - unsafe { - elements_indices.push_unchecked(E::from_usize(i).vortex_expect("i < list_end")) - }; + let indices_array = indices.clone().execute::(ctx)?; + match_each_integer_ptype!(indices_array.ptype(), |I| { + take_non_empty_non_degenerate_fsl::(array, indices, indices_array.as_view(), ctx) + }) +} + +/// Take for [`PiecewiseSequence`] indices with unit multipliers: a piece of `length` consecutive +/// lists gathers `length * list_size` consecutive elements, so each piece scales by `list_size` +/// instead of expanding into one element run per list. +fn take_piecewise_fsl( + array: ArrayView<'_, FixedSizeList>, + indices: ArrayView<'_, PiecewiseSequence>, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { + return Ok(None); + }; + + let list_size = array.list_size() as usize; + let array_len = array.as_ref().len(); + let new_len = indices.as_ref().len(); + let elements_len = new_len.checked_mul(list_size).ok_or_else(|| { + vortex_err!( + "FixedSizeList take output length overflow: {new_len} lists of size {list_size}" + ) + })?; + + let starts: Vec = match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + starts + .as_slice::() + .iter() + .map(|&start| start.as_()) + .collect() + }); + let lengths: Vec = match lengths { + Columnar::Constant(lengths) => vec![constant_unsigned_usize(&lengths); starts.len()], + Columnar::Canonical(lengths) => { + let lengths = lengths.into_primitive(); + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + lengths + .as_slice::() + .iter() + .map(|&length| length.as_()) + .collect() + }) } + }; + + let mut element_starts = BufferMut::::with_capacity(starts.len()); + let mut element_lengths = BufferMut::::with_capacity(starts.len()); + let mut total_len = 0usize; + for (&start, &length) in starts.iter().zip_eq(&lengths) { + if length == 0 { + continue; + } + let end = start + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + if end > array_len { + vortex_bail!(OutOfBounds: end - 1, 0, array_len); + } + total_len = total_len + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + element_starts.push(u64::try_from(start * list_size)?); + element_lengths.push(u64::try_from(length * list_size)?); } + vortex_ensure!( + total_len == new_len, + "PiecewiseSequenceArray expanded length {total_len} does not match declared length {new_len}" + ); - let elements_indices = elements_indices.freeze(); - debug_assert_eq!(elements_indices.len(), new_len * list_size); + let new_elements = + if let ([start], [length]) = (element_starts.as_slice(), element_lengths.as_slice()) { + let start = usize::try_from(*start)?; + array + .elements() + .slice(start..start + usize::try_from(*length)?)? + } else { + let run_count = element_starts.len(); + let starts = + PrimitiveArray::new(element_starts.freeze(), Validity::NonNullable).into_array(); + let lengths = + PrimitiveArray::new(element_lengths.freeze(), Validity::NonNullable).into_array(); + let multipliers = ConstantArray::new(1u64, run_count).into_array(); + // SAFETY: starts, lengths, and multipliers are non-nullable u64 arrays of equal length, + // and the piece lengths were validated to sum to `new_len`, so the element runs sum to + // `elements_len`. + let element_indices = unsafe { + PiecewiseSequenceArray::new_unchecked(starts, lengths, multipliers, elements_len) + }; + array.elements().take(element_indices.into_array())? + }; + let new_validity = array.validity()?.take(indices.as_ref())?; + + // SAFETY: `new_elements` has `new_len * list_size` elements, and `Validity::take` produces + // validity for `new_len`. + unsafe { + FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len) + } + .into_array() + .optimize_ctx(ctx.session()) + .map(Some) +} - let elements_indices_array = PrimitiveArray::new(elements_indices, Validity::NonNullable); - let new_elements = array.elements().take(elements_indices_array.into_array())?; - debug_assert_eq!(new_elements.len(), new_len * list_size); +fn take_non_empty_degenerate_fsl( + array: ArrayView<'_, FixedSizeList>, + indices: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + debug_assert!(!array.is_empty()); + debug_assert_eq!(array.list_size(), 0); + vortex_ensure!( + array.elements().is_empty(), + "degenerate list must have empty elements" + ); + + let indices_array = indices.clone().execute::(ctx)?; + match_each_integer_ptype!(indices_array.ptype(), |I| { + bounds_check_valid_indices::(&indices_array.as_view(), array.as_ref().len(), ctx) + })?; + let new_validity = array.validity()?.take(indices)?; + let new_len = indices_array.len(); - // Both inputs are non-nullable, so the result is non-nullable. - Ok(unsafe { + // SAFETY: degenerate FSL inputs have no elements, valid index payloads were checked against + // the source length, and `Validity::take` produces validity for `new_len`. + unsafe { FixedSizeListArray::new_unchecked( - new_elements, + array.elements().clone(), array.list_size(), - Validity::NonNullable, + new_validity, new_len, ) } - .into_array()) + .into_array() + .optimize_ctx(ctx.session()) } -/// Takes from an array when either the array or indices are nullable. -fn take_nullable_fsl( +fn take_non_empty_non_degenerate_fsl( array: ArrayView<'_, FixedSizeList>, + indices: &ArrayRef, indices_array: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> VortexResult { + debug_assert!(!array.is_empty()); + debug_assert_ne!(array.list_size(), 0); + + let (new_elements, new_len) = + take_non_empty_non_degenerate_elements::(array, indices_array, ctx)?; + let new_validity = array.validity()?.take(indices)?; + + // SAFETY: `new_elements` has `new_len * list_size` elements, and `Validity::take` produces + // validity for `new_len`. + unsafe { + FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len) + } + .into_array() + .optimize_ctx(ctx.session()) +} + +fn take_non_empty_non_degenerate_elements( + array: ArrayView<'_, FixedSizeList>, + indices_array: ArrayView<'_, Primitive>, + ctx: &mut ExecutionCtx, +) -> VortexResult<(ArrayRef, usize)> { + debug_assert!(!array.is_empty()); + debug_assert_ne!(array.list_size(), 0); + let list_size = array.list_size() as usize; + let array_len = array.as_ref().len(); let indices: &[I] = indices_array.as_slice::(); let new_len = indices.len(); + let elements_len = new_len.checked_mul(list_size).ok_or_else(|| { + vortex_err!( + "FixedSizeList take output length overflow: {new_len} lists of size {list_size}" + ) + })?; - let array_validity = array - .fixed_size_list_validity() - .execute_mask(array.as_ref().len(), ctx) - .vortex_expect("Failed to compute validity mask"); - let indices_len = indices_array.as_ref().len(); - let indices_validity = match indices_array - .validity() - .vortex_expect("Failed to compute validity mask") - { - Validity::NonNullable | Validity::AllValid => Mask::new_true(indices_len), - Validity::AllInvalid => Mask::new_false(indices_len), - Validity::Array(a) => a.execute::(ctx)?.execute_mask(ctx), - }; + let indices_validity = indices_array.validity()?.execute_mask(new_len, ctx)?; + let starts = indices + .iter() + .zip_eq(indices_validity.iter()) + .map(|(&data_idx, is_index_valid)| { + if !is_index_valid { + return Ok(0); + } - // We must use placeholder zeros for null lists to maintain the array length without - // propagating nullability to the element array's take operation. - let mut elements_indices = BufferMut::::with_capacity(new_len * list_size); - let mut new_validity_builder = BitBufferMut::with_capacity(new_len); - - // Build the element indices while tracking which lists are null. - for (data_idx, is_index_valid) in indices.iter().zip(indices_validity.iter()) { - let data_idx = data_idx - .to_usize() - .unwrap_or_else(|| vortex_panic!("Failed to convert index to usize: {}", data_idx)); - - // The list is null if the index is null or the indexed element is null. - if !is_index_valid || !array_validity.value(data_idx) { - // Append placeholder zeros for null lists. These will be masked by the validity array. - // We cannot use append_nulls here as explained above. - unsafe { elements_indices.push_n_unchecked(E::zero(), list_size) }; - new_validity_builder.append(false); - } else { - // Append the actual element indices for this list. - let list_start = data_idx * list_size; - let list_end = (data_idx + 1) * list_size; - - // Expand the list into individual element indices. - for i in list_start..list_end { - // SAFETY: We've allocated enough space for enough indices for all `new_len` lists (that each consist of `list_size = list_end - list_start` elements), so we know we have enough capacity. - unsafe { - elements_indices.push_unchecked(E::from_usize(i).vortex_expect("i < list_end")) - }; + let data_idx: usize = data_idx.as_(); + if data_idx >= array_len { + vortex_bail!(OutOfBounds: data_idx, 0, array_len); } + Ok((data_idx * list_size) as u64) + }) + .process_results(|iter| iter.collect::>())?; + + let new_elements = + take_element_runs(array.elements(), starts.freeze(), list_size, elements_len)?; - new_validity_builder.append(true); + Ok((new_elements, new_len)) +} + +fn bounds_check_valid_indices( + indices_array: &ArrayView<'_, Primitive>, + array_len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let indices: &[I] = indices_array.as_slice::(); + let indices_validity = indices_array.validity()?.execute_mask(indices.len(), ctx)?; + + for (&data_idx, is_index_valid) in indices.iter().zip_eq(indices_validity.iter()) { + if is_index_valid { + let data_idx = data_idx.as_(); + if data_idx >= array_len { + vortex_bail!(OutOfBounds: data_idx, 0, array_len); + } } } + Ok(()) +} - let elements_indices = elements_indices.freeze(); - debug_assert_eq!(elements_indices.len(), new_len * list_size); - - let elements_indices_array = PrimitiveArray::new(elements_indices, Validity::NonNullable); - let new_elements = array.elements().take(elements_indices_array.into_array())?; - debug_assert_eq!(new_elements.len(), new_len * list_size); +fn default_elements( + array: ArrayView<'_, FixedSizeList>, + len: usize, + allocator: &vortex_buffer::BufferAllocatorRef, +) -> ArrayRef { + let mut builder = builder_with_capacity_in(array.elements().dtype(), len, allocator); + builder.append_defaults(len); + builder.finish() +} - // At least one input was nullable, so the result is nullable. - let new_validity = Validity::from(new_validity_builder.freeze()); - debug_assert!(new_validity.maybe_len().is_none_or(|vl| vl == new_len)); +fn take_element_runs( + elements: &ArrayRef, + starts: Buffer, + length: usize, + output_len: usize, +) -> VortexResult { + let run_count = starts.len(); + let starts = PrimitiveArray::new(starts, Validity::NonNullable).into_array(); + let lengths = ConstantArray::new(length as u64, run_count).into_array(); + let multipliers = ConstantArray::new(1u64, run_count).into_array(); - Ok(unsafe { - FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len) - } - .into_array()) + // SAFETY: callers produced one start per output row after validating list indices against the + // source FSL length. `length` and multiplier 1 are represented as non-nullable unsigned + // constant arrays, and `output_len` was computed as `run_count * length`. + let indices = + unsafe { PiecewiseSequenceArray::new_unchecked(starts, lengths, multipliers, output_len) }; + elements.take(indices.into_array()) } diff --git a/vortex-array/src/arrays/fixed_size_list/mod.rs b/vortex-array/src/arrays/fixed_size_list/mod.rs index 3d90d38e6e9..adf5c9bf167 100644 --- a/vortex-array/src/arrays/fixed_size_list/mod.rs +++ b/vortex-array/src/arrays/fixed_size_list/mod.rs @@ -3,8 +3,11 @@ mod array; pub use array::FixedSizeListArrayExt; +pub use array::FixedSizeListArraySlotsExt; pub use array::FixedSizeListData; pub use array::FixedSizeListDataParts; +pub use array::FixedSizeListSlots; +pub use array::FixedSizeListSlotsView; pub use vtable::FixedSizeListArray; pub(crate) mod compute; diff --git a/vortex-array/src/arrays/fixed_size_list/tests/basic.rs b/vortex-array/src/arrays/fixed_size_list/tests/basic.rs index 961801094d8..b471ebb07a0 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/basic.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/basic.rs @@ -10,6 +10,7 @@ use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; diff --git a/vortex-array/src/arrays/fixed_size_list/tests/degenerate.rs b/vortex-array/src/arrays/fixed_size_list/tests/degenerate.rs index f609fd72405..73085b8c8a8 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/degenerate.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/degenerate.rs @@ -11,6 +11,7 @@ use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; diff --git a/vortex-array/src/arrays/fixed_size_list/tests/nested.rs b/vortex-array/src/arrays/fixed_size_list/tests/nested.rs index 40fd2bc066a..b754229ff90 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/nested.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/nested.rs @@ -408,11 +408,12 @@ fn test_fsl_of_list() { let fsl_size = 3; // Create a ListBuilder for i32 values. - let mut list_builder = ListBuilder::::with_capacity( + let mut list_builder = ListBuilder::::with_capacity_in( Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), Nullability::NonNullable, 12, 6, + vortex_buffer::BufferAllocatorRef::static_ref(), ); // Add 6 lists (2 FSL * 3 lists each). @@ -489,11 +490,12 @@ fn test_fsl_of_nullable_list() { let fsl_size = 2; // Create a ListBuilder with nullable lists. - let mut list_builder = ListBuilder::::with_capacity( + let mut list_builder = ListBuilder::::with_capacity_in( Arc::new(DType::Primitive(PType::U16, Nullability::NonNullable)), Nullability::Nullable, 8, 4, + vortex_buffer::BufferAllocatorRef::static_ref(), ); // Add 4 lists (2 FSL * 2 lists each). diff --git a/vortex-array/src/arrays/fixed_size_list/tests/take.rs b/vortex-array/src/arrays/fixed_size_list/tests/take.rs index af8f94f2d5f..3904527abfd 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/take.rs @@ -2,7 +2,9 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use rstest::rstest; +use vortex_buffer::Buffer; use vortex_buffer::buffer; +use vortex_error::VortexResult; use super::common::create_basic_fsl; use super::common::create_empty_fsl; @@ -10,10 +12,13 @@ use super::common::create_large_fsl; use super::common::create_nullable_fsl; use super::common::create_single_element_fsl; use crate::ArrayRef; +use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; +use crate::arrays::ConstantArray; use crate::arrays::FixedSizeListArray; +use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; @@ -147,6 +152,105 @@ fn test_take_fsl_with_null_indices_preserves_elements() { assert_arrays_eq!(expected, result, &mut ctx); } +#[test] +fn test_take_fsl_null_index_ignores_out_of_bounds_payload() { + let mut ctx = array_session().create_execution_ctx(); + let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); + let fsl = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 3); + + let indices = PrimitiveArray::new( + buffer![1u32, 99, 2], + Validity::from_iter([true, false, true]), + ) + .into_array(); + let result = fsl.take(indices).unwrap(); + + let expected = FixedSizeListArray::new( + buffer![3i32, 4, 0, 0, 5, 6].into_array(), + 2, + Validity::from_iter([true, false, true]), + 3, + ); + assert_arrays_eq!(expected, result, &mut ctx); +} + +fn contiguous_pieces(starts: &[u64], lengths: &[u64]) -> VortexResult { + let len = usize::try_from(lengths.iter().sum::())?; + Ok(PiecewiseSequenceArray::try_new( + starts.iter().copied().collect::>().into_array(), + lengths + .iter() + .copied() + .collect::>() + .into_array(), + ConstantArray::new(1u64, starts.len()).into_array(), + len, + )? + .into_array()) +} + +fn nullable_fsl_i32() -> FixedSizeListArray { + FixedSizeListArray::new( + PrimitiveArray::from_iter(0i32..12).into_array(), + 3, + Validity::from_iter([true, false, true, true]), + 4, + ) +} + +#[rstest] +#[case::multi_piece(&[2, 0], &[2, 1], &[2u64, 3, 0])] +#[case::single_piece_slices(&[1], &[3], &[1u64, 2, 3])] +#[case::zero_length_piece(&[3, 0, 1], &[0, 1, 2], &[0u64, 1, 2])] +fn test_take_piecewise( + #[case] starts: &[u64], + #[case] lengths: &[u64], + #[case] expanded: &[u64], +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let fsl = nullable_fsl_i32(); + + let result = fsl.take(contiguous_pieces(starts, lengths)?)?; + let expected = fsl.take(Buffer::from(expanded.to_vec()).into_array())?; + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_take_piecewise_out_of_bounds() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let fsl = nullable_fsl_i32(); + + let indices = contiguous_pieces(&[2], &[3])?; + let result = fsl + .take(indices) + .and_then(|taken| taken.execute::(&mut ctx)); + + assert!(result.is_err(), "expected out-of-bounds error"); + Ok(()) +} + +#[test] +fn test_take_piecewise_non_unit_multiplier() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let fsl = nullable_fsl_i32(); + + // Multiplier 2 falls back to the generic path but must stay correct. + let indices = PiecewiseSequenceArray::try_new( + buffer![0u64].into_array(), + buffer![2u64].into_array(), + buffer![2u64].into_array(), + 2, + )? + .into_array(); + let result = fsl.take(indices)?; + let expected = fsl.take(buffer![0u64, 2].into_array())?; + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) +} + // Element index overflow: with u8 indices and list_size=16, data_idx=16 produces element index // 16*16=256 which overflows u8. The take kernel must widen the element index type. #[rstest] @@ -205,11 +309,12 @@ fn test_take_nullable_arrays_fsl_specific( 2 }; - let mut builder = FixedSizeListBuilder::with_capacity( + let mut builder = FixedSizeListBuilder::with_capacity_in( DType::Primitive(PType::I32, Nullability::NonNullable).into(), list_size, Nullability::Nullable, array_values.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), ); for value in array_values { diff --git a/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs b/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs index 9393b8563d4..656a08b0f66 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs +++ b/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs @@ -26,9 +26,7 @@ use crate::array::ArrayView; use crate::array::VTable; use crate::array::with_empty_buffers; use crate::arrays::fixed_size_list::FixedSizeListData; -use crate::arrays::fixed_size_list::array::ELEMENTS_SLOT; -use crate::arrays::fixed_size_list::array::NUM_SLOTS; -use crate::arrays::fixed_size_list::array::SLOT_NAMES; +use crate::arrays::fixed_size_list::FixedSizeListSlots; use crate::arrays::fixed_size_list::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; @@ -115,14 +113,15 @@ impl VTable for FixedSizeList { slots: &[Option], ) -> VortexResult<()> { vortex_ensure!( - slots.len() == NUM_SLOTS, - "FixedSizeListArray expected {NUM_SLOTS} slots, found {}", + slots.len() == FixedSizeListSlots::COUNT, + "FixedSizeListArray expected {} slots, found {}", + FixedSizeListSlots::COUNT, slots.len() ); let DType::FixedSizeList(_, list_size, nullability) = dtype else { vortex_bail!("Expected `DType::FixedSizeList`, got {dtype:?}"); }; - let elements = slots[ELEMENTS_SLOT] + let elements = slots[FixedSizeListSlots::ELEMENTS] .as_ref() .vortex_expect("FixedSizeListArray elements slot"); vortex_ensure!( @@ -197,7 +196,7 @@ impl VTable for FixedSizeList { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + FixedSizeListSlots::NAMES[idx].to_string() } fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { diff --git a/vortex-array/src/arrays/fixed_width/array.rs b/vortex-array/src/arrays/fixed_width/array.rs new file mode 100644 index 00000000000..96ee6e6255f --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/array.rs @@ -0,0 +1,62 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::ByteBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; + +use crate::array::Array; +use crate::array::ArrayView; +use crate::array::VTable; +use crate::validity::Validity; + +/// Type-specific access needed by shared fixed-width structural compute. +/// +/// The shared `take` and `filter` kernels view an implementing array as `array.len()` records of +/// [`byte_width`] bytes each, which lets them move whole records without knowing the logical +/// type. +/// +/// [`byte_width`]: FixedWidthArray::byte_width +pub(crate) trait FixedWidthArray: VTable { + /// Returns the number of bytes each record occupies. + fn byte_width(array: ArrayView<'_, Self>) -> usize; + + /// Returns the records of `array` as a single host-resident byte buffer. + /// + /// The returned buffer must contain exactly `array.len() * byte_width` bytes. + fn values(array: ArrayView<'_, Self>) -> ByteBuffer; + + /// Rebuilds an array of this encoding from a records buffer, preserving the logical type of + /// `array`. + /// + /// Callers must provide a `values` buffer of exactly `len * byte_width` bytes and a + /// `validity` with logical length `len`. The shared kernels enforce the buffer length through + /// the module-level `with_values` helper rather than in each implementation. + fn with_values( + array: ArrayView<'_, Self>, + values: ByteBuffer, + len: usize, + validity: Validity, + ) -> VortexResult>; +} + +/// Rebuilds a fixed-width array from `len` records in `values`, validating the buffer length +/// before dispatching to [`FixedWidthArray::with_values`]. +pub(crate) fn with_values( + array: ArrayView<'_, V>, + values: ByteBuffer, + len: usize, + validity: Validity, +) -> VortexResult> { + let expected_len = len + .checked_mul(V::byte_width(array)) + .ok_or_else(|| vortex_err!("Fixed-width values buffer length overflows usize"))?; + vortex_ensure_eq!( + values.len(), + expected_len, + "Fixed-width values buffer length {} does not match expected length {expected_len}", + values.len(), + ); + V::with_values(array, values, len, validity) +} diff --git a/vortex-array/src/arrays/fixed_width/filter.rs b/vortex-array/src/arrays/fixed_width/filter.rs new file mode 100644 index 00000000000..7d7eec9d635 --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/filter.rs @@ -0,0 +1,70 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_buffer::ByteBuffer; +use vortex_error::VortexExpect; +use vortex_mask::MaskValues; +use vortex_mask::MaskValuesRef; + +use super::FixedWidthArray; +use super::match_each_record_width; +use super::with_values; +use crate::array::Array; +use crate::arrays::filter::filter_buffer; +use crate::arrays::filter::filter_validity; + +#[cfg(test)] +#[expect(clippy::cast_possible_truncation)] +mod tests; + +pub(crate) fn filter(array: &Array, mask: &MaskValuesRef) -> Array { + let array = array.as_view(); + let values = filter_records(V::values(array), V::byte_width(array), mask.as_ref()); + let validity = filter_validity( + array + .validity() + .vortex_expect("validity is derivable for a valid fixed-width array"), + mask, + ); + with_values(array, values, mask.true_count(), validity) + .vortex_expect("filtering fixed-width values preserves array invariants") +} + +fn filter_records(values: ByteBuffer, byte_width: usize, mask: &MaskValues) -> ByteBuffer { + let alignment = values.alignment(); + + match_each_record_width!( + byte_width, + |W| { + let records = Buffer::<[u8; W]>::from_byte_buffer(values); + // `filter_buffer` picks between in-place compaction, cached indices/slices, + // byte-compress, and bitmap iteration based on record width and mask density. + let filtered = filter_buffer(records, mask); + filtered.into_byte_buffer().aligned(alignment) + }, + _ => { + match values.try_into_mut() { + Ok(mut values) => { + let mut destination = 0; + mask.bit_buffer().for_each_set_index(|index| { + let source = index * byte_width; + values.copy_within(source..source + byte_width, destination); + destination += byte_width; + }); + values.truncate(destination); + values.freeze().into_byte_buffer().aligned(alignment) + } + Err(values) => { + let mut filtered = BufferMut::with_capacity(mask.true_count() * byte_width); + mask.bit_buffer().for_each_set_index(|index| { + let start = index * byte_width; + filtered.extend_from_slice(&values[start..start + byte_width]); + }); + filtered.freeze().into_byte_buffer().aligned(alignment) + } + } + } + ) +} diff --git a/vortex-array/src/arrays/fixed_width/filter/tests.rs b/vortex-array/src/arrays/fixed_width/filter/tests.rs new file mode 100644 index 00000000000..cf01e5d39e5 --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/filter/tests.rs @@ -0,0 +1,85 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_buffer::Buffer; +use vortex_buffer::buffer; +use vortex_mask::Mask; + +use super::filter_records; +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::DecimalArray; +use crate::arrays::PrimitiveArray; +use crate::compute::conformance::filter::LARGE_SIZE; +use crate::compute::conformance::filter::MEDIUM_SIZE; +use crate::compute::conformance::filter::test_filter_conformance; +use crate::dtype::DecimalDType; +use crate::dtype::i256; +use crate::validity::Validity; + +#[test] +fn filter_fallback_width_records() { + let Mask::Values(mask) = Mask::from_iter([true, false, true, false]) else { + panic!("a mixed mask must have mask values"); + }; + let expected = [0u8, 1, 2, 6, 7, 8]; + + // A uniquely owned buffer takes the in-place `copy_within` path. + let owned = Buffer::from_iter(0u8..12); + let filtered = filter_records(owned, 3, &mask); + assert_eq!(filtered.as_slice(), &expected); + + // Retaining a second reference forces the copying path instead. + let shared = Buffer::from_iter(0u8..12); + let _retained = shared.clone(); + let filtered = filter_records(shared, 3, &mask); + assert_eq!(filtered.as_slice(), &expected); +} + +#[rstest] +#[case::primitive_i8(PrimitiveArray::from_iter([-2i8, -1, 0, 1, 2]).into_array())] +#[case::primitive_u16(PrimitiveArray::from_iter([1u16, 2, 3, 4, 5]).into_array())] +#[case::primitive_i32(PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]).into_array())] +#[case::primitive_f32(PrimitiveArray::from_iter([0.1f32, 0.2, 0.3, 0.4, 0.5]).into_array())] +#[case::primitive_nullable(PrimitiveArray::from_option_iter( + [Some(1i64), None, Some(3), Some(4), None], +).into_array())] +#[case::primitive_large(PrimitiveArray::from_iter(0..LARGE_SIZE as u32).into_array())] +#[case::primitive_medium(PrimitiveArray::from_iter(0..MEDIUM_SIZE as i64).into_array())] +#[case::decimal_i8(DecimalArray::new( + buffer![1i8, 2, 3, 4, 5], + DecimalDType::new(2, 0), + Validity::NonNullable, +).into_array())] +#[case::decimal_i32(DecimalArray::new( + buffer![123i32, 456, -123, 0, 999], + DecimalDType::new(8, 2), + Validity::NonNullable, +).into_array())] +#[case::decimal_i64(DecimalArray::new( + buffer![12345i64, 67890, -12345, 0, 99999], + DecimalDType::new(18, 2), + Validity::NonNullable, +).into_array())] +#[case::decimal_i128(DecimalArray::new( + buffer![12345i128, 67890, -12345, 0, 99999], + DecimalDType::new(38, 4), + Validity::from_iter([true, false, true, true, false]), +).into_array())] +#[case::decimal_i256(DecimalArray::new( + buffer![ + i256::from_i128(12345), + i256::from_i128(67890), + i256::from_i128(-12345), + i256::ZERO, + i256::from_i128(99999), + ], + DecimalDType::new(76, 4), + Validity::NonNullable, +).into_array())] +fn fixed_width_filter_conformance(#[case] array: ArrayRef) { + test_filter_conformance(&array, &mut array_session().create_execution_ctx()); +} diff --git a/vortex-array/src/arrays/fixed_width/mod.rs b/vortex-array/src/arrays/fixed_width/mod.rs new file mode 100644 index 00000000000..ec19f05413c --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/mod.rs @@ -0,0 +1,47 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Shared structural operations for fixed-width canonical arrays. + +mod array; +pub(crate) mod filter; +pub(crate) mod take; +pub(crate) mod vtable; + +pub(crate) use self::array::FixedWidthArray; +pub(crate) use self::array::with_values; + +/// Dispatches a runtime byte width to a compile-time `const $W: usize` for every record width +/// with a dedicated fixed-width kernel, falling back to `$fallback` for any other width. +macro_rules! match_each_record_width { + ($byte_width:expr, | $W:ident | $body:block,_ => $fallback:block) => { + match $byte_width { + 1 => { + const $W: usize = 1; + $body + } + 2 => { + const $W: usize = 2; + $body + } + 4 => { + const $W: usize = 4; + $body + } + 8 => { + const $W: usize = 8; + $body + } + 16 => { + const $W: usize = 16; + $body + } + 32 => { + const $W: usize = 32; + $body + } + _ => $fallback, + } + }; +} +pub(crate) use match_each_record_width; diff --git a/vortex-array/src/arrays/fixed_width/take/avx2/gather.rs b/vortex-array/src/arrays/fixed_width/take/avx2/gather.rs new file mode 100644 index 00000000000..79e0a43358f --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/avx2/gather.rs @@ -0,0 +1,293 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::arch::x86_64::__m256i; +use std::arch::x86_64::_mm_loadu_si128; +use std::arch::x86_64::_mm_setzero_si128; +use std::arch::x86_64::_mm_shuffle_epi32; +use std::arch::x86_64::_mm_storeu_si128; +use std::arch::x86_64::_mm_unpacklo_epi64; +use std::arch::x86_64::_mm256_cmpgt_epi32; +use std::arch::x86_64::_mm256_cmpgt_epi64; +use std::arch::x86_64::_mm256_cvtepu8_epi32; +use std::arch::x86_64::_mm256_cvtepu8_epi64; +use std::arch::x86_64::_mm256_cvtepu16_epi32; +use std::arch::x86_64::_mm256_cvtepu16_epi64; +use std::arch::x86_64::_mm256_cvtepu32_epi64; +use std::arch::x86_64::_mm256_extracti128_si256; +use std::arch::x86_64::_mm256_loadu_si256; +use std::arch::x86_64::_mm256_mask_i32gather_epi32; +use std::arch::x86_64::_mm256_mask_i64gather_epi32; +use std::arch::x86_64::_mm256_mask_i64gather_epi64; +use std::arch::x86_64::_mm256_set1_epi32; +use std::arch::x86_64::_mm256_set1_epi64x; +use std::arch::x86_64::_mm256_setzero_si256; +use std::arch::x86_64::_mm256_storeu_si256; +use std::arch::x86_64::_mm256_xor_si256; +use std::convert::identity; + +/// Moves one SIMD block of fixed-width values into an output buffer. +pub(super) trait GatherFn { + /// The number of data elements written on each iteration. + const WIDTH: usize; + /// The number of indices read on each iteration. + const STRIDE: usize = Self::WIDTH; + + /// Gather values from `src` into `dst`. + /// + /// `max_idx` is the exclusive upper bound on valid indices. `None` means the bound does not + /// fit in the index type, so every representable index is in-bounds. + /// + /// # Safety + /// + /// `indices` must be readable for `STRIDE` elements. `src` must contain `max_idx` elements, + /// or more than `Index::MAX` elements when `max_idx` is `None`, and `dst` must be writable + /// for `WIDTH` elements. + /// + /// Returns a vector mask with all lanes set when every gathered index was in-bounds. + unsafe fn gather( + indices: *const Index, + max_idx: Option, + src: *const Lane, + dst: *mut Lane, + ) -> __m256i; +} + +/// AVX2 gather implementations for 32- and 64-bit value lanes. +pub(super) enum Avx2Gather {} + +macro_rules! cmpgt_epu32 { + ($lhs:expr, $rhs:expr) => {{ + // AVX2 only supplies a signed integer comparison. XORing each lane with its sign bit + // maps the unsigned ordering into the signed ordering without changing the relative + // order. + let sign_bit = _mm256_set1_epi32(i32::MIN); + _mm256_cmpgt_epi32( + _mm256_xor_si256($lhs, sign_bit), + _mm256_xor_si256($rhs, sign_bit), + ) + }}; +} + +macro_rules! cmpgt_epu64 { + ($lhs:expr, $rhs:expr) => {{ + // AVX2 only supplies a signed integer comparison. XORing each lane with its sign bit + // maps the unsigned ordering into the signed ordering without changing the relative + // order. + let sign_bit = _mm256_set1_epi64x(i64::MIN); + _mm256_cmpgt_epi64( + _mm256_xor_si256($lhs, sign_bit), + _mm256_xor_si256($rhs, sign_bit), + ) + }}; +} + +macro_rules! pack_i64_mask_for_i32_gather { + ($mask:expr) => {{ + let lo_bits = _mm256_extracti128_si256::<0>($mask); + let hi_bits = _mm256_extracti128_si256::<1>($mask); + let lo_packed = pack_i64_mask_half_for_i32_gather!(lo_bits); + let hi_packed = pack_i64_mask_half_for_i32_gather!(hi_bits); + _mm_unpacklo_epi64(lo_packed, hi_packed) + }}; +} + +macro_rules! pack_i64_mask_half_for_i32_gather { + ($mask:expr) => { + _mm_shuffle_epi32::<0b11_01_11_01>($mask) + }; +} + +macro_rules! impl_gather { + ($idx:ty, $({$value:ty => load: $load:ident, extend: $extend:ident, splat: $splat:ident, zero_vec: $zero_vec:ident, mask_indices: $mask_indices:ident, mask_cvt: |$mask_var:ident| $mask_cvt:block, gather: $masked_gather:ident, store: $store:ident, WIDTH = $WIDTH:literal, STRIDE = $STRIDE:literal }),+) => { + $( + impl_gather!(single; $idx, $value, load: $load, extend: $extend, splat: $splat, zero_vec: $zero_vec, mask_indices: $mask_indices, mask_cvt: |$mask_var| $mask_cvt, gather: $masked_gather, store: $store, WIDTH = $WIDTH, STRIDE = $STRIDE); + )* + }; + (single; $idx:ty, $value:ty, load: $load:ident, extend: $extend:ident, splat: $splat:ident, zero_vec: $zero_vec:ident, mask_indices: $mask_indices:ident, mask_cvt: |$mask_var:ident| $mask_cvt:block, gather: $masked_gather:ident, store: $store:ident, WIDTH = $WIDTH:literal, STRIDE = $STRIDE:literal) => { + impl GatherFn<$idx, $value> for Avx2Gather { + const WIDTH: usize = $WIDTH; + const STRIDE: usize = $STRIDE; + + #[allow(unused_unsafe, clippy::cast_possible_truncation)] + #[allow(clippy::inline_always)] + #[inline(always)] + unsafe fn gather( + indices: *const $idx, + max_idx: Option<$idx>, + src: *const $value, + dst: *mut $value, + ) -> __m256i { + const { + assert!($WIDTH <= $STRIDE, "dst cannot advance by more than the stride"); + } + + const SCALE: i32 = std::mem::size_of::<$value>() as i32; + + let indices_vec = unsafe { $load(indices.cast()) }; + let indices_vec = unsafe { $extend(indices_vec) }; + // An absent exclusive bound means it does not fit in the index type, so every + // representable index is in-bounds. Passing the valid mask to the gather masks + // every invalid lane before it can access `src`. + let valid_mask = if let Some(max_idx) = max_idx { + let max_idx_vec = unsafe { $splat(max_idx as _) }; + unsafe { $mask_indices!(max_idx_vec, indices_vec) } + } else { + unsafe { $splat(-1) } + }; + let gather_mask = { + let $mask_var = valid_mask; + $mask_cvt + }; + let values_vec = unsafe { + $masked_gather::( + $zero_vec(), + src.cast(), + indices_vec, + gather_mask, + ) + }; + unsafe { $store(dst.cast(), values_vec) }; + valid_mask + } + } + }; +} + +impl_gather!(u8, + { u32 => + load: _mm_loadu_si128, + extend: _mm256_cvtepu8_epi32, + splat: _mm256_set1_epi32, + zero_vec: _mm256_setzero_si256, + mask_indices: cmpgt_epu32, + mask_cvt: |x| { x }, + gather: _mm256_mask_i32gather_epi32, + store: _mm256_storeu_si256, + WIDTH = 8, STRIDE = 16 + }, + { u64 => + load: _mm_loadu_si128, + extend: _mm256_cvtepu8_epi64, + splat: _mm256_set1_epi64x, + zero_vec: _mm256_setzero_si256, + mask_indices: cmpgt_epu64, + mask_cvt: |x| { x }, + gather: _mm256_mask_i64gather_epi64, + store: _mm256_storeu_si256, + WIDTH = 4, STRIDE = 16 + } +); + +impl_gather!(u16, + { u32 => + load: _mm_loadu_si128, + extend: _mm256_cvtepu16_epi32, + splat: _mm256_set1_epi32, + zero_vec: _mm256_setzero_si256, + mask_indices: cmpgt_epu32, + mask_cvt: |x| { x }, + gather: _mm256_mask_i32gather_epi32, + store: _mm256_storeu_si256, + WIDTH = 8, STRIDE = 8 + }, + { u64 => + load: _mm_loadu_si128, + extend: _mm256_cvtepu16_epi64, + splat: _mm256_set1_epi64x, + zero_vec: _mm256_setzero_si256, + mask_indices: cmpgt_epu64, + mask_cvt: |x| { x }, + gather: _mm256_mask_i64gather_epi64, + store: _mm256_storeu_si256, + WIDTH = 4, STRIDE = 8 + } +); + +impl_gather!(u32, + { u32 => + load: _mm256_loadu_si256, + extend: identity, + splat: _mm256_set1_epi32, + zero_vec: _mm256_setzero_si256, + mask_indices: cmpgt_epu32, + mask_cvt: |x| { x }, + gather: _mm256_mask_i32gather_epi32, + store: _mm256_storeu_si256, + WIDTH = 8, STRIDE = 8 + }, + { u64 => + load: _mm_loadu_si128, + extend: _mm256_cvtepu32_epi64, + splat: _mm256_set1_epi64x, + zero_vec: _mm256_setzero_si256, + mask_indices: cmpgt_epu64, + mask_cvt: |x| { x }, + gather: _mm256_mask_i64gather_epi64, + store: _mm256_storeu_si256, + WIDTH = 4, STRIDE = 4 + } +); + +impl_gather!(u64, + { u32 => + load: _mm256_loadu_si256, + extend: identity, + splat: _mm256_set1_epi64x, + zero_vec: _mm_setzero_si128, + mask_indices: cmpgt_epu64, + mask_cvt: |mask| { unsafe { pack_i64_mask_for_i32_gather!(mask) } }, + gather: _mm256_mask_i64gather_epi32, + store: _mm_storeu_si128, + WIDTH = 4, STRIDE = 4 + }, + { u64 => + load: _mm256_loadu_si256, + extend: identity, + splat: _mm256_set1_epi64x, + zero_vec: _mm256_setzero_si256, + mask_indices: cmpgt_epu64, + mask_cvt: |x| { x }, + gather: _mm256_mask_i64gather_epi64, + store: _mm256_storeu_si256, + WIDTH = 4, STRIDE = 4 + } +); + +#[cfg(test)] +mod tests { + use std::arch::x86_64::_mm_movemask_epi8; + use std::arch::x86_64::_mm_set_epi64x; + + use super::*; + + #[test] + fn pack_i64_mask_for_i32_gather_preserves_lanes() { + if !is_x86_feature_detected!("sse2") { + return; + } + + for lane_bits in 0u8..16 { + let lane_mask = |lane: u32| { + if lane_bits & (1u8 << lane) == 0 { + 0 + } else { + -1 + } + }; + let actual = unsafe { + let lo_bits = _mm_set_epi64x(lane_mask(1), lane_mask(0)); + let hi_bits = _mm_set_epi64x(lane_mask(3), lane_mask(2)); + let lo_packed = pack_i64_mask_half_for_i32_gather!(lo_bits); + let hi_packed = pack_i64_mask_half_for_i32_gather!(hi_bits); + let packed = _mm_unpacklo_epi64(lo_packed, hi_packed); + _mm_movemask_epi8(packed) + }; + let expected = (0u32..4).fold(0, |bits, lane| { + bits | (((lane_bits >> lane) & 1) as i32 * (0b1111 << (lane * 4))) + }); + + assert_eq!(actual, expected, "lane mask {lane_bits:04b}"); + } + } +} diff --git a/vortex-array/src/arrays/fixed_width/take/avx2/mod.rs b/vortex-array/src/arrays/fixed_width/take/avx2/mod.rs new file mode 100644 index 00000000000..cba089f1011 --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/avx2/mod.rs @@ -0,0 +1,168 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! An AVX2 implementation of take operation using gather instructions. +//! +//! Only enabled for x86_64 hosts and it is gated at runtime behind feature detection to +//! ensure AVX2 instructions are available. + +#![cfg(any(target_arch = "x86_64", target_arch = "x86"))] + +mod gather; +#[cfg(test)] +mod tests; + +use std::arch::x86_64::__m256i; +use std::arch::x86_64::_mm256_and_si256; +use std::arch::x86_64::_mm256_movemask_epi8; +use std::arch::x86_64::_mm256_set1_epi32; + +use vortex_buffer::Alignment; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; + +use self::gather::Avx2Gather; +use self::gather::GatherFn; +use super::FixedWidthTakeValue; +use super::take_values_scalar; +use crate::dtype::PType; +use crate::dtype::UnsignedPType; +use crate::match_each_unsigned_integer_ptype; + +/// Takes the specified indices into a new [`Buffer`] using AVX2 SIMD. +/// +/// An AVX2 gather only moves raw bytes, so signedness and float-ness are irrelevant — only the +/// byte width of `V` matters. Any 4-byte value rides the gather through the `u32` lane and any +/// 8-byte value through the `u64` lane, regardless of its actual type. Values 1 or 2 bytes wide +/// (AVX2 has no sub-32-bit gather) and wider than 8 bytes fall back to the scalar kernel. +/// +/// [`FixedWidthTakeValue`] guarantees that the complete representation is initialized before it is +/// read through an integer lane. +/// +/// # Safety +/// +/// The caller must ensure the `avx2` feature is enabled. +#[target_feature(enable = "avx2")] +pub(super) unsafe fn take_avx2( + buffer: &[V], + indices: &[I], +) -> Buffer { + if buffer.is_empty() { + assert!( + indices.is_empty(), + "cannot take a non-empty set of indices from an empty buffer" + ); + return Buffer::empty(); + } + + // Dispatch on the gather lane width. The index type must still be concretized to select the + // right `GatherFn` impl, so re-dispatch it with `match_each_unsigned_integer_ptype!`. + macro_rules! dispatch { + ($lane:ty) => {{ + match_each_unsigned_integer_ptype!(I::PTYPE, |Idx| { + // SAFETY: `Idx` has the same `PTYPE` as `I`, so this is a no-op reinterpret of the + // index slice into the concrete type the gather impl is keyed on. + let indices = unsafe { std::mem::transmute::<&[I], &[Idx]>(indices) }; + exec_take::(buffer, indices) + }) + }}; + } + + match size_of::() { + // The i32 gather interprets u32 lanes as signed offsets. A valid high u32 index needs the + // scalar path when the values slice exceeds the non-negative i32 addressable range. + 4 if I::PTYPE == PType::U32 && !i32_gather_can_address(buffer.len()) => { + take_values_scalar(buffer, indices) + } + 4 => dispatch!(u32), + 8 => dispatch!(u64), + // 1/2-byte and >8-byte values have no AVX2 gather lane, so fall back to scalar. + _ => take_values_scalar(buffer, indices), + } +} + +const fn i32_gather_can_address(values_len: usize) -> bool { + values_len <= i32::MAX as usize + 1 +} + +/// AVX2 core inner loop for a given index type `Idx`, output element type `Out`, and gather +/// `Lane` type. +/// +/// `Out` is the element type written to the output buffer; `Lane` (`u32` or `u64`) is the +/// integer type the gather intrinsics operate on. The caller must pair them so that +/// `size_of::() == size_of::()` (the only caller, [`take_avx2`], picks `Lane` from +/// `size_of::()`). Each valid lane copies the initialized representation of an existing +/// `Out`; invalid lanes are masked and cause a panic before the output buffer is initialized. +/// Gather instructions tolerate the source's potentially weaker alignment. +#[allow(clippy::inline_always)] +#[inline(always)] +fn exec_take(values: &[Out], indices: &[Idx]) -> Buffer +where + Out: FixedWidthTakeValue, + Idx: UnsignedPType, + Gather: GatherFn, +{ + assert_eq!( + size_of::(), + size_of::(), + "gather lane and output element must have the same size" + ); + + let indices_len = indices.len(); + // The length is an exclusive upper bound on valid indices. `None` means the bound does not + // fit in the index type, so every representable index is in-bounds. + let max_index = Idx::from(values.len()); + let mut buffer = + BufferMut::::with_capacity_aligned(indices_len, Alignment::of::<__m256i>()); + let buf_uninit = buffer.spare_capacity_mut(); + + let mut offset = 0; + // SAFETY: `exec_take` is only called by `take_avx2`, whose caller guarantees AVX2 support. + let mut all_indices_valid = unsafe { _mm256_set1_epi32(-1) }; + // Loop terminates STRIDE elements before end of the indices array because the `GatherFn` + // might read up to STRIDE src elements at a time, even though it only advances WIDTH elements + // in the dst. + while offset + Gather::STRIDE < indices_len { + // SAFETY: `gather` preconditions satisfied: + // 1. `(indices + offset)..(indices + offset + STRIDE)` is in-bounds for indices + // allocation. + // 2. `buffer` has same len as indices so `buffer + offset + WIDTH` is always valid. + // 3. `size_of::() == size_of::()` (asserted above), so the `Lane`-typed + // pointers address the same bytes as the `Out`-typed `values`/`buffer` allocations. + let valid_mask = unsafe { + Gather::gather( + indices.as_ptr().add(offset), + max_index, + values.as_ptr().cast::(), + buf_uninit.as_mut_ptr().add(offset).cast::(), + ) + }; + // SAFETY: `exec_take` is only called by `take_avx2`, whose caller guarantees AVX2 support. + all_indices_valid = unsafe { _mm256_and_si256(all_indices_valid, valid_mask) }; + offset += Gather::WIDTH; + } + + // Invalid lanes were masked before gathering, so it is safe to defer the bounds failure until + // after the SIMD loop and avoid a conditional branch on every iteration. + assert!( + // SAFETY: `exec_take` is only called by `take_avx2`, whose caller guarantees AVX2 support. + unsafe { _mm256_movemask_epi8(all_indices_valid) } == -1, + "take index out of bounds" + ); + + // Remainder. + while offset < indices_len { + buf_uninit[offset].write(values[indices[offset].as_()]); + offset += 1; + } + + assert_eq!(offset, indices_len); + + // SAFETY: All elements have been initialized. + unsafe { buffer.set_len(indices_len) }; + + // Do not expose the temporary SIMD over-alignment as part of the returned buffer. + buffer = buffer.aligned(Alignment::of::()); + + buffer.freeze() +} diff --git a/vortex-array/src/arrays/fixed_width/take/avx2/tests.rs b/vortex-array/src/arrays/fixed_width/take/avx2/tests.rs new file mode 100644 index 00000000000..2ec97bcc38c --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/avx2/tests.rs @@ -0,0 +1,212 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![cfg(all(target_arch = "x86_64", not(miri)))] + +use std::panic::RefUnwindSafe; +use std::panic::catch_unwind; + +use vortex_buffer::Buffer; + +use super::super::FixedWidthTakeValue; +use super::take_avx2; +use crate::dtype::UnsignedPType; + +fn take_avx2_if_supported( + values: &[V], + indices: &[I], +) -> Option> { + if !is_x86_feature_detected!("avx2") { + return None; + } + + // SAFETY: AVX2 support was detected above, and `FixedWidthTakeValue` guarantees that every + // byte in the values is initialized. + Some(unsafe { take_avx2(values, indices) }) +} + +fn assert_avx2_take_panics(values: &[V], indices: &[I], expected: &str) +where + V: FixedWidthTakeValue + RefUnwindSafe, + I: UnsignedPType + RefUnwindSafe, +{ + if !is_x86_feature_detected!("avx2") { + return; + } + + // SAFETY: AVX2 support was detected above, and `FixedWidthTakeValue` guarantees that every + // byte in the values is initialized. + let result = catch_unwind(|| unsafe { take_avx2(values, indices) }); + let Err(payload) = result else { + panic!("take should panic for an invalid index"); + }; + let message = payload + .downcast_ref::<&str>() + .copied() + .or_else(|| payload.downcast_ref::().map(String::as_str)); + assert_eq!(message, Some(expected)); +} + +macro_rules! test_cases { + (index_type => $IDX:ty, value_types => $($VAL:ty),+) => { + paste::paste! { + $( + #[test] + #[allow(clippy::cast_possible_truncation)] + fn []() { + let values: Vec<$VAL> = (1..=127).map(|x| x as $VAL).collect(); + let indices: Vec<$IDX> = (0..127).collect(); + + let Some(result) = take_avx2_if_supported(&values, &indices) else { + return; + }; + assert_eq!(&values, result.as_slice()); + } + + #[test] + #[allow(clippy::cast_possible_truncation)] + fn []() { + let values: Vec<$VAL> = vec![]; + let indices: Vec<$IDX> = (0..127).collect(); + + assert_avx2_take_panics( + &values, + &indices, + "cannot take a non-empty set of indices from an empty buffer", + ); + } + + #[test] + #[allow(clippy::cast_possible_truncation)] + fn []() { + let values: Vec<$VAL> = (1..=127).map(|x| x as $VAL).collect(); + let indices: Vec<$IDX> = (127..=254).collect(); + + assert_avx2_take_panics(&values, &indices, "take index out of bounds"); + } + )+ + } + }; +} + +test_cases!( + index_type => u8, + value_types => u32, i32, u64, i64, f32, f64 +); +test_cases!( + index_type => u16, + value_types => u32, i32, u64, i64, f32, f64 +); +test_cases!( + index_type => u32, + value_types => u32, i32, u64, i64, f32, f64 +); +test_cases!( + index_type => u64, + value_types => u32, i32, u64, i64, f32, f64 +); + +#[test] +fn last_valid_u8_index() { + let values: Vec = (0..=255).collect(); + let indices: Vec = vec![255; 20]; + + let Some(result) = take_avx2_if_supported(&values, &indices) else { + return; + }; + assert_eq!(&[255; 20], result.as_slice()); +} + +#[test] +fn last_valid_u16_index() { + let values: Vec = (0..=65535).collect(); + let indices: Vec = vec![65535; 20]; + + let Some(result) = take_avx2_if_supported(&values, &indices) else { + return; + }; + assert_eq!(&[65535; 20], result.as_slice()); +} + +#[test] +fn empty_values_and_indices() { + let Some(result) = take_avx2_if_supported::(&[], &[]) else { + return; + }; + + assert!(result.is_empty()); +} + +#[test] +fn i32_gather_addressable_length_boundary() { + assert!(super::i32_gather_can_address(i32::MAX as usize + 1)); + assert!(!super::i32_gather_can_address(i32::MAX as usize + 2)); +} + +#[test] +fn invalid_index_only_in_simd_block() { + let values = vec![10u32, 20, 30]; + let indices = vec![3u32, 0, 1, 2, 0, 1, 2, 0, 1]; + + assert_avx2_take_panics(&values, &indices, "take index out of bounds"); +} + +#[test] +fn simd_array_u8x4() { + let values: Vec<[u8; 4]> = (1u32..=200).map(u32::to_le_bytes).collect(); + let indices: Vec = (0..200).collect(); + + let Some(result) = take_avx2_if_supported(&values, &indices) else { + return; + }; + assert_eq!(values.as_slice(), result.as_slice()); +} + +#[test] +fn scalar_fallback_u16() { + let values: Vec = (1..=300).collect(); + let indices: Vec = (0..300).collect(); + + let Some(result) = take_avx2_if_supported(&values, &indices) else { + return; + }; + assert_eq!(values.as_slice(), result.as_slice()); +} + +#[test] +fn scalar_fallback_array_u8x16() { + let values: Vec<[u8; 16]> = (0u128..200).map(u128::to_le_bytes).collect(); + let indices: Vec = (0..200).collect(); + + let Some(result) = take_avx2_if_supported(&values, &indices) else { + return; + }; + assert_eq!(values.as_slice(), result.as_slice()); +} + +#[test] +fn u32_max_index_in_u32_lane() { + let values = vec![0u32; 8]; + // The first eight indices execute in the SIMD loop; the scalar remainder is valid. + let indices = vec![0, u32::MAX, 2, 3, 4, 5, 6, 7, 0]; + + assert_avx2_take_panics(&values, &indices, "take index out of bounds"); +} + +#[test] +fn u64_max_index_in_u32_lane() { + let values = vec![0u32; 8]; + // The first four indices execute in the SIMD loop; the scalar remainder is valid. + let indices = vec![0, u64::MAX, 2, 3, 0]; + + assert_avx2_take_panics(&values, &indices, "take index out of bounds"); +} + +#[test] +fn u64_max_index_in_u64_lane() { + let values = vec![0u64; 8]; + // The first four indices execute in the SIMD loop; the scalar remainder is valid. + let indices = vec![0, u64::MAX, 2, 3, 0]; + + assert_avx2_take_panics(&values, &indices, "take index out of bounds"); +} diff --git a/vortex-array/src/arrays/fixed_width/take/mod.rs b/vortex-array/src/arrays/fixed_width/take/mod.rs new file mode 100644 index 00000000000..d6e2f744bdd --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/mod.rs @@ -0,0 +1,199 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#[cfg(any(target_arch = "x86_64", target_arch = "x86"))] +mod avx2; +mod records; +mod scalar; +mod slices; +#[cfg(test)] +mod tests; + +#[cfg(any(target_arch = "x86_64", target_arch = "x86"))] +use std::sync::LazyLock; + +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_mask::Mask; + +use self::records::take_byte_records; +use self::scalar::take_values_scalar; +use self::slices::take_slices; +use self::slices::take_slices_constant_length; +use super::FixedWidthArray; +use super::with_values; +use crate::ArrayRef; +use crate::Columnar; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::ConstantArray; +use crate::arrays::PiecewiseSequence; +use crate::arrays::PrimitiveArray; +use crate::arrays::dict::TakeExecute; +use crate::arrays::piecewise_sequence::constant_unsigned_usize; +use crate::arrays::piecewise_sequence::maybe_contiguous_slices; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::UnsignedPType; +use crate::dtype::half::f16; +use crate::match_each_unsigned_integer_ptype; +use crate::scalar::Scalar; + +#[cfg(any(target_arch = "x86_64", target_arch = "x86"))] +static HAS_AVX2: LazyLock = LazyLock::new(|| is_x86_feature_detected!("avx2")); + +impl TakeExecute for V { + fn take( + array: ArrayView<'_, Self>, + indices: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + take(array, indices, ctx) + } +} + +/// A fixed-width value whose initialized bytes may be moved through integer SIMD lanes. +/// +/// # Safety +/// +/// Implementors must have no uninitialized bytes. The shared AVX2 gather reads the complete +/// representation through a same-width integer lane before writing those bytes back unchanged. +pub(crate) unsafe trait FixedWidthTakeValue: Copy {} + +macro_rules! impl_fixed_width_take_value { + ($($ty:ty),+ $(,)?) => { + $( + // SAFETY: These scalar representations contain no padding or uninitialized bytes. + unsafe impl FixedWidthTakeValue for $ty {} + )+ + }; +} + +impl_fixed_width_take_value!(u8, u16, u32, u64, i8, i16, i32, i64, f16, f32, f64,); + +// SAFETY: Byte arrays have no padding and every byte is initialized. +unsafe impl FixedWidthTakeValue for [u8; N] {} + +pub(crate) fn take_values( + values: &[T], + indices: &[I], +) -> Buffer { + #[cfg(any(target_arch = "x86_64", target_arch = "x86"))] + if *HAS_AVX2 { + // SAFETY: AVX2 was detected above and `FixedWidthTakeValue` guarantees an initialized byte + // representation. The AVX2 dispatcher retains Primitive's existing scalar fallbacks and + // out-of-bounds behavior for every value width. + return unsafe { avx2::take_avx2(values, indices) }; + } + + take_values_scalar(values, indices) +} + +pub(crate) fn take( + array: ArrayView<'_, V>, + indices: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)? + { + return Ok(Some(taken)); + } + + let DType::Primitive(ptype, nullability) = indices.dtype() else { + vortex_bail!("Invalid indices dtype: {}", indices.dtype()) + }; + if !ptype.is_int() { + vortex_bail!("Invalid indices dtype: {}", indices.dtype()) + } + + let indices_validity = indices.validity()?; + let indices_nulls_zeroed = match indices_validity.execute_mask(indices.len(), ctx)? { + Mask::AllTrue(_) => indices.clone(), + Mask::AllFalse(_) => { + return Ok(Some( + ConstantArray::new(Scalar::null(array.dtype().as_nullable()), indices.len()) + .into_array(), + )); + } + Mask::Values(_) => indices + .clone() + .fill_null(Scalar::from(0).cast(indices.dtype())?)?, + }; + + let indices = if ptype.is_unsigned_int() { + indices_nulls_zeroed.execute::(ctx)? + } else { + indices_nulls_zeroed + .cast(DType::Primitive(ptype.to_unsigned(), *nullability))? + .execute::(ctx)? + }; + let validity = array + .validity()? + .take(&indices.clone().into_array())? + .and(indices_validity)?; + + let source = V::values(array); + let values = match_each_unsigned_integer_ptype!(indices.ptype(), |I| { + take_byte_records( + &source, + V::byte_width(array), + array.len(), + indices.as_slice::(), + ) + })?; + Ok(Some( + with_values(array, values, indices.len(), validity)?.into_array(), + )) +} + +fn take_contiguous_ranges( + array: ArrayView<'_, V>, + indices: ArrayView<'_, PiecewiseSequence>, + indices_ref: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { + return Ok(None); + }; + + let values = V::values(array); + let byte_width = V::byte_width(array); + let output_len = indices_ref.len(); + let taken = match lengths { + Columnar::Constant(lengths) => { + let length = constant_unsigned_usize(&lengths); + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + take_slices_constant_length( + &values, + byte_width, + array.len(), + starts.as_slice::(), + length, + output_len, + ) + }) + } + Columnar::Canonical(lengths) => { + let lengths = lengths.into_primitive(); + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + take_slices( + &values, + byte_width, + array.len(), + starts.as_slice::(), + lengths.as_slice::(), + output_len, + ) + }) + }) + } + }?; + let validity = array.validity()?.take(indices_ref)?; + Ok(Some( + with_values(array, taken, output_len, validity)?.into_array(), + )) +} diff --git a/vortex-array/src/arrays/fixed_width/take/records.rs b/vortex-array/src/arrays/fixed_width/take/records.rs new file mode 100644 index 00000000000..6cb47957991 --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/records.rs @@ -0,0 +1,49 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_buffer::ByteBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use super::take_values; +use crate::arrays::fixed_width::match_each_record_width; +use crate::dtype::UnsignedPType; + +pub(super) fn take_byte_records( + values: &ByteBuffer, + byte_width: usize, + record_count: usize, + indices: &[I], +) -> VortexResult { + let alignment = values.alignment(); + + match_each_record_width!( + byte_width, + |W| { + let records = Buffer::<[u8; W]>::from_byte_buffer(values.clone()); + debug_assert_eq!(records.len(), record_count); + Ok(take_values(records.as_slice(), indices) + .into_byte_buffer() + .aligned(alignment)) + }, + _ => { + let output_len = indices + .len() + .checked_mul(byte_width) + .ok_or_else(|| vortex_err!("Fixed-width take output length overflows usize"))?; + let mut result = BufferMut::::with_capacity(output_len); + for index in indices { + let index = index.as_(); + assert!( + index < record_count, + "take index {index} out of bounds for length {record_count}" + ); + let start = index * byte_width; + result.extend_from_slice(&values[start..start + byte_width]); + } + Ok(result.freeze().into_byte_buffer().aligned(alignment)) + } + ) +} diff --git a/vortex-array/src/arrays/fixed_width/take/scalar.rs b/vortex-array/src/arrays/fixed_width/take/scalar.rs new file mode 100644 index 00000000000..5a3510efc2d --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/scalar.rs @@ -0,0 +1,28 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; + +use crate::dtype::IntegerPType; + +#[allow(clippy::inline_always)] +#[inline(always)] +pub(crate) fn take_values_scalar( + values: &[T], + indices: &[I], +) -> Buffer { + // The explicit pointer loop keeps the source length in a register and avoids a capacity check + // for every output value. + let mut result = BufferMut::with_capacity(indices.len()); + let result_ptr = result.spare_capacity_mut().as_mut_ptr().cast::(); + + for (output_index, index) in indices.iter().enumerate() { + // SAFETY: `indices.len()` elements were reserved and each output position is written once. + unsafe { result_ptr.add(output_index).write(values[index.as_()]) }; + } + + // SAFETY: the loop initialized every element in the reserved output range. + unsafe { result.set_len(indices.len()) }; + result.freeze() +} diff --git a/vortex-array/src/arrays/fixed_width/take/slices.rs b/vortex-array/src/arrays/fixed_width/take/slices.rs new file mode 100644 index 00000000000..a4d9ceda84a --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/slices.rs @@ -0,0 +1,110 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ptr; + +use itertools::Itertools as _; +use vortex_buffer::BufferMut; +use vortex_buffer::ByteBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; + +use crate::dtype::UnsignedPType; + +pub(super) fn take_slices( + values: &ByteBuffer, + byte_width: usize, + record_count: usize, + starts: &[S], + lengths: &[L], + output_len: usize, +) -> VortexResult { + let slices = starts + .iter() + .zip_eq(lengths) + .map(|(&start, &length)| (start.as_(), length.as_())); + copy_slices(values, byte_width, record_count, slices, output_len) +} + +pub(super) fn take_slices_constant_length( + values: &ByteBuffer, + byte_width: usize, + record_count: usize, + starts: &[S], + length: usize, + output_len: usize, +) -> VortexResult { + let computed_len = starts + .len() + .checked_mul(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); + copy_slices( + values, + byte_width, + record_count, + starts.iter().map(|start| (start.as_(), length)), + output_len, + ) +} + +fn copy_slices( + values: &ByteBuffer, + byte_width: usize, + record_count: usize, + slices: impl IntoIterator, + output_len: usize, +) -> VortexResult { + let input_byte_len = record_count + .checked_mul(byte_width) + .ok_or_else(|| vortex_err!("Fixed-width values buffer length overflows usize"))?; + vortex_ensure!( + values.len() == input_byte_len, + "Fixed-width values buffer length does not match record count" + ); + + let output_byte_len = output_len + .checked_mul(byte_width) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + let mut result = BufferMut::::with_capacity_aligned(output_byte_len, values.alignment()); + let spare = &mut result.spare_capacity_mut()[..output_byte_len]; + let mut cursor = 0usize; + + for (start, length) in slices { + let end = start + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray slice end overflows usize"))?; + vortex_ensure!( + end <= record_count, + "PiecewiseSequenceArray slice {start}..{end} exceeds array length {record_count}" + ); + // These multiplications cannot overflow because `end <= record_count` and the complete + // values buffer length was checked above. + let byte_start = start * byte_width; + let byte_length = length * byte_width; + let source = &values[byte_start..][..byte_length]; + // SAFETY: `source` and the checked spare-capacity range have equal lengths and do not + // overlap. + unsafe { + ptr::copy_nonoverlapping( + source.as_ptr(), + spare[cursor..][..source.len()].as_mut_ptr().cast::(), + source.len(), + ); + } + cursor += source.len(); + } + + // SAFETY: The loop initialized the prefix `0..cursor` of the spare capacity. + unsafe { result.set_len(cursor) }; + vortex_ensure!( + result.len() == output_byte_len, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_byte_len}", + result.len() + ); + Ok(result.freeze().into_byte_buffer()) +} diff --git a/vortex-array/src/arrays/fixed_width/take/tests.rs b/vortex-array/src/arrays/fixed_width/take/tests.rs new file mode 100644 index 00000000000..ff167d27499 --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/tests.rs @@ -0,0 +1,205 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_buffer::Buffer; +use vortex_buffer::buffer; +use vortex_error::VortexResult; + +use super::records::take_byte_records; +use super::slices::take_slices; +use super::slices::take_slices_constant_length; +use super::take_values; +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::arrays::DecimalArray; +use crate::arrays::PiecewiseSequenceArray; +use crate::arrays::PrimitiveArray; +use crate::assert_arrays_eq; +use crate::compute::conformance::take::test_take_conformance; +use crate::dtype::DecimalDType; +use crate::dtype::i256; +use crate::validity::Validity; + +#[test] +fn take_four_byte_records() { + let values = [[1u8, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]; + let taken = take_values(&values, &[2u32, 0]); + assert_eq!(taken.as_slice(), &[[9, 10, 11, 12], [1, 2, 3, 4]]); +} + +#[test] +fn take_eight_byte_values() { + let taken = take_values(&[10i64, 20, 30], &[1u16, 2, 0]); + assert_eq!(taken.as_slice(), &[20, 30, 10]); +} + +#[rstest] +#[case(1)] +#[case(2)] +#[case(4)] +#[case(8)] +#[case(16)] +#[case(32)] +#[case::fallback(3)] +#[case::fallback_wide(12)] +fn take_runtime_width_records(#[case] byte_width: usize) -> VortexResult<()> { + let values = Buffer::from_iter((0u8..).take(3 * byte_width)); + let expected = values[2 * byte_width..3 * byte_width] + .iter() + .chain(&values[..byte_width]) + .copied() + .collect::>(); + let taken = take_byte_records(&values.into_byte_buffer(), byte_width, 3, &[2u32, 0])?; + assert_eq!(taken.as_slice(), expected); + Ok(()) +} + +#[test] +#[should_panic(expected = "take index 3 out of bounds for length 3")] +fn fallback_take_rejects_out_of_bounds_index() { + let values = Buffer::from_iter((0u8..).take(9)).into_byte_buffer(); + drop(take_byte_records(&values, 3, 3, &[3u32])); +} + +#[test] +fn take_variable_length_slices() -> VortexResult<()> { + let values = buffer![10u8, 11, 12, 13, 14].into_byte_buffer(); + let taken = take_slices(&values, 1, 5, &[1u32, 3], &[2u32, 1], 3)?; + assert_eq!(taken.as_slice(), &[11, 12, 13]); + Ok(()) +} + +#[test] +fn variable_length_slices_validate_output_length() { + let values = buffer![10u8, 11, 12, 13].into_byte_buffer(); + assert!(take_slices(&values, 1, 4, &[0u32, 2], &[1u32, 1], 3).is_err()); +} + +#[test] +fn take_constant_length_slices() -> VortexResult<()> { + let values = buffer![10u8, 11, 12, 13, 14].into_byte_buffer(); + let taken = take_slices_constant_length(&values, 1, 5, &[0u32, 3], 2, 4)?; + assert_eq!(taken.as_slice(), &[10, 11, 13, 14]); + Ok(()) +} + +#[test] +fn null_index_skips_out_of_bounds_primitive_value() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let values = PrimitiveArray::from_iter([10i32, 20, 30]); + let indices = PrimitiveArray::new( + buffer![1u64, 3], + Validity::Array(BoolArray::from_iter([true, false]).into_array()), + ); + + let taken = values.take(indices.into_array())?; + + assert_arrays_eq!( + taken, + PrimitiveArray::from_option_iter([Some(20i32), None]).into_array(), + &mut ctx + ); + Ok(()) +} + +#[test] +fn null_index_skips_out_of_bounds_decimal_value() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal_dtype = DecimalDType::new(19, 1); + let values = DecimalArray::new( + buffer![10i128, 20, 30], + decimal_dtype, + Validity::NonNullable, + ); + let indices = PrimitiveArray::new( + buffer![1u64, 3], + Validity::Array(BoolArray::from_iter([true, false]).into_array()), + ); + + let taken = values.take(indices.into_array())?; + + assert_arrays_eq!( + taken, + DecimalArray::from_option_iter([Some(20i128), None], decimal_dtype).into_array(), + &mut ctx + ); + Ok(()) +} + +#[test] +fn decimal_i256_take_consumes_piecewise_indices() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal_dtype = DecimalDType::new(76, 2); + let values = DecimalArray::new( + buffer![ + i256::from_i128(100), + i256::from_i128(200), + i256::from_i128(300), + i256::from_i128(400), + i256::from_i128(500), + ], + decimal_dtype, + Validity::NonNullable, + ); + let starts = PrimitiveArray::from_iter([1u64, 3]).into_array(); + let lengths = PrimitiveArray::from_iter([2u64, 1]).into_array(); + let multipliers = ConstantArray::new(1u64, 2).into_array(); + let indices = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 3)?.into_array(); + + let taken = values.take(indices)?; + + let expected = DecimalArray::new( + buffer![ + i256::from_i128(200), + i256::from_i128(300), + i256::from_i128(400), + ], + decimal_dtype, + Validity::NonNullable, + ); + assert_arrays_eq!(taken, expected, &mut ctx); + Ok(()) +} + +#[rstest] +#[case::primitive(PrimitiveArray::new( + buffer![0i32, 1, 2, 3, 4], + Validity::NonNullable, +).into_array())] +#[case::primitive_nullable(PrimitiveArray::from_option_iter( + [Some(1i64), None, Some(3), Some(4), None], +).into_array())] +#[case::decimal_i32(DecimalArray::new( + buffer![1i32, 2, 3, 4, 5], + DecimalDType::new(5, 0), + Validity::NonNullable, +).into_array())] +#[case::decimal_i64(DecimalArray::new( + buffer![10i64, 20, 30, 40, 50], + DecimalDType::new(10, 1), + Validity::NonNullable, +).into_array())] +#[case::decimal_i128(DecimalArray::new( + buffer![100i128, 200, 300, 400, 500], + DecimalDType::new(19, 2), + Validity::from_iter([true, false, true, true, false]), +).into_array())] +#[case::decimal_i256(DecimalArray::new( + buffer![ + i256::from_i128(100), + i256::from_i128(200), + i256::from_i128(300), + i256::from_i128(400), + i256::from_i128(500), + ], + DecimalDType::new(76, 2), + Validity::NonNullable, +).into_array())] +fn fixed_width_take_conformance(#[case] array: ArrayRef) { + test_take_conformance(&array, &mut array_session().create_execution_ctx()); +} diff --git a/vortex-array/src/arrays/fixed_width/vtable.rs b/vortex-array/src/arrays/fixed_width/vtable.rs new file mode 100644 index 00000000000..79d6ed20072 --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/vtable.rs @@ -0,0 +1,47 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_panic; + +use crate::buffer::BufferHandle; +use crate::dtype::Nullability; +use crate::serde::ArrayChildren; +use crate::validity::Validity; + +pub(crate) fn buffer(array_name: &str, values: &BufferHandle, idx: usize) -> BufferHandle { + match idx { + 0 => values.clone(), + _ => vortex_panic!("{array_name} buffer index {idx} out of bounds"), + } +} + +pub(crate) fn buffer_name(idx: usize) -> Option { + match idx { + 0 => Some("values".to_string()), + _ => None, + } +} + +pub(crate) fn single_buffer(buffers: &[BufferHandle]) -> VortexResult { + vortex_ensure!( + buffers.len() == 1, + "Expected 1 buffer, got {}", + buffers.len() + ); + Ok(buffers[0].clone()) +} + +pub(crate) fn deserialize_validity( + nullability: Nullability, + len: usize, + children: &dyn ArrayChildren, +) -> VortexResult { + match children.len() { + 0 => Ok(Validity::from(nullability)), + 1 => Ok(Validity::Array(children.get(0, &Validity::DTYPE, len)?)), + child_count => vortex_bail!("Expected 0 or 1 child, got {child_count}"), + } +} diff --git a/vortex-array/src/arrays/interleave/execute/mod.rs b/vortex-array/src/arrays/interleave/execute/mod.rs index 05dcd161f62..558db9e3309 100644 --- a/vortex-array/src/arrays/interleave/execute/mod.rs +++ b/vortex-array/src/arrays/interleave/execute/mod.rs @@ -4,19 +4,18 @@ //! Execution logic for [`Interleave`], dispatched on the value type. //! //! All values share a type (validated in [`Interleave::check`]), so the -//! physical gather kernel is chosen from the first value. The selector types are an orthogonal -//! concern handled within each kernel. Only boolean values are implemented today (see the [`bool`] module). +//! physical gather kernel is chosen from the array's dtype. The selector types are an orthogonal +//! concern handled within each kernel. //! //! [`Interleave::check`]: super::Interleave::check -//! [`bool`]: module@crate::arrays::interleave::execute::bool mod bool; +mod primitive; use vortex_error::VortexResult; use vortex_error::vortex_panic; use super::Interleave; -use super::InterleaveArrayExt; use crate::array::Array; use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; @@ -26,14 +25,14 @@ pub(super) fn execute( array: Array, ctx: &mut ExecutionCtx, ) -> VortexResult { - if array.value(0).dtype().is_boolean() { + if array.dtype().is_boolean() { bool::execute(array, ctx) + } else if array.dtype().is_primitive() { + primitive::execute(array, ctx) } else { - let value_dtype = array.value(0).dtype().clone(); vortex_panic!( - "interleave execution is only implemented for boolean values; value dtype {} is not \ - yet supported", - value_dtype + "interleave execution is not implemented for value dtype {}", + array.dtype() ) } } diff --git a/vortex-array/src/arrays/interleave/execute/primitive.rs b/vortex-array/src/arrays/interleave/execute/primitive.rs new file mode 100644 index 00000000000..e9af70a6e92 --- /dev/null +++ b/vortex-array/src/arrays/interleave/execute/primitive.rs @@ -0,0 +1,138 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Execution for primitive [`Interleave`] values. + +use num_traits::AsPrimitive; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; + +use super::super::Interleave; +use super::super::InterleaveArrayExt; +use crate::AnyColumnar; +use crate::array::Array; +use crate::arrays::Constant; +use crate::arrays::Primitive; +use crate::arrays::PrimitiveArray; +use crate::arrays::primitive::PrimitiveArrayExt; +use crate::dtype::NativePType; +use crate::executor::ExecutionCtx; +use crate::executor::ExecutionResult; +use crate::match_each_native_ptype; +use crate::match_each_unsigned_integer_ptype; +use crate::require_child; + +pub(super) fn execute( + mut array: Array, + _ctx: &mut ExecutionCtx, +) -> VortexResult { + let num_values = array.num_values(); + array = require_child!(array, array.array_indices(), 0 => Primitive); + array = require_child!(array, array.row_indices(), 1 => Primitive); + for i in 0..num_values { + array = require_child!(array, array.value(i), i + 2 => AnyColumnar); + } + + let validity = array.as_ref().validity()?; + let output = match_each_native_ptype!(array.dtype().as_ptype(), |T| { + let values = gather_values::(&array)?; + VortexResult::Ok(PrimitiveArray::new(values, validity)) + })?; + + Ok(ExecutionResult::done(output)) +} + +/// Physical storage for one primitive source array. +struct PrimitiveSource { + data: Buffer, + len: usize, + // Maps logical rows to `data`: zero for constants and identity for decoded buffers. + row_mask: usize, +} + +fn gather_values(array: &Array) -> VortexResult> { + let values = (0..array.num_values()) + .map(|i| { + let value = array.value(i); + let len = value.len(); + if let Some(constant) = value.as_opt::() { + // Validity carries nullness; a null constant's payload is never observed. + let payload = constant + .scalar() + .as_primitive() + .typed_value::() + .unwrap_or_default(); + PrimitiveSource { + data: buffer![payload], + len, + row_mask: 0, + } + } else { + PrimitiveSource { + data: value.as_::().to_buffer::(), + len, + row_mask: usize::MAX, + } + } + }) + .collect::>(); + let branches = array.array_indices().as_::(); + let rows = array.row_indices().as_::(); + + match_each_unsigned_integer_ptype!(branches.ptype(), |A| { + match_each_unsigned_integer_ptype!(rows.ptype(), |R| { + gather(&values, branches.as_slice::(), rows.as_slice::()) + }) + }) +} + +fn gather( + values: &[PrimitiveSource], + branches: &[A], + rows: &[R], +) -> VortexResult> +where + T: NativePType, + A: AsPrimitive, + R: AsPrimitive, +{ + // `zip` truncates to the shorter input. + vortex_ensure!( + rows.len() == branches.len(), + "interleave selectors differ in length: array_indices {}, row_indices {}", + branches.len(), + rows.len() + ); + + let output = + BufferMut::try_from_trusted_len_iter(branches.iter().zip(rows).map(|(branch, row)| { + let Some(source) = values.get((*branch).as_()) else { + vortex_bail!("interleave array index out of bounds"); + }; + let row = (*row).as_(); + vortex_ensure!(row < source.len, "interleave row index out of bounds"); + Ok(source.data[row & source.row_mask]) + }))?; + Ok(output.freeze()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rejects_out_of_bounds_selectors() { + let values = [PrimitiveSource { + data: buffer![1u32], + len: 1, + row_mask: 0, + }]; + + assert!(gather(&values, &[1u8], &[0u8]).is_err()); + assert!(gather(&values, &[0u8], &[1u8]).is_err()); + } +} diff --git a/vortex-array/src/arrays/interleave/mod.rs b/vortex-array/src/arrays/interleave/mod.rs index bff03ab055f..d29981249d4 100644 --- a/vortex-array/src/arrays/interleave/mod.rs +++ b/vortex-array/src/arrays/interleave/mod.rs @@ -463,6 +463,7 @@ mod tests { use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; + use crate::dtype::PType; /// Reference (oracle) implementation of the interleave spec, used only to validate the optimized /// [execute](super::execute) path. It is intentionally simple and slow: it pulls each output @@ -719,17 +720,55 @@ mod tests { } #[test] - #[should_panic(expected = "only implemented for boolean values")] - fn non_boolean_value_execution_panics() { - // Execution dispatches on the value type: primitive values have no kernel yet. - let v0 = PrimitiveArray::from_iter([1u32]).into_array(); - let v1 = PrimitiveArray::from_iter([2u32]).into_array(); - let array_indices = PrimitiveArray::from_iter([0u32, 1]).into_array(); - let row_indices = PrimitiveArray::from_iter([0u32, 0]).into_array(); - let interleaved = InterleaveArray::try_new(vec![v0, v1], array_indices, row_indices) - .vortex_expect("primitive values should construct") - .into_array(); + fn executes_primitive_values() -> VortexResult<()> { + let v0 = PrimitiveArray::from_iter([1.0f64, 2.0]).into_array(); + let v1 = PrimitiveArray::from_option_iter([Some(10.0f64), None]).into_array(); + let array_indices = PrimitiveArray::from_iter([0u8, 1, 0, 1]).into_array(); + let row_indices = PrimitiveArray::from_iter([0u32, 0, 1, 1]).into_array(); + let interleaved = + InterleaveArray::try_new(vec![v0, v1], array_indices, row_indices)?.into_array(); + let expected = + PrimitiveArray::from_option_iter([Some(1.0f64), Some(10.0), Some(2.0), None]) + .into_array(); + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(interleaved, expected, &mut ctx); + Ok(()) + } + + #[test] + fn executes_primitive_constant_values() -> VortexResult<()> { + let constant = ConstantArray::new(1.0f64, 2).into_array(); + let column = PrimitiveArray::from_iter([10.0f64, 20.0]).into_array(); + let array_indices = PrimitiveArray::from_iter([0u8, 1, 0, 1]).into_array(); + let row_indices = PrimitiveArray::from_iter([0u32, 0, 1, 1]).into_array(); + let interleaved = + InterleaveArray::try_new(vec![constant, column], array_indices, row_indices)? + .into_array(); + let expected = PrimitiveArray::from_iter([1.0f64, 10.0, 1.0, 20.0]).into_array(); let mut ctx = array_session().create_execution_ctx(); - interleaved.execute::(&mut ctx).ok(); + + assert_arrays_eq!(interleaved, expected, &mut ctx); + Ok(()) + } + + #[test] + fn executes_null_primitive_constant_values() -> VortexResult<()> { + let constant = ConstantArray::new( + Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable)), + 2, + ) + .into_array(); + let column = PrimitiveArray::from_iter([10.0f64, 20.0]).into_array(); + let array_indices = PrimitiveArray::from_iter([0u8, 1, 0, 1]).into_array(); + let row_indices = PrimitiveArray::from_iter([0u32, 0, 1, 1]).into_array(); + let interleaved = + InterleaveArray::try_new(vec![constant, column], array_indices, row_indices)? + .into_array(); + let expected = + PrimitiveArray::from_option_iter([None, Some(10.0f64), None, Some(20.0)]).into_array(); + let mut ctx = array_session().create_execution_ctx(); + + assert_arrays_eq!(interleaved, expected, &mut ctx); + Ok(()) } } diff --git a/vortex-array/src/arrays/list/array.rs b/vortex-array/src/arrays/list/array.rs index 8350b1a3558..419617c073c 100644 --- a/vortex-array/src/arrays/list/array.rs +++ b/vortex-array/src/arrays/list/array.rs @@ -6,7 +6,6 @@ use std::fmt::Formatter; use std::sync::Arc; use num_traits::AsPrimitive; -use smallvec::smallvec; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -26,6 +25,7 @@ use crate::array::ArrayParts; use crate::array::TypedArrayRef; use crate::array::child_to_validity; use crate::array::validity_to_child; +use crate::array_slots; use crate::arrays::ConstantArray; use crate::arrays::List; use crate::arrays::ListArray; @@ -39,14 +39,18 @@ use crate::match_each_native_ptype; use crate::scalar_fn::fns::operators::Operator; use crate::validity::Validity; -/// The elements data array containing all list elements concatenated together. -pub(super) const ELEMENTS_SLOT: usize = 0; -/// The offsets array defining the start/end of each list within the elements array. -pub(super) const OFFSETS_SLOT: usize = 1; -/// The validity bitmap indicating which list elements are non-null. -pub(super) const VALIDITY_SLOT: usize = 2; -pub(super) const NUM_SLOTS: usize = 3; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["elements", "offsets", "validity"]; +#[array_slots(List)] +pub struct ListSlots { + /// The elements data array containing all list elements concatenated together. + #[slot(0)] + pub elements: ArrayRef, + /// The offsets array defining the start/end of each list within the elements array. + #[slot(1)] + pub offsets: ArrayRef, + /// The validity bitmap indicating which list elements are non-null. + #[slot(2)] + pub validity: Option, +} /// A list array that stores variable-length lists of elements, similar to `Vec>`. /// @@ -124,11 +128,12 @@ impl ListData { validity: &Validity, len: usize, ) -> ArraySlots { - smallvec![ - Some(elements.clone()), - Some(offsets.clone()), - validity_to_child(validity, len), - ] + ListSlots { + elements: elements.clone(), + offsets: offsets.clone(), + validity: validity_to_child(validity, len), + } + .into_slots() } /// Creates a new `ListArray`. @@ -271,7 +276,7 @@ impl ListData { // the offset type and manual subtraction and fast path where `offsets[0] == 0`. } -pub trait ListArrayExt: TypedArrayRef { +pub trait ListArrayExt: ListArraySlotsExt { fn nullability(&self) -> crate::dtype::Nullability { match self.as_ref().dtype() { DType::List(_, nullability) => *nullability, @@ -279,21 +284,9 @@ pub trait ListArrayExt: TypedArrayRef { } } - fn elements(&self) -> &ArrayRef { - self.as_ref().slots()[ELEMENTS_SLOT] - .as_ref() - .vortex_expect("ListArray elements slot") - } - - fn offsets(&self) -> &ArrayRef { - self.as_ref().slots()[OFFSETS_SLOT] - .as_ref() - .vortex_expect("ListArray offsets slot") - } - fn list_validity(&self) -> Validity { child_to_validity( - self.as_ref().slots()[VALIDITY_SLOT].as_ref(), + self.as_ref().slots()[ListSlots::VALIDITY].as_ref(), self.nullability(), ) } @@ -407,10 +400,10 @@ impl Array { pub fn into_data_parts(self) -> ListDataParts { let dtype = self.dtype().clone(); - let elements = self.slots()[ELEMENTS_SLOT] + let elements = self.slots()[ListSlots::ELEMENTS] .clone() .vortex_expect("ListArray elements slot"); - let offsets = self.slots()[OFFSETS_SLOT] + let offsets = self.slots()[ListSlots::OFFSETS] .clone() .vortex_expect("ListArray offsets slot"); let validity = self.list_validity(); diff --git a/vortex-array/src/arrays/list/compute/cast.rs b/vortex-array/src/arrays/list/compute/cast.rs index 9fd27f731f1..930b6e3b8f4 100644 --- a/vortex-array/src/arrays/list/compute/cast.rs +++ b/vortex-array/src/arrays/list/compute/cast.rs @@ -9,7 +9,7 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::List; use crate::arrays::ListArray; -use crate::arrays::list::ListArrayExt; +use crate::arrays::list::ListArraySlotsExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::scalar_fn::fns::cast::CastKernel; diff --git a/vortex-array/src/arrays/list/compute/filter.rs b/vortex-array/src/arrays/list/compute/filter.rs index 2b98e3c36d1..fbf91cd0c32 100644 --- a/vortex-array/src/arrays/list/compute/filter.rs +++ b/vortex-array/src/arrays/list/compute/filter.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::sync::Arc; +use std::ops::Range; use num_traits::Zero; use vortex_buffer::BitBufferMut; @@ -10,7 +10,7 @@ use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_mask::MaskIter; -use vortex_mask::MaskValues; +use vortex_mask::MaskValuesRef; use crate::ArrayRef; use crate::Canonical; @@ -22,6 +22,7 @@ use crate::arrays::List; use crate::arrays::ListArray; use crate::arrays::filter::FilterKernel; use crate::arrays::list::ListArrayExt; +use crate::arrays::list::ListArraySlotsExt; use crate::dtype::IntegerPType; use crate::match_each_integer_ptype; use crate::validity::Validity; @@ -33,14 +34,54 @@ use crate::validity::Validity; /// Note that this is somewhat arbitrarily chosen... const MASK_EXPANSION_DENSITY_THRESHOLD: f64 = 0.05; -/// Construct an element mask from contiguous list offsets and a selection mask. +/// Crop when at least `1 / threshold` of referenced elements are unselected. +const PERCENTAGE_REFERENCED_UNSELECTED_ELEMENTS_THRESHOLD: usize = 20; + +/// Minimum percentage of referenced-but-unselected prefix and suffix elements required before +/// cropping. +const N_REFERENCED_UNSELECTED_ELEMENTS_THRESHOLD: usize = 1024; + +/// Return the element range to construct new mask over and to subsequently filter. In the general case this will be the range +/// of elements referenced by a sublist. +/// +/// If there are enough elements in head or tail that are referenced but not selected, it is more efficient +/// to bound the element range to the first and last selected indices. We can then expand the +/// mask only over this subset of elements, slice the elements array, and then filter. This avoids the +/// overhead of potentially wasteful mask reconstruction. +/// +/// Returns the range and a flag indicating whether the range is a subinterval of the referenced element range. +fn element_range_from_offsets( + offsets: &[O], + selection: &MaskValuesRef, +) -> (Range, bool) { + let referenced_elements_range = offsets[0].as_()..offsets[offsets.len() - 1].as_(); + + let selected_indices = selection.indices(); + let first_selected_sublist_index = selected_indices[0]; + let last_selected_sublist_index = selected_indices[selected_indices.len() - 1]; + let selected_elements_range = + offsets[first_selected_sublist_index].as_()..offsets[last_selected_sublist_index + 1].as_(); + let trimmed_element_count = referenced_elements_range.len() - selected_elements_range.len(); + + if trimmed_element_count >= N_REFERENCED_UNSELECTED_ELEMENTS_THRESHOLD + && trimmed_element_count.saturating_mul(PERCENTAGE_REFERENCED_UNSELECTED_ELEMENTS_THRESHOLD) + >= referenced_elements_range.len() + { + (selected_elements_range, true) + } else { + (referenced_elements_range, false) + } +} + +/// Construct an element mask relative to `element_range` from contiguous list offsets and an +/// outer-row selection mask. pub fn element_mask_from_offsets( offsets: &[O], - selection: &Arc, + selection: &MaskValuesRef, + element_range: &Range, ) -> Mask { - let first_offset = offsets.first().map_or(0, |first_offset| first_offset.as_()); - let last_offset = offsets.last().map_or(0, |last_offset| last_offset.as_()); - let len = last_offset - first_offset; + let first_offset = element_range.start; + let len = element_range.end - first_offset; let mut mask_builder = BitBufferMut::with_capacity(len); @@ -101,7 +142,7 @@ impl FilterKernel for List { ) -> VortexResult> { let selection = match mask { Mask::AllTrue(_) | Mask::AllFalse(_) => return Ok(None), - Mask::Values(v) => v, + Mask::Values(values) => values, }; let new_validity = match array.validity()? { @@ -120,7 +161,7 @@ impl FilterKernel for List { // TODO(ngates): for ultra-sparse masks, we don't need to optimize the entire offsets. let offsets = array.offsets().clone(); - let (new_offsets, element_mask) = + let (new_offsets, element_range, range_is_subinterval, element_mask) = match_each_integer_ptype!(offsets.dtype().as_ptype(), |O| { let offsets_buffer = offsets.execute::>(ctx)?; let offsets = offsets_buffer.as_slice(); @@ -136,12 +177,27 @@ impl FilterKernel for List { // TODO(ngates): for very dense masks, there may be no point in filtering the elements, // and instead we should construct a view against the unfiltered elements. - let element_mask = element_mask_from_offsets::(offsets, selection); + let (element_range, range_is_subinterval) = + element_range_from_offsets::(offsets, selection); + let element_mask = + element_mask_from_offsets::(offsets, selection, &element_range); - (new_offsets.freeze().into_array(), element_mask) + ( + new_offsets.freeze().into_array(), + element_range, + range_is_subinterval, + element_mask, + ) }); - let new_elements = array.sliced_elements()?.filter(element_mask)?; + let new_elements = if range_is_subinterval { + array + .elements() + .slice(element_range)? + .filter(element_mask)? + } else { + array.sliced_elements()?.filter(element_mask)? + }; // SAFETY: new_offsets are monotonically increasing starting from 0 with length // true_count + 1, and the elements have been filtered to match. @@ -150,3 +206,123 @@ impl FilterKernel for List { })) } } + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + use vortex_error::vortex_bail; + use vortex_mask::Mask; + + use super::element_mask_from_offsets; + use super::element_range_from_offsets; + + #[test] + fn element_mask_excludes_unselected_prefix_and_suffix() -> VortexResult<()> { + let Mask::Values(selection) = Mask::from_indices(5, [2]) else { + vortex_bail!("a partially selective mask uses Mask::Values") + }; + + let offsets = [10u32, 20010, 40010, 60010, 80010, 100010]; + let (range, range_is_subinterval) = element_range_from_offsets(&offsets, &selection); + let element_mask = element_mask_from_offsets(&offsets, &selection, &range); + + assert_eq!(range, 40010..60010); + assert!(range_is_subinterval); + assert!(element_mask.all_true()); + assert_eq!(element_mask.len(), 20000); + Ok(()) + } + + #[test] + fn element_mask_retains_gaps_between_selected_lists() -> VortexResult<()> { + let Mask::Values(selection) = Mask::from_indices(5, [1, 3]) else { + vortex_bail!("a partially selective mask uses Mask::Values") + }; + + let offsets = [10u32, 20010, 40010, 60010, 80010, 100010]; + let (range, range_is_subinterval) = element_range_from_offsets(&offsets, &selection); + let element_mask = element_mask_from_offsets(&offsets, &selection, &range); + + assert_eq!(range, 20010..80010); + assert!(range_is_subinterval); + assert_eq!(element_mask.len(), 60000); + assert_eq!(element_mask.true_count(), 40000); + Ok(()) + } + + #[test] + fn element_range_preserves_complete_range_for_short_lists() -> VortexResult<()> { + let Mask::Values(selection) = Mask::from_indices(5, [2]) else { + vortex_bail!("a partially selective mask uses Mask::Values") + }; + + let offsets = [10u32, 20, 30, 40, 50, 60]; + let (range, range_is_subinterval) = element_range_from_offsets(&offsets, &selection); + let element_mask = element_mask_from_offsets(&offsets, &selection, &range); + + assert_eq!(range, 10..60); + assert!(!range_is_subinterval); + assert_eq!(element_mask.len(), 50); + assert_eq!(element_mask.true_count(), 10); + Ok(()) + } + + #[test] + fn element_range_requires_minimum_savings() -> VortexResult<()> { + let Mask::Values(selection) = Mask::from_indices(2, [0]) else { + vortex_bail!("a partially selective mask uses Mask::Values") + }; + + let offsets = [0u32, 512, 1024]; + assert_eq!( + element_range_from_offsets(&offsets, &selection), + (0..1024, false) + ); + Ok(()) + } + + #[test] + fn element_range_requires_sufficient_savings_ratio() -> VortexResult<()> { + let Mask::Values(selection) = Mask::from_slices(100, vec![(0, 96)]) else { + vortex_bail!("a partially selective mask uses Mask::Values") + }; + let offsets = (0..=100).map(|index| index * 1000).collect::>(); + + assert_eq!( + element_range_from_offsets(&offsets, &selection), + (0..100_000, false) + ); + Ok(()) + } + + #[test] + fn element_range_crops_at_sufficient_savings_ratio() -> VortexResult<()> { + let Mask::Values(selection) = Mask::from_slices(100, vec![(0, 95)]) else { + vortex_bail!("a partially selective mask uses Mask::Values") + }; + let offsets = (0..=100).map(|index| index * 1000).collect::>(); + + assert_eq!( + element_range_from_offsets(&offsets, &selection), + (0..95_000, true) + ); + Ok(()) + } + + #[test] + fn element_mask_preserves_complete_range_for_edge_spanning_selection() -> VortexResult<()> { + let Mask::Values(selection) = Mask::from_indices(5, [0, 4]) else { + vortex_bail!("a partially selective mask uses Mask::Values") + }; + + let offsets = [10u32, 20010, 40010, 60010, 80010, 100010]; + let (range, range_is_subinterval) = element_range_from_offsets(&offsets, &selection); + let element_mask = element_mask_from_offsets(&offsets, &selection, &range); + + assert_eq!(range, 10..100010); + assert!(!range_is_subinterval); + assert_eq!(element_mask.len(), 100000); + assert_eq!(element_mask.true_count(), 40000); + Ok(()) + } +} diff --git a/vortex-array/src/arrays/list/compute/mask.rs b/vortex-array/src/arrays/list/compute/mask.rs index e9e895d08f1..f2727247646 100644 --- a/vortex-array/src/arrays/list/compute/mask.rs +++ b/vortex-array/src/arrays/list/compute/mask.rs @@ -8,7 +8,7 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::List; use crate::arrays::ListArray; -use crate::arrays::list::ListArrayExt; +use crate::arrays::list::ListArraySlotsExt; use crate::scalar_fn::fns::mask::MaskReduce; use crate::validity::Validity; diff --git a/vortex-array/src/arrays/list/compute/slice.rs b/vortex-array/src/arrays/list/compute/slice.rs index d29eaa5184c..f42d490237e 100644 --- a/vortex-array/src/arrays/list/compute/slice.rs +++ b/vortex-array/src/arrays/list/compute/slice.rs @@ -10,7 +10,7 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::List; use crate::arrays::ListArray; -use crate::arrays::list::ListArrayExt; +use crate::arrays::list::ListArraySlotsExt; use crate::arrays::slice::SliceReduce; impl SliceReduce for List { diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index 14e89d4c238..c010b0d2732 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -1,26 +1,37 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use itertools::Itertools as _; +use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_mask::Mask; use crate::ArrayRef; +use crate::Columnar; use crate::IntoArray; use crate::array::ArrayView; +use crate::arrays::ConstantArray; use crate::arrays::List; use crate::arrays::ListArray; +use crate::arrays::PiecewiseSequence; +use crate::arrays::PiecewiseSequenceArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; use crate::arrays::list::ListArrayExt; +use crate::arrays::list::ListArraySlotsExt; +use crate::arrays::piecewise_sequence::constant_unsigned_usize; +use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::primitive::PrimitiveArrayExt; -use crate::builders::ArrayBuilder; -use crate::builders::PrimitiveBuilder; use crate::dtype::IntegerPType; -use crate::dtype::Nullability; +use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_unsigned_integer_ptype; use crate::match_smallest_offset_type; +use crate::validity::Validity; // TODO(connor)[ListView]: Re-revert to the version where we simply convert to a `ListView` and call // the `ListView::take` compute function once `ListView` is more stable. @@ -37,21 +48,30 @@ impl TakeExecute for List { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_slices(array, piecewise_indices, indices, ctx)? + { + return Ok(Some(taken)); + } + + let new_validity = array.validity()?.take(indices)?; let indices = indices.clone().execute::(ctx)?; let indices = indices.reinterpret_cast(indices.ptype().to_unsigned()); let offsets = array.offsets().clone().execute::(ctx)?; let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); + let validity_mask = new_validity.execute_mask(indices.len(), ctx)?; // This is an over-approximation of the total number of elements in the resulting array. let total_approx = array.elements().len().saturating_mul(indices.len()); match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { match_each_unsigned_integer_ptype!(indices.ptype(), |I| { - match_smallest_offset_type!(total_approx, |OutputOffsetType| { - _take::( + match_smallest_offset_type!(total_approx, |OutOffset| { + take_with_piecewise_elements::( array, offsets.as_view(), indices.as_view(), - ctx, + new_validity, + &validity_mask, ) .map(Some) }) @@ -60,148 +80,766 @@ impl TakeExecute for List { } } -fn _take( +fn take_with_piecewise_elements( array: ArrayView<'_, List>, offsets_array: ArrayView<'_, Primitive>, indices_array: ArrayView<'_, Primitive>, - ctx: &mut ExecutionCtx, + new_validity: Validity, + validity_mask: &Mask, ) -> VortexResult { - let data_validity = array - .list_validity() - .execute_mask(array.as_ref().len(), ctx)?; - let indices_validity = indices_array - .validity() - .vortex_expect("Failed to compute validity mask") - .execute_mask(indices_array.as_ref().len(), ctx)?; - - if !indices_validity.all_true() || !data_validity.all_true() { - return _take_nullable::(array, offsets_array, indices_array, ctx); - } - let offsets: &[O] = offsets_array.as_slice(); let indices: &[I] = indices_array.as_slice(); - let mut new_offsets = PrimitiveBuilder::::with_capacity( - Nullability::NonNullable, - indices.len(), - ); - let mut elements_to_take = - PrimitiveBuilder::with_capacity(Nullability::NonNullable, 2 * indices.len()); - - let mut current_offset = OutputOffsetType::zero(); - new_offsets.append_zero(); + let offsets_capacity = indices + .len() + .checked_add(1) + .ok_or_else(|| vortex_err!("List take offsets length overflow"))?; + let mut new_offsets = BufferMut::::with_capacity(offsets_capacity); + let mut element_starts = BufferMut::::with_capacity(indices.len()); + let mut element_lengths = BufferMut::::with_capacity(indices.len()); + + let mut current_offset = 0usize; + new_offsets.push(OutOffset::zero()); + + for (&data_idx, is_valid) in indices.iter().zip_eq(validity_mask.iter()) { + if !is_valid { + new_offsets.push(new_offset_value::(current_offset)); + element_starts.push(0); + element_lengths.push(0); + continue; + } - for &data_idx in indices { let data_idx: usize = data_idx.as_(); let start = offsets[data_idx]; let stop = offsets[data_idx + 1]; - - // Annoyingly, we can't turn (start..end) into a range, so we're doing that manually. - // - // We could convert start and end to usize, but that would impose a potentially - // harder constraint - now we don't care if they fit into usize as long as their - // difference does. - let additional: usize = (stop - start).as_(); - - // TODO(0ax1): optimize this - elements_to_take.reserve_exact(additional); - for i in 0..additional { - elements_to_take.append_value(start + O::from_usize(i).vortex_expect("i < additional")); - } - current_offset += - OutputOffsetType::from_usize((stop - start).as_()).vortex_expect("offset conversion"); - new_offsets.append_value(current_offset); + let start: usize = start.as_(); + let stop: usize = stop.as_(); + let length = stop - start; + + current_offset = current_offset + .checked_add(length) + .ok_or_else(|| vortex_err!("List take output elements length overflow"))?; + new_offsets.push(new_offset_value::(current_offset)); + element_starts.push(start as u64); + element_lengths.push(length as u64); } - let elements_to_take = elements_to_take.finish(); - let new_offsets = new_offsets.finish(); - - let new_elements = array.elements().take(elements_to_take)?; + let new_offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable).into_array(); + let multipliers = ConstantArray::new(1u64, element_starts.len()).into_array(); + + // SAFETY: valid source rows contribute ranges derived from list offsets; null index/source + // rows contribute zero-length placeholder ranges. `current_offset` is the sum of all generated + // element lengths, and multiplier 1 preserves contiguous ranges. + let element_indices = unsafe { + PiecewiseSequenceArray::new_unchecked( + element_starts.into_array(), + element_lengths.into_array(), + multipliers, + current_offset, + ) + }; + let new_elements = array.elements().take(element_indices.into_array())?; - Ok(ListArray::try_new( - new_elements, - new_offsets, - array.validity()?.take(indices_array.array())?, - )? - .into_array()) + // SAFETY: offsets are rebuilt from the gathered element ranges and have one entry per output + // row plus the leading zero; validity is produced by the usual take-validity path. + Ok(unsafe { ListArray::new_unchecked(new_elements, new_offsets, new_validity) }.into_array()) } -// Kept out-of-line: as a single-callsite generic helper it would otherwise be inlined into every -// monomorphization of `_take`, duplicating the entire nullable path across all specializations. -#[inline(never)] -fn _take_nullable( +fn take_slices( array: ArrayView<'_, List>, - offsets_array: ArrayView<'_, Primitive>, - indices_array: ArrayView<'_, Primitive>, + indices: ArrayView<'_, PiecewiseSequence>, + indices_ref: &ArrayRef, ctx: &mut ExecutionCtx, -) -> VortexResult { - let offsets: &[O] = offsets_array.as_slice(); - let indices: &[I] = indices_array.as_slice(); +) -> VortexResult> { + let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { + return Ok(None); + }; let data_validity = array .list_validity() .execute_mask(array.as_ref().len(), ctx)?; - let indices_validity = indices_array - .validity() - .vortex_expect("Failed to compute validity mask") - .execute_mask(indices_array.as_ref().len(), ctx)?; - - let mut new_offsets = PrimitiveBuilder::::with_capacity( - Nullability::NonNullable, - indices.len(), + let offsets = array.offsets().clone().execute::(ctx)?; + let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); + let output_len = indices_ref.len(); + + let taken = match lengths { + Columnar::Constant(lengths) => { + let length = constant_unsigned_usize(&lengths); + take_slices_constant_start_dispatch( + array, + &starts, + length, + &offsets, + indices_ref, + output_len, + &data_validity, + )? + } + Columnar::Canonical(lengths) => { + let lengths = lengths.into_primitive(); + take_slices_start_dispatch( + array, + &starts, + &lengths, + &offsets, + indices_ref, + output_len, + &data_validity, + )? + } + }; + Ok(Some(taken)) +} + +fn take_slices_constant_start_dispatch( + array: ArrayView<'_, List>, + starts: &PrimitiveArray, + length: usize, + offsets: &PrimitiveArray, + indices_ref: &ArrayRef, + output_len: usize, + data_validity: &Mask, +) -> VortexResult { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + take_slices_constant_offset_dispatch::( + array, + starts, + length, + offsets, + indices_ref, + output_len, + data_validity, + ) + }) +} + +fn take_slices_constant_offset_dispatch( + array: ArrayView<'_, List>, + starts: &PrimitiveArray, + length: usize, + offsets: &PrimitiveArray, + indices_ref: &ArrayRef, + output_len: usize, + data_validity: &Mask, +) -> VortexResult +where + S: UnsignedPType, +{ + match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { + take_slices_constant_length::( + array, + starts.as_slice::(), + length, + offsets.as_slice::(), + indices_ref, + output_len, + data_validity, + ) + }) +} + +fn take_slices_start_dispatch( + array: ArrayView<'_, List>, + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + offsets: &PrimitiveArray, + indices_ref: &ArrayRef, + output_len: usize, + data_validity: &Mask, +) -> VortexResult { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + take_slices_length_dispatch::( + array, + starts, + lengths, + offsets, + indices_ref, + output_len, + data_validity, + ) + }) +} + +fn take_slices_length_dispatch( + array: ArrayView<'_, List>, + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + offsets: &PrimitiveArray, + indices_ref: &ArrayRef, + output_len: usize, + data_validity: &Mask, +) -> VortexResult +where + S: UnsignedPType, +{ + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + take_slices_offset_dispatch::( + array, + starts, + lengths, + offsets, + indices_ref, + output_len, + data_validity, + ) + }) +} + +fn take_slices_offset_dispatch( + array: ArrayView<'_, List>, + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + offsets: &PrimitiveArray, + indices_ref: &ArrayRef, + output_len: usize, + data_validity: &Mask, +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, +{ + match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { + take_slices_typed::( + array, + starts.as_slice::(), + lengths.as_slice::(), + offsets.as_slice::(), + indices_ref, + output_len, + data_validity, + ) + }) +} + +fn take_slices_constant_length( + array: ArrayView<'_, List>, + starts: &[S], + length: usize, + offsets: &[Offset], + indices_ref: &ArrayRef, + output_len: usize, + data_validity: &Mask, +) -> VortexResult +where + S: UnsignedPType, + Offset: UnsignedPType, +{ + let computed_len = starts + .len() + .checked_mul(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" ); + let all_valid = data_validity.all_true(); + let total_elements = if all_valid { + piecewise_list_elements_len_constant(offsets, starts, length)? + } else { + piecewise_list_elements_len_constant_validity(offsets, starts, length, data_validity)? + }; + let validity = array.validity()?.take(indices_ref)?; + + match_smallest_offset_type!(total_elements, |OutOffset| { + let gathered = if all_valid { + gather_piecewise_list_constant_length::( + array.elements(), + offsets, + starts, + length, + output_len, + total_elements, + )? + } else { + gather_piecewise_list_constant_length_validity::( + array.elements(), + offsets, + starts, + length, + output_len, + total_elements, + data_validity, + )? + }; + + // SAFETY: output offsets are rebuilt from valid monotonic source offsets; output elements + // are exactly the gathered child ranges referenced by those offsets; validity has one bit + // per output row. + Ok( + unsafe { ListArray::new_unchecked(gathered.elements, gathered.offsets, validity) } + .into_array(), + ) + }) +} - // This will be the indices we push down to the child array to call `take` with. - // - // There are 2 things to note here: - // - We do not know how many elements we need to take from our child since lists are variable - // size: thus we arbitrarily choose a capacity of `2 * # of indices`. - // - The type of the primitive builder needs to fit the largest offset of the (parent) - // `ListArray`, so we make this `PrimitiveBuilder` generic over `O` (instead of `I`). - let mut elements_to_take = - PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 2 * indices.len()); - - let mut current_offset = OutputOffsetType::zero(); - new_offsets.append_zero(); - - for (data_idx, index_valid) in indices.iter().zip(indices_validity.iter()) { - if !index_valid { - new_offsets.append_value(current_offset); +fn take_slices_typed( + array: ArrayView<'_, List>, + starts: &[S], + lengths: &[L], + offsets: &[Offset], + indices_ref: &ArrayRef, + output_len: usize, + data_validity: &Mask, +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, + Offset: UnsignedPType, +{ + let mut computed_len = 0usize; + for &length in lengths { + let length: usize = length.as_(); + computed_len = computed_len + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + } + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); + let all_valid = data_validity.all_true(); + let total_elements = if all_valid { + piecewise_list_elements_len(offsets, starts, lengths)? + } else { + piecewise_list_elements_len_validity(offsets, starts, lengths, data_validity)? + }; + + match_smallest_offset_type!(total_elements, |OutOffset| { + let gathered = if all_valid { + gather_piecewise_list::( + array.elements(), + offsets, + starts, + lengths, + output_len, + total_elements, + )? + } else { + gather_piecewise_list_validity::( + array.elements(), + offsets, + starts, + lengths, + output_len, + total_elements, + data_validity, + )? + }; + let validity = array.validity()?.take(indices_ref)?; + + // SAFETY: output offsets are rebuilt from valid monotonic source offsets; output elements + // are exactly the gathered child ranges referenced by those offsets; validity has one bit + // per output row. + Ok( + unsafe { ListArray::new_unchecked(gathered.elements, gathered.offsets, validity) } + .into_array(), + ) + }) +} + +struct GatheredList { + elements: ArrayRef, + offsets: ArrayRef, +} + +struct ValidPieceGather { + new_offsets: BufferMut, + element_starts: BufferMut, + element_lengths: BufferMut, + output_elements: usize, +} + +fn piecewise_list_elements_len_constant( + offsets: &[Offset], + starts: &[S], + length: usize, +) -> VortexResult +where + S: UnsignedPType, + Offset: UnsignedPType, +{ + if length == 0 { + return Ok(0); + } + + let mut total = 0usize; + for start in starts { + let start: usize = start.as_(); + let offset_range = &offsets[start..][..=length]; + let element_start: usize = offset_range[0].as_(); + let element_end: usize = offset_range[length].as_(); + total = total + .checked_add(element_end - element_start) + .ok_or_else(|| vortex_err!("List take output elements length overflow"))?; + } + Ok(total) +} + +fn piecewise_list_elements_len_constant_validity( + offsets: &[Offset], + starts: &[S], + length: usize, + data_validity: &Mask, +) -> VortexResult +where + S: UnsignedPType, + Offset: UnsignedPType, +{ + if length == 0 { + return Ok(0); + } + + let mut total = 0usize; + for start in starts { + let start: usize = start.as_(); + let additional = valid_piece_elements_len(offsets, data_validity, start, length)?; + total = total + .checked_add(additional) + .ok_or_else(|| vortex_err!("List take output elements length overflow"))?; + } + Ok(total) +} + +fn piecewise_list_elements_len( + offsets: &[Offset], + starts: &[S], + lengths: &[L], +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, + Offset: UnsignedPType, +{ + let mut total = 0usize; + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start: usize = start.as_(); + let length: usize = length.as_(); + let offset_range = &offsets[start..][..=length]; + let element_start: usize = offset_range[0].as_(); + let element_end: usize = offset_range[length].as_(); + total = total + .checked_add(element_end - element_start) + .ok_or_else(|| vortex_err!("List take output elements length overflow"))?; + } + Ok(total) +} + +fn piecewise_list_elements_len_validity( + offsets: &[Offset], + starts: &[S], + lengths: &[L], + data_validity: &Mask, +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, + Offset: UnsignedPType, +{ + let mut total = 0usize; + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start: usize = start.as_(); + let length: usize = length.as_(); + let additional = valid_piece_elements_len(offsets, data_validity, start, length)?; + total = total + .checked_add(additional) + .ok_or_else(|| vortex_err!("List take output elements length overflow"))?; + } + Ok(total) +} + +fn valid_piece_elements_len( + offsets: &[Offset], + data_validity: &Mask, + start: usize, + length: usize, +) -> VortexResult +where + Offset: UnsignedPType, +{ + let offset_range = &offsets[start..][..=length]; + let mut total = 0usize; + for (data_idx, window) in (start..).zip(offset_range.windows(2)) { + if !data_validity.value(data_idx) { continue; } + let element_start: usize = window[0].as_(); + let element_end: usize = window[1].as_(); + total = total + .checked_add(element_end - element_start) + .ok_or_else(|| vortex_err!("List take output elements length overflow"))?; + } + Ok(total) +} - let data_idx: usize = data_idx.as_(); +fn gather_piecewise_list_constant_length( + elements: &ArrayRef, + offsets: &[Offset], + starts: &[S], + length: usize, + output_len: usize, + total_elements: usize, +) -> VortexResult +where + S: UnsignedPType, + Offset: UnsignedPType, + OutOffset: IntegerPType, +{ + let offsets_capacity = output_len + .checked_add(1) + .ok_or_else(|| vortex_err!("List take offsets length overflow"))?; + let mut new_offsets = BufferMut::::with_capacity(offsets_capacity); + let mut element_starts = BufferMut::::with_capacity(starts.len()); + let mut element_lengths = BufferMut::::with_capacity(starts.len()); + let mut output_elements = 0usize; + + new_offsets.push(OutOffset::zero()); + for start in starts { + let start: usize = start.as_(); + if length == 0 { + continue; + } - if !data_validity.value(data_idx) { - new_offsets.append_value(current_offset); + let offset_range = &offsets[start..][..=length]; + let element_start: usize = offset_range[0].as_(); + let element_end: usize = offset_range[length].as_(); + for &offset in &offset_range[1..] { + let offset: usize = offset.as_(); + let relative = offset - element_start; + let output_offset = output_elements + relative; + new_offsets.push(new_offset_value::(output_offset)); + } + + let element_length = element_end - element_start; + element_starts.push(element_start as u64); + element_lengths.push(element_length as u64); + output_elements += element_length; + } + debug_assert_eq!(output_elements, total_elements); + + let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable).into_array(); + let multipliers = ConstantArray::new(1u64, element_starts.len()).into_array(); + // SAFETY: element ranges are derived from validated source list offsets, and total_elements is + // the sum of the gathered element range lengths. Multiplier 1 preserves contiguous ranges. + let element_indices = unsafe { + PiecewiseSequenceArray::new_unchecked( + element_starts.into_array(), + element_lengths.into_array(), + multipliers, + total_elements, + ) + }; + let elements = elements.take(element_indices.into_array())?; + + Ok(GatheredList { elements, offsets }) +} + +fn gather_piecewise_list_constant_length_validity( + elements: &ArrayRef, + offsets: &[Offset], + starts: &[S], + length: usize, + output_len: usize, + total_elements: usize, + data_validity: &Mask, +) -> VortexResult +where + S: UnsignedPType, + Offset: UnsignedPType, + OutOffset: IntegerPType, +{ + let offsets_capacity = output_len + .checked_add(1) + .ok_or_else(|| vortex_err!("List take offsets length overflow"))?; + let mut gather = ValidPieceGather { + new_offsets: BufferMut::::with_capacity(offsets_capacity), + element_starts: BufferMut::::with_capacity(output_len), + element_lengths: BufferMut::::with_capacity(output_len), + output_elements: 0, + }; + + gather.new_offsets.push(OutOffset::zero()); + for start in starts { + let start: usize = start.as_(); + if length == 0 { continue; } - let start = offsets[data_idx]; - let stop = offsets[data_idx + 1]; + gather_valid_piece(offsets, data_validity, start, length, &mut gather); + } + debug_assert_eq!(gather.output_elements, total_elements); + + let offsets = + PrimitiveArray::new(gather.new_offsets.freeze(), Validity::NonNullable).into_array(); + let multipliers = ConstantArray::new(1u64, gather.element_starts.len()).into_array(); + // SAFETY: element ranges come only from valid source list rows. Source list construction + // validated those row offsets, and null source rows produce no element range. + let element_indices = unsafe { + PiecewiseSequenceArray::new_unchecked( + gather.element_starts.into_array(), + gather.element_lengths.into_array(), + multipliers, + total_elements, + ) + }; + let elements = elements.take(element_indices.into_array())?; - // See the note in `_take` on the reasoning. - let additional: usize = (stop - start).as_(); + Ok(GatheredList { elements, offsets }) +} - elements_to_take.reserve_exact(additional); - for i in 0..additional { - elements_to_take.append_value(start + O::from_usize(i).vortex_expect("i < additional")); +fn gather_piecewise_list( + elements: &ArrayRef, + offsets: &[Offset], + starts: &[S], + lengths: &[L], + output_len: usize, + total_elements: usize, +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, + Offset: UnsignedPType, + OutOffset: IntegerPType, +{ + let offsets_capacity = output_len + .checked_add(1) + .ok_or_else(|| vortex_err!("List take offsets length overflow"))?; + let mut new_offsets = BufferMut::::with_capacity(offsets_capacity); + let mut element_starts = BufferMut::::with_capacity(starts.len()); + let mut element_lengths = BufferMut::::with_capacity(lengths.len()); + let mut output_elements = 0usize; + + new_offsets.push(OutOffset::zero()); + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start: usize = start.as_(); + let length: usize = length.as_(); + if length == 0 { + continue; + } + + let offset_range = &offsets[start..][..=length]; + let element_start: usize = offset_range[0].as_(); + let element_end: usize = offset_range[length].as_(); + for &offset in &offset_range[1..] { + let offset: usize = offset.as_(); + let relative = offset - element_start; + let output_offset = output_elements + relative; + new_offsets.push(new_offset_value::(output_offset)); } - current_offset += - OutputOffsetType::from_usize((stop - start).as_()).vortex_expect("offset conversion"); - new_offsets.append_value(current_offset); + + let element_length = element_end - element_start; + element_starts.push(element_start as u64); + element_lengths.push(element_length as u64); + output_elements += element_length; } + debug_assert_eq!(output_elements, total_elements); + + let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable).into_array(); + let multipliers = ConstantArray::new(1u64, element_starts.len()).into_array(); + // SAFETY: element ranges are derived from validated source list offsets, and total_elements is + // the sum of the gathered element range lengths. Multiplier 1 preserves contiguous ranges. + let element_indices = unsafe { + PiecewiseSequenceArray::new_unchecked( + element_starts.into_array(), + element_lengths.into_array(), + multipliers, + total_elements, + ) + }; + let elements = elements.take(element_indices.into_array())?; + + Ok(GatheredList { elements, offsets }) +} + +fn gather_piecewise_list_validity( + elements: &ArrayRef, + offsets: &[Offset], + starts: &[S], + lengths: &[L], + output_len: usize, + total_elements: usize, + data_validity: &Mask, +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, + Offset: UnsignedPType, + OutOffset: IntegerPType, +{ + let offsets_capacity = output_len + .checked_add(1) + .ok_or_else(|| vortex_err!("List take offsets length overflow"))?; + let mut gather = ValidPieceGather { + new_offsets: BufferMut::::with_capacity(offsets_capacity), + element_starts: BufferMut::::with_capacity(output_len), + element_lengths: BufferMut::::with_capacity(output_len), + output_elements: 0, + }; + + gather.new_offsets.push(OutOffset::zero()); + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start: usize = start.as_(); + let length: usize = length.as_(); + if length == 0 { + continue; + } + + gather_valid_piece(offsets, data_validity, start, length, &mut gather); + } + debug_assert_eq!(gather.output_elements, total_elements); + + let offsets = + PrimitiveArray::new(gather.new_offsets.freeze(), Validity::NonNullable).into_array(); + let multipliers = ConstantArray::new(1u64, gather.element_starts.len()).into_array(); + // SAFETY: element ranges come only from valid source list rows. Source list construction + // validated those row offsets, and null source rows produce no element range. + let element_indices = unsafe { + PiecewiseSequenceArray::new_unchecked( + gather.element_starts.into_array(), + gather.element_lengths.into_array(), + multipliers, + total_elements, + ) + }; + let elements = elements.take(element_indices.into_array())?; - let elements_to_take = elements_to_take.finish(); - let new_offsets = new_offsets.finish(); - let new_elements = array.elements().take(elements_to_take)?; + Ok(GatheredList { elements, offsets }) +} + +fn gather_valid_piece( + offsets: &[Offset], + data_validity: &Mask, + start: usize, + length: usize, + gather: &mut ValidPieceGather, +) where + Offset: UnsignedPType, + OutOffset: IntegerPType, +{ + let offset_range = &offsets[start..][..=length]; + for (data_idx, window) in (start..).zip(offset_range.windows(2)) { + if !data_validity.value(data_idx) { + gather + .new_offsets + .push(new_offset_value::(gather.output_elements)); + continue; + } + + let element_start: usize = window[0].as_(); + let element_end: usize = window[1].as_(); + let element_length = element_end - element_start; + if element_length != 0 { + gather.element_starts.push(element_start as u64); + gather.element_lengths.push(element_length as u64); + gather.output_elements += element_length; + } + gather + .new_offsets + .push(new_offset_value::(gather.output_elements)); + } +} - Ok(ListArray::try_new( - new_elements, - new_offsets, - array.validity()?.take(indices_array.array())?, - )? - .into_array()) +fn new_offset_value(value: usize) -> T { + T::from_usize(value).vortex_expect("output offset fits selected offset type") } #[cfg(test)] @@ -210,14 +848,19 @@ mod test { use rstest::rstest; use vortex_buffer::buffer; + use vortex_error::VortexResult; use crate::IntoArray as _; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::BoolArray; + use crate::arrays::ConstantArray; use crate::arrays::ListArray; use crate::arrays::ListViewArray; + use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; + use crate::arrays::listview::ListViewArrayExt; + use crate::assert_arrays_eq; use crate::compute::conformance::take::test_take_conformance; use crate::dtype::DType; use crate::dtype::Nullability; @@ -306,6 +949,53 @@ mod test { ); } + #[test] + fn null_index_ignores_out_of_bounds_payload() { + let mut ctx = array_session().create_execution_ctx(); + let list = ListArray::try_new( + buffer![1i32, 2, 3, 4].into_array(), + buffer![0u32, 2, 4].into_array(), + Validity::NonNullable, + ) + .unwrap() + .into_array(); + + let idx = PrimitiveArray::new( + buffer![1u32, 99, 0], + Validity::from_iter([true, false, true]), + ) + .into_array(); + let result = list.take(idx).unwrap(); + + let expected = ListArray::new( + buffer![3i32, 4, 1, 2].into_array(), + buffer![0u32, 2, 2, 4].into_array(), + Validity::from_iter([true, false, true]), + ); + assert_arrays_eq!(expected, result, &mut ctx); + } + + #[test] + fn null_source_row_uses_valid_empty_output_range() { + let mut ctx = array_session().create_execution_ctx(); + let list = ListArray::new( + buffer![1i32, 2, 7, 8].into_array(), + buffer![0u32, 2, 4].into_array(), + Validity::from_iter([true, false]), + ) + .into_array(); + + let idx = buffer![0u32, 1].into_array(); + let result = list.take(idx).unwrap(); + + let expected = ListArray::new( + buffer![1i32, 2].into_array(), + buffer![0u32, 2, 2].into_array(), + Validity::from_iter([true, false]), + ); + assert_arrays_eq!(expected, result, &mut ctx); + } + #[test] fn change_validity() { let list = ListArray::try_new( @@ -403,6 +1093,152 @@ mod test { ); } + #[test] + fn piecewise_sequence_take() { + let mut ctx = array_session().create_execution_ctx(); + let list = ListArray::try_new( + buffer![0i32, 1, 2, 3, 4, 5, 6].into_array(), + buffer![0u32, 2, 5, 5, 7].into_array(), + Validity::NonNullable, + ) + .unwrap() + .into_array(); + let idx = PiecewiseSequenceArray::try_new( + buffer![1u64, 0].into_array(), + buffer![2u64, 1].into_array(), + ConstantArray::new(1u64, 2).into_array(), + 3, + ) + .unwrap() + .into_array(); + + let result = list + .take(idx) + .unwrap() + .execute::(&mut ctx) + .unwrap(); + + let element_dtype: Arc = Arc::new(I32.into()); + assert_eq!( + result.execute_scalar(0, &mut ctx).unwrap(), + Scalar::list( + Arc::clone(&element_dtype), + vec![2i32.into(), 3.into(), 4.into()], + Nullability::NonNullable + ) + ); + assert_eq!( + result.execute_scalar(1, &mut ctx).unwrap(), + Scalar::list(Arc::clone(&element_dtype), vec![], Nullability::NonNullable) + ); + assert_eq!( + result.execute_scalar(2, &mut ctx).unwrap(), + Scalar::list( + element_dtype, + vec![0i32.into(), 1.into()], + Nullability::NonNullable + ) + ); + } + + #[test] + fn piecewise_sequence_take_nullable_list_constant_lengths() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let list = ListArray::try_new( + buffer![0i32, 1, 99, 100, 2, 3, 4, 5].into_array(), + buffer![0u32, 2, 4, 7, 8].into_array(), + Validity::Array(BoolArray::from_iter([true, false, true, true]).into_array()), + )? + .into_array(); + let idx = PiecewiseSequenceArray::try_new( + buffer![0u64].into_array(), + ConstantArray::new(4u64, 1).into_array(), + ConstantArray::new(1u64, 1).into_array(), + 4, + )? + .into_array(); + + let result = list.take(idx)?.execute::(&mut ctx)?; + assert_eq!(result.offset_at(0), 0); + assert_eq!(result.size_at(0), 2); + assert_eq!(result.offset_at(1), 2); + assert_eq!(result.size_at(1), 0); + assert_eq!(result.offset_at(2), 2); + assert_eq!(result.size_at(2), 3); + assert_eq!(result.offset_at(3), 5); + assert_eq!(result.size_at(3), 1); + + let element_dtype: Arc = Arc::new(I32.into()); + assert_eq!( + result.execute_scalar(0, &mut ctx)?, + Scalar::list( + Arc::clone(&element_dtype), + vec![0i32.into(), 1.into()], + Nullability::Nullable + ) + ); + assert!(result.is_invalid(1, &mut ctx)?); + assert_eq!( + result.execute_scalar(2, &mut ctx)?, + Scalar::list( + Arc::clone(&element_dtype), + vec![2i32.into(), 3.into(), 4.into()], + Nullability::Nullable + ) + ); + assert_eq!( + result.execute_scalar(3, &mut ctx)?, + Scalar::list(element_dtype, vec![5i32.into()], Nullability::Nullable) + ); + Ok(()) + } + + #[test] + fn piecewise_sequence_take_nullable_list_array_lengths() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let list = ListArray::try_new( + buffer![0i32, 1, 99, 100, 2, 3, 4, 5].into_array(), + buffer![0u32, 2, 4, 7, 8].into_array(), + Validity::Array(BoolArray::from_iter([true, false, true, true]).into_array()), + )? + .into_array(); + let idx = PiecewiseSequenceArray::try_new( + buffer![1u64, 0].into_array(), + buffer![2u64, 1].into_array(), + ConstantArray::new(1u64, 2).into_array(), + 3, + )? + .into_array(); + + let result = list.take(idx)?.execute::(&mut ctx)?; + assert_eq!(result.offset_at(0), 0); + assert_eq!(result.size_at(0), 0); + assert_eq!(result.offset_at(1), 0); + assert_eq!(result.size_at(1), 3); + assert_eq!(result.offset_at(2), 3); + assert_eq!(result.size_at(2), 2); + + let element_dtype: Arc = Arc::new(I32.into()); + assert!(result.is_invalid(0, &mut ctx)?); + assert_eq!( + result.execute_scalar(1, &mut ctx)?, + Scalar::list( + Arc::clone(&element_dtype), + vec![2i32.into(), 3.into(), 4.into()], + Nullability::Nullable + ) + ); + assert_eq!( + result.execute_scalar(2, &mut ctx)?, + Scalar::list( + element_dtype, + vec![0i32.into(), 1.into()], + Nullability::Nullable + ) + ); + Ok(()) + } + #[test] fn test_take_empty_array() { let list = ListArray::try_new( diff --git a/vortex-array/src/arrays/list/mod.rs b/vortex-array/src/arrays/list/mod.rs index 422a9523ff9..bfd43cdb401 100644 --- a/vortex-array/src/arrays/list/mod.rs +++ b/vortex-array/src/arrays/list/mod.rs @@ -3,8 +3,11 @@ mod array; pub use array::ListArrayExt; +pub use array::ListArraySlotsExt; pub use array::ListData; pub use array::ListDataParts; +pub use array::ListSlots; +pub use array::ListSlotsView; pub use vtable::ListArray; pub(crate) mod compute; diff --git a/vortex-array/src/arrays/list/test_harness.rs b/vortex-array/src/arrays/list/test_harness.rs index ba2c9d335f1..573c4d35a3c 100644 --- a/vortex-array/src/arrays/list/test_harness.rs +++ b/vortex-array/src/arrays/list/test_harness.rs @@ -4,20 +4,21 @@ use std::sync::Arc; use itertools::Itertools; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use crate::arrays::ListArray; use crate::builders::ArrayBuilder; use crate::builders::ListBuilder; use crate::dtype::DType; -use crate::dtype::IntegerPType; +use crate::dtype::OffsetBuilderPType; use crate::scalar::Scalar; impl ListArray { /// This is a convenience method to create a list array from an iterator of iterators. /// This method is slow however since each element is first converted to a scalar and then /// appended to the array. - pub fn from_iter_slow( + pub fn from_iter_slow( iter: I, dtype: Arc, ) -> VortexResult @@ -26,11 +27,12 @@ impl ListArray { ::Item: Into, { let iter = iter.into_iter(); - let mut builder = ListBuilder::::with_capacity( + let mut builder = ListBuilder::::with_capacity_in( Arc::clone(&dtype), crate::dtype::Nullability::NonNullable, 2 * iter.size_hint().0, iter.size_hint().0, + BufferAllocatorRef::static_ref(), ); for v in iter { @@ -44,7 +46,7 @@ impl ListArray { Ok(builder.finish_into_list()) } - pub fn from_iter_opt_slow>, T>( + pub fn from_iter_opt_slow>, T>( iter: I, dtype: Arc, ) -> VortexResult @@ -53,11 +55,12 @@ impl ListArray { T::Item: Into, { let iter = iter.into_iter(); - let mut builder = ListBuilder::::with_capacity( + let mut builder = ListBuilder::::with_capacity_in( Arc::clone(&dtype), crate::dtype::Nullability::Nullable, 2 * iter.size_hint().0, iter.size_hint().0, + BufferAllocatorRef::static_ref(), ); for v in iter { diff --git a/vortex-array/src/arrays/list/tests.rs b/vortex-array/src/arrays/list/tests.rs index 906f4e03588..93daf3dc4a4 100644 --- a/vortex-array/src/arrays/list/tests.rs +++ b/vortex-array/src/arrays/list/tests.rs @@ -386,8 +386,13 @@ fn test_list_filter_variable_sizes() { #[test] fn test_offset_to_0() { - let mut builder = - ListBuilder::::with_capacity(Arc::new(I32.into()), Nullability::NonNullable, 10, 5); + let mut builder = ListBuilder::::with_capacity_in( + Arc::new(I32.into()), + Nullability::NonNullable, + 10, + 5, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); builder .append_value( Scalar::list( diff --git a/vortex-array/src/arrays/list/vtable/mod.rs b/vortex-array/src/arrays/list/vtable/mod.rs index 55f7753c995..c55e7050351 100644 --- a/vortex-array/src/arrays/list/vtable/mod.rs +++ b/vortex-array/src/arrays/list/vtable/mod.rs @@ -26,12 +26,9 @@ use crate::array::ArrayParts; use crate::array::ArrayView; use crate::array::VTable; use crate::array::with_empty_buffers; -use crate::arrays::list::ListArrayExt; +use crate::arrays::list::ListArraySlotsExt; use crate::arrays::list::ListData; -use crate::arrays::list::array::ELEMENTS_SLOT; -use crate::arrays::list::array::NUM_SLOTS; -use crate::arrays::list::array::OFFSETS_SLOT; -use crate::arrays::list::array::SLOT_NAMES; +use crate::arrays::list::ListSlots; use crate::arrays::list::compute::rules::PARENT_RULES; use crate::arrays::listview::list_view_from_list; use crate::buffer::BufferHandle; @@ -39,6 +36,7 @@ use crate::builders::ArrayBuilder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; +use crate::match_each_list_builder; use crate::serde::ArrayChildren; use crate::validity::Validity; mod operations; @@ -123,14 +121,15 @@ impl VTable for List { slots: &[Option], ) -> VortexResult<()> { vortex_ensure!( - slots.len() == NUM_SLOTS, - "ListArray expected {NUM_SLOTS} slots, found {}", + slots.len() == ListSlots::COUNT, + "ListArray expected {} slots, found {}", + ListSlots::COUNT, slots.len() ); - let elements = slots[ELEMENTS_SLOT] + let elements = slots[ListSlots::ELEMENTS] .as_ref() .vortex_expect("ListArray elements slot"); - let offsets = slots[OFFSETS_SLOT] + let offsets = slots[ListSlots::OFFSETS] .as_ref() .vortex_expect("ListArray offsets slot"); vortex_ensure!( @@ -192,7 +191,7 @@ impl VTable for List { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + ListSlots::NAMES[idx].to_string() } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { @@ -201,12 +200,21 @@ impl VTable for List { )) } + // The complexity comes from the expansion of `match_each_list_builder!`. + #[expect(clippy::cognitive_complexity)] fn append_to_builder( array: ArrayView<'_, Self>, builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - builder.append_list_array(array, ctx) + match match_each_list_builder!(&mut *builder, |b| b.append_list_array(array, ctx)) { + Some(result) => result, + None => vortex_bail!( + "cannot append a List array of dtype {} to a {} builder", + array.dtype(), + builder.dtype() + ), + } } } diff --git a/vortex-array/src/arrays/listview/array.rs b/vortex-array/src/arrays/listview/array.rs index daf7100dc59..43f51971b8f 100644 --- a/vortex-array/src/arrays/listview/array.rs +++ b/vortex-array/src/arrays/listview/array.rs @@ -6,7 +6,6 @@ use std::fmt::Formatter; use std::sync::Arc; use num_traits::AsPrimitive; -use smallvec::smallvec; use vortex_buffer::BitBufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -26,6 +25,7 @@ use crate::array::ArrayParts; use crate::array::TypedArrayRef; use crate::array::child_to_validity; use crate::array::validity_to_child; +use crate::array_slots; use crate::arrays::ListView; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; @@ -42,23 +42,28 @@ use crate::match_each_unsigned_integer_ptype; use crate::scalar_fn::fns::operators::Operator; use crate::validity::Validity; -/// The `elements` data array, where each list scalar is a _slice_ of the `elements` array, and -/// each inner list element is a _scalar_ of the `elements` array. -pub(super) const ELEMENTS_SLOT: usize = 0; -/// The `offsets` array indicating the start position of each list in elements. -/// -/// Since we also store `sizes`, this `offsets` field is allowed to be stored out-of-order -/// (which is different from [`ListArray`](crate::arrays::ListArray)). -pub(super) const OFFSETS_SLOT: usize = 1; -/// The `sizes` array indicating the length of each list. -/// -/// This field is intended to be paired with a corresponding offset to determine the list scalar -/// we want to access. -pub(super) const SIZES_SLOT: usize = 2; -/// The validity bitmap indicating which list elements are non-null. -pub(super) const VALIDITY_SLOT: usize = 3; -pub(super) const NUM_SLOTS: usize = 4; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["elements", "offsets", "sizes", "validity"]; +#[array_slots(ListView)] +pub struct ListViewSlots { + /// The `elements` data array, where each list scalar is a _slice_ of the `elements` array, + /// and each inner list element is a _scalar_ of the `elements` array. + #[slot(0)] + pub elements: ArrayRef, + /// The `offsets` array indicating the start position of each list in elements. + /// + /// Since we also store `sizes`, this `offsets` field is allowed to be stored out-of-order + /// (which is different from [`ListArray`](crate::arrays::ListArray)). + #[slot(1)] + pub offsets: ArrayRef, + /// The `sizes` array indicating the length of each list. + /// + /// This field is intended to be paired with a corresponding offset to determine the list + /// scalar we want to access. + #[slot(2)] + pub sizes: ArrayRef, + /// The validity bitmap indicating which list elements are non-null. + #[slot(3)] + pub validity: Option, +} /// The canonical encoding for variable-length list arrays. /// @@ -166,12 +171,13 @@ impl ListViewData { validity: &Validity, len: usize, ) -> ArraySlots { - smallvec![ - Some(elements.clone()), - Some(offsets.clone()), - Some(sizes.clone()), - validity_to_child(validity, len), - ] + ListViewSlots { + elements: elements.clone(), + offsets: offsets.clone(), + sizes: sizes.clone(), + validity: validity_to_child(validity, len), + } + .into_slots() } /// Creates a new `ListViewArray`. @@ -347,7 +353,7 @@ fn fill_referenced_mask( } } -pub trait ListViewArrayExt: TypedArrayRef { +pub trait ListViewArrayExt: ListViewArraySlotsExt { fn nullability(&self) -> crate::dtype::Nullability { match self.as_ref().dtype() { DType::List(_, nullability) => *nullability, @@ -355,27 +361,9 @@ pub trait ListViewArrayExt: TypedArrayRef { } } - fn elements(&self) -> &ArrayRef { - self.as_ref().slots()[ELEMENTS_SLOT] - .as_ref() - .vortex_expect("ListViewArray elements slot") - } - - fn offsets(&self) -> &ArrayRef { - self.as_ref().slots()[OFFSETS_SLOT] - .as_ref() - .vortex_expect("ListViewArray offsets slot") - } - - fn sizes(&self) -> &ArrayRef { - self.as_ref().slots()[SIZES_SLOT] - .as_ref() - .vortex_expect("ListViewArray sizes slot") - } - fn listview_validity(&self) -> Validity { child_to_validity( - self.as_ref().slots()[VALIDITY_SLOT].as_ref(), + self.as_ref().slots()[ListViewSlots::VALIDITY].as_ref(), self.nullability(), ) } @@ -667,13 +655,13 @@ impl Array { } pub fn into_data_parts(self) -> ListViewDataParts { - let elements = self.slots()[ELEMENTS_SLOT] + let elements = self.slots()[ListViewSlots::ELEMENTS] .clone() .vortex_expect("ListViewArray elements slot"); - let offsets = self.slots()[OFFSETS_SLOT] + let offsets = self.slots()[ListViewSlots::OFFSETS] .clone() .vortex_expect("ListViewArray offsets slot"); - let sizes = self.slots()[SIZES_SLOT] + let sizes = self.slots()[ListViewSlots::SIZES] .clone() .vortex_expect("ListViewArray sizes slot"); let validity = self.listview_validity(); diff --git a/vortex-array/src/arrays/listview/compute/cast.rs b/vortex-array/src/arrays/listview/compute/cast.rs index 7f4c972940f..1c5e1fdfd84 100644 --- a/vortex-array/src/arrays/listview/compute/cast.rs +++ b/vortex-array/src/arrays/listview/compute/cast.rs @@ -9,7 +9,7 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::ListView; use crate::arrays::ListViewArray; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::scalar_fn::fns::cast::CastKernel; diff --git a/vortex-array/src/arrays/listview/compute/mask.rs b/vortex-array/src/arrays/listview/compute/mask.rs index e14233a71d0..2b5a89132fb 100644 --- a/vortex-array/src/arrays/listview/compute/mask.rs +++ b/vortex-array/src/arrays/listview/compute/mask.rs @@ -8,7 +8,7 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::ListView; use crate::arrays::ListViewArray; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::scalar_fn::fns::mask::MaskReduce; use crate::validity::Validity; diff --git a/vortex-array/src/arrays/listview/compute/rules.rs b/vortex-array/src/arrays/listview/compute/rules.rs index 9b9396d7fe7..233533202a8 100644 --- a/vortex-array/src/arrays/listview/compute/rules.rs +++ b/vortex-array/src/arrays/listview/compute/rules.rs @@ -10,7 +10,7 @@ use crate::arrays::Filter; use crate::arrays::ListView; use crate::arrays::ListViewArray; use crate::arrays::dict::TakeReduceAdaptor; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::slice::SliceReduceAdaptor; use crate::optimizer::rules::ArrayParentReduceRule; use crate::optimizer::rules::ParentRuleSet; diff --git a/vortex-array/src/arrays/listview/compute/slice.rs b/vortex-array/src/arrays/listview/compute/slice.rs index 2dd29278392..8cda72b871b 100644 --- a/vortex-array/src/arrays/listview/compute/slice.rs +++ b/vortex-array/src/arrays/listview/compute/slice.rs @@ -10,7 +10,7 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::ListView; use crate::arrays::ListViewArray; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::slice::SliceReduce; impl SliceReduce for ListView { diff --git a/vortex-array/src/arrays/listview/compute/take.rs b/vortex-array/src/arrays/listview/compute/take.rs index 2b6c016d2c3..f5ffe321821 100644 --- a/vortex-array/src/arrays/listview/compute/take.rs +++ b/vortex-array/src/arrays/listview/compute/take.rs @@ -12,7 +12,7 @@ use crate::arrays::ListView; use crate::arrays::ListViewArray; use crate::arrays::dict::TakeExecute; use crate::arrays::dict::TakeReduce; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::builtins::ArrayBuiltins; use crate::dtype::Nullability; use crate::match_each_integer_ptype; diff --git a/vortex-array/src/arrays/listview/compute/zip.rs b/vortex-array/src/arrays/listview/compute/zip.rs index 71640d4fb21..d5e2a6bf406 100644 --- a/vortex-array/src/arrays/listview/compute/zip.rs +++ b/vortex-array/src/arrays/listview/compute/zip.rs @@ -21,7 +21,7 @@ use crate::arrays::ChunkedArray; use crate::arrays::ListView; use crate::arrays::ListViewArray; use crate::arrays::chunked::ChunkedArrayExt; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; @@ -251,7 +251,7 @@ mod tests { use crate::arrays::ListView; use crate::arrays::ListViewArray; use crate::arrays::chunked::ChunkedArrayExt; - use crate::arrays::listview::ListViewArrayExt; + use crate::arrays::listview::ListViewArraySlotsExt; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; diff --git a/vortex-array/src/arrays/listview/conversion.rs b/vortex-array/src/arrays/listview/conversion.rs index e85567a9ed9..78c0e81844d 100644 --- a/vortex-array/src/arrays/listview/conversion.rs +++ b/vortex-array/src/arrays/listview/conversion.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -16,8 +17,11 @@ use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::extension::ExtensionArrayExt; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::arrays::list::ListArrayExt; +use crate::arrays::list::ListArraySlotsExt; use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::ListViewRebuildMode; use crate::arrays::struct_::StructArrayExt; use crate::builders::PrimitiveBuilder; @@ -40,17 +44,20 @@ pub fn list_view_from_list(list: ListArray, ctx: &mut ExecutionCtx) -> VortexRes // function might not expect the output `ListViewArray` to have a bunch of leading and trailing // garbage data when they turn it back into a `ListArray`. let list = list.reset_offsets(false, ctx)?; - let list_offsets = list.offsets().clone(); + + // `reset_offsets` leaves a lazy subtraction in the offsets slot, which both the `sizes` below + // and the view's own offsets read. Execute it once here so that it does not run per reader. + let list_offsets = list.offsets().clone().execute::(ctx)?; // Create `sizes` array by computing differences between consecutive offsets. // We use the same `DType` for the sizes as the `offsets` array to ensure compatibility. - let sizes = match_each_integer_ptype!(list_offsets.dtype().as_ptype(), |O| { - build_sizes_from_offsets::(&list, ctx)? + let sizes = match_each_integer_ptype!(list_offsets.ptype(), |O| { + build_sizes_from_offsets::(&list_offsets, ctx.allocator())? }); // We need to slice the `offsets` to remove the last element (`ListArray` has `n + 1` offsets). debug_assert_eq!(list_offsets.len(), list.len() + 1); - let adjusted_offsets = list_offsets.slice(0..list.len())?; + let adjusted_offsets = list_offsets.into_array().slice(0..list.len())?; // SAFETY: Since everything came from an existing valid `ListArray`, and the `sizes` were // derived from valid and in-order `offsets`, we know these fields are valid. @@ -66,22 +73,23 @@ pub fn list_view_from_list(list: ListArray, ctx: &mut ExecutionCtx) -> VortexRes }) } -/// Builds a sizes array from a `ListArray` by computing differences between consecutive offsets. +/// Builds a sizes array by computing differences between consecutive offsets. +/// +/// `offsets` **must** be the `n + 1` sorted offsets of a non-empty `ListArray` of `n` rows. fn build_sizes_from_offsets( - list: &ListArray, - ctx: &mut ExecutionCtx, + offsets: &PrimitiveArray, + allocator: &BufferAllocatorRef, ) -> VortexResult { - let len = list.len(); - let mut sizes_builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, len); + let offsets_slice = offsets.as_slice::(); + debug_assert!(offsets_slice.is_sorted()); + + let len = offsets_slice.len() - 1; + let mut sizes_builder = + PrimitiveBuilder::::with_capacity_in(Nullability::NonNullable, len, allocator); // Create `UninitRange` for direct memory access. let mut sizes_range = sizes_builder.uninit_range(len); - let offsets = list.offsets().clone().execute::(ctx)?; - let offsets_slice = offsets.as_slice::(); - debug_assert_eq!(len + 1, offsets_slice.len()); - debug_assert!(offsets_slice.is_sorted()); - // Compute sizes as the difference between consecutive offsets. for i in 0..len { let size = offsets_slice[i + 1] - offsets_slice[i]; @@ -98,13 +106,12 @@ fn build_sizes_from_offsets( // TODO(connor)[ListView]: Note that it is not exactly zero-copy because we have to add a single // offset at the end, but it is fast enough. -/// Creates a `ListArray` from a `ListViewArray`. The resulting `ListArray` will not have any -/// leading or trailing garbage data. +/// Creates a `ListArray` from a `ListViewArray`. /// -/// If `ListViewArray::is_zero_copy_to_list` is `true`, then this operation is fast +/// The resulting `ListArray` will not have any leading or trailing garbage data. /// -/// Otherwise, this function fall back to the (very) expensive path and will rebuild the -/// `ListArray` from scratch. +/// This operation is fast when `ListViewArray::is_zero_copy_to_list` is `true`. Otherwise it +/// falls back to the (very) expensive path and rebuilds the `ListArray` from scratch. pub fn list_from_list_view( list_view: ListViewArray, ctx: &mut ExecutionCtx, @@ -133,8 +140,9 @@ pub fn list_from_list_view( // TODO(connor)[ListView]: We can optimize this by always keeping extra memory in `ListViewArray` // offsets for an `n+1`th offset. -/// Builds a `ListArray` offsets array from a `ListViewArray` by constructing `n+1` offsets. -/// The last offset is computed as `last_offset + last_size`. +/// Builds a `ListArray` offsets array from a `ListViewArray` by constructing `n + 1` offsets. +/// +/// The last offset is computed as `last_offset + last_size`, and is zero for an empty array. /// /// # Safety /// @@ -146,7 +154,7 @@ unsafe fn build_list_offsets_from_list_view( ) -> ArrayRef { let len = list_view.len(); let mut offsets_builder = - PrimitiveBuilder::::with_capacity(Nullability::NonNullable, len + 1); + PrimitiveBuilder::::with_capacity_in(Nullability::NonNullable, len + 1, ctx.allocator()); // Create uninit range for direct memory access. let mut offsets_range = offsets_builder.uninit_range(len + 1); @@ -247,11 +255,11 @@ pub fn recursive_list_from_list_view( } } Canonical::Struct(struct_array) => { - let fields = struct_array.unmasked_fields(); - let mut converted_fields = Vec::with_capacity(fields.len()); + let mut converted_fields = + Vec::with_capacity(struct_array.iter_unmasked_fields().len()); let mut any_changed = false; - for field in fields.iter() { + for field in struct_array.iter_unmasked_fields() { let converted_field = recursive_list_from_list_view(field.clone(), ctx)?; // Avoid cloning if elements didn't change. any_changed |= !ArrayRef::ptr_eq(&converted_field, field); @@ -312,7 +320,8 @@ mod tests { use crate::arrays::StructArray; use crate::arrays::VarBinViewArray; use crate::arrays::list::ListArrayExt; - use crate::arrays::listview::ListViewArrayExt; + use crate::arrays::list::ListArraySlotsExt; + use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::list_from_list_view; use crate::arrays::listview::list_view_from_list; use crate::assert_arrays_eq; @@ -346,6 +355,24 @@ mod tests { Ok(()) } + #[test] + fn test_list_to_listview_resets_nonzero_offsets() -> VortexResult<()> { + let elements = buffer![0i32, 1, 2, 3, 4].into_array(); + let offsets = buffer![2u16, 4, 5].into_array(); + let list = ListArray::try_new(elements, offsets, Validity::NonNullable)?; + + let mut ctx = SESSION.create_execution_ctx(); + let list_view = list_view_from_list(list.clone(), &mut ctx)?; + + assert_arrays_eq!( + buffer![0u16, 2].into_array(), + list_view.offsets().clone(), + &mut ctx + ); + assert_arrays_eq!(list, list_view, &mut ctx); + Ok(()) + } + #[test] fn test_listview_to_list_zero_copy() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); diff --git a/vortex-array/src/arrays/listview/mod.rs b/vortex-array/src/arrays/listview/mod.rs index b240854c6cf..9955ae0f5fd 100644 --- a/vortex-array/src/arrays/listview/mod.rs +++ b/vortex-array/src/arrays/listview/mod.rs @@ -3,8 +3,11 @@ mod array; pub use array::ListViewArrayExt; +pub use array::ListViewArraySlotsExt; pub use array::ListViewData; pub use array::ListViewDataParts; +pub use array::ListViewSlots; +pub use array::ListViewSlotsView; pub use vtable::ListViewArray; pub(crate) mod compute; diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 342c9eccf19..dfcfc4283c7 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -5,16 +5,18 @@ use num_traits::FromPrimitive; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_mask::Mask; -use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::ConstantArray; use crate::arrays::ListViewArray; +use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::primitive::PrimitiveArrayExt; -use crate::builders::builder_with_capacity; use crate::builtins::ArrayBuiltins; use crate::dtype::IntegerPType; use crate::dtype::Nullability; @@ -150,8 +152,7 @@ impl ListViewArray { } /// Picks between [`rebuild_with_take`](Self::rebuild_with_take) and - /// [`rebuild_list_by_list`](Self::rebuild_list_by_list) based on element dtype and average - /// list size. + /// [`rebuild_with_piecewise`](Self::rebuild_with_piecewise) based on average list size. fn naive_rebuild( &self, ctx: &mut ExecutionCtx, @@ -167,12 +168,12 @@ impl ListViewArray { if Self::should_use_take(total, self.len()) { self.rebuild_with_take::(ctx) } else { - self.rebuild_list_by_list::(ctx) + self.rebuild_with_piecewise::(ctx) } } /// Returns `true` when we are confident that `rebuild_with_take` will - /// outperform `rebuild_list_by_list`. + /// outperform `rebuild_with_piecewise`. /// /// Take is dramatically faster for small lists (often 10-100×) because it /// avoids per-list builder overhead. LBL is the safer default for larger @@ -243,86 +244,69 @@ impl ListViewArray { .reinterpret_cast(size_ptype) .into_array(); - // SAFETY: same invariants as `rebuild_list_by_list` — offsets are sequential and - // non-overlapping, all (offset, size) pairs reference valid elements, and the validity - // array is preserved from the original. + // SAFETY: offsets are sequential and non-overlapping, all (offset, size) pairs reference + // valid elements, and the validity array is preserved from the original. Ok(unsafe { ListViewArray::new_unchecked(elements, offsets, sizes, self.validity()?) .with_zero_copy_to_list(true) }) } - /// Rebuilds elements list-by-list: canonicalize elements upfront, then for each list `slice` - /// the relevant range and `append_to_builder` into a typed builder. - fn rebuild_list_by_list( + /// Rebuilds elements using one contiguous-run take over the element child. + fn rebuild_with_piecewise( &self, ctx: &mut ExecutionCtx, ) -> VortexResult { - let element_dtype = self - .dtype() - .as_list_element_opt() - .vortex_expect("somehow had a canonical list that was not a list"); - let new_offset_ptype = rebuilt_offset_ptype(self.offsets().dtype().as_ptype()); let size_ptype = self.sizes().dtype().as_ptype(); let offsets_canonical = self.offsets().clone().execute::(ctx)?; let offsets_canonical = offsets_canonical.reinterpret_cast(offsets_canonical.ptype().to_unsigned()); - let offsets_slice = offsets_canonical.as_slice::(); let sizes_canonical = self.sizes().clone().execute::(ctx)?; let sizes_canonical = sizes_canonical.reinterpret_cast(sizes_canonical.ptype().to_unsigned()); - let sizes_slice = sizes_canonical.as_slice::(); - - let len = offsets_slice.len(); - - let mut new_offsets = BufferMut::::with_capacity(len); - // TODO(connor)[ListView]: Do we really need to do this? - // The only reason we need to rebuild the sizes here is that the validity may indicate that - // a list is null even though it has a non-zero size. This rebuild will set the size of all - // null lists to 0. - let mut new_sizes = BufferMut::::with_capacity(len); - - // Canonicalize the elements up front as we will be slicing the elements quite a lot. - let elements_canonical = self - .elements() - .clone() - .execute::(ctx)? - .into_array(); - - // Note that we do not know what the exact capacity should be of the new elements since - // there could be overlaps in the existing `ListViewArray`. - let mut new_elements_builder = - builder_with_capacity(element_dtype.as_ref(), self.elements().len()); - - // Resolve validity to a mask once instead of probing it per row (see `rebuild_with_take`). - let validity = self.validity()?.execute_mask(len, ctx)?; - - let mut n_elements = NewOffset::zero(); - for index in 0..len { - if !validity.value(index) { - // For NULL lists, place them after the previous item's data to maintain the - // no-overlap invariant for zero-copy to `ListArray` arrays. - new_offsets.push(n_elements); - new_sizes.push(S::zero()); - continue; - } - let offset = offsets_slice[index]; - let size = sizes_slice[index]; - - let start = offset.as_(); - let stop = start + size.as_(); - - new_offsets.push(n_elements); - new_sizes.push(size); - elements_canonical - .slice(start..stop)? - .append_to_builder(new_elements_builder.as_mut(), ctx)?; + let len = offsets_canonical.len(); + let validity = self.validity()?; + let validity_mask = validity.execute_mask(len, ctx)?; - n_elements += num_traits::cast(size).vortex_expect("Cast failed"); - } + let ranges = match_each_unsigned_integer_ptype!(offsets_canonical.ptype(), |O| { + rebuild_ranges::( + offsets_canonical.as_slice::(), + sizes_canonical.as_slice::(), + &validity_mask, + ) + })?; + + let RebuildRanges { + new_offsets, + new_sizes, + starts, + lengths, + elements_len, + } = ranges; + let constant_length = lengths + .first() + .copied() + .filter(|first| lengths.iter().all(|length| *length == *first)); + let lengths = match constant_length { + Some(length) => ConstantArray::new(length, starts.len()).into_array(), + None => lengths.into_array(), + }; + + // SAFETY: range starts and lengths are derived from valid ListView metadata; elements_len + // is the sum of all generated range lengths. Multiplier 1 preserves contiguous ranges. + let multipliers = ConstantArray::new(1u64, starts.len()).into_array(); + let element_indices = unsafe { + PiecewiseSequenceArray::new_unchecked( + starts.into_array(), + lengths, + multipliers, + elements_len, + ) + }; + let elements = self.elements().take(element_indices.into_array())?; // Built unsigned; reinterpret back to the signed-preserving result types. let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) @@ -331,24 +315,11 @@ impl ListViewArray { let sizes = PrimitiveArray::new(new_sizes.freeze(), Validity::NonNullable) .reinterpret_cast(size_ptype) .into_array(); - let elements = new_elements_builder.finish(); - debug_assert_eq!( - n_elements.as_(), - elements.len(), - "The accumulated elements somehow had the wrong length" - ); - - // SAFETY: - // - All offsets are sequential and non-overlapping (`n_elements` tracks running total). - // - Each `offset[i] + size[i]` equals `offset[i+1]` for all valid indices (including null - // lists). - // - All elements referenced by (offset, size) pairs exist within the new `elements` array. - // - The validity array is preserved from the original array unchanged - // - The array satisfies the zero-copy-to-list property by having sorted offsets, no gaps, - // and no overlaps. + // SAFETY: offsets are sequential and non-overlapping, all (offset, size) pairs reference + // valid elements, and validity is preserved from the original array. Ok(unsafe { - ListViewArray::new_unchecked(elements, offsets, sizes, self.validity()?) + ListViewArray::new_unchecked(elements, offsets, sizes, validity) .with_zero_copy_to_list(true) }) } @@ -414,6 +385,69 @@ impl ListViewArray { } } +struct RebuildRanges { + new_offsets: BufferMut, + new_sizes: BufferMut, + starts: BufferMut, + lengths: BufferMut, + elements_len: usize, +} + +fn rebuild_ranges( + offsets: &[O], + sizes: &[S], + validity_mask: &Mask, +) -> VortexResult> +where + NewOffset: IntegerPType, + O: IntegerPType, + S: IntegerPType, +{ + let len = offsets.len(); + let mut new_offsets = BufferMut::::with_capacity(len); + let mut new_sizes = BufferMut::::with_capacity(len); + let mut starts = BufferMut::::with_capacity(len); + let mut lengths = BufferMut::::with_capacity(len); + let mut n_elements = NewOffset::zero(); + let mut elements_len = 0usize; + + for (index, is_valid) in validity_mask.iter().enumerate() { + if !is_valid { + new_offsets.push(n_elements); + new_sizes.push(S::zero()); + starts.push(0); + lengths.push(0); + continue; + } + + let size = sizes[index]; + let start: usize = offsets[index].as_(); + let length: usize = size.as_(); + start.checked_add(length).ok_or_else(|| { + vortex_err!( + "ListView rebuild element range overflow for start {start} and length {length}" + ) + })?; + elements_len = elements_len + .checked_add(length) + .ok_or_else(|| vortex_err!("ListView rebuild elements length overflow"))?; + + new_offsets.push(n_elements); + new_sizes.push(size); + starts.push(start as u64); + lengths.push(length as u64); + n_elements += num_traits::cast(size).vortex_expect("Cast failed"); + } + + Ok(RebuildRanges { + new_offsets, + new_sizes, + starts, + lengths, + elements_len, + }) +} + #[cfg(test)] mod tests { use vortex_buffer::BitBuffer; @@ -426,6 +460,7 @@ mod tests { use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::listview::ListViewArrayExt; + use crate::arrays::listview::ListViewArraySlotsExt; use crate::assert_arrays_eq; use crate::dtype::Nullability; use crate::validity::Validity; @@ -511,6 +546,42 @@ mod tests { Ok(()) } + #[test] + fn test_rebuild_flatten_null_row_uses_valid_empty_range() -> VortexResult<()> { + let elements = PrimitiveArray::from_iter(vec![1i32, 2, 3, 4]).into_array(); + let offsets = PrimitiveArray::from_iter(vec![0u32, 1, 2]).into_array(); + let sizes = PrimitiveArray::from_iter(vec![1u32, 2, 2]).into_array(); + let validity = Validity::from_iter([true, false, true]); + + // SAFETY: all source ranges are valid, including the null row's non-empty range. + let listview = unsafe { ListViewArray::new_unchecked(elements, offsets, sizes, validity) }; + + let mut ctx = SESSION.create_execution_ctx(); + let flattened = listview.rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx)?; + + assert_eq!(flattened.offset_at(0), 0); + assert_eq!(flattened.size_at(0), 1); + assert_eq!(flattened.offset_at(1), 1); + assert_eq!(flattened.size_at(1), 0); + assert_eq!(flattened.offset_at(2), 1); + assert_eq!(flattened.size_at(2), 2); + assert!(flattened.validity()?.execute_is_valid(0, &mut ctx)?); + assert!(!flattened.validity()?.execute_is_valid(1, &mut ctx)?); + assert!(flattened.validity()?.execute_is_valid(2, &mut ctx)?); + + assert_arrays_eq!( + flattened.list_elements_at(0)?, + PrimitiveArray::from_iter([1i32]), + &mut ctx + ); + assert_arrays_eq!( + flattened.list_elements_at(2)?, + PrimitiveArray::from_iter([3i32, 4]), + &mut ctx + ); + Ok(()) + } + #[test] fn test_rebuild_trim_elements_basic() -> VortexResult<()> { // Test trimming both leading and trailing unused elements while preserving gaps in the @@ -561,6 +632,35 @@ mod tests { Ok(()) } + #[test] + fn test_rebuild_flatten_large_lists_with_piecewise_indices() -> VortexResult<()> { + let elements = PrimitiveArray::from_iter(0i32..300).into_array(); + let offsets = PrimitiveArray::from_iter(vec![10u32, 0]).into_array(); + let sizes = PrimitiveArray::from_iter(vec![150u32, 130]).into_array(); + let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable); + + let mut ctx = SESSION.create_execution_ctx(); + let flattened = listview.rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx)?; + + assert_eq!(flattened.elements().len(), 280); + assert_eq!(flattened.offset_at(0), 0); + assert_eq!(flattened.size_at(0), 150); + assert_eq!(flattened.offset_at(1), 150); + assert_eq!(flattened.size_at(1), 130); + + assert_arrays_eq!( + flattened.list_elements_at(0)?, + PrimitiveArray::from_iter(10i32..160), + &mut ctx + ); + assert_arrays_eq!( + flattened.list_elements_at(1)?, + PrimitiveArray::from_iter(0i32..130), + &mut ctx + ); + Ok(()) + } + #[test] fn test_rebuild_with_trailing_nulls_regression() -> VortexResult<()> { // Regression test for issue #5412 @@ -668,6 +768,7 @@ mod tests { /// the result offset/size dtypes preserve signedness (widened to >=32-bit for offsets). #[test] fn test_rebuild_preserves_signed_offset_and_size_types() -> VortexResult<()> { + use crate::arrays::listview::ListViewArraySlotsExt; use crate::dtype::PType; // Overlapping lists force an actual rebuild rather than the zero-copy fast path. diff --git a/vortex-array/src/arrays/listview/tests/density.rs b/vortex-array/src/arrays/listview/tests/density.rs index 8538a9fe20c..6bff870a1b6 100644 --- a/vortex-array/src/arrays/listview/tests/density.rs +++ b/vortex-array/src/arrays/listview/tests/density.rs @@ -18,6 +18,7 @@ use crate::IntoArray; use crate::VortexSessionExecute; use crate::arrays::ListViewArray; use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::tests::common::create_empty_elements_listview; use crate::expr::stats::Precision; use crate::expr::stats::Stat; diff --git a/vortex-array/src/arrays/listview/tests/filter.rs b/vortex-array/src/arrays/listview/tests/filter.rs index e8cc7636b30..d8a6cb4f4db 100644 --- a/vortex-array/src/arrays/listview/tests/filter.rs +++ b/vortex-array/src/arrays/listview/tests/filter.rs @@ -19,6 +19,7 @@ use crate::arrays::ConstantArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::assert_arrays_eq; use crate::compute::conformance::filter::test_filter_conformance; use crate::validity::Validity; diff --git a/vortex-array/src/arrays/listview/tests/nullability.rs b/vortex-array/src/arrays/listview/tests/nullability.rs index 4889b73de21..b32e2052435 100644 --- a/vortex-array/src/arrays/listview/tests/nullability.rs +++ b/vortex-array/src/arrays/listview/tests/nullability.rs @@ -13,6 +13,7 @@ use crate::arrays::BoolArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; diff --git a/vortex-array/src/arrays/listview/tests/operations.rs b/vortex-array/src/arrays/listview/tests/operations.rs index 51f03b1abeb..38ac286b2d1 100644 --- a/vortex-array/src/arrays/listview/tests/operations.rs +++ b/vortex-array/src/arrays/listview/tests/operations.rs @@ -22,6 +22,7 @@ use crate::arrays::ListView; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; use crate::compute::conformance::mask::test_mask_conformance; diff --git a/vortex-array/src/arrays/listview/tests/take.rs b/vortex-array/src/arrays/listview/tests/take.rs index 2af397e262f..6f74081a8c5 100644 --- a/vortex-array/src/arrays/listview/tests/take.rs +++ b/vortex-array/src/arrays/listview/tests/take.rs @@ -18,6 +18,7 @@ use crate::arrays::ConstantArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::assert_arrays_eq; use crate::compute::conformance::take::test_take_conformance; use crate::validity::Validity; diff --git a/vortex-array/src/arrays/listview/vtable/mod.rs b/vortex-array/src/arrays/listview/vtable/mod.rs index 546e77898f8..806e7a6563c 100644 --- a/vortex-array/src/arrays/listview/vtable/mod.rs +++ b/vortex-array/src/arrays/listview/vtable/mod.rs @@ -26,19 +26,16 @@ use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::VTable; use crate::array::with_empty_buffers; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::ListViewData; -use crate::arrays::listview::array::ELEMENTS_SLOT; -use crate::arrays::listview::array::NUM_SLOTS; -use crate::arrays::listview::array::OFFSETS_SLOT; -use crate::arrays::listview::array::SIZES_SLOT; -use crate::arrays::listview::array::SLOT_NAMES; +use crate::arrays::listview::ListViewSlots; use crate::arrays::listview::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; +use crate::match_each_list_builder; use crate::serde::ArrayChildren; use crate::validity::Validity; mod kernel; @@ -128,17 +125,18 @@ impl VTable for ListView { slots: &[Option], ) -> VortexResult<()> { vortex_ensure!( - slots.len() == NUM_SLOTS, - "ListViewArray expected {NUM_SLOTS} slots, found {}", + slots.len() == ListViewSlots::COUNT, + "ListViewArray expected {} slots, found {}", + ListViewSlots::COUNT, slots.len() ); - let elements = slots[ELEMENTS_SLOT] + let elements = slots[ListViewSlots::ELEMENTS] .as_ref() .vortex_expect("ListViewArray elements slot"); - let offsets = slots[OFFSETS_SLOT] + let offsets = slots[ListViewSlots::OFFSETS] .as_ref() .vortex_expect("ListViewArray offsets slot"); - let sizes = slots[SIZES_SLOT] + let sizes = slots[ListViewSlots::SIZES] .as_ref() .vortex_expect("ListViewArray sizes slot"); vortex_ensure!( @@ -219,19 +217,28 @@ impl VTable for ListView { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + ListViewSlots::NAMES[idx].to_string() } fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { Ok(ExecutionResult::done(array)) } + // The complexity comes from the expansion of `match_each_list_builder!`. + #[expect(clippy::cognitive_complexity)] fn append_to_builder( array: ArrayView<'_, Self>, builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - builder.append_listview_array(array, ctx) + match match_each_list_builder!(&mut *builder, |b| b.append_listview_array(array, ctx)) { + Some(result) => result, + None => vortex_bail!( + "cannot append a ListView array of dtype {} to a {} builder", + array.dtype(), + builder.dtype() + ), + } } fn reduce_parent( diff --git a/vortex-array/src/arrays/map/array.rs b/vortex-array/src/arrays/map/array.rs new file mode 100644 index 00000000000..1cdc362b606 --- /dev/null +++ b/vortex-array/src/arrays/map/array.rs @@ -0,0 +1,190 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hasher; +use std::sync::Arc; + +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::ArrayEq; +use crate::ArrayHash; +use crate::ArrayRef; +use crate::ArraySlots; +use crate::EqMode; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayParts; +use crate::array::TypedArrayRef; +use crate::array_slots; +use crate::arrays::ListView; +use crate::arrays::ListViewArray; +use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::map::Map; +use crate::dtype::DType; +use crate::dtype::MapDType; +use crate::validity::Validity; + +#[array_slots(Map)] +pub struct MapSlots { + /// The list-view storage of non-null `{key, value}` entry structs. + #[slot(0)] + pub entries: ArrayRef, +} + +/// Encoding-specific metadata for [`crate::arrays::MapArray`]. +/// +/// All map metadata is represented by the outer [`DType::Map`] and the entries child, so this +/// value is intentionally empty. +#[derive(Clone, Debug, Default)] +pub struct MapData; + +impl Display for MapData { + fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result { + Ok(()) + } +} + +impl ArrayEq for MapData { + fn array_eq(&self, _other: &Self, _accuracy: EqMode) -> bool { + true + } +} + +impl ArrayHash for MapData { + fn array_hash(&self, _state: &mut H, _accuracy: EqMode) {} +} + +impl MapData { + pub(crate) fn make_slots(entries: ArrayRef) -> ArraySlots { + MapSlots { entries }.into_slots() + } +} + +/// The logical and physical inputs used to construct a [`crate::arrays::MapArray`]. +pub struct MapDataParts { + /// The key/value type and sortedness assertion for the map. + pub map_dtype: MapDType, + /// The physical list-view storage of `{key, value}` entry structs. + pub entries: ListViewArray, +} + +/// Accessors for the canonical map representation. +pub trait MapArrayExt: MapArraySlotsExt { + /// Returns the entry structs for one map row. + fn entries_at(&self, index: usize) -> VortexResult { + self.entries().as_::().list_elements_at(index) + } + + /// Returns the number of entries in one map row. + fn entry_count_at(&self, index: usize) -> usize { + self.entries().as_::().size_at(index) + } + + /// Returns the outer map validity delegated from the entries list-view. + fn map_validity(&self) -> Validity { + self.entries().as_::().listview_validity() + } + + /// Returns this map's key/value type information. + fn map_dtype(&self) -> &MapDType { + self.as_ref() + .dtype() + .as_map_opt() + .vortex_expect("MapArray requires a map dtype") + } + + /// Returns whether producers assert sorted keys within each map value. + fn keys_sorted(&self) -> bool { + self.map_dtype().keys_sorted() + } +} +impl> MapArrayExt for T {} + +impl Array { + /// Creates a canonical map array from its map dtype and list-view entry storage. + /// + /// # Panics + /// + /// Panics if `entries` is not a list of the map dtype's non-nullable `{key, value}` entry + /// struct with matching outer nullability. + pub fn new(map_dtype: MapDType, entries: ListViewArray) -> Self { + Self::try_new(map_dtype, entries).vortex_expect("MapArray construction failed") + } + + /// Constructs a canonical map array from its map dtype and list-view entry storage. + /// + /// # Errors + /// + /// Returns an error when the entry child is not `ListView>`, has a + /// different outer nullability, or has a different length than the outer map array. + pub fn try_new(map_dtype: MapDType, entries: ListViewArray) -> VortexResult { + let nullability = entries.nullability(); + let dtype = DType::Map(map_dtype, nullability); + let len = entries.len(); + let slots = MapData::make_slots(entries.into_array()); + let parts = ArrayParts::new(Map, dtype, len, MapData).with_slots(slots); + Self::try_from_parts(parts) + } + + /// Creates a canonical map array without validating its entry storage. + /// + /// # Safety + /// + /// The caller must ensure that `entries` has dtype + /// `List(Struct { key, value }, entries.nullability())`, where the struct exactly matches + /// `map_dtype.entries_dtype()`. + pub unsafe fn new_unchecked(map_dtype: MapDType, entries: ListViewArray) -> Self { + let nullability = entries.nullability(); + let dtype = DType::Map(map_dtype, nullability); + let len = entries.len(); + let slots = MapData::make_slots(entries.into_array()); + let parts = ArrayParts::new(Map, dtype, len, MapData).with_slots(slots); + unsafe { Self::from_parts_unchecked(parts) } + } + + /// Decomposes this map array into its logical dtype and physical entries child. + pub fn into_data_parts(self) -> MapDataParts { + let map_dtype = self + .dtype() + .as_map_opt() + .vortex_expect("MapArray requires a map dtype") + .clone(); + let entries = self.entries().clone().downcast::(); + MapDataParts { map_dtype, entries } + } +} + +fn expected_entries_dtype(map_dtype: &MapDType, nullability: crate::dtype::Nullability) -> DType { + DType::List(Arc::new(map_dtype.entries_dtype()), nullability) +} + +pub(super) fn validate_entries( + map_dtype: &MapDType, + nullability: crate::dtype::Nullability, + len: usize, + entries: &ArrayRef, +) -> VortexResult<()> { + vortex_ensure!( + entries.is::(), + "MapArray entries must use vortex.listview encoding, got {}", + entries.encoding_id() + ); + vortex_ensure!( + entries.len() == len, + "MapArray entries length {} does not match outer length {len}", + entries.len() + ); + + let expected_dtype = expected_entries_dtype(map_dtype, nullability); + vortex_ensure!( + entries.dtype() == &expected_dtype, + "MapArray entries dtype {} does not match expected {expected_dtype}", + entries.dtype() + ); + + Ok(()) +} diff --git a/vortex-array/src/arrays/map/compute/cast.rs b/vortex-array/src/arrays/map/compute/cast.rs new file mode 100644 index 00000000000..bfbab6e17e9 --- /dev/null +++ b/vortex-array/src/arrays/map/compute/cast.rs @@ -0,0 +1,85 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; + +use crate::ArrayRef; +use crate::array::ArrayView; +use crate::arrays::ListView; +use crate::arrays::map::Map; +use crate::arrays::map::MapArrayExt; +use crate::arrays::map::MapArraySlotsExt; +use crate::arrays::map::compute::rebuild_map_from_array; +use crate::dtype::DType; +use crate::dtype::MapDType; +use crate::executor::ExecutionCtx; +use crate::scalar_fn::fns::cast::CastKernel; +use crate::scalar_fn::fns::cast::CastReduce; + +fn prepare_map_cast_target( + array: ArrayView<'_, Map>, + dtype: &DType, +) -> VortexResult> { + let Some(target_map_dtype) = dtype.as_map_opt() else { + return Ok(None); + }; + + if target_map_dtype.keys_sorted() && !array.keys_sorted() { + vortex_bail!( + "Cannot cast {} to {dtype}: source does not assert sorted map keys", + array.dtype() + ); + } + + let target_entries_dtype = DType::List( + Arc::new(target_map_dtype.entries_dtype()), + dtype.nullability(), + ); + Ok(Some((target_map_dtype.clone(), target_entries_dtype))) +} + +impl CastReduce for Map { + fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { + let Some((target_map_dtype, target_entries_dtype)) = prepare_map_cast_target(array, dtype)? + else { + return Ok(None); + }; + + let Some(entries) = ::cast( + array.entries().as_::(), + &target_entries_dtype, + )? + else { + return Ok(None); + }; + + rebuild_map_from_array(target_map_dtype, entries).map(Some) + } +} + +impl CastKernel for Map { + fn cast( + array: ArrayView<'_, Self>, + dtype: &DType, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some((target_map_dtype, target_entries_dtype)) = prepare_map_cast_target(array, dtype)? + else { + return Ok(None); + }; + + let Some(entries) = ::cast( + array.entries().as_::(), + &target_entries_dtype, + ctx, + )? + else { + return Ok(None); + }; + + rebuild_map_from_array(target_map_dtype, entries).map(Some) + } +} diff --git a/vortex-array/src/arrays/map/compute/filter.rs b/vortex-array/src/arrays/map/compute/filter.rs new file mode 100644 index 00000000000..53a862bb35a --- /dev/null +++ b/vortex-array/src/arrays/map/compute/filter.rs @@ -0,0 +1,53 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::ListView; +use crate::arrays::ListViewArray; +use crate::arrays::MapArray; +use crate::arrays::filter::FilterKernel; +use crate::arrays::filter::FilterReduce; +use crate::arrays::listview::ListViewArraySlotsExt; +use crate::arrays::map::Map; +use crate::arrays::map::MapArrayExt; +use crate::arrays::map::MapArraySlotsExt; +use crate::executor::ExecutionCtx; + +impl FilterReduce for Map { + fn filter(array: ArrayView<'_, Self>, mask: &Mask) -> VortexResult> { + let entries = array.entries().as_::(); + + // SAFETY: filtering row metadata keeps offsets and sizes paired, preserves the original + // elements, and filters validity to the same output length. The zero-copy-to-list flag is + // not carried over: dropping a non-empty row leaves a gap in the referenced elements. + let filtered_entries = unsafe { + ListViewArray::new_unchecked( + entries.elements().clone(), + entries.offsets().filter(mask.clone())?, + entries.sizes().filter(mask.clone())?, + entries.validity()?.filter(mask)?, + ) + }; + + { + let map_dtype = array.map_dtype().clone(); + MapArray::try_new(map_dtype, filtered_entries).map(IntoArray::into_array) + } + .map(Some) + } +} + +impl FilterKernel for Map { + fn filter( + array: ArrayView<'_, Self>, + mask: &Mask, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + ::filter(array, mask) + } +} diff --git a/vortex-array/src/arrays/map/compute/mask.rs b/vortex-array/src/arrays/map/compute/mask.rs new file mode 100644 index 00000000000..41703a2a3ea --- /dev/null +++ b/vortex-array/src/arrays/map/compute/mask.rs @@ -0,0 +1,37 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::array::ArrayView; +use crate::arrays::ListView; +use crate::arrays::map::Map; +use crate::arrays::map::MapArrayExt; +use crate::arrays::map::MapArraySlotsExt; +use crate::arrays::map::compute::rebuild_map_from_array; +use crate::executor::ExecutionCtx; +use crate::scalar_fn::fns::mask::MaskKernel; +use crate::scalar_fn::fns::mask::MaskReduce; + +impl MaskReduce for Map { + fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult> { + let Some(entries) = + ::mask(array.entries().as_::(), mask)? + else { + return Ok(None); + }; + + rebuild_map_from_array(array.map_dtype().clone(), entries).map(Some) + } +} + +impl MaskKernel for Map { + fn mask( + array: ArrayView<'_, Self>, + mask: &ArrayRef, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + ::mask(array, mask) + } +} diff --git a/vortex-array/src/arrays/map/compute/mod.rs b/vortex-array/src/arrays/map/compute/mod.rs new file mode 100644 index 00000000000..ea0be66f0e4 --- /dev/null +++ b/vortex-array/src/arrays/map/compute/mod.rs @@ -0,0 +1,31 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Map array compute-kernel namespace. + +mod cast; +mod filter; +mod mask; +pub(crate) mod rules; +mod slice; +mod take; + +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::arrays::ListView; +use crate::arrays::MapArray; +use crate::dtype::MapDType; + +fn rebuild_map_from_array(map_dtype: MapDType, entries: ArrayRef) -> VortexResult { + let map_entries = entries.try_downcast::().map_err(|arr| { + vortex_err!( + "Map entries operation expected vortex.listview/ListView, got {}", + arr.encoding_id() + ) + })?; + + MapArray::try_new(map_dtype, map_entries).map(IntoArray::into_array) +} diff --git a/vortex-array/src/arrays/map/compute/rules.rs b/vortex-array/src/arrays/map/compute/rules.rs new file mode 100644 index 00000000000..27622ff1c5f --- /dev/null +++ b/vortex-array/src/arrays/map/compute/rules.rs @@ -0,0 +1,18 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use crate::arrays::Map; +use crate::arrays::dict::TakeReduceAdaptor; +use crate::arrays::filter::FilterReduceAdaptor; +use crate::arrays::slice::SliceReduceAdaptor; +use crate::optimizer::rules::ParentRuleSet; +use crate::scalar_fn::fns::cast::CastReduceAdaptor; +use crate::scalar_fn::fns::mask::MaskReduceAdaptor; + +pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ + ParentRuleSet::lift(&FilterReduceAdaptor(Map)), + ParentRuleSet::lift(&CastReduceAdaptor(Map)), + ParentRuleSet::lift(&MaskReduceAdaptor(Map)), + ParentRuleSet::lift(&SliceReduceAdaptor(Map)), + ParentRuleSet::lift(&TakeReduceAdaptor(Map)), +]); diff --git a/vortex-array/src/arrays/map/compute/slice.rs b/vortex-array/src/arrays/map/compute/slice.rs new file mode 100644 index 00000000000..463dee74f60 --- /dev/null +++ b/vortex-array/src/arrays/map/compute/slice.rs @@ -0,0 +1,39 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::array::ArrayView; +use crate::arrays::ListView; +use crate::arrays::map::Map; +use crate::arrays::map::MapArrayExt; +use crate::arrays::map::MapArraySlotsExt; +use crate::arrays::map::compute::rebuild_map_from_array; +use crate::arrays::slice::SliceKernel; +use crate::arrays::slice::SliceReduce; +use crate::executor::ExecutionCtx; + +impl SliceReduce for Map { + fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { + let Some(sliced_entries) = + ::slice(array.entries().as_::(), range)? + else { + return Ok(None); + }; + + rebuild_map_from_array(array.map_dtype().clone(), sliced_entries).map(Some) + } +} + +impl SliceKernel for Map { + fn slice( + array: ArrayView<'_, Self>, + range: Range, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + ::slice(array, range) + } +} diff --git a/vortex-array/src/arrays/map/compute/take.rs b/vortex-array/src/arrays/map/compute/take.rs new file mode 100644 index 00000000000..58be62e9645 --- /dev/null +++ b/vortex-array/src/arrays/map/compute/take.rs @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::array::ArrayView; +use crate::arrays::ListView; +use crate::arrays::dict::TakeExecute; +use crate::arrays::dict::TakeReduce; +use crate::arrays::map::Map; +use crate::arrays::map::MapArrayExt; +use crate::arrays::map::MapArraySlotsExt; +use crate::arrays::map::compute::rebuild_map_from_array; +use crate::executor::ExecutionCtx; + +impl TakeReduce for Map { + fn take(array: ArrayView<'_, Self>, indices: &ArrayRef) -> VortexResult> { + let Some(entries) = + ::take(array.entries().as_::(), indices)? + else { + return Ok(None); + }; + + rebuild_map_from_array(array.map_dtype().clone(), entries).map(Some) + } +} + +impl TakeExecute for Map { + fn take( + array: ArrayView<'_, Self>, + indices: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some(entries) = + ::take(array.entries().as_::(), indices, ctx)? + else { + return Ok(None); + }; + + rebuild_map_from_array(array.map_dtype().clone(), entries).map(Some) + } +} diff --git a/vortex-array/src/arrays/map/mod.rs b/vortex-array/src/arrays/map/mod.rs new file mode 100644 index 00000000000..ef1a04166bc --- /dev/null +++ b/vortex-array/src/arrays/map/mod.rs @@ -0,0 +1,25 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Canonical map arrays backed by [`ListView`](crate::arrays::ListView) entry storage. + +mod array; +pub use array::MapArrayExt; +pub use array::MapArraySlotsExt; +pub use array::MapData; +pub use array::MapDataParts; +pub use array::MapSlots; +pub use array::MapSlotsView; + +pub(crate) mod compute; + +mod vtable; +pub use vtable::Map; +pub use vtable::MapArray; + +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/arrays/map/tests.rs b/vortex-array/src/arrays/map/tests.rs new file mode 100644 index 00000000000..e1df51ce6d3 --- /dev/null +++ b/vortex-array/src/arrays/map/tests.rs @@ -0,0 +1,733 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::ByteBufferMut; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_session::registry::ReadContext; + +use crate::Array; +use crate::ArrayContext; +use crate::ArrayParts; +use crate::ArrayVTable; +use crate::Canonical; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::BoolArray; +use crate::arrays::ChunkedArray; +use crate::arrays::ConstantArray; +use crate::arrays::FilterArray; +use crate::arrays::ListView; +use crate::arrays::ListViewArray; +use crate::arrays::Map; +use crate::arrays::MapArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::map::MapArrayExt; +use crate::arrays::map::MapArraySlotsExt; +use crate::arrays::map::MapData; +use crate::arrays::map::MapDataParts; +use crate::arrays::map::MapSlots; +use crate::assert_arrays_eq; +use crate::builders::ArrayBuilder; +use crate::builders::MapBuilder; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::MapDType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::scalar::Scalar; +use crate::serde::SerializeOptions; +use crate::serde::SerializedArray; +use crate::session::ArraySessionExt; +use crate::validity::Validity; + +fn map_dtype() -> VortexResult { + MapDType::try_new( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + ) +} + +fn key(value: i32) -> Scalar { + Scalar::primitive(value, Nullability::NonNullable) +} + +fn value(value: Option<&str>) -> Scalar { + match value { + Some(value) => Scalar::utf8(value, Nullability::Nullable), + None => Scalar::null(DType::Utf8(Nullability::Nullable)), + } +} + +fn sample_scalar(dtype: DType) -> VortexResult { + Scalar::try_map( + dtype, + [ + (key(2), value(Some("two"))), + (key(1), value(None)), + (key(2), value(Some("duplicate"))), + ], + ) +} + +fn map_array_from_rows( + map_dtype: MapDType, + nullability: Nullability, + rows: impl IntoIterator>>, +) -> VortexResult { + let rows = rows.into_iter().collect::>(); + let dtype = DType::Map(map_dtype.clone(), nullability); + let mut builder = MapBuilder::::with_capacity_in( + map_dtype, + nullability, + rows.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + + for row in rows { + let scalar = match row { + Some(entries) => Scalar::try_map(dtype.clone(), entries)?, + None => Scalar::null(dtype.clone()), + }; + builder.append_scalar(&scalar)?; + } + + Ok(builder.finish_into_map()) +} + +fn sample_array() -> VortexResult { + let map_dtype = map_dtype()?; + map_array_from_rows( + map_dtype, + Nullability::Nullable, + [ + Some(vec![ + (key(2), value(Some("two"))), + (key(1), value(None)), + (key(2), value(Some("duplicate"))), + ]), + Some(vec![]), + None, + ], + ) +} + +#[test] +fn constructs_map_with_listview_entries() -> VortexResult<()> { + let MapDataParts { map_dtype, entries } = sample_array()?.into_data_parts(); + let array = MapArray::try_new(map_dtype.clone(), entries)?; + + assert_eq!( + array.dtype(), + &DType::Map(map_dtype.clone(), Nullability::Nullable) + ); + assert!(array.keys_sorted()); + assert_eq!(array.entry_count_at(0), 3); + assert_eq!(array.entry_count_at(1), 0); + assert_eq!(array.entries_at(0)?.dtype(), &map_dtype.entries_dtype()); + + let mut ctx = array_session().create_execution_ctx(); + assert_eq!( + array + .map_validity() + .execute_mask(array.len(), &mut ctx)? + .true_count(), + 2 + ); + + Ok(()) +} + +#[test] +fn accepts_duplicate_and_unsorted_keys() -> VortexResult<()> { + let array = sample_array()?; + let mut ctx = array_session().create_execution_ctx(); + + assert_eq!( + array.execute_scalar(0, &mut ctx)?, + sample_scalar(array.dtype().clone())? + ); + + Ok(()) +} + +#[test] +fn rejects_malformed_entry_storage() -> VortexResult<()> { + let map_dtype = map_dtype()?; + let offsets = PrimitiveArray::from_iter([0u64]).into_array(); + let sizes = PrimitiveArray::from_iter([1u64]).into_array(); + let non_struct_entries = ListViewArray::try_new( + PrimitiveArray::from_iter([1i32]).into_array(), + offsets, + sizes, + Validity::NonNullable, + )?; + assert!(MapArray::try_new(map_dtype.clone(), non_struct_entries).is_err()); + + let nullable_entry_struct = + ConstantArray::new(Scalar::null(map_dtype.entries_dtype().as_nullable()), 1).into_array(); + let null_entries = ListViewArray::try_new( + nullable_entry_struct, + PrimitiveArray::from_iter([0u64]).into_array(), + PrimitiveArray::from_iter([1u64]).into_array(), + Validity::NonNullable, + )?; + assert!(MapArray::try_new(map_dtype.clone(), null_entries).is_err()); + + let MapDataParts { entries, .. } = sample_array()?.into_data_parts(); + let parts = ArrayParts::new( + Map, + DType::Map(map_dtype, Nullability::NonNullable), + entries.len(), + MapData, + ) + .with_slots( + MapSlots { + entries: entries.into_array(), + } + .into_slots(), + ); + assert!(Array::::try_from_parts(parts).is_err()); + + assert!( + MapDType::try_new( + DType::Primitive(PType::I32, Nullability::Nullable), + DType::Utf8(Nullability::Nullable), + false, + ) + .is_err() + ); + + Ok(()) +} + +#[test] +fn scalar_access_preserves_null_and_empty_maps() -> VortexResult<()> { + let array = sample_array()?; + let mut ctx = array_session().create_execution_ctx(); + + assert_eq!( + array.execute_scalar(0, &mut ctx)?, + sample_scalar(array.dtype().clone())? + ); + assert!(array.execute_scalar(1, &mut ctx)?.as_map().is_empty()); + assert!(array.execute_scalar(2, &mut ctx)?.is_null()); + + Ok(()) +} + +#[test] +fn scalar_access_preserves_variable_entry_counts_and_utf8_pairs() -> VortexResult<()> { + let map_dtype = MapDType::try_new( + DType::Utf8(Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + false, + )?; + let dtype = DType::Map(map_dtype.clone(), Nullability::Nullable); + let mut builder = MapBuilder::::with_capacity_in( + map_dtype, + Nullability::Nullable, + 4, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + let rows = [ + Some(vec![ + ( + Scalar::utf8("alpha", Nullability::NonNullable), + Scalar::utf8("one", Nullability::Nullable), + ), + ( + Scalar::utf8("longer-key", Nullability::NonNullable), + Scalar::utf8("longer-value", Nullability::Nullable), + ), + ]), + Some(vec![( + Scalar::utf8("z", Nullability::NonNullable), + Scalar::utf8("last", Nullability::Nullable), + )]), + None, + Some(vec![ + ( + Scalar::utf8("repeated", Nullability::NonNullable), + Scalar::utf8("first", Nullability::Nullable), + ), + ( + Scalar::utf8("repeated", Nullability::NonNullable), + Scalar::null(DType::Utf8(Nullability::Nullable)), + ), + ( + Scalar::utf8("tail", Nullability::NonNullable), + Scalar::utf8("", Nullability::Nullable), + ), + ]), + ]; + let expected = rows + .into_iter() + .map(|row| { + row.map_or_else( + || Ok(Scalar::null(dtype.clone())), + |entries| Scalar::try_map(dtype.clone(), entries), + ) + }) + .collect::>>()?; + + for scalar in &expected { + builder.append_scalar(scalar)?; + } + let array = builder.finish_into_map(); + let mut ctx = array_session().create_execution_ctx(); + + assert_eq!(array.entry_count_at(0), 2); + assert_eq!(array.entry_count_at(1), 1); + assert_eq!(array.entry_count_at(2), 0); + assert_eq!(array.entry_count_at(3), 3); + for (index, expected) in expected.into_iter().enumerate() { + assert_eq!(array.execute_scalar(index, &mut ctx)?, expected); + } + + Ok(()) +} + +#[test] +fn slice_preserves_map_rows() -> VortexResult<()> { + let source = sample_array()?.into_array(); + let expected = map_array_from_rows( + map_dtype()?, + Nullability::Nullable, + [ + Some(vec![ + (key(2), value(Some("two"))), + (key(1), value(None)), + (key(2), value(Some("duplicate"))), + ]), + Some(vec![]), + ], + )?; + let sliced = source.slice(0..2)?; + let mut ctx = array_session().create_execution_ctx(); + + assert!(sliced.is::()); + assert_arrays_eq!(sliced, expected, &mut ctx); + + Ok(()) +} + +#[test] +fn filter_handles_all_none_and_mixed_maps() -> VortexResult<()> { + let source = sample_array()?.into_array(); + let map_dtype = map_dtype()?; + let mut ctx = array_session().create_execution_ctx(); + + let all = source.filter(Mask::from_iter([true, true, true]))?; + assert!(all.is::()); + assert_arrays_eq!(all, sample_array()?, &mut ctx); + + let none = source.filter(Mask::from_iter([false, false, false]))?; + let expected_none = map_array_from_rows( + map_dtype.clone(), + Nullability::Nullable, + Vec::>>::new(), + )?; + assert!(none.is::()); + assert_arrays_eq!(none, expected_none, &mut ctx); + + let mixed = source.filter(Mask::from_iter([true, false, true]))?; + let expected_mixed = map_array_from_rows( + map_dtype, + Nullability::Nullable, + [ + Some(vec![ + (key(2), value(Some("two"))), + (key(1), value(None)), + (key(2), value(Some("duplicate"))), + ]), + None, + ], + )?; + assert!(mixed.is::()); + assert_arrays_eq!(mixed, expected_mixed, &mut ctx); + + Ok(()) +} + +#[test] +fn filter_dropping_nonempty_middle_row_clears_zero_copy_flag() -> VortexResult<()> { + let source = map_array_from_rows( + map_dtype()?, + Nullability::Nullable, + [ + Some(vec![ + (key(1), value(Some("one"))), + (key(2), value(Some("two"))), + ]), + Some(vec![(key(3), value(Some("three")))]), + Some(vec![(key(4), value(Some("four")))]), + ], + )? + .into_array(); + let expected = map_array_from_rows( + map_dtype()?, + Nullability::Nullable, + [ + Some(vec![ + (key(1), value(Some("one"))), + (key(2), value(Some("two"))), + ]), + Some(vec![(key(4), value(Some("four")))]), + ], + )?; + let mask = Mask::from_iter([true, false, true]); + let mut ctx = array_session().create_execution_ctx(); + + let reduced = source.filter(mask.clone())?; + assert!(reduced.is::()); + // The dropped middle row leaves a gap in the referenced entry elements, so the filtered + // entries must not claim zero-copy-to-list convertibility. + assert!( + !reduced + .as_::() + .entries() + .as_::() + .into_owned() + .is_zero_copy_to_list() + ); + assert_arrays_eq!(reduced, expected, &mut ctx); + + let executed = FilterArray::new(source, mask) + .into_array() + .execute::(&mut ctx)?; + assert!( + !executed + .entries() + .as_::() + .into_owned() + .is_zero_copy_to_list() + ); + assert_arrays_eq!(executed, expected, &mut ctx); + + Ok(()) +} + +#[test] +fn take_supports_reordered_duplicate_and_nullable_indices() -> VortexResult<()> { + let source = sample_array()?.into_array(); + let map_dtype = map_dtype()?; + let mut ctx = array_session().create_execution_ctx(); + + let taken = source.take(PrimitiveArray::from_iter([2u64, 0, 0]).into_array())?; + let expected_taken = map_array_from_rows( + map_dtype.clone(), + Nullability::Nullable, + [ + None, + Some(vec![ + (key(2), value(Some("two"))), + (key(1), value(None)), + (key(2), value(Some("duplicate"))), + ]), + Some(vec![ + (key(2), value(Some("two"))), + (key(1), value(None)), + (key(2), value(Some("duplicate"))), + ]), + ], + )?; + assert!(taken.is::()); + assert_arrays_eq!(taken, expected_taken, &mut ctx); + + let nonnullable_source = map_array_from_rows( + map_dtype.clone(), + Nullability::NonNullable, + [ + Some(vec![(key(1), value(Some("one")))]), + Some(vec![]), + Some(vec![(key(3), value(None))]), + ], + )? + .into_array(); + let nullable_taken = nonnullable_source + .take(PrimitiveArray::from_option_iter([Some(1u64), None, Some(0)]).into_array())?; + let expected_nullable_taken = map_array_from_rows( + map_dtype, + Nullability::Nullable, + [Some(vec![]), None, Some(vec![(key(1), value(Some("one")))])], + )?; + assert!(nullable_taken.is::()); + assert_eq!(nullable_taken.dtype().nullability(), Nullability::Nullable); + assert_arrays_eq!(nullable_taken, expected_nullable_taken, &mut ctx); + + Ok(()) +} + +#[test] +fn mask_combines_with_existing_map_validity() -> VortexResult<()> { + let source = sample_array()?.into_array(); + let mask = BoolArray::from_iter([true, false, true]).into_array(); + let masked = source.mask(mask)?; + let expected = map_array_from_rows( + map_dtype()?, + Nullability::Nullable, + [ + Some(vec![ + (key(2), value(Some("two"))), + (key(1), value(None)), + (key(2), value(Some("duplicate"))), + ]), + None, + None, + ], + )?; + let mut ctx = array_session().create_execution_ctx(); + + assert!(masked.is::()); + assert_arrays_eq!(masked, expected, &mut ctx); + + Ok(()) +} + +#[test] +fn cast_widens_map_key_value_and_outer_nullability() -> VortexResult<()> { + let source_map_dtype = MapDType::try_new( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::NonNullable), + true, + )?; + let source = map_array_from_rows( + source_map_dtype, + Nullability::NonNullable, + [ + Some(vec![ + ( + Scalar::primitive(1i32, Nullability::NonNullable), + Scalar::utf8("one", Nullability::NonNullable), + ), + ( + Scalar::primitive(2i32, Nullability::NonNullable), + Scalar::utf8("two", Nullability::NonNullable), + ), + ]), + Some(vec![]), + ], + )? + .into_array(); + let target_dtype = DType::map( + DType::Primitive(PType::I64, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + Nullability::Nullable, + )?; + let target_map_dtype = target_dtype + .as_map_opt() + .vortex_expect("target dtype is map") + .clone(); + let cast = source.cast(target_dtype)?; + let expected = map_array_from_rows( + target_map_dtype, + Nullability::Nullable, + [ + Some(vec![ + ( + Scalar::primitive(1i64, Nullability::NonNullable), + Scalar::utf8("one", Nullability::Nullable), + ), + ( + Scalar::primitive(2i64, Nullability::NonNullable), + Scalar::utf8("two", Nullability::Nullable), + ), + ]), + Some(vec![]), + ], + )?; + let mut ctx = array_session().create_execution_ctx(); + + assert!(cast.is::()); + assert_arrays_eq!(cast, expected, &mut ctx); + + Ok(()) +} + +#[test] +fn cast_can_drop_but_not_create_sortedness_assertion() -> VortexResult<()> { + let sorted_source = sample_array()?.into_array(); + let unsorted_target = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + false, + Nullability::Nullable, + )?; + let cast = sorted_source.cast(unsorted_target)?; + let expected = map_array_from_rows( + MapDType::try_new( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + false, + )?, + Nullability::Nullable, + [ + Some(vec![ + (key(2), value(Some("two"))), + (key(1), value(None)), + (key(2), value(Some("duplicate"))), + ]), + Some(vec![]), + None, + ], + )?; + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(cast, expected, &mut ctx); + + let unsorted_map_dtype = MapDType::try_new( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + false, + )?; + let unsorted_source = map_array_from_rows( + unsorted_map_dtype, + Nullability::Nullable, + [Some(vec![(key(1), value(Some("one")))])], + )? + .into_array(); + let sorted_target = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + Nullability::Nullable, + )?; + assert!(unsorted_source.cast(sorted_target).is_err()); + + Ok(()) +} + +#[test] +fn null_map_cast_cannot_create_sortedness_assertion() -> VortexResult<()> { + let unsorted_map_dtype = MapDType::try_new( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + false, + )?; + let unsorted_dtype = DType::Map(unsorted_map_dtype.clone(), Nullability::Nullable); + let sorted_dtype = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + Nullability::Nullable, + )?; + let scalar = Scalar::null(unsorted_dtype); + + assert!(scalar.cast(&sorted_dtype).is_err()); + let constant_cast = ConstantArray::new(scalar, 2) + .into_array() + .cast(sorted_dtype.clone())?; + let mut ctx = array_session().create_execution_ctx(); + assert!(constant_cast.execute::(&mut ctx).is_err()); + + let all_null = + map_array_from_rows(unsorted_map_dtype, Nullability::Nullable, [None, None])?.into_array(); + assert!(all_null.cast(sorted_dtype).is_err()); + + Ok(()) +} + +#[test] +fn filter_preserves_duplicate_map_keys() -> VortexResult<()> { + let source = sample_array()?.into_array(); + let filtered = source.filter(Mask::from_iter([true, false, false]))?; + let expected = map_array_from_rows( + map_dtype()?, + Nullability::Nullable, + [Some(vec![ + (key(2), value(Some("two"))), + (key(1), value(None)), + (key(2), value(Some("duplicate"))), + ])], + )?; + let mut ctx = array_session().create_execution_ctx(); + + assert_arrays_eq!(filtered, expected, &mut ctx); + + Ok(()) +} + +#[test] +fn builder_appends_existing_map_arrays() -> VortexResult<()> { + let source = sample_array()?; + let mut builder = MapBuilder::::with_capacity_in( + source.map_dtype().clone(), + source.dtype().nullability(), + 0, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + let mut ctx = array_session().create_execution_ctx(); + builder.append_map_array(source.as_view(), &mut ctx)?; + builder.append_map_array(source.as_view(), &mut ctx)?; + let array = builder.finish_into_map(); + + assert_eq!(array.len(), 6); + assert_eq!( + array.execute_scalar(0, &mut ctx)?, + source.execute_scalar(0, &mut ctx)? + ); + assert!(array.execute_scalar(2, &mut ctx)?.is_null()); + assert_eq!( + array.execute_scalar(3, &mut ctx)?, + source.execute_scalar(0, &mut ctx)? + ); + + Ok(()) +} + +#[test] +fn canonicalizes_empty_constant_and_chunked_maps() -> VortexResult<()> { + let map_dtype = map_dtype()?; + let dtype = DType::Map(map_dtype, Nullability::Nullable); + let empty = Canonical::empty(&dtype); + assert!(empty.as_map().is_empty()); + + let mut ctx = array_session().create_execution_ctx(); + let constant = ConstantArray::new(sample_scalar(dtype.clone())?, 2) + .into_array() + .execute::(&mut ctx)?; + assert_eq!(constant.as_map().len(), 2); + + let first = sample_array()?.into_array(); + let second = sample_array()?.into_array(); + let chunked = ChunkedArray::try_new(vec![first, second], dtype)?.into_array(); + let canonical = chunked.execute::(&mut ctx)?; + assert_eq!(canonical.as_map().len(), 6); + + Ok(()) +} + +#[test] +fn serde_roundtrip_uses_registered_map_vtable() -> VortexResult<()> { + let session = array_session(); + assert!(session.arrays().registry().contains_key(&Map.id())); + + let array = sample_array()?.into_array(); + let dtype = array.dtype().clone(); + let len = array.len(); + let array_ctx = ArrayContext::empty(); + let serialized = array.serialize(&array_ctx, &session, &SerializeOptions::default())?; + let mut concat = ByteBufferMut::empty(); + for buffer in serialized { + concat.extend_from_slice(buffer.as_ref()); + } + + let serialized = SerializedArray::try_from(concat.freeze())?; + let decoded = + serialized.decode(&dtype, len, &ReadContext::new(array_ctx.to_ids()), &session)?; + assert!(decoded.is::()); + + let mut ctx = session.create_execution_ctx(); + for index in 0..len { + assert_eq!( + decoded.execute_scalar(index, &mut ctx)?, + array.execute_scalar(index, &mut ctx)? + ); + } + + Ok(()) +} diff --git a/vortex-array/src/arrays/map/vtable/kernel.rs b/vortex-array/src/arrays/map/vtable/kernel.rs new file mode 100644 index 00000000000..860496834d6 --- /dev/null +++ b/vortex-array/src/arrays/map/vtable/kernel.rs @@ -0,0 +1,28 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_session::VortexSession; + +use crate::ArrayVTable; +use crate::arrays::Dict; +use crate::arrays::Filter; +use crate::arrays::Map; +use crate::arrays::Slice; +use crate::arrays::dict::TakeExecuteAdaptor; +use crate::arrays::filter::FilterExecuteAdaptor; +use crate::arrays::slice::SliceExecuteAdaptor; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::cast::Cast; +use crate::scalar_fn::fns::cast::CastExecuteAdaptor; +use crate::scalar_fn::fns::mask::Mask; +use crate::scalar_fn::fns::mask::MaskExecuteAdaptor; + +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Cast.id(), Map, CastExecuteAdaptor(Map)); + kernels.register_execute_parent_kernel(Dict.id(), Map, TakeExecuteAdaptor(Map)); + kernels.register_execute_parent_kernel(Filter.id(), Map, FilterExecuteAdaptor(Map)); + kernels.register_execute_parent_kernel(Mask.id(), Map, MaskExecuteAdaptor(Map)); + kernels.register_execute_parent_kernel(Slice.id(), Map, SliceExecuteAdaptor(Map)); +} diff --git a/vortex-array/src/arrays/map/vtable/mod.rs b/vortex-array/src/arrays/map/vtable/mod.rs new file mode 100644 index 00000000000..f11e3e72a89 --- /dev/null +++ b/vortex-array/src/arrays/map/vtable/mod.rs @@ -0,0 +1,180 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ArrayParts; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::ExecutionResult; +use crate::array::Array; +use crate::array::ArrayId; +use crate::array::ArrayView; +use crate::array::VTable; +use crate::array::ValidityVTableFromChild; +use crate::array::with_empty_buffers; +use crate::arrays::ListView; +use crate::arrays::map::MapData; +use crate::arrays::map::MapSlots; +use crate::arrays::map::MapSlotsView; +use crate::arrays::map::array::validate_entries; +use crate::arrays::map::compute::rules::PARENT_RULES; +use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; +use crate::dtype::DType; +use crate::match_each_map_builder; +use crate::serde::ArrayChildren; + +mod kernel; +mod operations; +mod validity; + +/// A [`Map`]-encoded Vortex array. +pub type MapArray = Array; + +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + +/// The canonical encoding for [`DType::Map`]. +/// +/// A map array has one `ListView>` child. Its outer dtype retains map-specific +/// metadata such as the `keys_sorted` assertion. +#[derive(Clone, Debug, Default)] +pub struct Map; + +impl VTable for Map { + type TypedArrayData = MapData; + + type OperationsVTable = Self; + type ValidityVTable = ValidityVTableFromChild; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.map"); + *ID + } + + fn validate( + &self, + _data: &MapData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure!( + slots.len() == MapSlots::COUNT, + "MapArray expected {} slot, found {}", + MapSlots::COUNT, + slots.len() + ); + + let DType::Map(map_dtype, nullability) = dtype else { + vortex_bail!("Expected map dtype, got {dtype}"); + }; + let slots = MapSlotsView::from_slots(slots); + validate_entries(map_dtype, *nullability, len, slots.entries) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + vortex_panic!("MapArray buffer index {idx} out of bounds") + } + + fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option { + None + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + + fn serialize( + _array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + Ok(Some(vec![])) + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + buffers: &[BufferHandle], + children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + if !metadata.is_empty() { + vortex_bail!( + "MapArray expects empty metadata, got {} bytes", + metadata.len() + ); + } + vortex_ensure!(buffers.is_empty(), "MapArray expects no buffers"); + + let DType::Map(map_dtype, nullability) = dtype else { + vortex_bail!("Expected map dtype, got {dtype}"); + }; + vortex_ensure!( + children.len() == MapSlots::COUNT, + "MapArray expected {} child, found {}", + MapSlots::COUNT, + children.len() + ); + + let expected_entries_dtype = + DType::List(std::sync::Arc::new(map_dtype.entries_dtype()), *nullability); + let entries = children.get(MapSlots::ENTRIES, &expected_entries_dtype, len)?; + vortex_ensure!( + entries.is::(), + "MapArray entries must use vortex.listview encoding, got {}", + entries.encoding_id() + ); + + let slots = MapData::make_slots(entries); + Ok(ArrayParts::new(self.clone(), dtype.clone(), len, MapData).with_slots(slots)) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + MapSlots::NAMES[idx].to_string() + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done(array)) + } + + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + match match_each_map_builder!(&mut *builder, |b| b.append_map_array(array, ctx)) { + Some(result) => result, + None => vortex_bail!( + "cannot append a Map array of dtype {} to a {} builder", + array.dtype(), + builder.dtype() + ), + } + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + PARENT_RULES.evaluate(array, parent, child_idx) + } +} diff --git a/vortex-array/src/arrays/map/vtable/operations.rs b/vortex-array/src/arrays/map/vtable/operations.rs new file mode 100644 index 00000000000..d6e8fe87f12 --- /dev/null +++ b/vortex-array/src/arrays/map/vtable/operations.rs @@ -0,0 +1,35 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::ExecutionCtx; +use crate::array::ArrayView; +use crate::array::OperationsVTable; +use crate::arrays::Map; +use crate::arrays::StructArray; +use crate::arrays::map::MapArrayExt; +use crate::arrays::struct_::StructArrayExt; +use crate::scalar::Scalar; + +impl OperationsVTable for Map { + fn scalar_at( + array: ArrayView<'_, Map>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let entries = array.entries_at(index)?.execute::(ctx)?; + let keys = entries.unmasked_field(0); + let values = entries.unmasked_field(1); + let pairs = (0..entries.len()) + .map(|entry_index| { + Ok(( + keys.execute_scalar(entry_index, ctx)?, + values.execute_scalar(entry_index, ctx)?, + )) + }) + .collect::>>()?; + + Scalar::try_map(array.dtype().clone(), pairs) + } +} diff --git a/vortex-array/src/arrays/map/vtable/validity.rs b/vortex-array/src/arrays/map/vtable/validity.rs new file mode 100644 index 00000000000..e81c7602077 --- /dev/null +++ b/vortex-array/src/arrays/map/vtable/validity.rs @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use crate::ArrayRef; +use crate::array::ArrayView; +use crate::array::ValidityChild; +use crate::arrays::Map; +use crate::arrays::map::MapArraySlotsExt; + +impl ValidityChild for Map { + fn validity_child(array: ArrayView<'_, Map>) -> ArrayRef { + array.entries().clone() + } +} diff --git a/vortex-array/src/arrays/masked/array.rs b/vortex-array/src/arrays/masked/array.rs index 077446b3f02..6062fa4e0d7 100644 --- a/vortex-array/src/arrays/masked/array.rs +++ b/vortex-array/src/arrays/masked/array.rs @@ -23,8 +23,10 @@ use crate::validity::Validity; #[array_slots(Masked)] pub struct MaskedSlots { /// The underlying child array being masked. + #[slot(0)] pub child: ArrayRef, /// The validity bitmap defining which elements are non-null. + #[slot(1)] pub validity: Option, } diff --git a/vortex-array/src/arrays/masked/compute/mask.rs b/vortex-array/src/arrays/masked/compute/mask.rs index e830f15bfa0..ba901086459 100644 --- a/vortex-array/src/arrays/masked/compute/mask.rs +++ b/vortex-array/src/arrays/masked/compute/mask.rs @@ -4,11 +4,10 @@ use vortex_error::VortexResult; use crate::ArrayRef; +use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::Masked; use crate::arrays::masked::MaskedArraySlotsExt; -use crate::arrays::scalar_fn::ScalarFnFactoryExt; -use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::fns::mask::Mask as MaskExpr; use crate::scalar_fn::fns::mask::MaskReduce; use crate::validity::Validity; @@ -20,11 +19,7 @@ impl MaskReduce for Masked { .validity()? .and(Validity::Array(mask.clone()))? .to_array(array.len()); - let masked_child = MaskExpr.try_new_array( - array.child().len(), - EmptyOptions, - [array.child().clone(), combined_mask], - )?; + let masked_child = MaskExpr::try_new(array.child().clone(), combined_mask)?.into_array(); Ok(Some(masked_child)) } } diff --git a/vortex-array/src/arrays/masked/compute/rules.rs b/vortex-array/src/arrays/masked/compute/rules.rs index 3accb455c3f..ac2123aedc8 100644 --- a/vortex-array/src/arrays/masked/compute/rules.rs +++ b/vortex-array/src/arrays/masked/compute/rules.rs @@ -1,16 +1,47 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::array::ArrayView; use crate::arrays::Masked; use crate::arrays::dict::TakeReduceAdaptor; use crate::arrays::filter::FilterReduceAdaptor; +use crate::arrays::masked::MaskedArrayExt; +use crate::arrays::masked::MaskedArraySlotsExt; +use crate::arrays::scalar_fn::ExactScalarFn; +use crate::arrays::scalar_fn::ScalarFnArrayView; use crate::arrays::slice::SliceReduceAdaptor; +use crate::builtins::ArrayBuiltins; +use crate::optimizer::rules::ArrayParentReduceRule; use crate::optimizer::rules::ParentRuleSet; +use crate::scalar_fn::fns::get_item::GetItem; use crate::scalar_fn::fns::mask::MaskReduceAdaptor; pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&FilterReduceAdaptor(Masked)), + ParentRuleSet::lift(&MaskedGetItemRule), ParentRuleSet::lift(&MaskReduceAdaptor(Masked)), ParentRuleSet::lift(&SliceReduceAdaptor(Masked)), ParentRuleSet::lift(&TakeReduceAdaptor(Masked)), ]); + +#[derive(Debug)] +struct MaskedGetItemRule; + +impl ArrayParentReduceRule for MaskedGetItemRule { + type Parent = ExactScalarFn; + + fn reduce_parent( + &self, + array: ArrayView<'_, Masked>, + parent: ScalarFnArrayView<'_, GetItem>, + _child_idx: usize, + ) -> VortexResult> { + let field = array.child().get_item(parent.options.clone())?; + Ok(Some( + field.mask(array.masked_validity().to_array(array.len()))?, + )) + } +} diff --git a/vortex-array/src/arrays/masked/execute.rs b/vortex-array/src/arrays/masked/execute.rs index c4173dd0cf6..eebcf1a2335 100644 --- a/vortex-array/src/arrays/masked/execute.rs +++ b/vortex-array/src/arrays/masked/execute.rs @@ -13,18 +13,26 @@ use crate::arrays::BoolArray; use crate::arrays::DecimalArray; use crate::arrays::ExtensionArray; use crate::arrays::FixedSizeListArray; +use crate::arrays::ListView; use crate::arrays::ListViewArray; +use crate::arrays::MapArray; use crate::arrays::MaskedArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; +use crate::arrays::UnionArray; use crate::arrays::VarBinViewArray; use crate::arrays::VariantArray; use crate::arrays::bool::BoolArrayExt; use crate::arrays::extension::ExtensionArrayExt; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use crate::arrays::listview::ListViewArraySlotsExt; +use crate::arrays::map::MapArrayExt; +use crate::arrays::map::MapArraySlotsExt; use crate::arrays::struct_::StructArrayExt; -use crate::arrays::variant::VariantArrayExt; +use crate::arrays::union::UnionArrayExt; +use crate::arrays::union::UnionArraySlotsExt; +use crate::arrays::variant::VariantArraySlotsExt; use crate::builtins::ArrayBuiltins; use crate::executor::ExecutionCtx; use crate::validity::Validity; @@ -46,10 +54,12 @@ pub fn mask_validity_canonical( Canonical::Decimal(a) => Canonical::Decimal(mask_validity_decimal(a, validity)?), Canonical::VarBinView(a) => Canonical::VarBinView(mask_validity_varbinview(a, validity)?), Canonical::List(a) => Canonical::List(mask_validity_listview(a, validity)?), + Canonical::Map(a) => Canonical::Map(mask_validity_map(a, validity)?), Canonical::FixedSizeList(a) => { Canonical::FixedSizeList(mask_validity_fixed_size_list(a, validity)?) } Canonical::Struct(a) => Canonical::Struct(mask_validity_struct(a, validity)?), + Canonical::Union(a) => Canonical::Union(mask_validity_union(a, validity)?), Canonical::Extension(a) => Canonical::Extension(mask_validity_extension(a, validity, ctx)?), Canonical::Variant(a) => Canonical::Variant(mask_validity_variant(a, validity, ctx)?), }) @@ -66,7 +76,7 @@ fn mask_validity_primitive( ) -> VortexResult { let ptype = array.ptype(); let new_validity = Validity::and(array.validity()?, validity)?; - // SAFETY: validity has same length as values + // SAFETY: We're only changing validity, not the data structure. Ok(unsafe { PrimitiveArray::new_unchecked_from_handle( array.buffer_handle().clone(), @@ -78,7 +88,7 @@ fn mask_validity_primitive( fn mask_validity_decimal(array: DecimalArray, validity: Validity) -> VortexResult { let new_validity = Validity::and(array.validity()?, validity)?; - // SAFETY: We're only changing validity, not the data structure + // SAFETY: We're only changing validity, not the data structure. Ok(unsafe { DecimalArray::new_unchecked_handle( array.buffer_handle().clone(), @@ -96,7 +106,7 @@ fn mask_validity_varbinview( ) -> VortexResult { let dtype = array.dtype().as_nullable(); let new_validity = Validity::and(array.validity()?, validity)?; - // SAFETY: We're only changing validity, not the data structure + // SAFETY: We're only changing validity, not the data structure. Ok(unsafe { VarBinViewArray::new_handle_unchecked( array.views_handle().clone(), @@ -109,7 +119,8 @@ fn mask_validity_varbinview( fn mask_validity_listview(array: ListViewArray, validity: Validity) -> VortexResult { let new_validity = Validity::and(array.validity()?, validity)?; - // SAFETY: We're only changing validity, not the data structure + // SAFETY: We're only changing validity, not the data structure. + let is_zctl = array.is_zero_copy_to_list(); Ok(unsafe { ListViewArray::new_unchecked( array.elements().clone(), @@ -117,9 +128,15 @@ fn mask_validity_listview(array: ListViewArray, validity: Validity) -> VortexRes array.sizes().clone(), new_validity, ) + .with_zero_copy_to_list(is_zctl) }) } +fn mask_validity_map(array: MapArray, validity: Validity) -> VortexResult { + let entries = mask_validity_listview(array.entries().as_::().into_owned(), validity)?; + MapArray::try_new(array.map_dtype().clone(), entries) +} + fn mask_validity_fixed_size_list( array: FixedSizeListArray, validity: Validity, @@ -127,7 +144,7 @@ fn mask_validity_fixed_size_list( let len = array.len(); let list_size = array.list_size(); let new_validity = Validity::and(array.validity()?, validity)?; - // SAFETY: We're only changing validity, not the data structure + // SAFETY: We're only changing validity, not the data structure. Ok(unsafe { FixedSizeListArray::new_unchecked(array.elements().clone(), list_size, new_validity, len) }) @@ -136,10 +153,26 @@ fn mask_validity_fixed_size_list( fn mask_validity_struct(array: StructArray, validity: Validity) -> VortexResult { let len = array.len(); let new_validity = Validity::and(array.validity()?, validity)?; - let fields = array.unmasked_fields(); - let struct_fields = array.struct_fields(); - // SAFETY: We're only changing validity, not the data structure - Ok(unsafe { StructArray::new_unchecked(fields, struct_fields.clone(), len, new_validity) }) + // SAFETY: We're only changing validity, not the data structure. + Ok(unsafe { + StructArray::new_unchecked( + array.iter_unmasked_fields().cloned(), + array.struct_fields().clone(), + len, + new_validity, + ) + }) +} + +fn mask_validity_union(array: UnionArray, validity: Validity) -> VortexResult { + let type_ids = array + .type_ids() + .clone() + .mask(validity.to_array(array.len()))?; + let variants = array.variants().clone(); + + // SAFETY: We're only changing validity, not the data structure. + Ok(unsafe { UnionArray::new_unchecked(type_ids, variants, array.iter_children().cloned()) }) } fn mask_validity_extension( @@ -147,7 +180,7 @@ fn mask_validity_extension( validity: Validity, ctx: &mut ExecutionCtx, ) -> VortexResult { - // For extension arrays, we need to mask the underlying storage + // For extension arrays, we need to mask the underlying storage. let storage = array.storage_array().clone().execute::(ctx)?; let masked_storage = mask_validity_canonical(storage, validity, ctx)?; let masked_storage = masked_storage.into_array(); diff --git a/vortex-array/src/arrays/masked/tests.rs b/vortex-array/src/arrays/masked/tests.rs index 842599bfa52..b326380bbf4 100644 --- a/vortex-array/src/arrays/masked/tests.rs +++ b/vortex-array/src/arrays/masked/tests.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use rstest::rstest; +use vortex_buffer::BitBuffer; +use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -10,6 +12,8 @@ use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; +use crate::arrays::BoolArray; +use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::dtype::DType; @@ -40,6 +44,24 @@ fn test_dtype_nullability_with_nullable_child() { assert!(child.dtype().is_nullable()); } +#[test] +fn test_empty_child_with_array_validity() -> VortexResult<()> { + let child_validity = + Validity::Array(BoolArray::new(BitBuffer::new_set(0), Validity::NonNullable).into_array()); + let child = PrimitiveArray::new(Buffer::::empty(), child_validity).into_array(); + let validity = + Validity::Array(BoolArray::new(BitBuffer::new_set(0), Validity::NonNullable).into_array()); + + let mut ctx = array_session().create_execution_ctx(); + assert!(child.all_valid(&mut ctx)?); + assert!(child.all_invalid(&mut ctx)?); + + let array = MaskedArray::try_new(child, validity)?; + + assert!(array.is_empty()); + Ok(()) +} + #[test] fn test_canonical_dtype_matches_array_dtype() -> VortexResult<()> { // The canonical form should have the same nullability as the array's dtype. @@ -140,3 +162,23 @@ fn test_masked_child_preserves_length(#[case] validity: Validity) { .unwrap(), ); } + +#[test] +fn masked_listview_execute_preserves_zctl_true() -> VortexResult<()> { + // Masking only intersects validity, so the zero-copy-to-list survives + // execution. + let elements = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); + let offsets = PrimitiveArray::from_iter([0i32, 2]).into_array(); + let sizes = PrimitiveArray::from_iter([2i32, 1]).into_array(); + let list_view = unsafe { + ListViewArray::new_unchecked(elements, offsets, sizes, Validity::NonNullable) + .with_zero_copy_to_list(true) + }; + + let masked = MaskedArray::try_new(list_view.into_array(), Validity::from_iter([true, false]))?; + let canonical = masked + .into_array() + .execute::(&mut array_session().create_execution_ctx())?; + assert!(canonical.into_listview().is_zero_copy_to_list()); + Ok(()) +} diff --git a/vortex-array/src/arrays/masked/vtable/mod.rs b/vortex-array/src/arrays/masked/vtable/mod.rs index c7e32a7ee3c..ccadee27e98 100644 --- a/vortex-array/src/arrays/masked/vtable/mod.rs +++ b/vortex-array/src/arrays/masked/vtable/mod.rs @@ -73,6 +73,7 @@ impl VTable for Masked { *ID } + #[expect(clippy::disallowed_methods)] fn validate( &self, _data: &MaskedData, @@ -92,6 +93,10 @@ impl VTable for Masked { child.dtype().as_nullable() == *dtype, "MaskedArray dtype does not match child and validity" ); + vortex_ensure!( + child.all_valid(&mut legacy_session().create_execution_ctx())?, + "MaskedArray children must not have nulls", + ); Ok(()) } diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index 15b6b5aa6be..f96eebead65 100644 --- a/vortex-array/src/arrays/mod.rs +++ b/vortex-array/src/arrays/mod.rs @@ -5,8 +5,8 @@ //! //! Canonical arrays are the default uncompressed representation for a logical dtype: //! [`NullArray`], [`BoolArray`], [`PrimitiveArray`], [`DecimalArray`], [`VarBinViewArray`], -//! [`ListViewArray`], [`FixedSizeListArray`], [`StructArray`], [`ExtensionArray`], and -//! [`VariantArray`]. +//! [`ListViewArray`], [`MapArray`], [`FixedSizeListArray`], [`StructArray`], [`UnionArray`], +//! [`ExtensionArray`], and [`VariantArray`]. //! //! Utility and lazy arrays represent common transformations without immediately materializing //! their result. Examples include [`ChunkedArray`] for concatenation, [`ConstantArray`] for repeated @@ -60,6 +60,8 @@ pub mod filter; pub use filter::Filter; pub use filter::FilterArray; +pub(crate) mod fixed_width; + pub mod fixed_size_list; pub use fixed_size_list::FixedSizeList; pub use fixed_size_list::FixedSizeListArray; @@ -76,6 +78,10 @@ pub mod listview; pub use listview::ListView; pub use listview::ListViewArray; +pub mod map; +pub use map::Map; +pub use map::MapArray; + pub mod masked; pub use masked::Masked; pub use masked::MaskedArray; @@ -88,6 +94,10 @@ pub mod patched; pub use patched::Patched; pub use patched::PatchedArray; +pub mod piecewise_sequence; +pub use piecewise_sequence::PiecewiseSequence; +pub use piecewise_sequence::PiecewiseSequenceArray; + pub mod primitive; pub use primitive::Primitive; pub use primitive::PrimitiveArray; @@ -108,6 +118,10 @@ pub mod struct_; pub use struct_::Struct; pub use struct_::StructArray; +pub mod union; +pub use union::Union; +pub use union::UnionArray; + pub mod varbin; pub use varbin::VarBin; pub use varbin::VarBinArray; @@ -131,6 +145,7 @@ pub(crate) fn initialize(session: &VortexSession) { fixed_size_list::initialize(session); list::initialize(session); listview::initialize(session); + map::initialize(session); patched::initialize(session); primitive::initialize(session); struct_::initialize(session); diff --git a/vortex-array/src/arrays/patched/array.rs b/vortex-array/src/arrays/patched/array.rs index 5e80545665e..4522c8223b8 100644 --- a/vortex-array/src/arrays/patched/array.rs +++ b/vortex-array/src/arrays/patched/array.rs @@ -54,12 +54,16 @@ pub struct PatchedData { #[array_slots(Patched)] pub struct PatchedSlots { /// The inner array containing the base unpatched values. + #[slot(0)] pub inner: ArrayRef, /// The lane offsets array for locating patches within lanes. + #[slot(1)] pub lane_offsets: ArrayRef, /// The indices of patched (exception) values. + #[slot(2)] pub patch_indices: ArrayRef, /// The patched (exception) values at the corresponding indices. + #[slot(3)] pub patch_values: ArrayRef, } @@ -369,16 +373,42 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; use crate::array_slots; + use crate::arrays::Chunked; use crate::arrays::Null; use crate::arrays::PrimitiveArray; + use crate::arrays::Union; use crate::validity::Validity; #[array_slots(Null)] struct OptionalPatchedSlots { + #[slot(0)] required: ArrayRef, + #[slot(1)] maybe: Option, } + #[array_slots(Chunked)] + struct VariadicSlots { + #[slot(0)] + offsets: ArrayRef, + #[slot(1)] + maybe_validity: Option, + #[slot(2..)] + chunks: Vec, + } + + /// The same layout as [`VariadicSlots`], but with every field declaration moved. The + /// `#[slot(..)]` annotations must keep the storage layout identical. + #[array_slots(Union)] + struct ShuffledVariadicSlots { + #[slot(2..)] + chunks: Vec, + #[slot(1)] + maybe_validity: Option, + #[slot(0)] + offsets: ArrayRef, + } + #[test] fn generated_slots_round_trip() { let required = PrimitiveArray::new(buffer![1u8, 2, 3], Validity::NonNullable).into_array(); @@ -405,4 +435,128 @@ mod tests { assert_eq!(rebuilt.inner.len(), required.len()); assert_eq!(rebuilt.patch_values.len(), optional.len()); } + + #[test] + fn variadic_slots_round_trip() { + let offsets = PrimitiveArray::new(buffer![0u64, 3, 5], Validity::NonNullable).into_array(); + let chunk0 = PrimitiveArray::new(buffer![1u8, 2, 3], Validity::NonNullable).into_array(); + let chunk1 = PrimitiveArray::new(buffer![4u8, 5], Validity::NonNullable).into_array(); + + assert_eq!(VariadicSlots::OFFSETS, 0); + assert_eq!(VariadicSlots::MAYBE_VALIDITY, 1); + assert_eq!(VariadicSlots::CHUNKS_OFFSET, 2); + assert_eq!(VariadicSlots::FIXED_COUNT, 2); + assert_eq!(VariadicSlots::slot_name(0), "offsets"); + assert_eq!(VariadicSlots::slot_name(3), "chunks[1]"); + + let slot_vec = vec![Some(offsets.clone()), None, Some(chunk0), Some(chunk1)]; + + let view = VariadicSlotsView::from_slots(&slot_vec); + assert_eq!(view.offsets.len(), 3); + assert!(view.maybe_validity.is_none()); + assert_eq!(view.chunks.len(), 2); + assert_eq!(view.chunks[0].len(), 3); + assert_eq!(view.chunks.get(1).map(|c| c.len()), Some(2)); + assert!(view.chunks.get(2).is_none()); + assert_eq!( + view.chunks.iter().map(|c| c.len()).collect::>(), + vec![3, 2] + ); + + let owned = view.to_owned(); + assert_eq!(owned.chunks.len(), 2); + + let owned = VariadicSlots::from_slots(slot_vec.into()); + assert_eq!(owned.offsets.len(), offsets.len()); + assert!(owned.maybe_validity.is_none()); + assert_eq!(owned.chunks.len(), 2); + + let slots = owned.into_slots(); + assert_eq!(slots.len(), 4); + assert!(slots[1].is_none()); + assert_eq!( + slots[VariadicSlots::CHUNKS_OFFSET] + .as_ref() + .map(|c| c.len()), + Some(3) + ); + } + + #[test] + fn variadic_slots_empty_tail() { + let offsets = PrimitiveArray::new(buffer![0u64], Validity::NonNullable).into_array(); + let slot_vec = vec![Some(offsets), None]; + + let view = VariadicSlotsView::from_slots(&slot_vec); + assert!(view.chunks.is_empty()); + + let owned = VariadicSlots::from_slots(slot_vec.into()); + assert!(owned.chunks.is_empty()); + assert_eq!(owned.into_slots().len(), 2); + } + + #[test] + fn slot_indices_follow_annotations_not_declaration_order() { + assert_eq!( + ShuffledVariadicSlots::OFFSETS, + VariadicSlots::OFFSETS, + "field declaration order must not move a slot" + ); + assert_eq!( + ShuffledVariadicSlots::MAYBE_VALIDITY, + VariadicSlots::MAYBE_VALIDITY + ); + assert_eq!( + ShuffledVariadicSlots::CHUNKS_OFFSET, + VariadicSlots::CHUNKS_OFFSET + ); + assert_eq!( + ShuffledVariadicSlots::FIXED_COUNT, + VariadicSlots::FIXED_COUNT + ); + assert_eq!(ShuffledVariadicSlots::slot_name(0), "offsets"); + assert_eq!(ShuffledVariadicSlots::slot_name(1), "maybe_validity"); + assert_eq!(ShuffledVariadicSlots::slot_name(3), "chunks[1]"); + } + + #[test] + fn shuffled_declaration_order_round_trips_through_storage() { + let offsets = PrimitiveArray::new(buffer![0u64, 3], Validity::NonNullable).into_array(); + let validity = PrimitiveArray::new(buffer![1u8, 1], Validity::NonNullable).into_array(); + let chunk = PrimitiveArray::new(buffer![1u8, 2, 3], Validity::NonNullable).into_array(); + + let slot_vec = vec![ + Some(offsets.clone()), + Some(validity.clone()), + Some(chunk.clone()), + ]; + + let view = ShuffledVariadicSlotsView::from_slots(&slot_vec); + assert_eq!(view.offsets.len(), offsets.len()); + assert_eq!(view.maybe_validity.map(|v| v.len()), Some(validity.len())); + assert_eq!(view.chunks.len(), 1); + assert_eq!(view.chunks[0].len(), chunk.len()); + + // `into_slots` must emit annotation order, not the shuffled declaration order. + let round_tripped = ShuffledVariadicSlots::from_slots(slot_vec.into()).into_slots(); + assert_eq!(round_tripped.len(), 3); + assert_eq!( + round_tripped[ShuffledVariadicSlots::OFFSETS] + .as_ref() + .map(|s| s.len()), + Some(offsets.len()) + ); + assert_eq!( + round_tripped[ShuffledVariadicSlots::MAYBE_VALIDITY] + .as_ref() + .map(|s| s.len()), + Some(validity.len()) + ); + assert_eq!( + round_tripped[ShuffledVariadicSlots::CHUNKS_OFFSET] + .as_ref() + .map(|s| s.len()), + Some(chunk.len()) + ); + } } diff --git a/vortex-array/src/arrays/patched/vtable/mod.rs b/vortex-array/src/arrays/patched/vtable/mod.rs index e36d426c0be..e7d0adae0dd 100644 --- a/vortex-array/src/arrays/patched/vtable/mod.rs +++ b/vortex-array/src/arrays/patched/vtable/mod.rs @@ -138,16 +138,6 @@ impl VTable for Patched { with_empty_buffers(self, array, buffers) } - fn child(array: ArrayView<'_, Self>, idx: usize) -> ArrayRef { - match idx { - PatchedSlots::INNER => array.inner().clone(), - PatchedSlots::LANE_OFFSETS => array.lane_offsets().clone(), - PatchedSlots::PATCH_INDICES => array.patch_indices().clone(), - PatchedSlots::PATCH_VALUES => array.patch_values().clone(), - _ => vortex_panic!("invalid child index for PatchedArray: {idx}"), - } - } - fn serialize( array: ArrayView<'_, Self>, _session: &VortexSession, @@ -375,7 +365,7 @@ mod tests { use crate::arrays::patched::PatchedSlots; use crate::arrays::patched::PatchedSlotsView; use crate::assert_arrays_eq; - use crate::builders::builder_with_capacity; + use crate::builders::builder_with_capacity_in; use crate::patches::Patches; use crate::serde::SerializeOptions; use crate::serde::SerializedArray; @@ -467,7 +457,11 @@ mod tests { .unwrap() .into_array(); - let mut builder = builder_with_capacity(array.dtype(), array.len()); + let mut builder = builder_with_capacity_in( + array.dtype(), + array.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); array.append_to_builder(builder.as_mut(), &mut ctx).unwrap(); let result = builder.finish(); @@ -502,7 +496,11 @@ mod tests { .slice(3..1024) .unwrap(); - let mut builder = builder_with_capacity(array.dtype(), array.len()); + let mut builder = builder_with_capacity_in( + array.dtype(), + array.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); array.append_to_builder(builder.as_mut(), &mut ctx).unwrap(); let result = builder.finish(); @@ -537,7 +535,11 @@ mod tests { .unwrap() .into_array(); - let mut builder = builder_with_capacity(array.dtype(), array.len()); + let mut builder = builder_with_capacity_in( + array.dtype(), + array.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); array.append_to_builder(builder.as_mut(), &mut ctx).unwrap(); let result = builder.finish(); @@ -598,7 +600,12 @@ mod tests { let session = array_session(); session.arrays().register(Patched); - let ctx = ArrayContext::empty().with_registry(session.arrays().registry().clone()); + let ctx = ArrayContext::empty().with_allowed_ids( + session + .arrays() + .registry() + .read(|map| map.keys().copied().collect()), + ); let serialized = array .serialize(&ctx, &session, &SerializeOptions::default()) .unwrap(); diff --git a/vortex-array/src/arrays/piecewise_sequence/array.rs b/vortex-array/src/arrays/piecewise_sequence/array.rs new file mode 100644 index 00000000000..389dc6df999 --- /dev/null +++ b/vortex-array/src/arrays/piecewise_sequence/array.rs @@ -0,0 +1,73 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use smallvec::smallvec; +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::array::Array; +use crate::array::ArrayParts; +use crate::array::EmptyArrayData; +use crate::array::TypedArrayRef; +use crate::array_slots; +use crate::arrays::PiecewiseSequence; +use crate::dtype::PType; + +#[array_slots(PiecewiseSequence)] +pub struct PiecewiseSequenceSlots { + /// The inclusive start index of each sequential piece. + #[slot(0)] + pub starts: ArrayRef, + /// The length of each sequential piece. + #[slot(1)] + pub lengths: ArrayRef, + /// The distance between consecutive indices in each piece. + #[slot(2)] + pub multipliers: ArrayRef, +} + +/// Extension methods for [`Array`] values using the [`PiecewiseSequence`] encoding. +pub trait PiecewiseSequenceArrayExt: + TypedArrayRef + PiecewiseSequenceArraySlotsExt +{ +} +impl> PiecewiseSequenceArrayExt for T {} + +impl Array { + /// Constructs a new `PiecewiseSequenceArray` from start, length, and multiplier arrays. + /// + /// This validates only structural invariants: all children must be non-nullable unsigned + /// integer arrays with matching lengths, and the outer array length is the declared expanded + /// index length. Individual ranges are checked when the index array is executed or consumed by + /// a take implementation. + pub fn try_new( + starts: ArrayRef, + lengths: ArrayRef, + multipliers: ArrayRef, + len: usize, + ) -> VortexResult { + Array::try_from_parts( + ArrayParts::new(PiecewiseSequence, PType::U64.into(), len, EmptyArrayData) + .with_slots(smallvec![Some(starts), Some(lengths), Some(multipliers)]), + ) + } + + /// Constructs a new `PiecewiseSequenceArray` without validation. + /// + /// # Safety + /// + /// The caller must guarantee the same structural invariants as [`Self::try_new`]. + pub unsafe fn new_unchecked( + starts: ArrayRef, + lengths: ArrayRef, + multipliers: ArrayRef, + len: usize, + ) -> Self { + unsafe { + Array::from_parts_unchecked( + ArrayParts::new(PiecewiseSequence, PType::U64.into(), len, EmptyArrayData) + .with_slots(smallvec![Some(starts), Some(lengths), Some(multipliers)]), + ) + } + } +} diff --git a/vortex-array/src/arrays/piecewise_sequence/mod.rs b/vortex-array/src/arrays/piecewise_sequence/mod.rs new file mode 100644 index 00000000000..05b7b13edf4 --- /dev/null +++ b/vortex-array/src/arrays/piecewise_sequence/mod.rs @@ -0,0 +1,168 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Index encoding for concatenated sequential ranges. +//! +//! A `PiecewiseSequenceArray` represents the expanded index sequence +//! `starts[i] + j * multipliers[i]` for `j` in `0..lengths[i]` for each piece `i`. It is +//! intended for take operations that can gather regular runs without materializing one index per +//! element. + +use itertools::Itertools; +use num_traits::AsPrimitive; +use vortex_buffer::BufferMut; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::Columnar; +use crate::array::ArrayView; +use crate::arrays::ConstantArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::piecewise_sequence::array::PiecewiseSequenceArraySlotsExt; +use crate::dtype::UnsignedPType; +use crate::executor::ExecutionCtx; +use crate::scalar::PValue; + +pub mod array; +mod vtable; + +#[cfg(test)] +mod tests; + +pub use array::PiecewiseSequenceArrayExt; +pub use vtable::*; + +pub(crate) fn check_index_arrays( + starts: &ArrayRef, + lengths: &ArrayRef, + multipliers: &ArrayRef, +) -> VortexResult<()> { + check_index_array("starts", starts)?; + check_index_array("lengths", lengths)?; + check_index_array("multipliers", multipliers)?; + vortex_ensure!( + starts.len() == lengths.len(), + "PiecewiseSequenceArray starts length {} does not match lengths length {}", + starts.len(), + lengths.len() + ); + vortex_ensure!( + starts.len() == multipliers.len(), + "PiecewiseSequenceArray starts length {} does not match multipliers length {}", + starts.len(), + multipliers.len() + ); + Ok(()) +} + +pub(crate) fn execute_index_arrays( + array: ArrayView<'_, PiecewiseSequence>, + ctx: &mut ExecutionCtx, +) -> VortexResult<(PrimitiveArray, PrimitiveArray, PrimitiveArray)> { + let starts = array.starts().clone().execute::(ctx)?; + let lengths = array.lengths().clone().execute::(ctx)?; + let multipliers = array.multipliers().clone().execute::(ctx)?; + Ok((starts, lengths, multipliers)) +} + +pub(crate) fn maybe_contiguous_slices( + array: ArrayView<'_, PiecewiseSequence>, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + if !is_constant_one(array.multipliers()) { + return Ok(None); + } + + let starts = array.starts().clone().execute::(ctx)?; + let lengths = array.lengths().clone().execute::(ctx)?; + Ok(Some((starts, lengths))) +} + +pub(crate) fn is_constant_one(multipliers: &ArrayRef) -> bool { + let Some(scalar) = multipliers.as_constant() else { + return false; + }; + matches!( + scalar.as_primitive_opt().and_then(|scalar| scalar.pvalue()), + Some(PValue::U8(1) | PValue::U16(1) | PValue::U32(1) | PValue::U64(1)) + ) +} + +pub(crate) fn constant_unsigned_usize(array: &ConstantArray) -> usize { + let pvalue = array + .scalar() + .as_primitive_opt() + .and_then(|scalar| scalar.pvalue()) + .vortex_expect("validated PiecewiseSequence length constants are primitive"); + + match pvalue { + PValue::U8(value) => value as usize, + PValue::U16(value) => value as usize, + PValue::U32(value) => value as usize, + PValue::U64(value) => value.as_(), + _ => unreachable!("validated PiecewiseSequence length constants are unsigned"), + } +} + +fn check_index_array(name: &str, array: &ArrayRef) -> VortexResult<()> { + vortex_ensure!( + array.dtype().is_unsigned_int(), + "PiecewiseSequenceArray {name} must have unsigned integer dtype, got {}", + array.dtype() + ); + vortex_ensure!( + !array.dtype().is_nullable(), + "PiecewiseSequenceArray {name} must be non-nullable, got {}", + array.dtype() + ); + Ok(()) +} + +pub(crate) fn materialize_ranges( + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + multipliers: &PrimitiveArray, + output_len: usize, +) -> VortexResult> +where + S: UnsignedPType, + L: UnsignedPType, + M: UnsignedPType, +{ + let starts = starts.as_slice::(); + let lengths = lengths.as_slice::(); + let multipliers = multipliers.as_slice::(); + let mut values = BufferMut::with_capacity(output_len); + let mut computed_len = 0usize; + + for ((&start, &length), &multiplier) in starts.iter().zip_eq(lengths).zip_eq(multipliers) { + let start: usize = start.as_(); + let length: usize = length.as_(); + let multiplier: usize = multiplier.as_(); + if length != 0 { + let last_offset = length - 1; + let last_delta = last_offset + .checked_mul(multiplier) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + start + .checked_add(last_delta) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + } + computed_len = computed_len + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + + values.extend((0..length).map(|offset| (start + offset * multiplier) as u64)); + } + + if computed_len != output_len { + vortex_bail!( + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); + } + Ok(values) +} diff --git a/vortex-array/src/arrays/piecewise_sequence/tests.rs b/vortex-array/src/arrays/piecewise_sequence/tests.rs new file mode 100644 index 00000000000..a087404fee2 --- /dev/null +++ b/vortex-array/src/arrays/piecewise_sequence/tests.rs @@ -0,0 +1,374 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use num_traits::AsPrimitive; +use vortex_buffer::buffer; +use vortex_error::VortexResult; + +use crate::ArrayContext; +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::arrays::DecimalArray; +use crate::arrays::FixedSizeListArray; +use crate::arrays::ListArray; +use crate::arrays::PiecewiseSequenceArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::VarBinArray; +use crate::arrays::VarBinViewArray; +use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::Nullability; +use crate::serde::SerializeOptions; +use crate::validity::Validity; + +fn piecewise_indices( + starts: impl IntoIterator, + lengths: &[u64], +) -> VortexResult { + let len = lengths + .iter() + .map(|&length| -> usize { length.as_() }) + .sum(); + let starts = PrimitiveArray::from_iter(starts).into_array(); + let lengths = PrimitiveArray::from_iter(lengths.iter().copied()).into_array(); + let multipliers = ConstantArray::new(1u64, lengths.len()).into_array(); + Ok(PiecewiseSequenceArray::try_new(starts, lengths, multipliers, len)?.into_array()) +} + +fn piecewise_indices_constant_length( + starts: impl IntoIterator, + length: u64, +) -> VortexResult { + let starts = starts.into_iter().collect::>(); + let length_usize: usize = length.as_(); + let len = starts.len() * length_usize; + let piece_count = starts.len(); + let starts = PrimitiveArray::from_iter(starts).into_array(); + let lengths = ConstantArray::new(length, piece_count).into_array(); + let multipliers = ConstantArray::new(1u64, piece_count).into_array(); + Ok(PiecewiseSequenceArray::try_new(starts, lengths, multipliers, len)?.into_array()) +} + +#[test] +fn materializes_piecewise_indices() -> VortexResult<()> { + let starts = buffer![3u64, 15, 21].into_array(); + let lengths = buffer![3u64, 3, 3].into_array(); + let multipliers = ConstantArray::new(1u64, 3).into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 9)?.into_array(); + + let expected = PrimitiveArray::from_iter([3u64, 4, 5, 15, 16, 17, 21, 22, 23]).into_array(); + assert_arrays_eq!(array, expected, &mut array_session().create_execution_ctx()); + Ok(()) +} + +#[test] +fn materializes_repeated_and_empty_ranges() -> VortexResult<()> { + let starts = buffer![5u64, 2, 5].into_array(); + let lengths = buffer![2u64, 0, 2].into_array(); + let multipliers = ConstantArray::new(1u64, 3).into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 4)?.into_array(); + + let expected = PrimitiveArray::from_iter([5u64, 6, 5, 6]).into_array(); + assert_arrays_eq!(array, expected, &mut array_session().create_execution_ctx()); + Ok(()) +} + +#[test] +fn materializes_multiplied_ranges() -> VortexResult<()> { + let starts = buffer![3u64, 15].into_array(); + let lengths = buffer![3u64, 2].into_array(); + let multipliers = buffer![2u64, 4].into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 5)?.into_array(); + + let expected = PrimitiveArray::from_iter([3u64, 5, 7, 15, 19]).into_array(); + assert_arrays_eq!(array, expected, &mut array_session().create_execution_ctx()); + Ok(()) +} + +#[test] +fn supports_constant_lengths() -> VortexResult<()> { + let starts = buffer![0u64, 10, 20].into_array(); + let lengths = ConstantArray::new(2u64, 3).into_array(); + let multipliers = ConstantArray::new(1u64, 3).into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 6)?.into_array(); + + let expected = PrimitiveArray::from_iter([0u64, 1, 10, 11, 20, 21]).into_array(); + assert_arrays_eq!(array, expected, &mut array_session().create_execution_ctx()); + Ok(()) +} + +#[test] +fn scalar_at_maps_into_piece() -> VortexResult<()> { + let starts = buffer![3u64, 15, 21].into_array(); + let lengths = buffer![3u64, 3, 3].into_array(); + let multipliers = buffer![1u64, 1, 1].into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 9)?.into_array(); + let mut ctx = array_session().create_execution_ctx(); + + assert_eq!(array.execute_scalar(0, &mut ctx)?, 3u64.into()); + assert_eq!(array.execute_scalar(4, &mut ctx)?, 16u64.into()); + assert_eq!(array.execute_scalar(8, &mut ctx)?, 23u64.into()); + Ok(()) +} + +#[test] +fn constructor_defers_range_value_validation() -> VortexResult<()> { + let starts = buffer![u64::MAX].into_array(); + let lengths = buffer![2u64].into_array(); + let multipliers = buffer![1u64].into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 2)?.into_array(); + + let err = array + .execute::(&mut array_session().create_execution_ctx()) + .unwrap_err(); + assert!( + err.to_string() + .contains("PiecewiseSequenceArray range overflows usize"), + "{err}" + ); + Ok(()) +} + +#[test] +fn execution_checks_declared_length() -> VortexResult<()> { + let starts = buffer![0u64, 3].into_array(); + let lengths = buffer![2u64, 2].into_array(); + let multipliers = ConstantArray::new(1u64, 2).into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 3)?.into_array(); + + let err = array + .execute::(&mut array_session().create_execution_ctx()) + .unwrap_err(); + assert!( + err.to_string() + .contains("PiecewiseSequenceArray expanded length 4 does not match declared length 3"), + "{err}" + ); + Ok(()) +} + +#[test] +fn serialization_is_not_supported() -> VortexResult<()> { + let array = PiecewiseSequenceArray::try_new( + buffer![3u32, 15, 21].into_array(), + buffer![2u16, 0, 2].into_array(), + buffer![2u8, 1, 3].into_array(), + 4, + )? + .into_array(); + + let err = array + .serialize( + &ArrayContext::empty(), + &array_session(), + &SerializeOptions::default(), + ) + .unwrap_err(); + assert!( + err.to_string() + .contains("Array vortex.piecewise-sequence does not support serialization"), + "{err}" + ); + Ok(()) +} + +#[test] +fn primitive_take_consumes_piecewise_indices() -> VortexResult<()> { + let values = PrimitiveArray::from_iter(0i32..20).into_array(); + let taken = values.take(piecewise_indices([3, 10], &[2, 3])?)?; + + assert_arrays_eq!( + taken, + PrimitiveArray::from_iter([3i32, 4, 10, 11, 12]).into_array(), + &mut array_session().create_execution_ctx() + ); + Ok(()) +} + +#[test] +fn primitive_take_handles_non_unit_multiplier() -> VortexResult<()> { + let values = PrimitiveArray::from_iter(0i32..20).into_array(); + let indices = PiecewiseSequenceArray::try_new( + buffer![3u64].into_array(), + buffer![3u64].into_array(), + buffer![2u64].into_array(), + 3, + )? + .into_array(); + let taken = values.take(indices)?; + + assert_arrays_eq!( + taken, + PrimitiveArray::from_iter([3i32, 5, 7]).into_array(), + &mut array_session().create_execution_ctx() + ); + Ok(()) +} + +#[test] +fn bool_take_consumes_piecewise_indices() -> VortexResult<()> { + let values = BoolArray::from_iter([true, false, true, true, false, false]).into_array(); + let taken = values.take(piecewise_indices([1, 4], &[2, 2])?)?; + + assert_arrays_eq!( + taken, + BoolArray::from_iter([false, true, false, false]).into_array(), + &mut array_session().create_execution_ctx() + ); + Ok(()) +} + +#[test] +fn decimal_take_consumes_piecewise_indices() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(19, 2); + let values = DecimalArray::from_iter([10i128, 20, 30, 40, 50, 60], decimal_dtype).into_array(); + let taken = values.take(piecewise_indices([1, 4], &[2, 1])?)?; + + assert_arrays_eq!( + taken, + DecimalArray::from_iter([20i128, 30, 50], decimal_dtype).into_array(), + &mut array_session().create_execution_ctx() + ); + Ok(()) +} + +#[test] +fn varbinview_take_consumes_piecewise_indices() -> VortexResult<()> { + let values = VarBinViewArray::from_iter( + ["a", "bb", "ccc", "dddd", "eeeee"].map(Some), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + let taken = values.take(piecewise_indices([1, 3], &[2, 1])?)?; + + assert_arrays_eq!( + taken, + VarBinViewArray::from_iter( + ["bb", "ccc", "dddd"].map(Some), + DType::Utf8(Nullability::NonNullable) + ) + .into_array(), + &mut array_session().create_execution_ctx() + ); + Ok(()) +} + +#[test] +fn varbin_take_consumes_piecewise_indices() -> VortexResult<()> { + let values = VarBinArray::from_iter( + ["a", "bb", "ccc", "dddd", "eeeee"].map(Some), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + let taken = values.take(piecewise_indices([1, 3], &[2, 1])?)?; + + assert_arrays_eq!( + taken, + VarBinArray::from_iter( + ["bb", "ccc", "dddd"].map(Some), + DType::Utf8(Nullability::NonNullable) + ) + .into_array(), + &mut array_session().create_execution_ctx() + ); + Ok(()) +} + +#[test] +fn contiguous_take_consumers_support_constant_piecewise_lengths() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let indices = piecewise_indices_constant_length([1, 3], 2)?; + + let primitive = PrimitiveArray::from_iter(0i32..10).into_array(); + assert_arrays_eq!( + primitive.take(indices.clone())?, + PrimitiveArray::from_iter([1i32, 2, 3, 4]).into_array(), + &mut ctx + ); + + let bools = BoolArray::from_iter([false, true, false, true, false]).into_array(); + assert_arrays_eq!( + bools.take(indices.clone())?, + BoolArray::from_iter([true, false, true, false]).into_array(), + &mut ctx + ); + + let decimal_dtype = DecimalDType::new(19, 2); + let decimals = DecimalArray::from_iter([10i128, 20, 30, 40, 50], decimal_dtype).into_array(); + assert_arrays_eq!( + decimals.take(indices.clone())?, + DecimalArray::from_iter([20i128, 30, 40, 50], decimal_dtype).into_array(), + &mut ctx + ); + + let varbinview = VarBinViewArray::from_iter( + ["a", "bb", "ccc", "dddd", "eeeee"].map(Some), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + assert_arrays_eq!( + varbinview.take(indices.clone())?, + VarBinViewArray::from_iter( + ["bb", "ccc", "dddd", "eeeee"].map(Some), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(), + &mut ctx + ); + + let varbin = VarBinArray::from_iter( + ["a", "bb", "ccc", "dddd", "eeeee"].map(Some), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + assert_arrays_eq!( + varbin.take(indices)?, + VarBinArray::from_iter( + ["bb", "ccc", "dddd", "eeeee"].map(Some), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(), + &mut ctx + ); + + Ok(()) +} + +#[test] +fn list_take_consumes_constant_piecewise_lengths() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let list = ListArray::try_new( + buffer![0i32, 1, 2, 3, 4, 5, 6].into_array(), + buffer![0u32, 2, 5, 5, 7].into_array(), + Validity::NonNullable, + )? + .into_array(); + let list_indices = piecewise_indices_constant_length([1, 0], 2)?; + let expected = ListArray::try_new( + buffer![2i32, 3, 4, 0, 1, 2, 3, 4].into_array(), + buffer![0u32, 3, 3, 5, 8].into_array(), + Validity::NonNullable, + )? + .into_array(); + assert_arrays_eq!(list.take(list_indices)?, expected, &mut ctx); + + Ok(()) +} + +#[test] +fn fixed_size_list_take_builds_piecewise_element_indices() -> VortexResult<()> { + let elements = PrimitiveArray::from_iter(0i32..12).into_array(); + let array = FixedSizeListArray::new(elements, 3, Validity::NonNullable, 4).into_array(); + let taken = array.take(buffer![1u64, 3].into_array())?; + + let expected_elements = PrimitiveArray::from_iter([3i32, 4, 5, 9, 10, 11]).into_array(); + let expected = + FixedSizeListArray::new(expected_elements, 3, Validity::NonNullable, 2).into_array(); + assert_arrays_eq!(taken, expected, &mut array_session().create_execution_ctx()); + Ok(()) +} diff --git a/vortex-array/src/arrays/piecewise_sequence/vtable.rs b/vortex-array/src/arrays/piecewise_sequence/vtable.rs new file mode 100644 index 00000000000..d68a00815f3 --- /dev/null +++ b/vortex-array/src/arrays/piecewise_sequence/vtable.rs @@ -0,0 +1,205 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use itertools::Itertools; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ArrayParts; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::ExecutionResult; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayId; +use crate::array::ArrayView; +use crate::array::EmptyArrayData; +use crate::array::OperationsVTable; +use crate::array::VTable; +use crate::array::ValidityVTable; +use crate::array::with_empty_buffers; +use crate::arrays::PrimitiveArray; +use crate::arrays::piecewise_sequence::array::PiecewiseSequenceSlots; +use crate::arrays::piecewise_sequence::check_index_arrays; +use crate::arrays::piecewise_sequence::execute_index_arrays; +use crate::arrays::piecewise_sequence::materialize_ranges; +use crate::arrays::primitive::PrimitiveArrayExt; +use crate::buffer::BufferHandle; +use crate::dtype::DType; +use crate::dtype::PType; +use crate::dtype::UnsignedPType; +use crate::match_each_unsigned_integer_ptype; +use crate::scalar::Scalar; +use crate::serde::ArrayChildren; +use crate::validity::Validity; + +/// A [`PiecewiseSequence`]-encoded Vortex index array. +pub type PiecewiseSequenceArray = Array; + +#[derive(Clone, Debug)] +pub struct PiecewiseSequence; + +impl VTable for PiecewiseSequence { + type TypedArrayData = EmptyArrayData; + type OperationsVTable = Self; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.piecewise-sequence"); + *ID + } + + fn validate( + &self, + _data: &Self::TypedArrayData, + dtype: &DType, + _len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure!( + dtype == &DType::from(PType::U64), + "PiecewiseSequenceArray dtype must be u64, got {dtype}" + ); + vortex_ensure!( + slots.len() == PiecewiseSequenceSlots::NAMES.len(), + "PiecewiseSequenceArray requires {} slots, got {}", + PiecewiseSequenceSlots::NAMES.len(), + slots.len() + ); + let starts = slots[PiecewiseSequenceSlots::STARTS] + .as_ref() + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray starts slot must be present"))?; + let lengths = slots[PiecewiseSequenceSlots::LENGTHS] + .as_ref() + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray lengths slot must be present"))?; + let multipliers = slots[PiecewiseSequenceSlots::MULTIPLIERS] + .as_ref() + .ok_or_else(|| { + vortex_err!("PiecewiseSequenceArray multipliers slot must be present") + })?; + check_index_arrays(starts, lengths, multipliers) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, _idx: usize) -> BufferHandle { + vortex_panic!("PiecewiseSequenceArray has no buffers") + } + + fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option { + None + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + PiecewiseSequenceSlots::NAMES[idx].to_string() + } + + fn serialize( + _array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + Ok(None) + } + + fn deserialize( + &self, + _dtype: &DType, + _len: usize, + _metadata: &[u8], + _buffers: &[BufferHandle], + _children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_bail!("PiecewiseSequenceArray is not serializable") + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let (starts, lengths, multipliers) = execute_index_arrays(array.as_view(), ctx)?; + + let values = match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + match_each_unsigned_integer_ptype!(multipliers.ptype(), |M| { + materialize_ranges::(&starts, &lengths, &multipliers, array.len())? + }) + }) + }); + Ok(ExecutionResult::done( + PrimitiveArray::new(values.freeze(), Validity::NonNullable).into_array(), + )) + } +} + +impl OperationsVTable for PiecewiseSequence { + fn scalar_at( + array: ArrayView<'_, PiecewiseSequence>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let (starts, lengths, multipliers) = execute_index_arrays(array, ctx)?; + + let value = match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + match_each_unsigned_integer_ptype!(multipliers.ptype(), |M| { + scalar_at::(&starts, &lengths, &multipliers, index)? + }) + }) + }); + Ok(value.into()) + } +} + +impl ValidityVTable for PiecewiseSequence { + fn validity(_array: ArrayView<'_, PiecewiseSequence>) -> VortexResult { + Ok(Validity::NonNullable) + } +} + +fn scalar_at( + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + multipliers: &PrimitiveArray, + index: usize, +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, + M: UnsignedPType, +{ + let mut remaining = index; + for ((&start, &length), &multiplier) in starts + .as_slice::() + .iter() + .zip_eq(lengths.as_slice::()) + .zip_eq(multipliers.as_slice::()) + { + let length: usize = length.as_(); + if remaining < length { + let start: usize = start.as_(); + let multiplier: usize = multiplier.as_(); + let offset = remaining + .checked_mul(multiplier) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + let value = start + .checked_add(offset) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + return Ok(value as u64); + } + remaining -= length; + } + vortex_bail!("PiecewiseSequenceArray index {index} out of bounds") +} diff --git a/vortex-array/src/arrays/primitive/array/mod.rs b/vortex-array/src/arrays/primitive/array/mod.rs index 867c7d0c783..39cdede9996 100644 --- a/vortex-array/src/arrays/primitive/array/mod.rs +++ b/vortex-array/src/arrays/primitive/array/mod.rs @@ -44,14 +44,17 @@ use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::min_max::min_max; use crate::array::child_to_validity; use crate::array::validity_to_child; +use crate::array_slots; use crate::arrays::bool::BoolArrayExt; use crate::buffer::BufferHandle; use crate::builtins::ArrayBuiltins; -/// The validity bitmap indicating which elements are non-null. -pub(super) const VALIDITY_SLOT: usize = 0; -pub(super) const NUM_SLOTS: usize = 1; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"]; +#[array_slots(Primitive)] +pub struct PrimitiveSlots { + /// The validity bitmap indicating which elements are non-null. + #[slot(0)] + pub validity: Option, +} /// A primitive array that stores [native types][crate::dtype::NativePType] in a contiguous buffer /// of memory, along with an optional validity child. @@ -119,12 +122,12 @@ pub trait PrimitiveArrayExt: TypedArrayRef { } fn validity_child(&self) -> Option<&ArrayRef> { - self.as_ref().slots()[VALIDITY_SLOT].as_ref() + self.as_ref().slots()[PrimitiveSlots::VALIDITY].as_ref() } fn validity(&self) -> Validity { child_to_validity( - self.as_ref().slots()[VALIDITY_SLOT].as_ref(), + self.as_ref().slots()[PrimitiveSlots::VALIDITY].as_ref(), self.nullability(), ) } @@ -144,7 +147,11 @@ pub trait PrimitiveArrayExt: TypedArrayRef { "can't reinterpret cast between integers of two different widths" ); - PrimitiveArray::from_buffer_handle(self.buffer_handle().clone(), ptype, self.validity()) + PrimitiveArray::from_buffer_handle( + self.buffer_handle().clone(), + ptype, + PrimitiveArrayExt::validity(self), + ) } /// Narrow the array to the smallest possible integer type that can represent all values. @@ -156,7 +163,7 @@ pub trait PrimitiveArrayExt: TypedArrayRef { let Some(min_max) = min_max(self.as_ref(), ctx, NumericalAggregateOpts::default())? else { return Ok(PrimitiveArray::new( Buffer::::zeroed(self.len()), - self.validity(), + PrimitiveArrayExt::validity(self), )); }; diff --git a/vortex-array/src/arrays/primitive/compute/cast.rs b/vortex-array/src/arrays/primitive/compute/cast.rs index 09defa442b8..93f10bdc5cf 100644 --- a/vortex-array/src/arrays/primitive/compute/cast.rs +++ b/vortex-array/src/arrays/primitive/compute/cast.rs @@ -2,12 +2,14 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use num_traits::AsPrimitive; +use num_traits::CheckedMul; use num_traits::NumCast; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_compute::lane_kernels::IndexedSinkExt; use vortex_compute::lane_kernels::IndexedSourceExt; use vortex_compute::lane_kernels::ReinterpretSink; +use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; @@ -18,16 +20,28 @@ use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn; use crate::array::ArrayView; +use crate::arrays::DecimalArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::primitive::PrimitiveArrayExt; +use crate::dtype::BigCast; use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::DecimalType; +use crate::dtype::IntegerPType; +use crate::dtype::NativeDecimalType; use crate::dtype::NativePType; use crate::dtype::Nullability; use crate::dtype::PType; +use crate::dtype::ToI256; +use crate::dtype::i256; use crate::expr::stats::Stat; use crate::expr::stats::StatsProvider; +use crate::match_each_decimal_value_type; +use crate::match_each_integer_ptype; use crate::match_each_native_ptype; +use crate::match_each_signed_integer_ptype; +use crate::scalar::DecimalValue; use crate::scalar_fn::fns::cast::CastKernel; use crate::scalar_fn::fns::cast::CastReduce; use crate::validity::Validity; @@ -68,6 +82,9 @@ impl CastKernel for Primitive { dtype: &DType, ctx: &mut ExecutionCtx, ) -> VortexResult> { + if let DType::Decimal(decimal_dtype, nullability) = dtype { + return cast_to_decimal(array, *decimal_dtype, *nullability, ctx).map(Some); + } let DType::Primitive(new_ptype, new_nullability) = dtype else { return Ok(None); }; @@ -105,6 +122,343 @@ impl CastKernel for Primitive { } } +fn cast_to_decimal( + array: ArrayView<'_, Primitive>, + decimal_dtype: DecimalDType, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if !array.ptype().is_int() { + vortex_bail!( + Compute: "Cannot cast floating primitive {} to decimal {}", + array.ptype(), decimal_dtype + ); + } + + let source_validity = array.validity()?; + let validity = source_validity + .clone() + .cast_nullability(nullability, array.len(), ctx)?; + let values_type = DecimalType::smallest_decimal_value_type(&decimal_dtype); + + if decimal_dtype.scale() == 0 + && signed_primitive_decimal_type(array.ptype()) == Some(values_type) + { + return match_each_signed_integer_ptype!(array.ptype(), |S| { + cast_unscaled_same_width_signed_integer_to_decimal::( + array, + decimal_dtype, + validity, + &source_validity, + ctx, + ) + }); + } + + let valid_values = source_validity.execute_mask(array.len(), ctx)?; + match_each_integer_ptype!(array.ptype(), |S| { + match_each_decimal_value_type!(values_type, |T| { + cast_integer_values_to_decimal::(array, decimal_dtype, validity, &valid_values) + }) + }) +} + +fn cast_unscaled_same_width_signed_integer_to_decimal( + array: ArrayView<'_, Primitive>, + decimal_dtype: DecimalDType, + validity: Validity, + source_validity: &Validity, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + S: IntegerPType + NativeDecimalType + ToI256, +{ + let values = array.as_slice::(); + let target_dtype = DType::Decimal(decimal_dtype, Nullability::NonNullable); + if !cached_values_fit_in(array, &target_dtype).unwrap_or(false) { + let valid_values = source_validity.execute_mask(array.len(), ctx)?; + validate_unscaled_signed_integer_values_to_decimal(values, decimal_dtype, &valid_values) + .map_err(|idx| primitive_to_decimal_cast_error(values[idx], decimal_dtype))?; + } + + // SAFETY: `S::DECIMAL_TYPE` has the same physical representation as the source ptype, and + // either exact min/max statistics or the validation above prove every valid value fits. + Ok(unsafe { + DecimalArray::new_unchecked_handle( + array.buffer_handle().clone(), + S::DECIMAL_TYPE, + decimal_dtype, + validity, + ) + .into_array() + }) +} + +fn validate_unscaled_signed_integer_values_to_decimal( + values: &[S], + decimal_dtype: DecimalDType, + valid_values: &Mask, +) -> Result<(), usize> +where + S: NativeDecimalType, +{ + let fits = |value| decimal_value_fits_precision(value, decimal_dtype); + match valid_values { + Mask::AllTrue(_) => values + .iter() + .position(|&value| !fits(value)) + .map_or(Ok(()), Err), + Mask::AllFalse(_) => Ok(()), + Mask::Values(mask) => { + let mut first_failure = None; + mask.bit_buffer().for_each_set_index(|idx| { + if first_failure.is_none() && !fits(values[idx]) { + first_failure = Some(idx); + } + }); + first_failure.map_or(Ok(()), Err) + } + } +} + +fn cast_integer_values_to_decimal( + array: ArrayView<'_, Primitive>, + decimal_dtype: DecimalDType, + validity: Validity, + valid_values: &Mask, +) -> VortexResult +where + S: IntegerPType + ToI256, + T: NativeDecimalType + CheckedMul, +{ + let scale = decimal_dtype.scale(); + let buffer = if scale == 0 { + cast_unscaled_integer_values_to_decimal::(array, decimal_dtype, valid_values)? + } else if scale > 0 { + cast_scaled_up_integer_values_to_decimal::(array, decimal_dtype, valid_values)? + } else { + cast_scaled_down_integer_values_to_decimal::(array, decimal_dtype, valid_values)? + }; + + Ok(DecimalArray::new(buffer, decimal_dtype, validity).into_array()) +} + +fn cast_unscaled_integer_values_to_decimal( + array: ArrayView<'_, Primitive>, + decimal_dtype: DecimalDType, + valid_values: &Mask, +) -> VortexResult> +where + S: IntegerPType + ToI256, + T: NativeDecimalType, +{ + let values = array.as_slice::(); + cast_integer_values_to_decimal_buffer(values, decimal_dtype, valid_values, |value| { + let value = ::from(value)?; + decimal_value_fits_precision(value, decimal_dtype).then_some(value) + }) +} + +fn cast_scaled_up_integer_values_to_decimal( + array: ArrayView<'_, Primitive>, + decimal_dtype: DecimalDType, + valid_values: &Mask, +) -> VortexResult> +where + S: IntegerPType + ToI256, + T: NativeDecimalType + CheckedMul, +{ + let values = array.as_slice::(); + let scale_factor = decimal_scale_factor::(decimal_dtype.scale())?; + cast_integer_values_to_decimal_buffer(values, decimal_dtype, valid_values, |value| { + let value = ::from(value)?; + let value = value.checked_mul(&scale_factor)?; + decimal_value_fits_precision(value, decimal_dtype).then_some(value) + }) +} + +fn cast_scaled_down_integer_values_to_decimal( + array: ArrayView<'_, Primitive>, + decimal_dtype: DecimalDType, + valid_values: &Mask, +) -> VortexResult> +where + S: IntegerPType + ToI256, + T: NativeDecimalType, +{ + let values = array.as_slice::(); + if decimal_dtype.scale().unsigned_abs() >= primitive_max_decimal_digits(array.ptype()) { + // The scale factor exceeds every source value, so only zero can be exactly rescaled. + return cast_integer_values_to_decimal_buffer( + values, + decimal_dtype, + valid_values, + |value| (value == S::default()).then_some(T::default()), + ); + } + + // Scaling down can shrink a value into a narrower target type, so first select the smallest + // signed carrier that can represent both the source and target. + let carrier_type = primitive_decimal_carrier_type(array.ptype()).max(T::DECIMAL_TYPE); + match_each_decimal_value_type!(carrier_type, |W| { + let scale_factor = decimal_scale_factor::(decimal_dtype.scale())?; + cast_integer_values_to_decimal_buffer(values, decimal_dtype, valid_values, |value| { + let value = ::from(value)?; + let quotient = value / scale_factor; + let value = (quotient * scale_factor == value).then_some(quotient)?; + let value = ::from(value)?; + decimal_value_fits_precision(value, decimal_dtype).then_some(value) + }) + }) +} + +fn cast_integer_values_to_decimal_buffer( + values: &[S], + decimal_dtype: DecimalDType, + valid_values: &Mask, + cast: impl Fn(S) -> Option, +) -> VortexResult> +where + S: IntegerPType + ToI256, + T: NativeDecimalType, +{ + cast_primitive_to_decimal_buffer(values, valid_values, cast) + .map_err(|idx| primitive_to_decimal_cast_error(values[idx], decimal_dtype)) +} + +/// Decimal physical type with the same representation, when `ptype` is signed. +fn signed_primitive_decimal_type(ptype: PType) -> Option { + match ptype { + PType::I8 => Some(DecimalType::I8), + PType::I16 => Some(DecimalType::I16), + PType::I32 => Some(DecimalType::I32), + PType::I64 => Some(DecimalType::I64), + PType::U8 | PType::U16 | PType::U32 | PType::U64 => None, + PType::F16 | PType::F32 | PType::F64 => { + unreachable!("floating primitives are rejected before selecting a decimal type") + } + } +} + +/// The smallest signed decimal physical type that represents every `ptype` value. +fn primitive_decimal_carrier_type(ptype: PType) -> DecimalType { + match ptype { + PType::I8 => DecimalType::I8, + PType::U8 | PType::I16 => DecimalType::I16, + PType::U16 | PType::I32 => DecimalType::I32, + PType::U32 | PType::I64 => DecimalType::I64, + PType::U64 => DecimalType::I128, + PType::F16 | PType::F32 | PType::F64 => { + unreachable!("floating primitives are rejected before selecting a decimal carrier") + } + } +} + +/// Maximum decimal digits in `ptype`; scales at least this large can exactly rescale only zero. +fn primitive_max_decimal_digits(ptype: PType) -> u8 { + match ptype { + PType::U8 | PType::I8 => 3, + PType::U16 | PType::I16 => 5, + PType::U32 | PType::I32 => 10, + PType::U64 => 20, + PType::I64 => 19, + PType::F16 | PType::F32 | PType::F64 => { + unreachable!("floating primitives are rejected before inspecting decimal digits") + } + } +} + +fn decimal_scale_factor(scale: i8) -> VortexResult +where + T: NativeDecimalType + CheckedMul, +{ + let exponent = if scale > 0 { + scale as u32 + } else { + (-(scale as i16)) as u32 + }; + let ten = ::from(10i8).ok_or_else( + || vortex_err!(Compute: "Cannot create decimal scale factor for scale {scale}"), + )?; + let mut factor = ::from(1i8).ok_or_else( + || vortex_err!(Compute: "Cannot create decimal scale factor for scale {scale}"), + )?; + for _ in 0..exponent { + factor = factor.checked_mul(&ten).ok_or_else( + || vortex_err!(Compute: "Cannot create decimal scale factor for scale {scale}"), + )?; + } + Ok(factor) +} + +fn decimal_value_fits_precision( + value: T, + decimal_dtype: DecimalDType, +) -> bool { + let precision = decimal_dtype.precision() as usize; + value >= T::MIN_BY_PRECISION[precision] && value <= T::MAX_BY_PRECISION[precision] +} + +fn cast_primitive_to_decimal_buffer( + values: &[S], + valid_values: &Mask, + cast: impl Fn(S) -> Option, +) -> Result, usize> +where + S: NativePType, + T: NativeDecimalType, +{ + if matches!(valid_values, Mask::AllFalse(_)) { + return Ok(BufferMut::::zeroed(values.len()).freeze()); + } + + let mut buffer = BufferMut::::with_capacity(values.len()); + match valid_values { + Mask::AllTrue(_) => { + values.try_map_into(&mut buffer.spare_capacity_mut()[..values.len()], &cast)?; + } + Mask::Values(mask) => { + values.try_map_masked_into( + mask.bit_buffer(), + &mut buffer.spare_capacity_mut()[..values.len()], + &cast, + )?; + } + Mask::AllFalse(_) => unreachable!("all-null values are handled before allocating"), + } + // SAFETY: the selected map kernel initialized every lane before returning Ok. + unsafe { buffer.set_len(values.len()) }; + Ok(buffer.freeze()) +} + +fn primitive_to_decimal_cast_error(value: S, decimal_dtype: DecimalDType) -> VortexError +where + S: IntegerPType + ToI256, +{ + let Some(value) = ::from(value) else { + return vortex_err!( + Compute: "primitive value cannot be represented while casting to {}", + decimal_dtype + ); + }; + + match DecimalValue::rescale_i256(value, 0, decimal_dtype.scale()) + .and_then(|value| DecimalValue::try_from_i256(value, decimal_dtype)) + { + Err(error) => error, + Ok(_) => { + debug_assert!( + false, + "primitive-to-decimal fast path rejected a value that the scalar cast accepts" + ); + vortex_err!( + Compute: "primitive value cannot be represented while casting to {}", + decimal_dtype + ) + } + } +} + /// Cast Primitive values from `F` to `T`. fn cast_values( array: ArrayView<'_, Primitive>, @@ -124,7 +478,28 @@ where // Returns `true` if every value of `from` is representable in `to` without loss. fn casts_losslessly_to(from: PType, to: PType) -> bool { - from.least_supertype(to) == Some(to) + if from == to { + return true; + } + if (from.is_unsigned_int() && to.is_unsigned_int()) + || (from.is_signed_int() && to.is_signed_int()) + || (from.is_float() && to.is_float()) + { + return from.byte_width() <= to.byte_width(); + } + if from.is_unsigned_int() && to.is_signed_int() { + return from.byte_width() < to.byte_width(); + } + if from.is_int() && to.is_float() { + let minimum_float_width = match from.byte_width() { + 1 => 2, + 2 => 4, + 4 => 8, + _ => return false, + }; + return to.byte_width() >= minimum_float_width; + } + false } // Skip the fallible kernel when type widening or (cached) min/max prove every value fits. @@ -268,19 +643,25 @@ mod test { use vortex_buffer::BitBuffer; use vortex_buffer::buffer; use vortex_error::VortexError; + use vortex_error::VortexResult; use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; + use crate::arrays::DecimalArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; use crate::compute::conformance::cast::test_cast_conformance; use crate::dtype::DType; + use crate::dtype::DecimalDType; + use crate::dtype::DecimalType; use crate::dtype::Nullability; use crate::dtype::PType; + use crate::dtype::i256; + use crate::expr::stats::Stat; use crate::validity::Validity; #[test] @@ -362,6 +743,270 @@ mod test { ); } + #[test] + fn cast_integer_to_decimal_rescales() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal_dtype = DecimalDType::new(5, 2); + let casted = PrimitiveArray::from_iter([42i32, -7]) + .into_array() + .cast(DType::Decimal(decimal_dtype, Nullability::NonNullable))? + .execute::(&mut ctx)?; + + assert_eq!( + casted.dtype(), + &DType::Decimal(decimal_dtype, Nullability::NonNullable) + ); + assert_eq!(casted.values_type(), DecimalType::I32); + assert_eq!(casted.buffer::().as_ref(), &[4_200, -700]); + Ok(()) + } + + #[test] + fn cast_same_width_signed_integer_to_decimal_reuses_buffer() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let source = PrimitiveArray::from_iter([42i32, -7]); + let source_ptr = source.as_slice::().as_ptr(); + let casted = source + .into_array() + .cast(DType::Decimal( + DecimalDType::new(9, 0), + Nullability::NonNullable, + ))? + .execute::(&mut ctx)?; + + assert_eq!(casted.buffer::().as_ptr(), source_ptr); + assert_eq!(casted.buffer::().as_ref(), &[42, -7]); + Ok(()) + } + + #[test] + fn cast_same_width_signed_integer_to_decimal_reuses_buffer_with_cached_bounds() + -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let source = PrimitiveArray::from_iter([42i32, -7]); + let source_ptr = source.as_slice::().as_ptr(); + let source = source.into_array(); + source + .statistics() + .compute_all(&[Stat::Min, Stat::Max], &mut ctx)?; + let casted = source + .cast(DType::Decimal( + DecimalDType::new(9, 0), + Nullability::NonNullable, + ))? + .execute::(&mut ctx)?; + + assert_eq!(casted.buffer::().as_ptr(), source_ptr); + Ok(()) + } + + #[test] + fn cast_same_width_signed_integer_to_decimal_ignores_out_of_range_nulls() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let source = PrimitiveArray::new(buffer![i32::MAX, 42], Validity::from_iter([false, true])); + let source_ptr = source.as_slice::().as_ptr(); + let casted = source + .into_array() + .cast(DType::Decimal( + DecimalDType::new(9, 0), + Nullability::Nullable, + ))? + .execute::(&mut ctx)?; + + assert_eq!(casted.buffer::().as_ptr(), source_ptr); + assert_eq!( + casted.validity()?.execute_mask(casted.len(), &mut ctx)?, + Mask::from(BitBuffer::from(vec![false, true])) + ); + Ok(()) + } + + #[test] + fn cast_same_width_signed_integer_to_decimal_checks_precision() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let casted = PrimitiveArray::from_iter([i32::MAX]) + .into_array() + .cast(DType::Decimal( + DecimalDType::new(9, 0), + Nullability::NonNullable, + ))?; + + let error = casted.execute::(&mut ctx).unwrap_err(); + assert!(error.to_string().contains("does not fit in precision")); + Ok(()) + } + + #[test] + fn cast_u64_to_decimal() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal_dtype = DecimalDType::new(20, 0); + let casted = PrimitiveArray::from_iter([u64::MAX]) + .into_array() + .cast(DType::Decimal(decimal_dtype, Nullability::NonNullable))? + .execute::(&mut ctx)?; + + assert_eq!(casted.values_type(), DecimalType::I128); + assert_eq!(casted.buffer::().as_ref(), &[i128::from(u64::MAX)]); + Ok(()) + } + + #[test] + fn cast_integer_to_decimal_reports_scale_up_overflow() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DType::Decimal(DecimalDType::new(38, 20), Nullability::NonNullable); + let casted = PrimitiveArray::from_iter([u64::MAX]) + .into_array() + .cast(dtype)?; + let actual = casted.execute::(&mut ctx).unwrap_err(); + + assert!( + actual + .to_string() + .contains("does not fit in precision of decimal(38,20)") + ); + Ok(()) + } + + #[test] + fn cast_u8_to_negative_scale_decimal() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal_dtype = DecimalDType::new(2, -2); + let casted = PrimitiveArray::from_iter([200u8]) + .into_array() + .cast(DType::Decimal(decimal_dtype, Nullability::NonNullable))? + .execute::(&mut ctx)?; + + assert_eq!(casted.values_type(), DecimalType::I8); + assert_eq!(casted.buffer::().as_ref(), &[2]); + Ok(()) + } + + #[test] + fn cast_u64_to_negative_scale_decimal() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal_dtype = DecimalDType::new(1, -19); + let casted = PrimitiveArray::from_iter([10_000_000_000_000_000_000u64]) + .into_array() + .cast(DType::Decimal(decimal_dtype, Nullability::NonNullable))? + .execute::(&mut ctx)?; + + assert_eq!(casted.values_type(), DecimalType::I8); + assert_eq!(casted.buffer::().as_ref(), &[1]); + Ok(()) + } + + #[test] + fn cast_integer_to_i256_decimal() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal_dtype = DecimalDType::new(39, 2); + let casted = PrimitiveArray::from_iter([42i64]) + .into_array() + .cast(DType::Decimal(decimal_dtype, Nullability::NonNullable))? + .execute::(&mut ctx)?; + + assert_eq!(casted.values_type(), DecimalType::I256); + assert_eq!(casted.buffer::().as_ref(), &[i256::from_i128(4_200)]); + Ok(()) + } + + #[test] + fn cast_all_null_integer_to_decimal() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal_dtype = DecimalDType::new(39, 2); + let casted = PrimitiveArray::new(buffer![i64::MAX, i64::MIN], Validity::AllInvalid) + .into_array() + .cast(DType::Decimal(decimal_dtype, Nullability::Nullable))? + .execute::(&mut ctx)?; + + assert!(matches!(casted.validity(), Ok(Validity::AllInvalid))); + assert_eq!(casted.buffer::().as_ref(), &[i256::ZERO, i256::ZERO]); + Ok(()) + } + + #[test] + fn cast_integer_to_negative_scale_decimal_requires_exact_rescale() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal_dtype = DecimalDType::new(3, -2); + let casted = PrimitiveArray::from_iter([1_200i32, -500]) + .into_array() + .cast(DType::Decimal(decimal_dtype, Nullability::NonNullable))? + .execute::(&mut ctx)?; + + assert_eq!(casted.buffer::().as_ref(), &[12, -5]); + + let error = PrimitiveArray::from_iter([42i32]) + .into_array() + .cast(DType::Decimal(decimal_dtype, Nullability::NonNullable))? + .execute::(&mut ctx) + .unwrap_err(); + assert!(error.to_string().contains("would lose precision")); + Ok(()) + } + + #[test] + fn cast_zero_to_large_negative_scale_decimal() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal_dtype = DecimalDType::new(3, -128); + let casted = PrimitiveArray::from_iter([0i32]) + .into_array() + .cast(DType::Decimal(decimal_dtype, Nullability::NonNullable))? + .execute::(&mut ctx)?; + + assert_eq!(casted.buffer::().as_ref(), &[0]); + Ok(()) + } + + #[test] + fn cast_integer_to_decimal_ignores_out_of_range_null_lanes() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal_dtype = DecimalDType::new(3, 1); + let casted = PrimitiveArray::new(buffer![999i32, 42], Validity::from_iter([false, true])) + .into_array() + .cast(DType::Decimal(decimal_dtype, Nullability::Nullable))? + .execute::(&mut ctx)?; + + assert_eq!(casted.buffer::().as_ref()[1], 420); + assert_eq!( + casted.validity()?.execute_mask(casted.len(), &mut ctx)?, + Mask::from(BitBuffer::from(vec![false, true])) + ); + Ok(()) + } + + #[test] + fn cast_integer_to_decimal_checks_precision() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let casted = PrimitiveArray::from_iter([100i32]) + .into_array() + .cast(DType::Decimal( + DecimalDType::new(2, 1), + Nullability::NonNullable, + ))?; + + let error = casted.execute::(&mut ctx).unwrap_err(); + assert!(error.to_string().contains("does not fit in precision")); + Ok(()) + } + + #[test] + fn cast_floating_primitive_to_decimal_fails() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let casted = PrimitiveArray::from_iter([1.0f64]) + .into_array() + .cast(DType::Decimal( + DecimalDType::new(3, 1), + Nullability::NonNullable, + ))?; + + let error = casted.execute::(&mut ctx).unwrap_err(); + assert!( + error + .to_string() + .contains("Cannot cast floating primitive f64 to decimal decimal(3,1)") + ); + Ok(()) + } + #[test] fn cast_i32_u32() { let arr = buffer![-1i32].into_array(); @@ -425,7 +1070,7 @@ mod test { /// Same-width integer cast where all values fit: should reinterpret the /// buffer without allocation (pointer identity). #[test] - fn cast_same_width_int_reinterprets_buffer() -> vortex_error::VortexResult<()> { + fn cast_same_width_int_reinterprets_buffer() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let src = PrimitiveArray::from_iter([0u32, 10, 100]); let src_ptr = src.as_slice::().as_ptr(); @@ -458,7 +1103,7 @@ mod test { /// All-null array cast between same-width types should succeed without /// touching the buffer contents. #[test] - fn cast_same_width_all_null() -> vortex_error::VortexResult<()> { + fn cast_same_width_all_null() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let arr = PrimitiveArray::new(buffer![0xFFu8, 0xFF], Validity::AllInvalid); let casted = arr @@ -473,7 +1118,7 @@ mod test { /// Same-width integer cast with nullable values: out-of-range nulls should /// not prevent the cast from succeeding. #[test] - fn cast_same_width_int_nullable_with_out_of_range_nulls() -> vortex_error::VortexResult<()> { + fn cast_same_width_int_nullable_with_out_of_range_nulls() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); // The null position holds u32::MAX which doesn't fit in i32, but it's // masked as invalid so the cast should still succeed via reinterpret. @@ -494,7 +1139,7 @@ mod test { } #[test] - fn cast_u32_to_u8_with_out_of_range_nulls() -> vortex_error::VortexResult<()> { + fn cast_u32_to_u8_with_out_of_range_nulls() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let arr = PrimitiveArray::new( buffer![1000u32, 10u32, 42u32], diff --git a/vortex-array/src/arrays/primitive/compute/fixed_width.rs b/vortex-array/src/arrays/primitive/compute/fixed_width.rs new file mode 100644 index 00000000000..6f3202ecc3c --- /dev/null +++ b/vortex-array/src/arrays/primitive/compute/fixed_width.rs @@ -0,0 +1,34 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::ByteBuffer; +use vortex_error::VortexResult; + +use crate::array::ArrayView; +use crate::arrays::Primitive; +use crate::arrays::PrimitiveArray; +use crate::arrays::fixed_width::FixedWidthArray; +use crate::validity::Validity; + +impl FixedWidthArray for Primitive { + fn byte_width(array: ArrayView<'_, Self>) -> usize { + array.ptype().byte_width() + } + + fn values(array: ArrayView<'_, Self>) -> ByteBuffer { + array.buffer_handle().to_host_sync() + } + + fn with_values( + array: ArrayView<'_, Self>, + values: ByteBuffer, + _len: usize, + validity: Validity, + ) -> VortexResult { + Ok(PrimitiveArray::from_byte_buffer( + values, + array.ptype(), + validity, + )) + } +} diff --git a/vortex-array/src/arrays/primitive/compute/mod.rs b/vortex-array/src/arrays/primitive/compute/mod.rs index 1ca4b17d3d5..382b42ee6e2 100644 --- a/vortex-array/src/arrays/primitive/compute/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/mod.rs @@ -4,10 +4,10 @@ mod between; mod cast; mod fill_null; +mod fixed_width; mod mask; pub(crate) mod rules; mod slice; -mod take; mod zip; #[cfg(test)] diff --git a/vortex-array/src/arrays/primitive/compute/take/avx2.rs b/vortex-array/src/arrays/primitive/compute/take/avx2.rs deleted file mode 100644 index fb64e6748b0..00000000000 --- a/vortex-array/src/arrays/primitive/compute/take/avx2.rs +++ /dev/null @@ -1,536 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! An AVX2 implementation of take operation using gather instructions. -//! -//! Only enabled for x86_64 hosts and it is gated at runtime behind feature detection to -//! ensure AVX2 instructions are available. - -#![cfg(any(target_arch = "x86_64", target_arch = "x86"))] - -use std::arch::x86_64::__m256i; -use std::arch::x86_64::_mm_loadu_si128; -use std::arch::x86_64::_mm_setzero_si128; -use std::arch::x86_64::_mm_shuffle_epi32; -use std::arch::x86_64::_mm_storeu_si128; -use std::arch::x86_64::_mm_unpacklo_epi64; -use std::arch::x86_64::_mm256_andnot_si256; -use std::arch::x86_64::_mm256_cmpgt_epi32; -use std::arch::x86_64::_mm256_cmpgt_epi64; -use std::arch::x86_64::_mm256_cvtepu8_epi32; -use std::arch::x86_64::_mm256_cvtepu8_epi64; -use std::arch::x86_64::_mm256_cvtepu16_epi32; -use std::arch::x86_64::_mm256_cvtepu16_epi64; -use std::arch::x86_64::_mm256_cvtepu32_epi64; -use std::arch::x86_64::_mm256_extracti128_si256; -use std::arch::x86_64::_mm256_loadu_si256; -use std::arch::x86_64::_mm256_mask_i32gather_epi32; -use std::arch::x86_64::_mm256_mask_i64gather_epi32; -use std::arch::x86_64::_mm256_mask_i64gather_epi64; -use std::arch::x86_64::_mm256_set1_epi32; -use std::arch::x86_64::_mm256_set1_epi64x; -use std::arch::x86_64::_mm256_setzero_si256; -use std::arch::x86_64::_mm256_storeu_si256; -use std::convert::identity; - -use vortex_buffer::Alignment; -use vortex_buffer::Buffer; -use vortex_buffer::BufferMut; -use vortex_error::VortexResult; - -use crate::ArrayRef; -use crate::IntoArray; -use crate::array::ArrayView; -use crate::arrays::PrimitiveArray; -use crate::arrays::primitive::compute::take::TakeImpl; -use crate::arrays::primitive::compute::take::take_primitive_scalar; -use crate::arrays::primitive::vtable::Primitive; -use crate::dtype::NativePType; -use crate::dtype::UnsignedPType; -use crate::match_each_native_ptype; -use crate::match_each_unsigned_integer_ptype; -use crate::validity::Validity; - -#[allow(unused)] -pub(super) struct TakeKernelAVX2; - -impl TakeImpl for TakeKernelAVX2 { - #[inline(always)] - fn take( - &self, - values: ArrayView<'_, Primitive>, - indices: ArrayView<'_, Primitive>, - validity: Validity, - ) -> VortexResult { - assert!(indices.ptype().is_unsigned_int()); - - Ok(match_each_unsigned_integer_ptype!(indices.ptype(), |I| { - match_each_native_ptype!(values.ptype(), |V| { - // SAFETY: This kernel is only selected when avx2 cpu-feature is detected. - unsafe { - take_primitive_avx2(values.as_slice::(), indices.as_slice::(), validity) - } - }) - }) - .into_array()) - } -} - -/// # Safety -/// -/// The caller must ensure that if the validity has a length, it is the same length as the indices, -/// and that the `avx2` feature is enabled. -#[target_feature(enable = "avx2")] -#[allow(unused)] -unsafe fn take_primitive_avx2( - values: &[V], - indices: &[I], - validity: Validity, -) -> PrimitiveArray -where - V: NativePType, - I: UnsignedPType, -{ - // SAFETY: The caller guarantees that the `avx2` feature is enabled. - let buffer = unsafe { take_avx2(values, indices) }; - - debug_assert!( - validity - .maybe_len() - .is_none_or(|validity_len| validity_len == buffer.len()) - ); - - // SAFETY: The caller ensures that the validity and indices have the same length, so the taken - // buffer and the validity must have the same length. - unsafe { PrimitiveArray::new_unchecked(buffer, validity) } -} - -// --------------------------------------------------------------------------- -// AVX2 SIMD take algorithm -// --------------------------------------------------------------------------- - -/// Takes the specified indices into a new [`Buffer`] using AVX2 SIMD. -/// -/// An AVX2 gather only moves raw bytes, so signedness and float-ness are irrelevant — only the -/// byte width of `V` matters. Any 4-byte value rides the gather through the `u32` lane and any -/// 8-byte value through the `u64` lane, regardless of its actual type. Values 1 or 2 bytes wide -/// (AVX2 has no sub-32-bit gather) and wider than 8 bytes (`i128`, decimals) fall back to the -/// scalar kernel. -/// -/// This treats `V` as plain-old-data: reinterpreting the gathered bytes as `V` is only sound -/// because every bit pattern is a valid `V`. All primitive and decimal-backing types satisfy -/// this, as does any `Copy` POD type the caller supplies. -/// -/// # Panics -/// -/// This function panics if any of the provided `indices` are out of bounds for `values`. -/// -/// # Safety -/// -/// The caller must ensure the `avx2` feature is enabled. -#[target_feature(enable = "avx2")] -#[doc(hidden)] -unsafe fn take_avx2(buffer: &[V], indices: &[I]) -> Buffer { - if buffer.is_empty() { - return Buffer::zeroed(indices.len()); - } - - // Dispatch on the gather lane width. The index type must still be concretized to select the - // right `GatherFn` impl, so re-dispatch it with `match_each_unsigned_integer_ptype!`. - macro_rules! dispatch { - ($lane:ty) => {{ - match_each_unsigned_integer_ptype!(I::PTYPE, |Idx| { - // SAFETY: `Idx` has the same `PTYPE` as `I`, so this is a no-op reinterpret of the - // index slice into the concrete type the gather impl is keyed on. - let indices = unsafe { std::mem::transmute::<&[I], &[Idx]>(indices) }; - exec_take::(buffer, indices) - }) - }}; - } - - match size_of::() { - 4 => dispatch!(u32), - 8 => dispatch!(u64), - // 1/2-byte and >8-byte values have no AVX2 gather lane, so fall back to scalar. - _ => take_primitive_scalar(buffer, indices), - } -} - -/// The main gather function that is used by the inner loop kernel for AVX2 gather. -trait GatherFn { - /// The number of data elements that are written to the `dst` on each loop iteration. - const WIDTH: usize; - /// The number of indices read from `indices` on each loop iteration. Depending on the - /// available instructions and bit-width we may stride by a larger amount than we actually - /// end up reading from `src` (governed by the `WIDTH` parameter). - const STRIDE: usize = Self::WIDTH; - - /// Gather values from `src` into the `dst` using the `indices`, optionally using SIMD - /// instructions. - /// - /// # Safety - /// - /// This function can read up to `STRIDE` elements through `indices`, and read/write up to - /// `WIDTH` elements through `src` and `dst` respectively. - unsafe fn gather(indices: *const Idx, max_idx: Idx, src: *const Values, dst: *mut Values); -} - -/// AVX2 version of [`GatherFn`] defined for 32- and 64-bit value types. -enum AVX2Gather {} - -macro_rules! impl_gather { - ($idx:ty, $({$value:ty => load: $load:ident, extend: $extend:ident, splat: $splat:ident, zero_vec: $zero_vec:ident, mask_indices: $mask_indices:ident, mask_cvt: |$mask_var:ident| $mask_cvt:block, gather: $masked_gather:ident, store: $store:ident, WIDTH = $WIDTH:literal, STRIDE = $STRIDE:literal }),+) => { - $( - impl_gather!(single; $idx, $value, load: $load, extend: $extend, splat: $splat, zero_vec: $zero_vec, mask_indices: $mask_indices, mask_cvt: |$mask_var| $mask_cvt, gather: $masked_gather, store: $store, WIDTH = $WIDTH, STRIDE = $STRIDE); - )* - }; - (single; $idx:ty, $value:ty, load: $load:ident, extend: $extend:ident, splat: $splat:ident, zero_vec: $zero_vec:ident, mask_indices: $mask_indices:ident, mask_cvt: |$mask_var:ident| $mask_cvt:block, gather: $masked_gather:ident, store: $store:ident, WIDTH = $WIDTH:literal, STRIDE = $STRIDE:literal) => { - impl GatherFn<$idx, $value> for AVX2Gather { - const WIDTH: usize = $WIDTH; - const STRIDE: usize = $STRIDE; - - #[allow(unused_unsafe, clippy::cast_possible_truncation)] - #[inline(always)] - unsafe fn gather(indices: *const $idx, max_idx: $idx, src: *const $value, dst: *mut $value) { - const { - assert!($WIDTH <= $STRIDE, "dst cannot advance by more than the stride"); - } - - const SCALE: i32 = std::mem::size_of::<$value>() as i32; - - let indices_vec = unsafe { $load(indices.cast()) }; - // Extend indices to fill vector register. - let indices_vec = unsafe { $extend(indices_vec) }; - - // Create a vec of the max idx. - let max_idx_vec = unsafe { $splat(max_idx as _) }; - // Create a mask for valid indices (where the max_idx > provided index). - let invalid_mask = unsafe { _mm256_andnot_si256($mask_indices(indices_vec, max_idx_vec), $splat(-1)) }; - let invalid_mask = { - let $mask_var = invalid_mask; - $mask_cvt - }; - let zero_vec = unsafe { $zero_vec() }; - - // Gather the values into new vector register, for masked positions - // it substitutes zero instead of accessing the src. - let values_vec = unsafe { $masked_gather::(zero_vec, src.cast(), indices_vec, invalid_mask) }; - - // Write the vec out to dst. - unsafe { $store(dst.cast(), values_vec) }; - } - } - }; -} - -// kernels for u8 indices -impl_gather!(u8, - // 32-bit values, loaded 8 at a time - { u32 => - load: _mm_loadu_si128, - extend: _mm256_cvtepu8_epi32, - splat: _mm256_set1_epi32, - zero_vec: _mm256_setzero_si256, - mask_indices: _mm256_cmpgt_epi32, - mask_cvt: |x| { x }, - gather: _mm256_mask_i32gather_epi32, - store: _mm256_storeu_si256, - WIDTH = 8, STRIDE = 16 - }, - - // 64-bit values, loaded 4 at a time - { u64 => - load: _mm_loadu_si128, - extend: _mm256_cvtepu8_epi64, - splat: _mm256_set1_epi64x, - zero_vec: _mm256_setzero_si256, - mask_indices: _mm256_cmpgt_epi64, - mask_cvt: |x| { x }, - gather: _mm256_mask_i64gather_epi64, - store: _mm256_storeu_si256, - WIDTH = 4, STRIDE = 16 - } -); - -// kernels for u16 indices -impl_gather!(u16, - // 32-bit values. 8x indices loaded at a time and 8x values written at a time. - { u32 => - load: _mm_loadu_si128, - extend: _mm256_cvtepu16_epi32, - splat: _mm256_set1_epi32, - zero_vec: _mm256_setzero_si256, - mask_indices: _mm256_cmpgt_epi32, - mask_cvt: |x| { x }, - gather: _mm256_mask_i32gather_epi32, - store: _mm256_storeu_si256, - WIDTH = 8, STRIDE = 8 - }, - - // 64-bit values. 8x indices loaded at a time and 4x values loaded at a time. - { u64 => - load: _mm_loadu_si128, - extend: _mm256_cvtepu16_epi64, - splat: _mm256_set1_epi64x, - zero_vec: _mm256_setzero_si256, - mask_indices: _mm256_cmpgt_epi64, - mask_cvt: |x| { x }, - gather: _mm256_mask_i64gather_epi64, - store: _mm256_storeu_si256, - WIDTH = 4, STRIDE = 8 - } -); - -// kernels for u32 indices -impl_gather!(u32, - // 32-bit values. 8x indices loaded at a time and 8x values written. - { u32 => - load: _mm256_loadu_si256, - extend: identity, - splat: _mm256_set1_epi32, - zero_vec: _mm256_setzero_si256, - mask_indices: _mm256_cmpgt_epi32, - mask_cvt: |x| { x }, - gather: _mm256_mask_i32gather_epi32, - store: _mm256_storeu_si256, - WIDTH = 8, STRIDE = 8 - }, - - // 64-bit values. - { u64 => - load: _mm_loadu_si128, - extend: _mm256_cvtepu32_epi64, - splat: _mm256_set1_epi64x, - zero_vec: _mm256_setzero_si256, - mask_indices: _mm256_cmpgt_epi64, - mask_cvt: |x| { x }, - gather: _mm256_mask_i64gather_epi64, - store: _mm256_storeu_si256, - WIDTH = 4, STRIDE = 4 - } -); - -// kernels for u64 indices -impl_gather!(u64, - { u32 => - load: _mm256_loadu_si256, - extend: identity, - splat: _mm256_set1_epi64x, - zero_vec: _mm_setzero_si128, - mask_indices: _mm256_cmpgt_epi64, - mask_cvt: |m| { - unsafe { - let lo_bits = _mm256_extracti128_si256::<0>(m); // lower half - let hi_bits = _mm256_extracti128_si256::<1>(m); // upper half - let lo_packed = _mm_shuffle_epi32::<0b01_01_01_01>(lo_bits); - let hi_packed = _mm_shuffle_epi32::<0b01_01_01_01>(hi_bits); - _mm_unpacklo_epi64(lo_packed, hi_packed) - } - }, - gather: _mm256_mask_i64gather_epi32, - store: _mm_storeu_si128, - WIDTH = 4, STRIDE = 4 - }, - - // 64-bit values. - { u64 => - load: _mm256_loadu_si256, - extend: identity, - splat: _mm256_set1_epi64x, - zero_vec: _mm256_setzero_si256, - mask_indices: _mm256_cmpgt_epi64, - mask_cvt: |x| { x }, - gather: _mm256_mask_i64gather_epi64, - store: _mm256_storeu_si256, - WIDTH = 4, STRIDE = 4 - } -); - -/// AVX2 core inner loop for a given index type `Idx`, output element type `Out`, and gather -/// `Lane` type. -/// -/// `Out` is the element type written to the output buffer; `Lane` (`u32` or `u64`) is the -/// integer type the gather intrinsics operate on. The caller must pair them so that -/// `size_of::() == size_of::()` (the only caller, [`take_avx2`], picks `Lane` from -/// `size_of::()`). The gather moves `size_of::()` raw bytes per element, which only -/// yields a valid `Out` because `Out` is plain-old-data (every bit pattern is a valid `Out`). -/// Pointers into the `Out`-typed slices are cast to `*const Lane`/`*mut Lane`; gather tolerates -/// the (possibly weaker) `Out` alignment. -#[inline(always)] -fn exec_take(values: &[Out], indices: &[Idx]) -> Buffer -where - Out: Copy, - Idx: UnsignedPType, - Gather: GatherFn, -{ - debug_assert_eq!( - size_of::(), - size_of::(), - "gather lane and output element must have the same size" - ); - - let indices_len = indices.len(); - let max_index = Idx::from(values.len()).unwrap_or_else(|| Idx::max_value()); - let mut buffer = - BufferMut::::with_capacity_aligned(indices_len, Alignment::of::<__m256i>()); - let buf_uninit = buffer.spare_capacity_mut(); - - let mut offset = 0; - // Loop terminates STRIDE elements before end of the indices array because the `GatherFn` - // might read up to STRIDE src elements at a time, even though it only advances WIDTH elements - // in the dst. - while offset + Gather::STRIDE < indices_len { - // SAFETY: `gather` preconditions satisfied: - // 1. `(indices + offset)..(indices + offset + STRIDE)` is in-bounds for indices - // allocation. - // 2. `buffer` has same len as indices so `buffer + offset + WIDTH` is always valid. - // 3. `size_of::() == size_of::()` (asserted above), so the `Lane`-typed - // pointers address the same bytes as the `Out`-typed `values`/`buffer` allocations. - unsafe { - Gather::gather( - indices.as_ptr().add(offset), - max_index, - values.as_ptr().cast::(), - buf_uninit.as_mut_ptr().add(offset).cast::(), - ) - }; - offset += Gather::WIDTH; - } - - // Remainder. - while offset < indices_len { - buf_uninit[offset].write(values[indices[offset].as_()]); - offset += 1; - } - - assert_eq!(offset, indices_len); - - // SAFETY: All elements have been initialized. - unsafe { buffer.set_len(indices_len) }; - - // Reset the buffer alignment to the output type. - // NOTE: if we don't do this, we pass back a Buffer which is over-aligned to the SIMD - // register width. The caller expects that this memory should be aligned to the value type - // so that we can slice it at value boundaries. - buffer = buffer.aligned(Alignment::of::()); - - buffer.freeze() -} - -#[cfg(test)] -#[cfg_attr(miri, ignore)] -#[cfg(target_arch = "x86_64")] -mod avx2_tests { - use super::*; - - macro_rules! test_cases { - (index_type => $IDX:ty, value_types => $($VAL:ty),+) => { - paste::paste! { - $( - // Test "happy path" take, valid indices on valid array. - #[test] - #[allow(clippy::cast_possible_truncation)] - fn []() { - let values: Vec<$VAL> = (1..=127).map(|x| x as $VAL).collect(); - let indices: Vec<$IDX> = (0..127).collect(); - - let result = unsafe { take_avx2(&values, &indices) }; - assert_eq!(&values, result.as_slice()); - } - - // Test take on empty array. - #[test] - #[should_panic] - #[allow(clippy::cast_possible_truncation)] - fn []() { - let values: Vec<$VAL> = vec![]; - let indices: Vec<$IDX> = (0..127).collect(); - let result = unsafe { take_avx2(&values, &indices) }; - assert!(result.is_empty()); - } - - // Test all invalid take indices mapping to zeros. - #[test] - #[should_panic] - #[allow(clippy::cast_possible_truncation)] - fn []() { - let values: Vec<$VAL> = (1..=127).map(|x| x as $VAL).collect(); - // All out of bounds indices. - let indices: Vec<$IDX> = (127..=254).collect(); - - let result = unsafe { take_avx2(&values, &indices) }; - assert_eq!(&[0 as $VAL; 127], result.as_slice()); - } - )+ - } - }; - } - - test_cases!( - index_type => u8, - value_types => u32, i32, u64, i64, f32, f64 - ); - test_cases!( - index_type => u16, - value_types => u32, i32, u64, i64, f32, f64 - ); - test_cases!( - index_type => u32, - value_types => u32, i32, u64, i64, f32, f64 - ); - test_cases!( - index_type => u64, - value_types => u32, i32, u64, i64, f32, f64 - ); - - #[test] - fn test_avx2_take_last_valid_index_u8() { - let values: Vec = (0..(255 + 1)).collect(); - let indices: Vec = vec![255; 20]; - - let result = unsafe { take_avx2(&values, &indices) }; - assert_eq!(&vec![255; indices.len()], result.as_slice()); - } - - #[test] - fn test_avx2_take_last_valid_index_u16() { - let values: Vec = (0..(65535 + 1)).collect(); - let indices: Vec = vec![65535; 20]; - - let result = unsafe { take_avx2(&values, &indices) }; - assert_eq!(&vec![65535; indices.len()], result.as_slice()); - } - - /// A `[u8; 4]` is a 4-byte `Copy` POD that is not a `NativePType`. This proves the kernel - /// gathers an arbitrary 4-byte value type through the `u32` SIMD lane. - #[test] - fn test_avx2_take_simd_array_u8x4() { - let values: Vec<[u8; 4]> = (1u32..=200).map(u32::to_le_bytes).collect(); - let indices: Vec = (0..200).collect(); - - let result = unsafe { take_avx2(&values, &indices) }; - assert_eq!(values.as_slice(), result.as_slice()); - } - - /// 2-byte values have no AVX2 gather, so they take the scalar fallback path and must still be - /// correct. - #[test] - fn test_avx2_take_scalar_fallback_u16() { - let values: Vec = (1..=300).collect(); - let indices: Vec = (0..300).collect(); - - let result = unsafe { take_avx2(&values, &indices) }; - assert_eq!(values.as_slice(), result.as_slice()); - } - - /// Values wider than 8 bytes (e.g. `i128`/decimal backing) exceed the gather lane and fall - /// back to the scalar kernel. - #[test] - fn test_avx2_take_scalar_fallback_array_u8x16() { - let values: Vec<[u8; 16]> = (0u128..200).map(u128::to_le_bytes).collect(); - let indices: Vec = (0..200).collect(); - - let result = unsafe { take_avx2(&values, &indices) }; - assert_eq!(values.as_slice(), result.as_slice()); - } -} diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs deleted file mode 100644 index dd562281608..00000000000 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ /dev/null @@ -1,247 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#[cfg(any(target_arch = "x86_64", target_arch = "x86"))] -mod avx2; - -use std::sync::LazyLock; - -use vortex_buffer::Buffer; -use vortex_buffer::BufferMut; -use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_mask::Mask; - -use crate::ArrayRef; -use crate::IntoArray; -use crate::array::ArrayView; -use crate::arrays::ConstantArray; -use crate::arrays::Primitive; -use crate::arrays::PrimitiveArray; -use crate::arrays::dict::TakeExecute; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -use crate::dtype::IntegerPType; -use crate::executor::ExecutionCtx; -use crate::match_each_integer_ptype; -use crate::match_each_native_ptype; -use crate::scalar::Scalar; -use crate::validity::Validity; - -// Kernel selection happens on the first call to `take` and uses a combination of compile-time -// and runtime feature detection to infer the best kernel for the platform. -static PRIMITIVE_TAKE_KERNEL: LazyLock<&'static dyn TakeImpl> = LazyLock::new(|| { - #[cfg(any(target_arch = "x86_64", target_arch = "x86"))] - { - if is_x86_feature_detected!("avx2") { - &avx2::TakeKernelAVX2 - } else { - &TakeKernelScalar - } - } - - #[cfg(not(any(target_arch = "x86_64", target_arch = "x86")))] - { - &TakeKernelScalar - } -}); - -trait TakeImpl: Send + Sync { - fn take( - &self, - array: ArrayView<'_, Primitive>, - indices: ArrayView<'_, Primitive>, - validity: Validity, - ) -> VortexResult; -} - -struct TakeKernelScalar; - -impl TakeImpl for TakeKernelScalar { - fn take( - &self, - array: ArrayView<'_, Primitive>, - indices: ArrayView<'_, Primitive>, - validity: Validity, - ) -> VortexResult { - match_each_native_ptype!(array.ptype(), |T| { - match_each_integer_ptype!(indices.ptype(), |I| { - let values = take_primitive_scalar(array.as_slice::(), indices.as_slice::()); - Ok(PrimitiveArray::new(values, validity).into_array()) - }) - }) - } -} - -impl TakeExecute for Primitive { - fn take( - array: ArrayView<'_, Primitive>, - indices: &ArrayRef, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - let DType::Primitive(ptype, null) = indices.dtype() else { - vortex_bail!("Invalid indices dtype: {}", indices.dtype()) - }; - - let indices_validity = indices.validity()?; - // Null index lanes are semantically ignored, but their physical values may be out of - // bounds. Redirect those lanes to zero for the cast/gather, then restore the original index - // validity below. - let indices_nulls_zeroed = match indices_validity.execute_mask(indices.len(), ctx)? { - Mask::AllTrue(_) => indices.clone(), - Mask::AllFalse(_) => { - return Ok(Some( - ConstantArray::new(Scalar::null(array.dtype().as_nullable()), indices.len()) - .into_array(), - )); - } - Mask::Values(_) => indices - .clone() - .fill_null(Scalar::from(0).cast(indices.dtype())?)?, - }; - - let unsigned_indices = if ptype.is_unsigned_int() { - indices_nulls_zeroed.execute::(ctx)? - } else { - // This will fail if all values cannot be converted to unsigned - indices_nulls_zeroed - .cast(DType::Primitive(ptype.to_unsigned(), *null))? - .execute::(ctx)? - }; - - let validity = array - .validity()? - .take(&unsigned_indices.clone().into_array())? - .and(indices_validity)?; - // Delegate to the best kernel based on the target CPU - { - let unsigned_indices = unsigned_indices.as_view(); - PRIMITIVE_TAKE_KERNEL - .take(array, unsigned_indices, validity) - .map(Some) - } - } -} - -// Compiler may see this as unused based on enabled features -#[inline(always)] -fn take_primitive_scalar(buffer: &[T], indices: &[I]) -> Buffer { - // NB: The simpler `indices.iter().map(|idx| buffer[idx.as_()]).collect()` generates suboptimal - // assembly where the buffer length is repeatedly loaded from the stack on each iteration. - - let mut result = BufferMut::with_capacity(indices.len()); - let ptr = result.spare_capacity_mut().as_mut_ptr().cast::(); - - // This explicit loop with pointer writes keeps the length in a register and avoids per-element - // capacity checks from `push()`. - for (i, idx) in indices.iter().enumerate() { - // SAFETY: We reserved `indices.len()` capacity, so `ptr.add(i)` is valid. - unsafe { ptr.add(i).write(buffer[idx.as_()]) }; - } - - // SAFETY: We just wrote exactly `indices.len()` elements. - unsafe { result.set_len(indices.len()) }; - result.freeze() -} - -#[cfg(any(target_arch = "x86_64", target_arch = "x86"))] -#[cfg(test)] -mod test { - use rstest::rstest; - use vortex_buffer::buffer; - use vortex_error::VortexExpect; - - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::array_session; - use crate::arrays::BoolArray; - use crate::arrays::PrimitiveArray; - use crate::arrays::primitive::compute::take::take_primitive_scalar; - use crate::compute::conformance::take::test_take_conformance; - use crate::scalar::Scalar; - use crate::validity::Validity; - - #[test] - fn test_take() { - let a = vec![1i32, 2, 3, 4, 5]; - let result = take_primitive_scalar(&a, &[0, 0, 4, 2]); - assert_eq!(result.as_slice(), &[1i32, 1, 5, 3]); - } - - #[test] - fn test_take_with_null_indices() { - let mut ctx = array_session().create_execution_ctx(); - let values = PrimitiveArray::new( - buffer![1i32, 2, 3, 4, 5], - Validity::Array(BoolArray::from_iter([true, true, false, false, true]).into_array()), - ); - let indices = PrimitiveArray::new( - buffer![0, 3, 4], - Validity::Array(BoolArray::from_iter([true, true, false]).into_array()), - ); - let actual = values.take(indices.into_array()).unwrap(); - assert_eq!( - actual.execute_scalar(0, &mut ctx).vortex_expect("no fail"), - Scalar::from(Some(1)) - ); - // position 3 is null - assert_eq!( - actual.execute_scalar(1, &mut ctx).vortex_expect("no fail"), - Scalar::null_native::() - ); - // the third index is null - assert_eq!( - actual.execute_scalar(2, &mut ctx).vortex_expect("no fail"), - Scalar::null_native::() - ); - } - - #[rstest] - #[case(PrimitiveArray::new(buffer![42i32], Validity::NonNullable))] - #[case(PrimitiveArray::new(buffer![0, 1], Validity::NonNullable))] - #[case(PrimitiveArray::new(buffer![0, 1, 2, 3, 4], Validity::NonNullable))] - #[case(PrimitiveArray::new(buffer![0, 1, 2, 3, 4, 5, 6, 7], Validity::NonNullable))] - #[case(PrimitiveArray::new(buffer![0, 1, 2, 3, 4], Validity::AllValid))] - #[case(PrimitiveArray::new( - buffer![0, 1, 2, 3, 4, 5], - Validity::Array(BoolArray::from_iter([true, false, true, false, true, true]).into_array()), - ))] - #[case(PrimitiveArray::from_option_iter([Some(1), None, Some(3), Some(4), None]))] - fn test_take_primitive_conformance(#[case] array: PrimitiveArray) { - test_take_conformance( - &array.into_array(), - &mut array_session().create_execution_ctx(), - ); - } -} - -#[cfg(test)] -mod tests { - use vortex_buffer::buffer; - - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::array_session; - use crate::arrays::BoolArray; - use crate::arrays::PrimitiveArray; - use crate::assert_arrays_eq; - use crate::validity::Validity; - - #[test] - fn take_null_index_skips_out_of_bounds_value() { - let mut ctx = array_session().create_execution_ctx(); - let values = PrimitiveArray::from_iter([10i32, 20, 30]); - let indices = PrimitiveArray::new( - buffer![1u64, 3], - Validity::Array(BoolArray::from_iter([true, false]).into_array()), - ); - - let taken = values.take(indices.into_array()).unwrap(); - - assert_arrays_eq!( - taken, - PrimitiveArray::from_option_iter([Some(20i32), None]).into_array(), - &mut ctx - ); - } -} diff --git a/vortex-array/src/arrays/primitive/mod.rs b/vortex-array/src/arrays/primitive/mod.rs index f1566f8174f..748beda8339 100644 --- a/vortex-array/src/arrays/primitive/mod.rs +++ b/vortex-array/src/arrays/primitive/mod.rs @@ -3,8 +3,10 @@ mod array; pub use array::PrimitiveArrayExt; +pub use array::PrimitiveArraySlotsExt; pub use array::PrimitiveData; pub use array::PrimitiveDataParts; +pub use array::PrimitiveSlots; pub use array::chunk_range; pub use array::patch_chunk; pub use vtable::PrimitiveArray; diff --git a/vortex-array/src/arrays/primitive/vtable/mod.rs b/vortex-array/src/arrays/primitive/vtable/mod.rs index 50189bf0037..45256a7a498 100644 --- a/vortex-array/src/arrays/primitive/vtable/mod.rs +++ b/vortex-array/src/arrays/primitive/vtable/mod.rs @@ -4,7 +4,6 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; -use vortex_error::vortex_panic; use crate::ArrayParts; use crate::ArrayRef; @@ -13,6 +12,7 @@ use crate::ExecutionResult; use crate::array::Array; use crate::array::ArrayView; use crate::array::VTable; +use crate::arrays::fixed_width::vtable as fixed_width; use crate::arrays::primitive::PrimitiveData; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; @@ -21,7 +21,6 @@ use crate::dtype::DType; use crate::dtype::PType; use crate::match_each_native_ptype; use crate::serde::ArrayChildren; -use crate::validity::Validity; mod kernel; mod operations; mod validity; @@ -34,7 +33,7 @@ use vortex_session::registry::CachedId; use crate::EqMode; use crate::array::ArrayId; -use crate::arrays::primitive::array::SLOT_NAMES; +use crate::arrays::primitive::array::PrimitiveSlots; use crate::arrays::primitive::compute::rules::RULES; use crate::hash::ArrayEq; use crate::hash::ArrayHash; @@ -74,17 +73,11 @@ impl VTable for Primitive { } fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { - match idx { - 0 => array.buffer_handle().clone(), - _ => vortex_panic!("PrimitiveArray buffer index {idx} out of bounds"), - } + fixed_width::buffer("PrimitiveArray", array.buffer_handle(), idx) } fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { - match idx { - 0 => Some("values".to_string()), - _ => None, - } + fixed_width::buffer_name(idx) } fn with_buffers( @@ -92,13 +85,8 @@ impl VTable for Primitive { array: ArrayView<'_, Self>, buffers: &[BufferHandle], ) -> VortexResult> { - vortex_ensure!( - buffers.len() == 1, - "Expected 1 buffer, got {}", - buffers.len() - ); let mut data = array.data().clone(); - data.buffer = buffers[0].clone(); + data.buffer = fixed_width::single_buffer(buffers)?; Ok( ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) .with_slots(array.slots().iter().cloned().collect()), @@ -128,7 +116,8 @@ impl VTable for Primitive { data.len(), len ); - let validity = crate::array::child_to_validity(slots[0].as_ref(), *nullability); + let validity = + crate::array::child_to_validity(slots[PrimitiveSlots::VALIDITY].as_ref(), *nullability); if let Some(validity_len) = validity.maybe_len() { vortex_ensure!( validity_len == len, @@ -157,19 +146,9 @@ impl VTable for Primitive { metadata.len() ); } - if buffers.len() != 1 { - vortex_bail!("Expected 1 buffer, got {}", buffers.len()); - } - let buffer = buffers[0].clone(); - - let validity = if children.is_empty() { - Validity::from(dtype.nullability()) - } else if children.len() == 1 { - let validity = children.get(0, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - vortex_bail!("Expected 0 or 1 child, got {}", children.len()); - }; + let buffer = fixed_width::single_buffer(buffers)?; + + let validity = fixed_width::deserialize_validity(dtype.nullability(), len, children)?; let ptype = PType::try_from(dtype)?; @@ -188,21 +167,14 @@ impl VTable for Primitive { ); } - vortex_ensure!( - buffer.is_aligned_to(Alignment::new(ptype.byte_width())), - "PrimitiveArray::build: Buffer (align={}) must be aligned to {}", - buffer.alignment(), - ptype.byte_width() - ); - - // SAFETY: checked ahead of time + // SAFETY: the buffer length and alignment are checked above. let slots = PrimitiveData::make_slots(&validity, len); let data = unsafe { PrimitiveData::new_unchecked_from_handle(buffer, ptype, validity) }; Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + PrimitiveSlots::NAMES[idx].to_string() } fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { diff --git a/vortex-array/src/arrays/scalar_fn/array.rs b/vortex-array/src/arrays/scalar_fn/array.rs index a999a25cd51..7f3a436996f 100644 --- a/vortex-array/src/arrays/scalar_fn/array.rs +++ b/vortex-array/src/arrays/scalar_fn/array.rs @@ -32,6 +32,7 @@ impl Display for ScalarFnData { impl ScalarFnData { /// Get the scalar function bound to this array. + #[allow(clippy::inline_always)] #[inline(always)] pub fn scalar_fn(&self) -> &ScalarFnRef { &self.scalar_fn @@ -87,13 +88,16 @@ impl Array { children: Vec, len: usize, ) -> VortexResult { + Self::validate_arity(&scalar_fn, children.len())?; Self::validate_children_len(&children, len)?; + let arg_dtypes: Vec<_> = children.iter().map(|c| c.dtype().clone()).collect(); let dtype = scalar_fn.return_dtype(&arg_dtypes)?; let data = ScalarFnData { scalar_fn: scalar_fn.clone(), }; let vtable = ScalarFn { id: scalar_fn.id() }; + Ok(unsafe { Array::from_parts_unchecked( ArrayParts::new(vtable, dtype, len, data) @@ -109,6 +113,15 @@ impl Array { Ok(child.len()) } + fn validate_arity(scalar_fn: &ScalarFnRef, child_count: usize) -> VortexResult<()> { + let arity = scalar_fn.signature().arity(); + vortex_ensure!( + arity.matches(child_count), + "ScalarFnArray requires {arity} children, got {child_count}" + ); + Ok(()) + } + fn validate_children_len(children: &[ArrayRef], len: usize) -> VortexResult<()> { vortex_ensure!( children.iter().all(|c| c.len() == len), diff --git a/vortex-array/src/arrays/scalar_fn/mod.rs b/vortex-array/src/arrays/scalar_fn/mod.rs index 6c508dcad96..681ff03b60a 100644 --- a/vortex-array/src/arrays/scalar_fn/mod.rs +++ b/vortex-array/src/arrays/scalar_fn/mod.rs @@ -7,5 +7,4 @@ mod rules; mod vtable; pub use array::ScalarFnArrayExt; -pub use vtable::ScalarFnFactoryExt; pub use vtable::*; diff --git a/vortex-array/src/arrays/scalar_fn/plugin.rs b/vortex-array/src/arrays/scalar_fn/plugin.rs index 044f79362cd..46c1a597f9a 100644 --- a/vortex-array/src/arrays/scalar_fn/plugin.rs +++ b/vortex-array/src/arrays/scalar_fn/plugin.rs @@ -2,16 +2,18 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_session::VortexSession; +use crate::ArrayDeserialization; use crate::ArrayId; use crate::ArrayPlugin; use crate::ArrayRef; +use crate::ArraySerialization; use crate::IntoArray; use crate::arrays::ScalarFnArray; use crate::arrays::scalar_fn::ExactScalarFn; use crate::arrays::scalar_fn::ScalarFnArrayView; -use crate::buffer::BufferHandle; use crate::dtype::DType; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::TypedScalarFnInstance; @@ -64,27 +66,38 @@ impl ArrayPlugin for ScalarFnArrayPlugi &self, array: &ArrayRef, session: &VortexSession, - ) -> VortexResult>> { + ) -> VortexResult> { // We serialize the scalar function options, along with any scalar function array data. let scalar_fn = array.as_::>(); - ::serialize(&self.0, &scalar_fn, session) + Ok( + ::serialize(&self.0, &scalar_fn, session)? + .map(|metadata| ArraySerialization::from_array(self.id(), array, metadata)), + ) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, session: &VortexSession, ) -> VortexResult { - let parts = ::deserialize( - &self.0, dtype, len, metadata, children, session, + vortex_ensure!( + parts.serialized_id == self.id(), + "scalar function array plugin {} does not recognize serialized ID {}", + self.id(), + parts.serialized_id, + ); + let len = parts.len; + let scalar_parts = ::deserialize( + &self.0, + parts.dtype, + parts.len, + parts.metadata, + parts.children, + session, )?; Ok(ScalarFnArray::try_new_with_len( - TypedScalarFnInstance::new(self.0.clone(), parts.options).erased(), - parts.children, + TypedScalarFnInstance::new(self.0.clone(), scalar_parts.options).erased(), + scalar_parts.children, len, )? .into_array()) diff --git a/vortex-array/src/arrays/scalar_fn/rules.rs b/vortex-array/src/arrays/scalar_fn/rules.rs index abc8e18e762..71f1a628eb3 100644 --- a/vortex-array/src/arrays/scalar_fn/rules.rs +++ b/vortex-array/src/arrays/scalar_fn/rules.rs @@ -1,11 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::any::Any; -use std::sync::Arc; - use itertools::Itertools; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::ArrayRef; @@ -18,16 +14,13 @@ use crate::arrays::ScalarFn; use crate::arrays::ScalarFnArray; use crate::arrays::Slice; use crate::arrays::StructArray; +use crate::arrays::filter::prepare_mask_for_reuse; use crate::arrays::scalar_fn::ScalarFnArrayExt; -use crate::dtype::DType; use crate::optimizer::rules::ArrayParentReduceRule; use crate::optimizer::rules::ArrayReduceRule; use crate::optimizer::rules::ParentRuleSet; use crate::optimizer::rules::ReduceRuleSet; -use crate::scalar_fn::ReduceCtx; -use crate::scalar_fn::ReduceNode; -use crate::scalar_fn::ReduceNodeRef; -use crate::scalar_fn::ScalarFnRef; +use crate::scalar_fn::ArrayReduceNode; use crate::scalar_fn::fns::pack::Pack; use crate::validity::Validity; @@ -35,7 +28,7 @@ pub(super) const RULES: ReduceRuleSet = ReduceRuleSet::new(&[&ScalarFnPackToStructRule, &ScalarFnAbstractReduceRule]); pub(super) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&ScalarFnUnaryFilterPushDownRule), + ParentRuleSet::lift(&ScalarFilterPushdownRule), ParentRuleSet::lift(&ScalarFnSliceReduceRule), ]); @@ -94,77 +87,18 @@ impl ArrayParentReduceRule for ScalarFnSliceReduceRule { struct ScalarFnAbstractReduceRule; impl ArrayReduceRule for ScalarFnAbstractReduceRule { fn reduce(&self, array: ArrayView<'_, ScalarFn>) -> VortexResult> { - if let Some(reduced) = array - .scalar_fn() - .reduce(array.as_ref(), &ArrayReduceCtx { len: array.len() })? - { - return Ok(Some( - reduced - .as_any() - .downcast_ref::() - .vortex_expect("ReduceNode is not an ArrayRef") - .clone(), - )); + let node = ArrayReduceNode::new(array.as_ref()); + if let Some(reduced) = array.scalar_fn().reduce_array(&node)? { + return Ok(Some(reduced.into_array())); } Ok(None) } } -impl ReduceNode for ArrayRef { - fn as_any(&self) -> &dyn Any { - self - } - - fn node_dtype(&self) -> VortexResult { - Ok(self.dtype().clone()) - } - - fn scalar_fn(&self) -> Option<&ScalarFnRef> { - self.as_opt::().map(|a| a.data().scalar_fn()) - } - - fn child(&self, idx: usize) -> ReduceNodeRef { - Arc::new(self.nth_child(idx).vortex_expect("child idx out of bounds")) - } - - fn child_count(&self) -> usize { - self.nchildren() - } -} - -struct ArrayReduceCtx { - // The length of the array being reduced - len: usize, -} -impl ReduceCtx for ArrayReduceCtx { - fn new_node( - &self, - scalar_fn: ScalarFnRef, - children: &[ReduceNodeRef], - ) -> VortexResult { - Ok(Arc::new( - ScalarFnArray::try_new_with_len( - scalar_fn, - children - .iter() - .map(|c| { - c.as_any() - .downcast_ref::() - .vortex_expect("ReduceNode is not an ArrayRef") - .clone() - }) - .collect(), - self.len, - )? - .into_array(), - )) - } -} - #[derive(Debug)] -struct ScalarFnUnaryFilterPushDownRule; +struct ScalarFilterPushdownRule; -impl ArrayParentReduceRule for ScalarFnUnaryFilterPushDownRule { +impl ArrayParentReduceRule for ScalarFilterPushdownRule { type Parent = Filter; fn reduce_parent( @@ -173,48 +107,145 @@ impl ArrayParentReduceRule for ScalarFnUnaryFilterPushDownRule { parent: ArrayView<'_, Filter>, _child_idx: usize, ) -> VortexResult> { - // If we only have one non-constant child, then it is _always_ cheaper to push down the - // filter over the children of the scalar function array. - if child + let nchildren = child .iter_children() .filter(|c| !c.is::()) - .count() - == 1 + .count(); + if nchildren > 1 + && let Some(values) = parent.filter_mask().values() { - let new_children: Vec<_> = child - .iter_children() - .map(|c| match c.as_opt::() { - Some(array) => { - Ok(ConstantArray::new(array.scalar().clone(), parent.len()).into_array()) - } - None => c.filter(parent.filter_mask().clone()), - }) - .try_collect()?; - - let new_array = - ScalarFnArray::try_new(child.scalar_fn().clone(), new_children)?.into_array(); - - return Ok(Some(new_array)); + prepare_mask_for_reuse(values, nchildren); } - Ok(None) + let new_children: Vec<_> = child + .iter_children() + .map(|c| match c.as_opt::() { + Some(array) => { + Ok(ConstantArray::new(array.scalar().clone(), parent.len()).into_array()) + } + None => c.filter(parent.filter_mask().clone()), + }) + .try_collect()?; + + Ok(Some( + ScalarFnArray::try_new_with_len(child.scalar_fn().clone(), new_children, parent.len())? + .into_array(), + )) } } #[cfg(test)] mod tests { + use rstest::rstest; use vortex_error::VortexExpect; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + use vortex_mask::Mask; + use super::ScalarFilterPushdownRule; + use crate::VortexSessionExecute; use crate::array::IntoArray; + use crate::array_session; use crate::arrays::ChunkedArray; + use crate::arrays::Constant; + use crate::arrays::Filter; + use crate::arrays::FilterArray; use crate::arrays::PrimitiveArray; + use crate::arrays::ScalarFn; + use crate::arrays::ScalarFnArray; + use crate::arrays::scalar_fn::ScalarFnArrayExt; use crate::arrays::scalar_fn::rules::ConstantArray; + use crate::assert_arrays_eq; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; use crate::expr::cast; use crate::expr::is_null; use crate::expr::root; + use crate::optimizer::rules::ArrayParentReduceRule; + use crate::scalar::Scalar; + use crate::scalar_fn::TypedScalarFnInstance; + use crate::scalar_fn::fns::binary::Binary; + use crate::scalar_fn::fns::literal::Literal; + use crate::scalar_fn::fns::operators::Operator; + + #[rstest] + #[case(0, [14, 14, 14, 14, 14])] + #[case(1, [7, 27, 47, 67, 87])] + #[case(2, [0, 40, 80, 120, 160])] + fn test_filter_pushdown( + #[case] nchildren: usize, + #[case] expected: [i32; 5], + ) -> VortexResult<()> { + let children = (0..2) + .map(|index| { + if index < nchildren { + PrimitiveArray::from_iter(0i32..100).into_array() + } else { + ConstantArray::new(7i32, 100).into_array() + } + }) + .collect(); + let array = ScalarFnArray::try_new( + TypedScalarFnInstance::new(Binary, Operator::Add).erased(), + children, + )?; + let mask = Mask::from_iter((0..100).map(|index| index % 20 == 0)); + let values = mask + .values() + .ok_or_else(|| vortex_err!("expected mask values"))?; + assert!(values.cached_indices().is_none()); + let parent = FilterArray::try_new(array.clone().into_array(), mask.clone())?; + + let result = ScalarFilterPushdownRule + .reduce_parent(array.as_view(), parent.as_view(), 0)? + .ok_or_else(|| vortex_err!("expected filter pushdown"))?; + + assert_eq!(values.cached_indices().is_some(), nchildren > 1); + let scalar_fn = result.as_::(); + assert_eq!(scalar_fn.len(), 5); + for (index, child) in scalar_fn.iter_children().enumerate() { + assert_eq!(child.len(), 5); + if index < nchildren { + assert!(child.is::()); + } else { + assert!(child.is::()); + } + } + let expected = PrimitiveArray::from_iter(expected); + assert_arrays_eq!( + result, + expected, + &mut array_session().create_execution_ctx() + ); + Ok(()) + } + + #[test] + fn test_filter_pushdown_without_children() -> VortexResult<()> { + let array = ScalarFnArray::try_new_with_len( + TypedScalarFnInstance::new(Literal, Scalar::from(7i32)).erased(), + vec![], + 3, + )?; + let parent = FilterArray::try_new( + array.clone().into_array(), + Mask::from_iter([true, false, true]), + )?; + + let result = ScalarFilterPushdownRule + .reduce_parent(array.as_view(), parent.as_view(), 0)? + .ok_or_else(|| vortex_err!("expected filter pushdown"))?; + + assert_eq!(result.as_::().nchildren(), 0); + assert_eq!(result.len(), 2); + assert_arrays_eq!( + result, + ConstantArray::new(7i32, 2), + &mut array_session().create_execution_ctx() + ); + Ok(()) + } #[test] fn test_empty_constants() { diff --git a/vortex-array/src/arrays/scalar_fn/vtable/mod.rs b/vortex-array/src/arrays/scalar_fn/vtable/mod.rs index d98e00975ea..c84a36a3169 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/mod.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/mod.rs @@ -20,9 +20,7 @@ use vortex_session::registry::CachedId; use crate::ArrayEq; use crate::ArrayHash; use crate::ArrayRef; -use crate::ArraySlots; use crate::EqMode; -use crate::IntoArray; use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayParts; @@ -38,6 +36,7 @@ use crate::dtype::DType; use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; use crate::expr::Expression; +use crate::expr::display::ExprDisplay; use crate::matcher::Matcher; use crate::scalar_fn; use crate::scalar_fn::Arity; @@ -84,10 +83,24 @@ impl VTable for ScalarFn { len: usize, slots: &[Option], ) -> VortexResult<()> { + let scalar_fn = data.scalar_fn(); vortex_ensure!( - data.scalar_fn.id() == self.id, + scalar_fn.id() == self.id, "ScalarFnArray data scalar_fn does not match vtable" ); + + let missing_children = slots.iter().filter(|slot| slot.is_none()).count(); + vortex_ensure!( + missing_children == 0, + "ScalarFnArray requires every child slot to be present, got {missing_children} missing" + ); + + let arity = scalar_fn.signature().arity(); + vortex_ensure!( + arity.matches(slots.len()), + "ScalarFnArray requires {arity} children, got {}", + slots.len() + ); vortex_ensure!( slots.iter().flatten().all(|c| c.len() == len), "All child arrays must have the same length as the scalar function array" @@ -99,7 +112,7 @@ impl VTable for ScalarFn { .map(|c| c.dtype().clone()) .collect_vec(); vortex_ensure!( - data.scalar_fn.return_dtype(&child_dtypes)? == *dtype, + scalar_fn.return_dtype(&child_dtypes)? == *dtype, "ScalarFnArray dtype does not match scalar function return dtype" ); Ok(()) @@ -155,7 +168,6 @@ impl VTable for ScalarFn { } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - ctx.log(format_args!("scalar_fn({}): executing", array.scalar_fn())); let args = VecExecutionArgs::new(array.children(), array.len()); array .scalar_fn() @@ -176,40 +188,6 @@ impl VTable for ScalarFn { } } -/// Array factory functions for scalar functions. -pub trait ScalarFnFactoryExt: scalar_fn::ScalarFnVTable { - fn try_new_array( - &self, - len: usize, - options: Self::Options, - children: impl Into>, - ) -> VortexResult { - let scalar_fn = scalar_fn::TypedScalarFnInstance::new(self.clone(), options).erased(); - - let children = children.into(); - vortex_ensure!( - children.iter().all(|c| c.len() == len), - "All child arrays must have the same length as the scalar function array" - ); - - let child_dtypes = children.iter().map(|c| c.dtype().clone()).collect_vec(); - let dtype = scalar_fn.return_dtype(&child_dtypes)?; - - let data = ScalarFnData { - scalar_fn: scalar_fn.clone(), - }; - let vtable = ScalarFn { id: scalar_fn.id() }; - Ok(unsafe { - Array::from_parts_unchecked( - ArrayParts::new(vtable, dtype, len, data) - .with_slots(children.into_iter().map(Some).collect::()), - ) - } - .into_array()) - } -} -impl ScalarFnFactoryExt for V {} - /// A matcher that matches any scalar function expression. #[derive(Debug)] pub struct AnyScalarFn; @@ -310,7 +288,7 @@ impl scalar_fn::ScalarFnVTable for ArrayExpr { fn fmt_sql( &self, options: &Self::Options, - _expr: &Expression, + _expr: &dyn ExprDisplay, f: &mut Formatter<'_>, ) -> std::fmt::Result { write!(f, "{}", options.0.encoding_id()) @@ -337,4 +315,8 @@ impl scalar_fn::ScalarFnVTable for ArrayExpr { let validity_array = options.0.validity()?.to_array(options.0.len()); Ok(Some(ArrayExpr.new_expr(FakeEq(validity_array), []))) } + + fn is_strict(&self, _options: &Self::Options) -> bool { + true + } } diff --git a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs index c19e3e2288c..40d75906356 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs @@ -58,14 +58,19 @@ mod tests { use vortex_buffer::buffer; use vortex_error::VortexResult; + use crate::ArraySlots; use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array::Array; + use crate::array::ArrayParts; use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::ScalarFnArray; use crate::arrays::scalar_fn::ScalarFnArrayExt; + use crate::arrays::scalar_fn::array::ScalarFnData; + use crate::arrays::scalar_fn::vtable::ScalarFn; use crate::assert_arrays_eq; use crate::scalar::Scalar; use crate::scalar_fn::TypedScalarFnInstance; @@ -110,6 +115,50 @@ mod tests { ); } + #[test] + fn rejects_wrong_arity() { + let child = buffer![1i32, 2, 3].into_array(); + let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Add).erased(); + + let Err(err) = ScalarFnArray::try_new(scalar_fn.clone(), vec![child.clone()]) else { + panic!("Binary must reject one child"); + }; + assert!( + err.to_string() + .contains("ScalarFnArray requires 2 children, got 1") + ); + + let Err(err) = ScalarFnArray::try_new(scalar_fn, vec![child.clone(), child.clone(), child]) + else { + panic!("Binary must reject three children"); + }; + assert!( + err.to_string() + .contains("ScalarFnArray requires 2 children, got 3") + ); + } + + #[test] + fn rejects_missing_child_slot() { + let lhs = buffer![1i32, 2, 3].into_array(); + let dtype = lhs.dtype().clone(); + let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Add).erased(); + let vtable = ScalarFn { id: scalar_fn.id() }; + let data = ScalarFnData { scalar_fn }; + let slots = [Some(lhs), None].into_iter().collect::(); + + let Err(err) = Array::::try_from_parts( + ArrayParts::new(vtable, dtype, 3, data).with_slots(slots), + ) else { + panic!("ScalarFnArray must reject missing child slots"); + }; + + assert!( + err.to_string() + .contains("ScalarFnArray requires every child slot to be present, got 1 missing") + ); + } + #[test] fn test_scalar_fn_without_children_requires_explicit_len() -> VortexResult<()> { let scalar_fn = TypedScalarFnInstance::new(Literal, Scalar::from(1i32)).erased(); diff --git a/vortex-array/src/arrays/scalar_fn/vtable/validity.rs b/vortex-array/src/arrays/scalar_fn/vtable/validity.rs index 39ab8c2c466..4e8ab0c95c0 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/validity.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/validity.rs @@ -21,7 +21,6 @@ use crate::legacy_session; use crate::scalar_fn::TypedScalarFnInstance; use crate::scalar_fn::VecExecutionArgs; use crate::scalar_fn::fns::literal::Literal; -use crate::scalar_fn::fns::root::Root; use crate::validity::Validity; /// Execute an expression tree recursively. @@ -32,10 +31,10 @@ fn execute_expr( row_count: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - // Handle Root expression - this should not happen in validity expressions - if expr.is::() { - vortex_bail!("Root expression cannot be executed in validity context"); - } + // Only Expression::Scalar is executable + let Some(scalar_fn) = expr.as_scalar() else { + vortex_bail!("Only Expression::Scalar is executable"); + }; // Handle Literal expression - create a constant array if expr.is::() { @@ -52,7 +51,7 @@ fn execute_expr( let args = VecExecutionArgs::new(inputs, row_count); - Ok(expr.scalar_fn().execute(&args, ctx)?.into_array()) + Ok(scalar_fn.execute(&args, ctx)?.into_array()) } impl ValidityVTable for ScalarFn { diff --git a/vortex-array/src/arrays/shared/array.rs b/vortex-array/src/arrays/shared/array.rs index 9619018ff89..fb588e7646e 100644 --- a/vortex-array/src/arrays/shared/array.rs +++ b/vortex-array/src/arrays/shared/array.rs @@ -8,9 +8,7 @@ use std::sync::Arc; use std::sync::OnceLock; use async_lock::Mutex as AsyncMutex; -use smallvec::smallvec; use vortex_error::SharedVortexResult; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::ArrayRef; @@ -19,12 +17,15 @@ use crate::IntoArray; use crate::array::Array; use crate::array::ArrayParts; use crate::array::TypedArrayRef; +use crate::array_slots; use crate::arrays::Shared; -/// The source array that is shared and lazily computed. -pub(super) const SOURCE_SLOT: usize = 0; -pub(super) const NUM_SLOTS: usize = 1; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["source"]; +#[array_slots(Shared)] +pub struct SharedSlots { + /// The source array that is shared and lazily computed. + #[slot(0)] + pub source: ArrayRef, +} /// A lazily-executing array wrapper with a one-way transition from source to cached form. /// @@ -43,13 +44,7 @@ impl Display for SharedData { } #[expect(async_fn_in_trait)] -pub trait SharedArrayExt: TypedArrayRef { - fn source(&self) -> &ArrayRef { - self.as_ref().slots()[SOURCE_SLOT] - .as_ref() - .vortex_expect("validated shared source slot") - } - +pub trait SharedArrayExt: TypedArrayRef + SharedArraySlotsExt { fn current_array_ref(&self) -> &ArrayRef { match self.cached.get() { Some(Ok(arr)) => arr, @@ -116,7 +111,7 @@ impl Array { unsafe { Array::from_parts_unchecked( ArrayParts::new(Shared, dtype, len, SharedData::new()) - .with_slots(smallvec![Some(source)]), + .with_slots(SharedSlots { source }.into_slots()), ) } } diff --git a/vortex-array/src/arrays/shared/mod.rs b/vortex-array/src/arrays/shared/mod.rs index 5d947388a5b..7f3e3cc9268 100644 --- a/vortex-array/src/arrays/shared/mod.rs +++ b/vortex-array/src/arrays/shared/mod.rs @@ -5,7 +5,10 @@ mod array; mod vtable; pub use array::SharedArrayExt; +pub use array::SharedArraySlotsExt; pub use array::SharedData; +pub use array::SharedSlots; +pub use array::SharedSlotsView; pub use vtable::Shared; pub use vtable::SharedArray; diff --git a/vortex-array/src/arrays/shared/vtable.rs b/vortex-array/src/arrays/shared/vtable.rs index fb2c67d1d65..758d091b557 100644 --- a/vortex-array/src/arrays/shared/vtable.rs +++ b/vortex-array/src/arrays/shared/vtable.rs @@ -26,7 +26,7 @@ use crate::array::ValidityVTable; use crate::array::with_empty_buffers; use crate::arrays::shared::SharedArrayExt; use crate::arrays::shared::SharedData; -use crate::arrays::shared::array::SLOT_NAMES; +use crate::arrays::shared::SharedSlots; use crate::buffer::BufferHandle; use crate::dtype::DType; use crate::scalar::Scalar; @@ -66,7 +66,7 @@ impl VTable for Shared { len: usize, slots: &[Option], ) -> VortexResult<()> { - let source = slots[0] + let source = slots[SharedSlots::SOURCE] .as_ref() .vortex_expect("SharedArray source slot must be present"); vortex_error::vortex_ensure!(source.dtype() == dtype, "SharedArray dtype mismatch"); @@ -95,7 +95,7 @@ impl VTable for Shared { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + SharedSlots::NAMES[idx].to_string() } fn serialize( diff --git a/vortex-array/src/arrays/slice/array.rs b/vortex-array/src/arrays/slice/array.rs index 9d7e0bb13ec..405b128fedc 100644 --- a/vortex-array/src/arrays/slice/array.rs +++ b/vortex-array/src/arrays/slice/array.rs @@ -5,21 +5,21 @@ use std::fmt::Display; use std::fmt::Formatter; use std::ops::Range; -use smallvec::smallvec; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_panic; use crate::ArrayRef; use crate::array::Array; use crate::array::ArrayParts; -use crate::array::TypedArrayRef; +use crate::array_slots; use crate::arrays::Slice; -/// The underlying child array being sliced. -pub(super) const CHILD_SLOT: usize = 0; -pub(super) const NUM_SLOTS: usize = 1; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["child"]; +#[array_slots(Slice)] +pub struct SliceSlots { + /// The underlying child array being sliced. + #[slot(0)] + pub child: ArrayRef, +} #[derive(Clone, Debug)] pub struct SliceData { @@ -36,15 +36,6 @@ pub struct SliceDataParts { pub range: Range, } -pub trait SliceArrayExt: TypedArrayRef { - fn child(&self) -> &ArrayRef { - self.as_ref().slots()[CHILD_SLOT] - .as_ref() - .vortex_expect("validated slice child slot") - } -} -impl> SliceArrayExt for T {} - impl SliceData { fn try_new(child_len: usize, range: Range) -> VortexResult { if range.end > child_len { @@ -89,7 +80,8 @@ impl Array { let data = SliceData::try_new(child.len(), range)?; Ok(unsafe { Array::from_parts_unchecked( - ArrayParts::new(Slice, dtype, len, data).with_slots(smallvec![Some(child)]), + ArrayParts::new(Slice, dtype, len, data) + .with_slots(SliceSlots { child }.into_slots()), ) }) } @@ -101,7 +93,8 @@ impl Array { let data = SliceData::new(range); unsafe { Array::from_parts_unchecked( - ArrayParts::new(Slice, dtype, len, data).with_slots(smallvec![Some(child)]), + ArrayParts::new(Slice, dtype, len, data) + .with_slots(SliceSlots { child }.into_slots()), ) } } diff --git a/vortex-array/src/arrays/slice/mod.rs b/vortex-array/src/arrays/slice/mod.rs index 8eec9d85310..bdbd75f45b6 100644 --- a/vortex-array/src/arrays/slice/mod.rs +++ b/vortex-array/src/arrays/slice/mod.rs @@ -15,9 +15,11 @@ mod vtable; use std::ops::Range; -pub use array::SliceArrayExt; +pub use array::SliceArraySlotsExt; pub use array::SliceData; pub use array::SliceDataParts; +pub use array::SliceSlots; +pub use array::SliceSlotsView; use vortex_error::VortexResult; pub use vtable::*; @@ -64,10 +66,14 @@ pub trait SliceKernel: VTable { ) -> VortexResult>; } -fn precondition(array: ArrayView<'_, V>, range: &Range) -> Option { +/// Short-circuits slice for the ranges that need no encoding-specific work. +/// +/// Returns `Some(result)` when the answer is already known, or `None` when the slice must proceed +/// normally. +fn short_circuit(array: ArrayView<'_, V>, range: &Range) -> Option { if range.start == 0 && range.end == array.len() { return Some(array.array().clone()); - }; + } if range.start == range.end { return Some(Canonical::empty(array.dtype()).into_array()); } @@ -91,7 +97,7 @@ where child_idx: usize, ) -> VortexResult> { assert_eq!(child_idx, 0); - if let Some(result) = precondition::(array, &parent.range) { + if let Some(result) = short_circuit::(array, &parent.range) { return Ok(Some(result)); } ::slice(array, parent.range.clone()) @@ -116,7 +122,7 @@ where ctx: &mut ExecutionCtx, ) -> VortexResult> { assert_eq!(child_idx, 0); - if let Some(result) = precondition::(array, &parent.range) { + if let Some(result) = short_circuit::(array, &parent.range) { return Ok(Some(result)); } ::slice(array, parent.range.clone(), ctx) diff --git a/vortex-array/src/arrays/slice/rules.rs b/vortex-array/src/arrays/slice/rules.rs index 3bc0a70f204..62ed8892541 100644 --- a/vortex-array/src/arrays/slice/rules.rs +++ b/vortex-array/src/arrays/slice/rules.rs @@ -1,9 +1,38 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::array::ArrayView; use crate::arrays::Slice; +use crate::arrays::scalar_fn::ExactScalarFn; +use crate::arrays::scalar_fn::ScalarFnArrayView; +use crate::arrays::slice::SliceArraySlotsExt; use crate::arrays::slice::SliceReduceAdaptor; +use crate::builtins::ArrayBuiltins; +use crate::optimizer::rules::ArrayParentReduceRule; use crate::optimizer::rules::ParentRuleSet; +use crate::scalar_fn::fns::get_item::GetItem; + +pub(super) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ + ParentRuleSet::lift(&SliceReduceAdaptor(Slice)), + ParentRuleSet::lift(&SliceGetItemRule), +]); + +#[derive(Debug)] +struct SliceGetItemRule; + +impl ArrayParentReduceRule for SliceGetItemRule { + type Parent = ExactScalarFn; -pub(super) const PARENT_RULES: ParentRuleSet = - ParentRuleSet::new(&[ParentRuleSet::lift(&SliceReduceAdaptor(Slice))]); + fn reduce_parent( + &self, + array: ArrayView<'_, Slice>, + parent: ScalarFnArrayView<'_, GetItem>, + _child_idx: usize, + ) -> VortexResult> { + let field = array.child().get_item(parent.options.clone())?; + Ok(Some(field.slice(array.slice_range().clone())?)) + } +} diff --git a/vortex-array/src/arrays/slice/slice_.rs b/vortex-array/src/arrays/slice/slice_.rs index 0f5fba2d80b..e30ddcc6a10 100644 --- a/vortex-array/src/arrays/slice/slice_.rs +++ b/vortex-array/src/arrays/slice/slice_.rs @@ -10,7 +10,7 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::Slice; use crate::arrays::SliceArray; -use crate::arrays::slice::SliceArrayExt; +use crate::arrays::slice::SliceArraySlotsExt; use crate::arrays::slice::SliceReduce; impl SliceReduce for Slice { diff --git a/vortex-array/src/arrays/slice/vtable.rs b/vortex-array/src/arrays/slice/vtable.rs index 1ff2e5bc84d..c88c28a9a4d 100644 --- a/vortex-array/src/arrays/slice/vtable.rs +++ b/vortex-array/src/arrays/slice/vtable.rs @@ -28,10 +28,9 @@ use crate::array::OperationsVTable; use crate::array::VTable; use crate::array::ValidityVTable; use crate::array::with_empty_buffers; -use crate::arrays::slice::SliceArrayExt; -use crate::arrays::slice::array::CHILD_SLOT; -use crate::arrays::slice::array::SLOT_NAMES; +use crate::arrays::slice::SliceArraySlotsExt; use crate::arrays::slice::array::SliceData; +use crate::arrays::slice::array::SliceSlots; use crate::arrays::slice::rules::PARENT_RULES; use crate::buffer::BufferHandle; use crate::dtype::DType; @@ -78,10 +77,10 @@ impl VTable for Slice { slots: &[Option], ) -> VortexResult<()> { vortex_ensure!( - slots[CHILD_SLOT].is_some(), + slots[SliceSlots::CHILD].is_some(), "SliceArray child slot must be present" ); - let child = slots[CHILD_SLOT] + let child = slots[SliceSlots::CHILD] .as_ref() .vortex_expect("validated child slot"); vortex_ensure!( @@ -126,7 +125,7 @@ impl VTable for Slice { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + SliceSlots::NAMES[idx].to_string() } fn serialize( @@ -151,7 +150,7 @@ impl VTable for Slice { } fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { - let array = require_child!(array, array.child(), CHILD_SLOT => AnyCanonical); + let array = require_child!(array, array.child(), SliceSlots::CHILD => AnyCanonical); debug_assert!(array.child().is_canonical()); // TODO(ngates): we should inline canonical slice logic here. diff --git a/vortex-array/src/arrays/struct_/array.rs b/vortex-array/src/arrays/struct_/array.rs index 8b939deb911..5f2a184bfb7 100644 --- a/vortex-array/src/arrays/struct_/array.rs +++ b/vortex-array/src/arrays/struct_/array.rs @@ -3,9 +3,9 @@ use std::borrow::Borrow; use std::iter::once; -use std::sync::Arc; use itertools::Itertools; +use vortex_array_macros::array_slots; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -30,10 +30,16 @@ use crate::dtype::StructFields; use crate::validity::Validity; // StructArray has a variable number of slots: [validity?, field_0, ..., field_N] -/// The validity bitmap indicating which struct elements are non-null. -pub(super) const VALIDITY_SLOT: usize = 0; -/// The offset at which the struct field arrays begin in the slots vector. -pub(super) const FIELDS_OFFSET: usize = 1; +/// Slot layout of a [`Struct`] array: `[validity?, fields...]`. +#[array_slots(Struct)] +pub struct StructSlots { + /// The optional row-level validity child. + #[slot(0)] + pub validity: Option, + /// The field arrays, one per struct field, all sharing the outer length. + #[slot(1..)] + pub fields: Vec, +} /// A struct array that stores multiple named fields as columns, similar to a database row. /// @@ -159,21 +165,43 @@ pub(super) const FIELDS_OFFSET: usize = 1; /// ``` pub struct StructDataParts { pub struct_fields: StructFields, - pub fields: Arc<[ArrayRef]>, + pub fields: Vec, pub validity: Validity, } +/// Allocate the slots of a [`Struct`] array holding only its validity, with room for `nfields` +/// field slots to be pushed. +/// +/// Callers that build their field arrays fallibly use this directly, so they don't have to stage +/// them in a `Vec` just to hand them over as an infallible iterator. +pub(super) fn struct_slots_with_capacity( + validity: &Validity, + length: usize, + nfields: usize, +) -> ArraySlots { + let mut slots = ArraySlots::with_capacity(StructSlots::FIELDS_OFFSET + nfields); + slots.push(validity_to_child(validity, length)); + slots +} + pub(super) fn make_struct_slots( - fields: &[ArrayRef], + fields: impl IntoIterator, validity: &Validity, length: usize, ) -> ArraySlots { - once(validity_to_child(validity, length)) - .chain(fields.iter().cloned().map(Some)) - .collect() + // Take the fields by value so callers that already own them move their `ArrayRef`s straight + // into the `SmallVec` instead of paying a refcount bump per field. + let fields = fields.into_iter(); + let mut slots = struct_slots_with_capacity(validity, length, fields.size_hint().0); + slots.extend(fields.map(Some)); + slots } -pub trait StructArrayExt: TypedArrayRef { +/// Struct-specific accessors. +/// +/// Slot accessors (`validity`, `fields`, `slots_view`) live on the generated +/// [`StructArraySlotsExt`] supertrait; this trait layers struct-specific lookups on top. +pub trait StructArrayExt: StructArraySlotsExt { fn nullability(&self) -> crate::dtype::Nullability { match self.as_ref().dtype() { DType::Struct(_, nullability) => *nullability, @@ -186,25 +214,29 @@ pub trait StructArrayExt: TypedArrayRef { } fn struct_validity(&self) -> Validity { - child_to_validity( - self.as_ref().slots()[VALIDITY_SLOT].as_ref(), - self.nullability(), - ) + child_to_validity(self.validity(), self.nullability()) } - fn iter_unmasked_fields(&self) -> impl Iterator + '_ { - self.as_ref().slots()[FIELDS_OFFSET..] - .iter() - .map(|s| s.as_ref().vortex_expect("StructArray field slot")) + /// Iterate over the field arrays in declaration order. + fn iter_unmasked_fields(&self) -> impl ExactSizeIterator + '_ { + self.fields().iter() } - fn unmasked_fields(&self) -> Arc<[ArrayRef]> { - self.iter_unmasked_fields().cloned().collect() + /// The field array at `idx`, or `None` if `idx` is out of bounds. + /// + /// Use this over [`unmasked_field`](Self::unmasked_field) when the index comes from outside + /// the library and an out-of-bounds value is an error to report rather than a bug to panic on. + fn unmasked_field_opt(&self, idx: usize) -> Option<&ArrayRef> { + self.fields().get(idx) } + /// The field array at `idx`. + /// + /// # Panics + /// + /// If `idx` is out of bounds. fn unmasked_field(&self, idx: usize) -> &ArrayRef { - self.as_ref().slots()[FIELDS_OFFSET + idx] - .as_ref() + self.unmasked_field_opt(idx) .vortex_expect("StructArray field slot") } @@ -235,7 +267,7 @@ impl Array { /// Creates a new `StructArray`. pub fn new( names: FieldNames, - fields: impl Into>, + fields: impl IntoIterator, length: usize, validity: Validity, ) -> Self { @@ -246,14 +278,20 @@ impl Array { /// Constructs a new `StructArray`. pub fn try_new( names: FieldNames, - fields: impl Into>, + fields: impl IntoIterator, length: usize, validity: Validity, ) -> VortexResult { - let fields = fields.into(); - let field_dtypes: Vec<_> = fields.iter().map(|d| d.dtype().clone()).collect(); + let fields = fields.into_iter(); + let (lower, _) = fields.size_hint(); + let mut field_dtypes = Vec::with_capacity(lower); + let mut slots = ArraySlots::with_capacity(StructSlots::FIELDS_OFFSET + lower); + slots.push(validity_to_child(&validity, length)); + for field in fields { + field_dtypes.push(field.dtype().clone()); + slots.push(Some(field)); + } let dtype = StructFields::new(names, field_dtypes); - let slots = make_struct_slots(&fields, &validity, length); Array::try_from_parts( ArrayParts::new( Struct, @@ -271,14 +309,13 @@ impl Array { /// /// Caller must ensure the field arrays match the supplied dtype, length, and validity. pub unsafe fn new_unchecked( - fields: impl Into>, + fields: impl IntoIterator, dtype: StructFields, length: usize, validity: Validity, ) -> Self { - let fields = fields.into(); let outer_dtype = DType::Struct(dtype, validity.nullability()); - let slots = make_struct_slots(&fields, &validity, length); + let slots = make_struct_slots(fields, &validity, length); unsafe { Array::from_parts_unchecked( ArrayParts::new(Struct, outer_dtype, length, EmptyArrayData).with_slots(slots), @@ -288,14 +325,13 @@ impl Array { /// Constructs a new `StructArray` with an explicit dtype. pub fn try_new_with_dtype( - fields: impl Into>, + fields: impl IntoIterator, dtype: StructFields, length: usize, validity: Validity, ) -> VortexResult { - let fields = fields.into(); let outer_dtype = DType::Struct(dtype, validity.nullability()); - let slots = make_struct_slots(&fields, &validity, length); + let slots = make_struct_slots(fields, &validity, length); Array::try_from_parts( ArrayParts::new(Struct, outer_dtype, length, EmptyArrayData).with_slots(slots), ) @@ -383,7 +419,7 @@ impl Array { StructFields::new(FieldNames::default(), Vec::new()), crate::dtype::Nullability::NonNullable, ); - let slots = make_struct_slots(&[], &Validity::NonNullable, len); + let slots = make_struct_slots([], &Validity::NonNullable, len); unsafe { Array::from_parts_unchecked( ArrayParts::new(Struct, dtype, len, EmptyArrayData).with_slots(slots), @@ -393,10 +429,7 @@ impl Array { // TODO(ngates): remove this... it doesn't help to consume self. pub fn into_data_parts(self) -> StructDataParts { - let fields: Arc<[ArrayRef]> = self.slots()[FIELDS_OFFSET..] - .iter() - .map(|s| s.as_ref().vortex_expect("StructArray field slot").clone()) - .collect(); + let fields = self.fields().to_vec(); let validity = self.validity().vortex_expect("StructArray validity"); StructDataParts { struct_fields: self.struct_fields().clone(), @@ -412,18 +445,14 @@ impl Array { let position = struct_dtype.find(name.as_ref())?; - let slot_position = FIELDS_OFFSET + position; - let field = self.slots()[slot_position] - .as_ref() - .vortex_expect("StructArray field slot") - .clone(); - let new_slots: ArraySlots = self - .slots() - .iter() - .enumerate() - .filter(|(i, _)| *i != slot_position) - .map(|(_, s)| s.clone()) - .collect(); + let slot_position = StructSlots::FIELDS_OFFSET + position; + let field = self.unmasked_field(position).clone(); + // `Filter` has a zero lower bound, so build the slots with an exact capacity instead of + // letting `collect` grow them. + let slots = self.slots(); + let mut new_slots = ArraySlots::with_capacity(slots.len() - 1); + new_slots.extend(slots[..slot_position].iter().cloned()); + new_slots.extend(slots[slot_position + 1..].iter().cloned()); let new_dtype = struct_dtype.without_field(position).ok()?; let new_array = unsafe { @@ -448,11 +477,7 @@ impl Array { let types = struct_dtype.fields().chain(once(array.dtype().clone())); let new_fields = StructFields::new(names.collect(), types.collect()); - let children: Arc<[ArrayRef]> = self.slots()[FIELDS_OFFSET..] - .iter() - .map(|s| s.as_ref().vortex_expect("StructArray field slot").clone()) - .chain(once(array)) - .collect(); + let children = self.iter_unmasked_fields().cloned().chain(once(array)); Self::try_new_with_dtype(children, new_fields, self.len(), self.validity()?) } @@ -505,10 +530,7 @@ impl Array { .enumerate() .map(|(i, dtype)| { // SAFETY: We establish above that every array has the same type. - let chunks = field_arrays_per_chunk - .iter() - .map(|x| x[i].clone()) - .collect(); + let chunks = field_arrays_per_chunk.iter().map(|x| x[i].clone()); unsafe { ChunkedArray::new_unchecked(chunks, dtype) }.into_array() }) .collect::>(); @@ -541,11 +563,12 @@ impl Array { struct_validity.clone() }; - // Nothing to push down. + // Nothing to push down. The fields are unchanged, so reuse the existing `StructFields` + // rather than re-deriving it field by field. if struct_validity.definitely_no_nulls() { - return Self::try_new( - self.names().clone(), - self.unmasked_fields(), + return Self::try_new_with_dtype( + self.iter_unmasked_fields().cloned(), + self.struct_fields().clone(), self.len(), new_validity, ); diff --git a/vortex-array/src/arrays/struct_/compute/cast.rs b/vortex-array/src/arrays/struct_/compute/cast.rs index bfae815ceb5..4e0640d6807 100644 --- a/vortex-array/src/arrays/struct_/compute/cast.rs +++ b/vortex-array/src/arrays/struct_/compute/cast.rs @@ -139,7 +139,6 @@ mod tests { use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::VarBinArray; - use crate::arrays::scalar_fn::ScalarFnFactoryExt; use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; @@ -209,9 +208,7 @@ mod tests { Nullability::NonNullable, ); - let cast = Cast - .try_new_array(source.len(), target.clone(), [source]) - .unwrap(); + let cast = Cast::new(source, target.clone()).into_array(); let parent_id = cast.encoding_id(); let session = VortexSession::empty().with_some(KernelSession::empty()); let kernels = session.kernels(); diff --git a/vortex-array/src/arrays/struct_/compute/mask.rs b/vortex-array/src/arrays/struct_/compute/mask.rs index a36fa1743e1..1b5b526e69d 100644 --- a/vortex-array/src/arrays/struct_/compute/mask.rs +++ b/vortex-array/src/arrays/struct_/compute/mask.rs @@ -15,7 +15,7 @@ use crate::validity::Validity; impl MaskReduce for Struct { fn mask(array: ArrayView<'_, Struct>, mask: &ArrayRef) -> VortexResult> { StructArray::try_new_with_dtype( - array.unmasked_fields().iter().cloned().collect::>(), + array.iter_unmasked_fields().cloned(), array.struct_fields().clone(), array.len(), array.validity()?.and(Validity::Array(mask.clone()))?, diff --git a/vortex-array/src/arrays/struct_/compute/rules.rs b/vortex-array/src/arrays/struct_/compute/rules.rs index 01981803239..48c70d9446a 100644 --- a/vortex-array/src/arrays/struct_/compute/rules.rs +++ b/vortex-array/src/arrays/struct_/compute/rules.rs @@ -102,10 +102,14 @@ impl ArrayParentReduceRule for StructGetItemRule { })?; match child.validity()? { - Validity::NonNullable | Validity::AllValid => { - // If the struct is non-nullable or all valid, the field's validity is unchanged + Validity::NonNullable => { + // If the struct is non-nullable, the field's validity is unchanged Ok(Some(field.clone())) } + Validity::AllValid => { + // If struct is nullable, field must also be nullable + field.clone().cast(field.dtype().as_nullable()).map(Some) + } Validity::AllInvalid => { // If everything is invalid, the field is also all invalid Ok(Some( @@ -333,7 +337,7 @@ mod tests { .unwrap() .execute::(&mut SESSION.create_execution_ctx()) .unwrap(); - assert_eq!(result.unmasked_fields().len(), 2); + assert_eq!(result.iter_unmasked_fields().len(), 2); assert_arrays_eq!( result.unmasked_field_by_name("a").unwrap(), buffer![1i32, 2, 3].into_array(), diff --git a/vortex-array/src/arrays/struct_/compute/take.rs b/vortex-array/src/arrays/struct_/compute/take.rs index a63c6c85cf6..b98cbcd2be8 100644 --- a/vortex-array/src/arrays/struct_/compute/take.rs +++ b/vortex-array/src/arrays/struct_/compute/take.rs @@ -20,7 +20,7 @@ impl TakeReduce for Struct { // an out of bounds element. if array.is_empty() { return StructArray::try_new_with_dtype( - array.iter_unmasked_fields().cloned().collect::>(), + array.iter_unmasked_fields().cloned(), array.struct_fields().clone(), indices.len(), Validity::AllInvalid, diff --git a/vortex-array/src/arrays/struct_/mod.rs b/vortex-array/src/arrays/struct_/mod.rs index da49411c3bc..ab533dd8cdb 100644 --- a/vortex-array/src/arrays/struct_/mod.rs +++ b/vortex-array/src/arrays/struct_/mod.rs @@ -3,7 +3,10 @@ mod array; pub use array::StructArrayExt; +pub use array::StructArraySlotsExt; pub use array::StructDataParts; +pub use array::StructSlots; +pub use array::StructSlotsView; pub use vtable::StructArray; pub(crate) mod compute; diff --git a/vortex-array/src/arrays/struct_/vtable/mod.rs b/vortex-array/src/arrays/struct_/vtable/mod.rs index 0222a3c4712..c1639501288 100644 --- a/vortex-array/src/arrays/struct_/vtable/mod.rs +++ b/vortex-array/src/arrays/struct_/vtable/mod.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use itertools::Itertools; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -18,9 +17,8 @@ use crate::array::EmptyArrayData; use crate::array::VTable; use crate::array::child_to_validity; use crate::array::with_empty_buffers; -use crate::arrays::struct_::array::FIELDS_OFFSET; -use crate::arrays::struct_::array::VALIDITY_SLOT; -use crate::arrays::struct_::array::make_struct_slots; +use crate::arrays::struct_::array::StructSlots; +use crate::arrays::struct_::array::struct_slots_with_capacity; use crate::arrays::struct_::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; @@ -77,7 +75,7 @@ impl VTable for Struct { ); } - let validity = child_to_validity(slots[VALIDITY_SLOT].as_ref(), *nullability); + let validity = child_to_validity(slots[StructSlots::VALIDITY].as_ref(), *nullability); if let Some(validity_len) = validity.maybe_len() && validity_len != len { @@ -88,7 +86,7 @@ impl VTable for Struct { ); } - let field_slots = &slots[FIELDS_OFFSET..]; + let field_slots = &slots[StructSlots::FIELDS_OFFSET..]; if field_slots.is_empty() { return Ok(()); } @@ -174,24 +172,26 @@ impl VTable for Struct { ); }; - let field_children: Vec<_> = (0..struct_dtype.nfields()) - .map(|i| { - let child_dtype = struct_dtype - .field_by_index(i) - .vortex_expect("no out of bounds"); - children.get(non_data_children + i, &child_dtype, len) - }) - .try_collect()?; + let mut slots = struct_slots_with_capacity(&validity, len, struct_dtype.nfields()); + for i in 0..struct_dtype.nfields() { + let child_dtype = struct_dtype + .field_by_index(i) + .vortex_expect("no out of bounds"); + slots.push(Some(children.get( + non_data_children + i, + &child_dtype, + len, + )?)); + } - let slots = make_struct_slots(&field_children, &validity, len); Ok(ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData).with_slots(slots)) } fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String { - if idx == VALIDITY_SLOT { + if idx == StructSlots::VALIDITY { "validity".to_string() } else { - array.dtype().as_struct_fields().names()[idx - FIELDS_OFFSET].to_string() + array.dtype().as_struct_fields().names()[idx - StructSlots::FIELDS_OFFSET].to_string() } } diff --git a/vortex-array/src/arrays/struct_/vtable/operations.rs b/vortex-array/src/arrays/struct_/vtable/operations.rs index d59bf7eeca3..b491e231520 100644 --- a/vortex-array/src/arrays/struct_/vtable/operations.rs +++ b/vortex-array/src/arrays/struct_/vtable/operations.rs @@ -9,6 +9,7 @@ use crate::array::OperationsVTable; use crate::arrays::Struct; use crate::arrays::struct_::StructArrayExt; use crate::scalar::Scalar; +use crate::scalar::ScalarValue; impl OperationsVTable for Struct { fn scalar_at( @@ -16,12 +17,17 @@ impl OperationsVTable for Struct { index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let field_scalars: VortexResult> = array + let field_values = array .iter_unmasked_fields() - .map(|field| field.execute_scalar(index, ctx)) - .collect(); + .map(|field| field.execute_scalar(index, ctx).map(Scalar::into_value)) + .collect::>>()?; // SAFETY: The vtable guarantees index is in-bounds and non-null before this is called. - // Each field's scalar_at returns a scalar with the field's own dtype. - Ok(unsafe { Scalar::struct_unchecked(array.dtype().clone(), field_scalars?) }) + // Each field's scalar_at returns a value with the field's own dtype. + Ok(unsafe { + Scalar::new_unchecked( + array.dtype().clone(), + Some(ScalarValue::Tuple(field_values)), + ) + }) } } diff --git a/vortex-array/src/arrays/union/array.rs b/vortex-array/src/arrays/union/array.rs new file mode 100644 index 00000000000..bd5a8c5196c --- /dev/null +++ b/vortex-array/src/arrays/union/array.rs @@ -0,0 +1,232 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::ArraySlots; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayParts; +use crate::array::EmptyArrayData; +use crate::array::TypedArrayRef; +use crate::array_slots; +use crate::arrays::ConstantArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::Union; +use crate::arrays::union::union_type_ids_dtype; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::dtype::UnionVariants; +use crate::scalar::Scalar; + +/// Slot layout of a canonical sparse union array. +#[array_slots(Union)] +pub struct UnionSlots { + /// The row-aligned array of type IDs selecting a union child. + #[slot(0)] + pub type_ids: ArrayRef, + /// The row-aligned sparse children in variant order. + #[slot(1..)] + pub children: Vec, +} + +pub(super) fn make_union_parts( + type_ids: ArrayRef, + variants: UnionVariants, + children: impl IntoIterator, +) -> ArrayParts { + let len = type_ids.len(); + let nullability = type_ids.dtype().nullability(); + let children = children.into_iter(); + let (lower, _) = children.size_hint(); + let mut slots = ArraySlots::with_capacity(UnionSlots::CHILDREN_OFFSET + lower); + slots.push(Some(type_ids)); + slots.extend(children.map(Some)); + + ArrayParts::new( + Union, + DType::Union(variants, nullability), + len, + EmptyArrayData, + ) + .with_slots(slots) +} + +/// Concrete parts of a [`UnionArray`](super::UnionArray). +pub struct UnionDataParts { + /// The union variant schema. + pub variants: UnionVariants, + /// The row-aligned type IDs. + pub type_ids: ArrayRef, + /// The row-aligned sparse children in variant order. + pub children: Vec, +} + +/// Accessors for a canonical sparse union array. +/// +/// Slot accessors (`type_ids`, `children`, `slots_view`) live on the generated +/// [`UnionArraySlotsExt`] supertrait; this trait layers union-specific lookups on top. +pub trait UnionArrayExt: UnionArraySlotsExt { + /// The union's variant schema. + fn variants(&self) -> &UnionVariants { + match self.as_ref().dtype() { + DType::Union(variants, _) => variants, + _ => unreachable!("UnionArrayExt requires a union dtype"), + } + } + + /// Iterate over sparse children in variant order. + fn iter_children(&self) -> impl ExactSizeIterator + '_ { + self.children().iter() + } + + /// Return a sparse child by its variant index. + fn child(&self, index: usize) -> Option<&ArrayRef> { + self.children().get(index) + } + + /// Return a sparse child selected by a data-level type ID. + fn child_by_type_id(&self, type_id: u8) -> Option<&ArrayRef> { + self.child(self.variants().tag_to_child_index(type_id)?) + } + + /// Return a sparse child selected by its variant name, if present. + fn child_by_name_opt(&self, name: impl AsRef) -> Option<&ArrayRef> { + self.child(self.variants().find(name)?) + } + + /// Return a sparse child selected by its variant name. + fn child_by_name(&self, name: impl AsRef) -> VortexResult<&ArrayRef> { + let name = name.as_ref(); + self.child_by_name_opt(name).ok_or_else(|| { + vortex_err!( + "Variant {name} not found in union array with names {:?}", + self.variants().names() + ) + }) + } +} +impl> UnionArrayExt for T {} + +impl Array { + /// Construct a canonical sparse union array. + /// + /// # Panics + /// + /// Panics if the components do not satisfy the invariants documented on + /// [`Self::new_unchecked`]. + pub fn new( + type_ids: ArrayRef, + variants: UnionVariants, + children: impl IntoIterator, + ) -> Self { + Self::try_new(type_ids, variants, children).vortex_expect("UnionArray construction failed") + } + + /// Try to construct a canonical sparse union array. + /// + /// Type ID values are not validated during construction. Accessing a non-null row whose type + /// ID is not declared by `variants` will panic. + /// + /// # Errors + /// + /// Returns an error if `type_ids` is not a `u8` array, or if the sparse children do not match + /// the variant schema and outer array length. + pub fn try_new( + type_ids: ArrayRef, + variants: UnionVariants, + children: impl IntoIterator, + ) -> VortexResult { + vortex_ensure!( + matches!(type_ids.dtype(), DType::Primitive(PType::U8, _)), + "UnionArray type_ids must be u8, got {}", + type_ids.dtype() + ); + + Array::try_from_parts(make_union_parts(type_ids, variants, children)) + } + + /// Construct a canonical sparse union array without validation. + /// + /// # Safety + /// + /// The caller must ensure `type_ids` is a `u8` array, every child has the corresponding variant + /// dtype, and all arrays have the same length. Null type IDs represent outer union nulls. + pub unsafe fn new_unchecked( + type_ids: ArrayRef, + variants: UnionVariants, + children: impl IntoIterator, + ) -> Self { + unsafe { Array::from_parts_unchecked(make_union_parts(type_ids, variants, children)) } + } + + /// Deconstruct this array into its type IDs, variant schema, and sparse children. + pub fn into_data_parts(self) -> UnionDataParts { + let variants = self.variants().clone(); + let type_ids = self.type_ids().clone(); + let children = self.iter_children().cloned().collect(); + UnionDataParts { + variants, + type_ids, + children, + } + } + + /// Construct a `len`-row union in which every row holds `scalar`. + /// + /// Unselected children are filled with their variant's default value, a null for a nullable + /// variant and a zero for a non-nullable one. An outer null `scalar` selects no child at all. + /// + /// # Errors + /// + /// Returns an error if `scalar` does not have a union dtype. + pub fn constant(scalar: &Scalar, len: usize) -> VortexResult { + let union = scalar + .as_union_opt() + .ok_or_else(|| vortex_err!("Expected a union scalar, got {}", scalar.dtype()))?; + let variants = union.variants().clone(); + let nullability = union.nullability(); + + let type_ids = match union.type_id() { + Some(type_id) => Scalar::primitive(type_id, nullability), + None => Scalar::null(union_type_ids_dtype(nullability)), + }; + + let selected = union.child_index().zip(union.child()); + + let children = variants + .variants() + .enumerate() + .map(|(index, dtype)| { + let value = match &selected { + Some((selected_index, child)) if *selected_index == index => child.clone(), + _ => Scalar::default_value(&dtype), + }; + + ConstantArray::new(value, len).into_array() + }) + .collect::>(); + + Self::try_new( + ConstantArray::new(type_ids, len).into_array(), + variants, + children, + ) + } + + /// Create an empty array for a union dtype. + pub(crate) fn empty(variants: UnionVariants, nullability: Nullability) -> Self { + let type_ids = PrimitiveArray::empty::(nullability).into_array(); + let children: Vec<_> = variants + .variants() + .map(|dtype| crate::Canonical::empty(&dtype).into_array()) + .collect(); + + Self::new(type_ids, variants, children) + } +} diff --git a/vortex-array/src/arrays/union/compute/mask.rs b/vortex-array/src/arrays/union/compute/mask.rs new file mode 100644 index 00000000000..63bb9ae60f7 --- /dev/null +++ b/vortex-array/src/arrays/union/compute/mask.rs @@ -0,0 +1,25 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Union; +use crate::arrays::UnionArray; +use crate::arrays::union::UnionArrayExt; +use crate::arrays::union::UnionArraySlotsExt; +use crate::builtins::ArrayBuiltins; +use crate::scalar_fn::fns::mask::MaskReduce; + +impl MaskReduce for Union { + fn mask(array: ArrayView<'_, Union>, mask: &ArrayRef) -> VortexResult> { + UnionArray::try_new( + array.type_ids().clone().mask(mask.clone())?, + array.variants().clone(), + array.iter_children().cloned(), + ) + .map(|a| Some(a.into_array())) + } +} diff --git a/vortex-array/src/arrays/union/compute/mod.rs b/vortex-array/src/arrays/union/compute/mod.rs new file mode 100644 index 00000000000..b4143ad873f --- /dev/null +++ b/vortex-array/src/arrays/union/compute/mod.rs @@ -0,0 +1,8 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +pub(crate) mod rules; + +mod mask; +mod slice; +mod take; diff --git a/vortex-array/src/arrays/union/compute/rules.rs b/vortex-array/src/arrays/union/compute/rules.rs new file mode 100644 index 00000000000..921d66a24d3 --- /dev/null +++ b/vortex-array/src/arrays/union/compute/rules.rs @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use crate::arrays::Union; +use crate::arrays::dict::TakeReduceAdaptor; +use crate::arrays::slice::SliceReduceAdaptor; +use crate::optimizer::rules::ParentRuleSet; +use crate::scalar_fn::fns::mask::MaskReduceAdaptor; + +pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ + ParentRuleSet::lift(&MaskReduceAdaptor(Union)), + ParentRuleSet::lift(&SliceReduceAdaptor(Union)), + ParentRuleSet::lift(&TakeReduceAdaptor(Union)), +]); diff --git a/vortex-array/src/arrays/union/compute/slice.rs b/vortex-array/src/arrays/union/compute/slice.rs new file mode 100644 index 00000000000..be1336f5154 --- /dev/null +++ b/vortex-array/src/arrays/union/compute/slice.rs @@ -0,0 +1,30 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; + +use itertools::Itertools; +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Union; +use crate::arrays::UnionArray; +use crate::arrays::slice::SliceReduce; +use crate::arrays::union::UnionArrayExt; +use crate::arrays::union::UnionArraySlotsExt; + +impl SliceReduce for Union { + fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { + let type_ids = array.type_ids().slice(range.clone())?; + let children: Vec = array + .iter_children() + .map(|child| child.slice(range.clone())) + .try_collect()?; + + Ok(Some( + UnionArray::try_new(type_ids, array.variants().clone(), children)?.into_array(), + )) + } +} diff --git a/vortex-array/src/arrays/union/compute/take.rs b/vortex-array/src/arrays/union/compute/take.rs new file mode 100644 index 00000000000..5ad0a6f1287 --- /dev/null +++ b/vortex-array/src/arrays/union/compute/take.rs @@ -0,0 +1,52 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use itertools::Itertools; +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Union; +use crate::arrays::UnionArray; +use crate::arrays::dict::TakeReduce; +use crate::arrays::union::UnionArrayExt; +use crate::arrays::union::UnionArraySlotsExt; +use crate::builtins::ArrayBuiltins; +use crate::scalar::Scalar; + +/// Gathers the type IDs and every sparse child at `indices`. +/// +/// Sparse children are row-aligned with the union, so a gather must visit all of them. Take costs +/// `O(variants * indices)`, which only the dense encoding fixes. +/// +/// The type IDs carry the union's validity, so gathering them with the original `indices` turns a +/// null index into an outer union null. The children are gathered with the nulls filled in, which +/// keeps their declared variant dtypes. +impl TakeReduce for Union { + fn take(array: ArrayView<'_, Union>, indices: &ArrayRef) -> VortexResult> { + // An empty union has no row to point at, so the indices must be all null. + if array.is_empty() { + return UnionArray::constant(&Scalar::null(array.dtype().as_nullable()), indices.len()) + .map(UnionArray::into_array) + .map(Some); + } + + let type_ids = array.type_ids().take(indices.clone())?; + + // This stays a lazy node, so the fill runs once per child. `TakeReduce` has no + // `ExecutionCtx` to materialize it with, and the cost scales with the indices, not the + // data behind them. + let fill_scalar = Scalar::zero_value(&indices.dtype().as_nonnullable()); + let child_indices = indices.clone().fill_null(fill_scalar)?; + + let children: Vec = array + .iter_children() + .map(|child| child.take(child_indices.clone())) + .try_collect()?; + + UnionArray::try_new(type_ids, array.variants().clone(), children) + .map(UnionArray::into_array) + .map(Some) + } +} diff --git a/vortex-array/src/arrays/union/mod.rs b/vortex-array/src/arrays/union/mod.rs new file mode 100644 index 00000000000..0dda06e2766 --- /dev/null +++ b/vortex-array/src/arrays/union/mod.rs @@ -0,0 +1,39 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Canonical sparse union arrays. +//! +//! A [`UnionArray`] stores one `u8` type ID per row followed by one row-aligned child for each +//! variant. The type ID selects which child's value is active for a row; values in all other +//! children at that row are placeholders. Outer union nulls are stored as nulls in the type IDs +//! child, independently of nulls in the selected variant child. +//! +//! Type ID values are not validated during construction. Accessing a non-null row whose type ID +//! is not declared by the union variants will panic. +//! +//! Mask rewrites only the type IDs. Slice and take rewrite every child to keep them row-aligned, +//! so each costs `O(variants)` child operations. + +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; + +mod array; +pub use array::UnionArrayExt; +pub use array::UnionArraySlotsExt; +pub use array::UnionDataParts; +pub use array::UnionSlots; +pub use array::UnionSlotsView; +pub use vtable::UnionArray; + +pub(crate) mod compute; + +mod vtable; +pub use vtable::Union; + +pub(crate) fn union_type_ids_dtype(nullability: Nullability) -> DType { + DType::Primitive(PType::U8, nullability) +} + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/arrays/union/tests/mod.rs b/vortex-array/src/arrays/union/tests/mod.rs new file mode 100644 index 00000000000..fc0abec37b8 --- /dev/null +++ b/vortex-array/src/arrays/union/tests/mod.rs @@ -0,0 +1,340 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::ByteBufferMut; +use vortex_buffer::buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_mask::Mask; +use vortex_session::registry::ReadContext; + +use crate::ArrayContext; +use crate::Canonical; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; +use crate::array_session; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::Union; +use crate::arrays::UnionArray; +use crate::arrays::union::UnionArrayExt; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::dtype::UnionVariants; +use crate::scalar::Scalar; +use crate::serde::SerializeOptions; +use crate::serde::SerializedArray; +use crate::validity::Validity; + +mod take; + +fn variants() -> VortexResult { + UnionVariants::try_new( + ["number", "flag"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Bool(Nullability::NonNullable), + ], + vec![5, 9], + ) +} + +fn union_array() -> VortexResult { + UnionArray::try_new( + PrimitiveArray::from_iter([5u8, 9, 5]).into_array(), + variants()?, + vec![ + PrimitiveArray::from_iter([10i32, 0, 30]).into_array(), + BoolArray::from_iter([false, true, false]).into_array(), + ], + ) +} + +fn nullable_variants() -> VortexResult { + UnionVariants::try_new( + ["number", "optional"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Primitive(PType::I64, Nullability::Nullable), + ], + vec![5, 9], + ) +} + +fn nullable_union_array() -> VortexResult { + UnionArray::try_new( + PrimitiveArray::from_option_iter([Some(5u8), None, Some(9), Some(9)]).into_array(), + nullable_variants()?, + vec![ + PrimitiveArray::from_iter([10i32, 0, 0, 0]).into_array(), + PrimitiveArray::from_option_iter([Some(0i64), Some(0), None, Some(40)]).into_array(), + ], + ) +} + +#[test] +fn scalar_at_uses_type_id_indirection() -> VortexResult<()> { + let array = union_array()?; + let mut ctx = array_session().create_execution_ctx(); + + assert_eq!( + array.child_by_name("flag")?.dtype(), + &DType::Bool(Nullability::NonNullable) + ); + assert!(array.child_by_name_opt("missing").is_none()); + assert_eq!( + array.execute_scalar(0, &mut ctx)?, + Scalar::union(variants()?, 5, 10i32.into(), Nullability::NonNullable)? + ); + assert_eq!( + array.execute_scalar(1, &mut ctx)?, + Scalar::union(variants()?, 9, true.into(), Nullability::NonNullable)? + ); + assert!(matches!(array.validity()?, Validity::NonNullable)); + + Ok(()) +} + +#[test] +fn validates_sparse_components() -> VortexResult<()> { + let children = vec![ + PrimitiveArray::from_iter([10i32, 0, 30]).into_array(), + BoolArray::from_iter([false, true, false]).into_array(), + ]; + + assert!( + UnionArray::try_new( + PrimitiveArray::from_iter([5u8, 9]).into_array(), + variants()?, + children, + ) + .is_err() + ); + assert!( + UnionArray::try_new( + PrimitiveArray::from_iter([5u8, 9, 5]).into_array(), + variants()?, + vec![PrimitiveArray::from_iter([10i32, 0, 30]).into_array()], + ) + .is_err() + ); + assert!( + UnionArray::try_new( + PrimitiveArray::from_iter([5u8, 9, 5]).into_array(), + variants()?, + vec![ + PrimitiveArray::new(buffer![10i32, 0, 30], Validity::AllValid).into_array(), + BoolArray::from_iter([false, true, false]).into_array(), + ], + ) + .is_err() + ); + + Ok(()) +} + +#[test] +#[should_panic(expected = "Unknown UnionArray type ID 7")] +fn invalid_type_id_panics_when_accessed() { + let array = UnionArray::try_new( + PrimitiveArray::from_iter([5u8, 7, 5]).into_array(), + variants().vortex_expect("valid Union variants"), + vec![ + PrimitiveArray::from_iter([10i32, 0, 30]).into_array(), + BoolArray::from_iter([false, true, false]).into_array(), + ], + ) + .vortex_expect("structurally valid UnionArray"); + let mut ctx = array_session().create_execution_ctx(); + + let _scalar = array + .execute_scalar(1, &mut ctx) + .vortex_expect("UnionArray scalar access"); +} + +#[test] +fn outer_nulls_are_independent_from_inner_nulls() -> VortexResult<()> { + let variants = nullable_variants()?; + let array = nullable_union_array()?; + let mut ctx = array_session().create_execution_ctx(); + + assert_eq!( + array.dtype(), + &DType::Union(variants.clone(), Nullability::Nullable) + ); + assert_eq!( + array.validity()?.execute_mask(array.len(), &mut ctx)?, + Mask::from_iter([true, false, true, true]) + ); + assert_eq!( + array.execute_scalar(0, &mut ctx)?, + Scalar::union(variants.clone(), 5, 10i32.into(), Nullability::Nullable)? + ); + assert_eq!( + array.execute_scalar(1, &mut ctx)?, + Scalar::null(DType::Union(variants.clone(), Nullability::Nullable)) + ); + assert_eq!( + array.execute_scalar(2, &mut ctx)?, + Scalar::union( + variants, + 9, + Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable)), + Nullability::Nullable, + )? + ); + + Ok(()) +} + +#[test] +fn masking_adds_outer_nulls_only() -> VortexResult<()> { + let masked = union_array()? + .into_array() + .mask(BoolArray::from_iter([true, false, true]).into_array())?; + let mut ctx = array_session().create_execution_ctx(); + let masked = masked.execute::(&mut ctx)?; + + assert_eq!( + masked.dtype(), + &DType::Union(variants()?, Nullability::Nullable) + ); + assert_eq!( + masked.validity()?.execute_mask(masked.len(), &mut ctx)?, + Mask::from_iter([true, false, true]) + ); + assert_eq!( + masked.execute_scalar(1, &mut ctx)?, + Scalar::null(DType::Union(variants()?, Nullability::Nullable)) + ); + assert_eq!( + masked.execute_scalar(2, &mut ctx)?, + Scalar::union(variants()?, 5, 30i32.into(), Nullability::Nullable,)? + ); + + Ok(()) +} + +#[test] +fn slice_and_filter_preserve_sparse_alignment() -> VortexResult<()> { + let array = union_array()?.into_array(); + let mut ctx = array_session().create_execution_ctx(); + + let sliced = array.slice(1..3)?; + let filtered = array.filter(Mask::from_iter([true, false, true]))?; + + assert_eq!( + sliced.execute_scalar(0, &mut ctx)?, + Scalar::union(variants()?, 9, true.into(), Nullability::NonNullable,)? + ); + assert_eq!( + filtered.execute_scalar(1, &mut ctx)?, + Scalar::union(variants()?, 5, 30i32.into(), Nullability::NonNullable,)? + ); + + Ok(()) +} + +/// A constant union canonicalizes into a sparse union whose selected child repeats the scalar's +/// value and whose unselected children hold placeholders. +#[test] +fn constant_union_canonicalizes_to_sparse_union() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let i32_variant = DType::Primitive(PType::I32, Nullability::NonNullable); + let i64_variant = DType::Primitive(PType::I64, Nullability::Nullable); + + // Each case pairs a constant scalar with the dtype of a variant it does not select. + let cases = [ + ( + Scalar::union(variants()?, 9, true.into(), Nullability::NonNullable)?, + i32_variant.clone(), + ), + ( + Scalar::union(nullable_variants()?, 5, 7i32.into(), Nullability::Nullable)?, + i64_variant, + ), + ( + Scalar::null(DType::Union(variants()?, Nullability::Nullable)), + i32_variant, + ), + ]; + + for (scalar, unselected_dtype) in cases { + let canonical = ConstantArray::new(scalar.clone(), 3) + .into_array() + .execute::(&mut ctx)? + .into_union(); + + assert_eq!(canonical.dtype(), scalar.dtype()); + + let unselected = canonical + .iter_children() + .find(|child| child.dtype() == &unselected_dtype) + .ok_or_else(|| vortex_err!("No child with dtype {unselected_dtype}"))?; + assert_eq!( + unselected.execute_scalar(0, &mut ctx)?, + Scalar::default_value(&unselected_dtype) + ); + + for index in 0..canonical.len() { + assert_eq!(canonical.execute_scalar(index, &mut ctx)?, scalar); + } + } + + Ok(()) +} + +#[test] +fn constant_union_reports_uncompressed_size() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let scalar = Scalar::union(variants()?, 5, 10i32.into(), Nullability::NonNullable)?; + + // Four `u8` type IDs, four `i32` rows, and four `bool` placeholder bits rounded up to a byte. + let size = uncompressed_size_in_bytes(&ConstantArray::new(scalar, 4).into_array(), &mut ctx)?; + + assert_eq!(size, 4 + 4 * 4 + 1); + + Ok(()) +} + +#[test] +fn serde_roundtrip() -> VortexResult<()> { + let session = array_session(); + let mut execution_ctx = session.create_execution_ctx(); + for array in [union_array()?, nullable_union_array()?] { + let dtype = array.dtype().clone(); + let len = array.len(); + let array_ctx = ArrayContext::empty(); + let serialized = array.clone().into_array().serialize( + &array_ctx, + &session, + &SerializeOptions::default(), + )?; + let mut concat = ByteBufferMut::empty(); + for buffer in serialized { + concat.extend_from_slice(buffer.as_ref()); + } + let decoded = SerializedArray::try_from(concat.freeze())?.decode( + &dtype, + len, + &ReadContext::new(array_ctx.to_ids()), + &session, + )?; + let decoded = decoded.as_::(); + + assert_eq!(decoded.variants(), array.variants()); + for index in 0..len { + assert_eq!( + decoded.array().execute_scalar(index, &mut execution_ctx)?, + array.execute_scalar(index, &mut execution_ctx)? + ); + } + } + + Ok(()) +} diff --git a/vortex-array/src/arrays/union/tests/take.rs b/vortex-array/src/arrays/union/tests/take.rs new file mode 100644 index 00000000000..0088ace43b2 --- /dev/null +++ b/vortex-array/src/arrays/union/tests/take.rs @@ -0,0 +1,171 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use super::nullable_union_array; +use super::nullable_variants; +use super::union_array; +use super::variants; +use crate::ArrayRef; +use crate::Canonical; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::PrimitiveArray; +use crate::arrays::Union; +use crate::arrays::UnionArray; +use crate::arrays::dict::TakeReduce; +use crate::arrays::union::UnionArrayExt; +use crate::compute::conformance::take::test_take_conformance; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::scalar::Scalar; + +/// Take `indices`, asserting the result reduced back to a union instead of staying a dictionary. +#[track_caller] +fn take(array: &UnionArray, indices: ArrayRef) -> VortexResult { + Ok(array + .clone() + .into_array() + .take(indices)? + .as_::() + .into_owned()) +} + +/// Assert that `array` holds exactly `expected`, row for row. +#[track_caller] +fn assert_rows(array: &UnionArray, expected: Vec) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + assert_eq!(array.len(), expected.len()); + for (index, expected) in expected.into_iter().enumerate() { + assert_eq!(array.execute_scalar(index, &mut ctx)?, expected); + } + + Ok(()) +} + +#[test] +fn take_reorders_and_repeats_variant_selection() -> VortexResult<()> { + let variants = variants()?; + let nullability = Nullability::NonNullable; + + let taken = take( + &union_array()?, + PrimitiveArray::from_iter([2u64, 1, 0, 1]).into_array(), + )?; + + assert_eq!(taken.dtype(), &DType::Union(variants.clone(), nullability)); + assert_rows( + &taken, + vec![ + Scalar::union(variants.clone(), 5, 30i32.into(), nullability)?, // + Scalar::union(variants.clone(), 9, true.into(), nullability)?, // + Scalar::union(variants.clone(), 5, 10i32.into(), nullability)?, // + Scalar::union(variants, 9, true.into(), nullability)?, // + ], + ) +} + +#[test] +fn null_indices_become_outer_nulls_and_leave_children_alone() -> VortexResult<()> { + let variants = variants()?; + let nullability = Nullability::Nullable; + + let taken = take( + &union_array()?, + PrimitiveArray::from_option_iter([Some(1u64), None, Some(0)]).into_array(), + )?; + + assert_eq!(taken.dtype(), &DType::Union(variants.clone(), nullability)); + + // A nullable index widens the union but never its variants. + assert_eq!( + taken.child_by_name("number")?.dtype(), + &DType::Primitive(PType::I32, Nullability::NonNullable) + ); + assert_eq!( + taken.child_by_name("flag")?.dtype(), + &DType::Bool(Nullability::NonNullable) + ); + + assert_rows( + &taken, + vec![ + Scalar::union(variants.clone(), 9, true.into(), nullability)?, // + Scalar::null(DType::Union(variants.clone(), nullability)), // + Scalar::union(variants, 5, 10i32.into(), nullability)?, // + ], + ) +} + +#[test] +fn take_keeps_outer_and_inner_nulls_distinct() -> VortexResult<()> { + let variants = nullable_variants()?; + let nullability = Nullability::Nullable; + + // Row 1 is an outer null and row 2 is a present union selecting a null child. + let taken = take( + &nullable_union_array()?, + PrimitiveArray::from_iter([1u64, 2, 3]).into_array(), + )?; + + assert_rows( + &taken, + vec![ + Scalar::null(DType::Union(variants.clone(), nullability)), // + Scalar::union( + variants.clone(), + 9, + Scalar::null(DType::Primitive(PType::I64, nullability)), + nullability, + )?, // + Scalar::union( + variants, + 9, + Scalar::primitive(40i64, nullability), + nullability, + )?, // + ], + ) +} + +/// `take` short-circuits an empty source into a constant before the union sees it, so this covers +/// that path and the union's own gather. +#[test] +fn take_from_empty_union_is_all_null() -> VortexResult<()> { + let variants = variants()?; + let nullability = Nullability::Nullable; + let empty = UnionArray::empty(variants.clone(), nullability).into_array(); + let indices = PrimitiveArray::from_option_iter([None::, None]).into_array(); + + let via_take = empty + .take(indices.clone())? + .execute::(&mut array_session().create_execution_ctx())? + .into_union(); + let via_reduce = ::take(empty.as_::(), &indices)? + .ok_or_else(|| vortex_err!("Union take must never decline"))? + .as_::() + .into_owned(); + + let expected = || vec![Scalar::null(DType::Union(variants.clone(), nullability)); 2]; + + assert_rows(&via_take, expected())?; + assert_rows(&via_reduce, expected()) +} + +#[rstest] +#[case::non_nullable(union_array())] +#[case::nullable(nullable_union_array())] +fn take_conformance(#[case] array: VortexResult) -> VortexResult<()> { + test_take_conformance( + &array?.into_array(), + &mut array_session().create_execution_ctx(), + ); + + Ok(()) +} diff --git a/vortex-array/src/arrays/union/vtable/mod.rs b/vortex-array/src/arrays/union/vtable/mod.rs new file mode 100644 index 00000000000..1f29b950b0a --- /dev/null +++ b/vortex-array/src/arrays/union/vtable/mod.rs @@ -0,0 +1,177 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::ExecutionResult; +use crate::array::Array; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayView; +use crate::array::EmptyArrayData; +use crate::array::VTable; +use crate::array::with_empty_buffers; +use crate::arrays::union::UnionArrayExt; +use crate::arrays::union::UnionSlots; +use crate::arrays::union::array::make_union_parts; +use crate::arrays::union::compute::rules::PARENT_RULES; +use crate::arrays::union::union_type_ids_dtype; +use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; +use crate::dtype::DType; +use crate::serde::ArrayChildren; + +mod operations; +mod validate; +mod validity; + +/// A canonical [`Union`]-encoded array. +/// +/// This is similar to Arrow's "spase" union type. +pub type UnionArray = Array; + +/// The canonical Union encoding. +/// +/// This is similar to Arrow's "spase" union type. +#[derive(Clone, Debug)] +pub struct Union; + +impl VTable for Union { + type TypedArrayData = EmptyArrayData; + type OperationsVTable = Self; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.union"); + *ID + } + + fn validate( + &self, + _data: &EmptyArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + let type_ids = slots + .get(UnionSlots::TYPE_IDS) + .and_then(Option::as_ref) + .ok_or_else(|| vortex_err!("UnionArray is missing its type_ids slot"))?; + let variant_arrays = slots + .get(UnionSlots::CHILDREN_OFFSET..) + .unwrap_or_default() + .iter() + .enumerate() + .map(|(index, slot)| { + slot.as_ref() + .ok_or_else(|| vortex_err!("UnionArray is missing child {index}")) + }) + .collect::>>()?; + + validate::validate_union_components(type_ids, &variant_arrays, dtype, len) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + vortex_panic!("UnionArray buffer index {idx} out of bounds") + } + + fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { + vortex_panic!("UnionArray buffer_name index {idx} out of bounds") + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + + fn serialize( + _array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + Ok(Some(vec![])) + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + buffers: &[BufferHandle], + children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!(metadata.is_empty(), "UnionArray expects empty metadata"); + vortex_ensure!(buffers.is_empty(), "UnionArray expects no buffers"); + let DType::Union(variants, nullability) = dtype else { + vortex_bail!("Expected union dtype, found {dtype}") + }; + vortex_ensure_eq!( + children.len(), + UnionSlots::CHILDREN_OFFSET + variants.len(), + "UnionArray expected {} children, found {}", + UnionSlots::CHILDREN_OFFSET + variants.len(), + children.len() + ); + + let type_ids = children.get( + UnionSlots::TYPE_IDS, + &union_type_ids_dtype(*nullability), + len, + )?; + let sparse_children = variants + .variants() + .enumerate() + .map(|(index, dtype)| children.get(UnionSlots::CHILDREN_OFFSET + index, &dtype, len)) + .collect::>>()?; + + Ok(make_union_parts( + type_ids, + variants.clone(), + sparse_children, + )) + } + + fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String { + if idx == UnionSlots::TYPE_IDS { + "type_ids".to_string() + } else { + array.variants().names()[idx - UnionSlots::CHILDREN_OFFSET].to_string() + } + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done(array)) + } + + fn append_to_builder( + _array: ArrayView<'_, Self>, + _builder: &mut dyn ArrayBuilder, + _ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + todo!("TODO(connor)[Union]: implement append_to_builder for Union arrays") + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + PARENT_RULES.evaluate(array, parent, child_idx) + } +} diff --git a/vortex-array/src/arrays/union/vtable/operations.rs b/vortex-array/src/arrays/union/vtable/operations.rs new file mode 100644 index 00000000000..39ba95ed4ee --- /dev/null +++ b/vortex-array/src/arrays/union/vtable/operations.rs @@ -0,0 +1,42 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; + +use crate::ExecutionCtx; +use crate::array::ArrayView; +use crate::array::OperationsVTable; +use crate::arrays::Union; +use crate::arrays::union::UnionArrayExt; +use crate::arrays::union::UnionArraySlotsExt; +use crate::scalar::Scalar; + +impl OperationsVTable for Union { + fn scalar_at( + array: ArrayView<'_, Union>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let type_id_scalar = array.type_ids().execute_scalar(index, ctx)?; + let Some(type_id) = type_id_scalar.as_primitive().typed_value::() else { + return Ok(Scalar::null(array.dtype().clone())); + }; + + let Some(child_index) = array.variants().tag_to_child_index(type_id) else { + vortex_panic!("Unknown UnionArray type ID {type_id}") + }; + let child = array + .child(child_index) + .ok_or_else(|| vortex_err!("UnionArray is missing child {child_index}"))?; + let child_scalar = child.execute_scalar(index, ctx)?; + + Scalar::union( + array.variants().clone(), + type_id, + child_scalar, + array.dtype().nullability(), + ) + } +} diff --git a/vortex-array/src/arrays/union/vtable/validate.rs b/vortex-array/src/arrays/union/vtable/validate.rs new file mode 100644 index 00000000000..e01c5486a94 --- /dev/null +++ b/vortex-array/src/arrays/union/vtable/validate.rs @@ -0,0 +1,62 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; + +use crate::ArrayRef; +use crate::arrays::union::UnionSlots; +use crate::arrays::union::union_type_ids_dtype; +use crate::dtype::DType; + +pub(super) fn validate_union_components( + type_ids: &ArrayRef, + variant_arrays: &[&ArrayRef], + dtype: &DType, + len: usize, +) -> VortexResult<()> { + let DType::Union(variants, nullability) = dtype else { + vortex_bail!("Expected union dtype, found {dtype}") + }; + vortex_ensure_eq!( + variant_arrays.len(), + variants.len(), + "UnionArray has {} slots but expected {}", + UnionSlots::CHILDREN_OFFSET + variant_arrays.len(), + UnionSlots::CHILDREN_OFFSET + variants.len() + ); + + let expected_union_type_ids_dtype = union_type_ids_dtype(*nullability); + vortex_ensure_eq!( + type_ids.dtype(), + &expected_union_type_ids_dtype, + "UnionArray type_ids must have dtype {expected_union_type_ids_dtype}, got {}", + type_ids.dtype() + ); + vortex_ensure_eq!( + type_ids.len(), + len, + "UnionArray type_ids length {} does not match outer length {len}", + type_ids.len() + ); + + for (index, (variant_dtype, child)) in + variants.variants().zip(variant_arrays.iter()).enumerate() + { + vortex_ensure_eq!( + child.len(), + len, + "UnionArray child {index} length {} does not match outer length {len}", + child.len() + ); + vortex_ensure_eq!( + child.dtype(), + &variant_dtype, + "UnionArray child {index} has dtype {} but expected {variant_dtype}", + child.dtype() + ); + } + + Ok(()) +} diff --git a/vortex-array/src/arrays/union/vtable/validity.rs b/vortex-array/src/arrays/union/vtable/validity.rs new file mode 100644 index 00000000000..1da64995b14 --- /dev/null +++ b/vortex-array/src/arrays/union/vtable/validity.rs @@ -0,0 +1,16 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::array::ArrayView; +use crate::array::ValidityVTable; +use crate::arrays::Union; +use crate::arrays::union::UnionArraySlotsExt; +use crate::validity::Validity; + +impl ValidityVTable for Union { + fn validity(array: ArrayView<'_, Union>) -> VortexResult { + array.type_ids().validity() + } +} diff --git a/vortex-array/src/arrays/validation_tests.rs b/vortex-array/src/arrays/validation_tests.rs index 30758a5c29e..fc7a6ee0a30 100644 --- a/vortex-array/src/arrays/validation_tests.rs +++ b/vortex-array/src/arrays/validation_tests.rs @@ -16,6 +16,8 @@ mod tests { use vortex_error::VortexError; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::DecimalArray; use crate::arrays::FixedSizeListArray; @@ -171,6 +173,7 @@ mod tests { Arc::new([]), DType::Utf8(Nullability::NonNullable), Validity::NonNullable, + &mut array_session().create_execution_ctx(), ); assert!(result.is_ok()); } @@ -190,6 +193,7 @@ mod tests { buffers, DType::Binary(Nullability::NonNullable), Validity::NonNullable, + &mut array_session().create_execution_ctx(), ); assert!(matches!(result, Err(VortexError::InvalidArgument(_, _)))); diff --git a/vortex-array/src/arrays/varbin/array.rs b/vortex-array/src/arrays/varbin/array.rs index 4cb74fe48b4..a763fb4f1de 100644 --- a/vortex-array/src/arrays/varbin/array.rs +++ b/vortex-array/src/arrays/varbin/array.rs @@ -5,8 +5,8 @@ use std::fmt::Display; use std::fmt::Formatter; use num_traits::AsPrimitive; -use smallvec::smallvec; use vortex_array::arrays::PrimitiveArray; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -21,22 +21,26 @@ use crate::array::ArrayParts; use crate::array::TypedArrayRef; use crate::array::child_to_validity; use crate::array::validity_to_child; +use crate::array_slots; use crate::arrays::VarBin; use crate::arrays::varbin::builder::VarBinBuilder; use crate::buffer::BufferHandle; use crate::dtype::DType; -use crate::dtype::IntegerPType; use crate::dtype::Nullability; +use crate::dtype::OffsetBuilderPType; use crate::legacy_session; use crate::match_each_integer_ptype; use crate::validity::Validity; -/// The offsets array defining the start/end of each variable-length binary element. -pub(super) const OFFSETS_SLOT: usize = 0; -/// The validity bitmap indicating which elements are non-null. -pub(super) const VALIDITY_SLOT: usize = 1; -pub(super) const NUM_SLOTS: usize = 2; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["offsets", "validity"]; +#[array_slots(VarBin)] +pub struct VarBinSlots { + /// The offsets array defining the start/end of each variable-length binary element. + #[slot(0)] + pub offsets: ArrayRef, + /// The validity bitmap indicating which elements are non-null. + #[slot(1)] + pub validity: Option, +} #[derive(Clone, Debug)] pub struct VarBinData { @@ -83,7 +87,11 @@ impl VarBinData { } pub(crate) fn make_slots(offsets: ArrayRef, validity: &Validity, len: usize) -> ArraySlots { - smallvec![Some(offsets), validity_to_child(validity, len)] + VarBinSlots { + offsets, + validity: validity_to_child(validity, len), + } + .into_slots() } /// Constructs a new `VarBinArray`. @@ -302,17 +310,7 @@ impl VarBinData { } } -pub trait VarBinArrayExt: TypedArrayRef { - fn offsets(&self) -> &ArrayRef { - self.as_ref().slots()[OFFSETS_SLOT] - .as_ref() - .vortex_expect("VarBinArray offsets slot") - } - - fn validity_child(&self) -> Option<&ArrayRef> { - self.as_ref().slots()[VALIDITY_SLOT].as_ref() - } - +pub trait VarBinArrayExt: VarBinArraySlotsExt { fn dtype_parts(&self) -> (bool, Nullability) { match self.as_ref().dtype() { DType::Utf8(nullability) => (true, *nullability), @@ -331,7 +329,7 @@ pub trait VarBinArrayExt: TypedArrayRef { fn varbin_validity(&self) -> Validity { child_to_validity( - self.as_ref().slots()[VALIDITY_SLOT].as_ref(), + self.as_ref().slots()[VarBinSlots::VALIDITY].as_ref(), self.nullability(), ) } @@ -386,11 +384,15 @@ impl Array { dtype: DType, ) -> Self { let iter = iter.into_iter(); - let mut builder = VarBinBuilder::::with_capacity(iter.size_hint().0); + let mut builder = VarBinBuilder::::with_capacity_in( + dtype, + iter.size_hint().0, + BufferAllocatorRef::static_ref(), + ); for v in iter { builder.append(v.as_ref().map(|o| o.as_ref())); } - builder.finish(dtype) + builder.finish_into_varbin() } pub fn from_iter_nonnull, I: IntoIterator>( @@ -398,23 +400,31 @@ impl Array { dtype: DType, ) -> Self { let iter = iter.into_iter(); - let mut builder = VarBinBuilder::::with_capacity(iter.size_hint().0); + let mut builder = VarBinBuilder::::with_capacity_in( + dtype, + iter.size_hint().0, + BufferAllocatorRef::static_ref(), + ); for v in iter { builder.append_value(v); } - builder.finish(dtype) + builder.finish_into_varbin() } fn from_vec_sized(vec: Vec, dtype: DType) -> Self where - O: IntegerPType, + O: OffsetBuilderPType, T: AsRef<[u8]>, { - let mut builder = VarBinBuilder::::with_capacity(vec.len()); + let mut builder = VarBinBuilder::::with_capacity_in( + dtype, + vec.len(), + BufferAllocatorRef::static_ref(), + ); for v in vec { builder.append_value(v.as_ref()); } - builder.finish(dtype) + builder.finish_into_varbin() } /// Create from a vector of string slices. @@ -457,7 +467,7 @@ impl Array { let len = offsets.len().saturating_sub(1); let slots = VarBinData::make_slots(offsets, &validity, len); let data = VarBinData::build( - slots[OFFSETS_SLOT] + slots[VarBinSlots::OFFSETS] .as_ref() .vortex_expect("VarBinArray offsets slot") .clone(), diff --git a/vortex-array/src/arrays/varbin/builder.rs b/vortex-array/src/arrays/varbin/builder.rs index 422139cc851..d915efdf2fd 100644 --- a/vortex-array/src/arrays/varbin/builder.rs +++ b/vortex-array/src/arrays/varbin/builder.rs @@ -1,109 +1,436 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::Any; +use std::mem::MaybeUninit; + use num_traits::AsPrimitive; +use vortex_buffer::Alignment; use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; +use vortex_buffer::ByteBuffer; +use vortex_buffer::ByteBufferMut; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; use vortex_error::vortex_panic; +use vortex_mask::AllOr; +use vortex_mask::Mask; +use crate::ArrayRef; +use crate::ArrayView; +use crate::Canonical; +use crate::ExecutionCtx; use crate::IntoArray; #[cfg(debug_assertions)] use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; +use crate::arrays::VarBin; use crate::arrays::VarBinArray; +use crate::arrays::VarBinView; +use crate::arrays::varbin::VarBinArrayExt; +use crate::arrays::varbin::VarBinArraySlotsExt; +use crate::arrays::varbinview::VarBinViewArrayExt; +use crate::builders::ArrayBuilder; use crate::dtype::DType; -use crate::dtype::IntegerPType; +use crate::dtype::OffsetBuilderPType; use crate::expr::stats::Precision; use crate::expr::stats::Stat; #[cfg(debug_assertions)] use crate::legacy_session; +use crate::match_each_integer_ptype; +use crate::scalar::Scalar; use crate::validity::Validity; -pub struct VarBinBuilder { +/// Builder for [`VarBinArray`] values with `O`-typed offsets. +/// +/// This is the offset-based counterpart to +/// [`VarBinViewBuilder`](crate::builders::VarBinViewBuilder): values are laid out contiguously in +/// a single byte buffer described by a monotonically increasing offsets buffer. An `i32` or `i64` +/// builder can therefore be handed straight to Arrow as a `Utf8`/`LargeUtf8`/`Binary`/ +/// `LargeBinary` array without re-laying out the bytes. +/// +/// Encodings that can decode into this layout should specialize +/// [`append_to_builder`](crate::vtable::VTable::append_to_builder) with +/// [`match_each_varbin_builder!`](crate::match_each_varbin_builder), which recovers the concrete +/// offset type from a `&mut dyn ArrayBuilder` so the decode loop is monomorphized over it. +pub struct VarBinBuilder { + dtype: DType, offsets: BufferMut, - data: BufferMut, + data: ByteBufferMut, validity: BitBufferMut, } -impl Default for VarBinBuilder { - fn default() -> Self { - Self::new() +impl VarBinBuilder { + /// Creates an empty builder for `dtype`. + #[deprecated(note = "use `new_in` with an explicit allocator")] + pub fn new(dtype: DType) -> Self { + Self::new_in(dtype, BufferAllocatorRef::static_ref()) + } + + /// Creates an empty builder for `dtype` using `allocator`. + pub fn new_in(dtype: DType, allocator: &BufferAllocatorRef) -> Self { + Self::with_capacity_in(dtype, 0, allocator) } -} -impl VarBinBuilder { - pub fn new() -> Self { - Self::with_capacity(0) + /// Creates a builder for `dtype` with room for `capacity` values. + #[deprecated(note = "use `with_capacity_in` with an explicit allocator")] + pub fn with_capacity(dtype: DType, capacity: usize) -> Self { + Self::with_capacity_in(dtype, capacity, BufferAllocatorRef::static_ref()) } - pub fn with_capacity(len: usize) -> Self { - let mut offsets = BufferMut::with_capacity(len + 1); + /// Creates a builder for `dtype` with room for `capacity` values using `allocator`. + pub fn with_capacity_in(dtype: DType, capacity: usize, allocator: &BufferAllocatorRef) -> Self { + assert!( + matches!(dtype, DType::Utf8(_) | DType::Binary(_)), + "VarBinBuilder dtype must be Utf8 or Binary, got {dtype}" + ); + let mut offsets = BufferMut::with_capacity_in(capacity + 1, allocator.clone()); offsets.push(O::zero()); Self { + dtype, offsets, - data: BufferMut::empty(), - validity: BitBufferMut::with_capacity(len), + data: BufferMut::empty_aligned_in(Alignment::of::(), allocator.clone()), + validity: BitBufferMut::with_capacity_in(capacity, allocator.clone()), } } + /// Creates a builder for `dtype` with room for `capacity` values totalling `bytes` bytes. + #[deprecated(note = "use `with_capacity_bytes_in` with an explicit allocator")] + pub fn with_capacity_bytes(dtype: DType, capacity: usize, bytes: usize) -> Self { + Self::with_capacity_bytes_in(dtype, capacity, bytes, BufferAllocatorRef::static_ref()) + } + + /// Creates a builder using `allocator` with room for `capacity` values and `bytes` bytes. + pub fn with_capacity_bytes_in( + dtype: DType, + capacity: usize, + bytes: usize, + allocator: &BufferAllocatorRef, + ) -> Self { + let mut builder = Self::with_capacity_in(dtype, capacity, allocator); + builder.reserve_data(bytes); + builder + } + + /// Reserves room for `additional` value bytes. + /// + /// [`reserve_exact`](ArrayBuilder::reserve_exact) takes a row count and so cannot size the + /// value bytes; callers that know the byte total should use this to size the buffer once + /// instead of letting it grow. + pub fn reserve_data(&mut self, additional: usize) { + self.data.reserve(additional); + } + + /// Appends one value, or a null when it is `None`. + /// + /// # Panics + /// + /// Panics if the resulting end offset does not fit in `O`. See + /// [`append_value`](Self::append_value). #[inline] pub fn append(&mut self, value: Option<&[u8]>) { match value { Some(v) => self.append_value(v), - None => self.append_null(), + None => self.push_null(), } } + /// Appends one non-null value. + /// + /// # Panics + /// + /// Panics if the resulting end offset does not fit in `O`. The bulk appends report that as an + /// error instead; use [`append_n_values`](Self::append_n_values) for a fallible single append, + /// or an `i64` builder for byte totals past `i32::MAX`. #[inline] pub fn append_value(&mut self, value: impl AsRef<[u8]>) { - let slice = value.as_ref(); - self.offsets - .push(O::from(self.data.len() + slice.len()).unwrap_or_else(|| { - vortex_panic!( - "Failed to convert sum of {} and {} to offset of type {}", - self.data.len(), - slice.len(), - std::any::type_name::() - ) - })); - self.data.extend_from_slice(slice); + self.push_value(value.as_ref()); self.validity.append_true(); } + /// Appends the same non-null value `n` times. + /// + /// # Errors + /// + /// Returns an error, leaving the builder unchanged, if the resulting end offsets do not fit + /// in `O`. + pub fn append_n_values(&mut self, value: impl AsRef<[u8]>, n: usize) -> VortexResult<()> { + let value = value.as_ref(); + let Some(num_bytes) = value.len().checked_mul(n) else { + vortex_bail!("Byte count overflow: {} values of {} bytes", n, value.len()); + }; + // Checking the total up front is what keeps `push_value` below from panicking. + self.check_offset_limit(self.data.len(), num_bytes)?; + self.offsets.reserve(n); + self.data.reserve(num_bytes); + for _ in 0..n { + self.push_value(value); + } + self.validity.append_n(true, n); + Ok(()) + } + + /// Appends a null value. + /// + /// Unlike [`append_null`](ArrayBuilder::append_null) this does not check that the builder is + /// nullable; the offsets and validity stay consistent either way, and a non-nullable `dtype` + /// discards the validity bits on [`finish_into_varbin`](Self::finish_into_varbin). #[inline] - pub fn append_null(&mut self) { - self.offsets.push(self.offsets[self.offsets.len() - 1]); - self.validity.append_false(); + pub fn push_null(&mut self) { + self.push_nulls(1) } + /// Appends `n` null values. See [`push_null`](Self::push_null). #[inline] - pub fn append_n_nulls(&mut self, n: usize) { - self.offsets.push_n(self.offsets[self.offsets.len() - 1], n); + pub fn push_nulls(&mut self, n: usize) { + self.offsets.push_n(self.last_offset(), n); self.validity.append_n(false, n); } + /// Appends the same UTF-8 or binary scalar `n` times. + /// + /// # Errors + /// + /// Returns an error if `scalar` has a different dtype than the builder, or if the resulting + /// end offsets do not fit in `O`. + pub fn append_scalar_repeated(&mut self, scalar: &Scalar, n: usize) -> VortexResult<()> { + vortex_ensure!( + scalar.dtype() == &self.dtype, + "VarBinBuilder expected scalar with dtype {}, got {}", + self.dtype, + scalar.dtype() + ); + match &self.dtype { + DType::Utf8(_) => match scalar.as_utf8().value() { + Some(value) => self.append_n_values(value, n)?, + None => self.push_nulls(n), + }, + DType::Binary(_) => match scalar.as_binary().value() { + Some(value) => self.append_n_values(value, n)?, + None => self.push_nulls(n), + }, + dtype => vortex_bail!("VarBinBuilder cannot append scalar of dtype {dtype}"), + } + Ok(()) + } + + /// Appends values from one contiguous byte buffer described by relative end offsets. + /// + /// Each entry of `end_offsets` marks the end of one value relative to the start of `values`, + /// and there must be exactly one per entry in `validity`. #[inline] - pub fn append_values(&mut self, values: &[u8], end_offsets: impl Iterator, num: usize) + pub fn append_values

( + &mut self, + values: &[u8], + end_offsets: impl Iterator, + validity: &Mask, + ) -> VortexResult<()> where - O: 'static, + P: AsPrimitive, usize: AsPrimitive, { - self.offsets - .extend(end_offsets.map(|offset| offset + self.data.len().as_())); + // Offsets are committed first: they are the only fallible part, and failing before the + // bytes are appended leaves the builder untouched. + self.extend_offsets(values.len(), validity.len(), end_offsets)?; self.data.extend_from_slice(values); - self.validity.append_n(true, num); + self.append_validity(validity); + Ok(()) + } + + /// Appends values whose bytes `decode` writes straight into the builder's byte storage. + /// + /// `decode` is handed at least `num_bytes + slack` bytes of uninitialized storage and returns + /// the number of bytes it initialized, which must be exactly `num_bytes`. Each entry of + /// `lengths` is the byte length of one value — including nulls, which are zero-length — so + /// there must be one per entry in `validity`. + /// + /// `slack` is spare headroom past the values themselves. It exists so that a decoder which + /// stores in wide fixed-size chunks can keep using them through the final value instead of + /// dropping into a byte-at-a-time tail; a decoder must still never write past the slice it is + /// handed, so `0` is correct for one that is already exact. + /// + /// Decoding in place saves staging the decoded heap in a temporary buffer and copying it in. + /// + /// `decode` is taken as a `dyn` reference so a caller that resolves the `lengths` type with + /// [`match_each_integer_ptype!`](crate::match_each_integer_ptype) can build the closure once + /// outside that match rather than inlining a whole decoder into each of its arms. + /// + /// # Errors + /// + /// Returns an error, leaving the builder unchanged, if `num_bytes` would push an offset past + /// what `O` can hold — checked before `decode` runs — or if `decode` reports a different byte + /// count than `num_bytes`, or if `lengths` does not describe those bytes. + /// + /// # Safety + /// + /// `decode` must initialize the first `n` bytes of the slice it is passed, where `n` is the + /// value it returns. Those bytes are published as initialized without ever being read first, + /// so a `decode` that over-reports leaves the builder holding uninitialized memory. + pub unsafe fn append_decoded

( + &mut self, + num_bytes: usize, + slack: usize, + lengths: &[P], + validity: &Mask, + decode: &mut dyn FnMut(&mut [MaybeUninit]) -> VortexResult, + ) -> VortexResult<()> + where + P: AsPrimitive, + usize: AsPrimitive, + { + let Some(capacity) = num_bytes.checked_add(slack) else { + vortex_bail!("Decoded size overflow: {num_bytes} + {slack}"); + }; + // Checked before decoding: an `i32` builder that the decoded bytes would overflow should + // not pay for the decompression first. + self.check_offset_limit(self.data.len(), num_bytes)?; + self.data.reserve(capacity); + + let data_len = self.data.len(); + let written = decode(self.data.spare_capacity_mut())?; + vortex_ensure!( + written == num_bytes, + "Decoded {written} bytes, expected {num_bytes}" + ); + + // The decoded bytes live in spare capacity until `set_len` below, so an invalid `lengths` + // still leaves the builder unchanged. + self.extend_offsets(num_bytes, validity.len(), prefix_sums(lengths))?; + + // SAFETY: `decode` reported initializing `written` spare bytes, and the caller guarantees + // that report is accurate; `written == num_bytes` was checked above. + unsafe { self.data.set_len(data_len + num_bytes) }; + self.append_validity(validity); + Ok(()) + } + + /// Appends `validity.len()` values, taking each valid row's bytes from `value`. + /// + /// `value` is called once per valid row in ascending row order and the slices it returns must + /// total `num_bytes` bytes. Null rows consume no bytes and repeat the preceding offset. + /// + /// The valid rows are walked a `u64` word at a time rather than pulled one at a time through + /// an iterator, and both buffers are sized once up front, so the copy loop costs one offset + /// store plus one `memcpy` per value. A caller whose values only exist as a stream can still + /// use this by pulling from the stream inside `value`: the rows are visited in ascending + /// order. + pub fn append_valid_slices<'a, F>( + &mut self, + num_bytes: usize, + validity: &Mask, + value: F, + ) -> VortexResult<()> + where + F: FnMut(usize) -> &'a [u8], + usize: AsPrimitive, + { + let data_start = self.data.len(); + match self.gather_valid_slices(num_bytes, validity, value) { + Ok(()) => { + self.append_validity(validity); + Ok(()) + } + Err(error) => { + // The offsets are never committed on failure, so only the copied bytes need to go. + self.data.truncate(data_start); + Err(error) + } + } + } + + /// Appends a [`VarBinArray`], reusing its offsets instead of walking its values. + pub fn append_varbin( + &mut self, + array: ArrayView<'_, VarBin>, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> + where + usize: AsPrimitive, + { + let offsets = array.offsets().clone().execute::(ctx)?; + let bytes: ByteBuffer = array.sliced_bytes(); + let validity = array + .varbin_validity() + .execute_mask(array.as_ref().len(), ctx)?; + match_each_integer_ptype!(offsets.ptype(), |P| { + let offsets = offsets.as_slice::

(); + let first: usize = offsets[0].as_(); + self.append_values( + bytes.as_slice(), + // Wrapping keeps a corrupt offsets child from panicking here; `append_values` + // rejects the resulting non-monotonic offsets. + offsets[1..] + .iter() + .map(|offset| AsPrimitive::::as_(*offset).wrapping_sub(first)), + &validity, + ) + }) + } + + /// Appends a [`VarBinViewArray`](crate::arrays::VarBinViewArray), compacting its values. + pub fn append_varbinview( + &mut self, + array: ArrayView<'_, VarBinView>, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> + where + usize: AsPrimitive, + { + let len = array.as_ref().len(); + let validity = array.varbinview_validity().execute_mask(len, ctx)?; + + // Resolve the views slice and the data buffers once. Reading them per row costs a buffer + // handle clone, and the byte total needs only the fixed-width view headers. + let views = array.views(); + let buffers = array + .data_buffers() + .iter() + .map(|buffer| buffer.as_host().as_slice()) + .collect::>(); + + let num_bytes = match validity.bit_buffer() { + AllOr::All => views.iter().map(|view| view.len() as usize).sum(), + AllOr::None => 0, + AllOr::Some(bits) => { + let mut total = 0; + bits.for_each_set_index(|index| total += views[index].len() as usize); + total + } + }; + + self.append_valid_slices(num_bytes, &validity, |index| views[index].bytes(&buffers)) } + /// Finishes the appended values into a [`VarBinArray`] and resets the builder. #[allow(clippy::disallowed_methods)] - pub fn finish(self, dtype: DType) -> VarBinArray { - let offsets = PrimitiveArray::new(self.offsets.freeze(), Validity::NonNullable); - let nulls = self.validity.freeze(); + pub fn finish_into_varbin(&mut self) -> VarBinArray { + assert_eq!( + self.offsets.len() - 1, + self.validity.len(), + "The offset count must be one more than the validity length" + ); + + let allocator = self.offsets.allocator().clone(); + let mut fresh_offsets = BufferMut::with_capacity_in(1, allocator.clone()); + fresh_offsets.push(O::zero()); + let offsets = PrimitiveArray::new( + std::mem::replace(&mut self.offsets, fresh_offsets).freeze(), + Validity::NonNullable, + ); + let data = std::mem::replace( + &mut self.data, + BufferMut::empty_aligned_in(Alignment::of::(), allocator.clone()), + ); + let nulls = + std::mem::replace(&mut self.validity, BitBufferMut::empty_in(allocator)).freeze(); - let validity = Validity::from_bit_buffer(nulls, dtype.nullability()); + let validity = Validity::from_bit_buffer(nulls, self.dtype.nullability()); - // The builder guarantees offsets are monotonically increasing, so we can set - // this stat eagerly. This avoids an O(n) recomputation when the array is - // deserialized and VarBinArray::validate checks sortedness. + // The builder adds offsets in monotonically increasing order. Store this statistic to + // prevent VarBinArray::validate from recomputing it after deserialization. #[cfg(debug_assertions)] { let offsets_are_sorted = offsets @@ -122,19 +449,302 @@ impl VarBinBuilder { // - Validity matches the dtype nullability. // - UTF-8 validity is ensured by the caller when using DType::Utf8. unsafe { - VarBinArray::new_unchecked(offsets.into_array(), self.data.freeze(), dtype, validity) + VarBinArray::new_unchecked( + offsets.into_array(), + data.freeze(), + self.dtype.clone(), + validity, + ) + } + } + + #[inline] + fn last_offset(&self) -> O { + self.offsets[self.offsets.len() - 1] + } + + /// Appends `value`'s bytes and its end offset, leaving the validity bits to the caller. + #[inline] + fn push_value(&mut self, value: &[u8]) { + self.offsets + .push(O::from(self.data.len() + value.len()).unwrap_or_else(|| { + vortex_panic!( + "Failed to convert sum of {} and {} to offset of type {}", + self.data.len(), + value.len(), + std::any::type_name::() + ) + })); + self.data.extend_from_slice(value); + } + + fn append_validity(&mut self, validity: &Mask) { + match validity { + Mask::AllTrue(len) => self.validity.append_n(true, *len), + Mask::AllFalse(len) => self.validity.append_n(false, *len), + Mask::Values(values) => self.validity.append_buffer(values.bit_buffer()), + } + } + + /// Appends `count` end offsets derived from `end_offsets`, shifted past the current data end. + /// + /// `end_offsets` must be monotonically non-decreasing and end at exactly `num_bytes`. Offsets + /// are written through the reserved spare capacity and committed with a single `set_len`, so a + /// rejected input leaves the buffer untouched and the caller can propagate the error. + fn extend_offsets

( + &mut self, + num_bytes: usize, + count: usize, + end_offsets: impl Iterator, + ) -> VortexResult<()> + where + P: AsPrimitive, + usize: AsPrimitive, + { + let data_start = self.data.len(); + let offsets_len = self.offsets.len(); + self.check_offset_limit(data_start, num_bytes)?; + self.offsets.reserve(count); + + // Writing into the spare capacity keeps the output cursor in a register: `push` rewrites + // the buffer length on every value, which the optimizer cannot hoist out of the loop. + let spare = &mut self.offsets.spare_capacity_mut()[..count]; + let mut end_offsets = end_offsets; + let mut previous = 0usize; + for slot in spare.iter_mut() { + let Some(end) = end_offsets.next() else { + vortex_bail!("End offset count is less than the validity length {count}"); + }; + let end = end.as_(); + vortex_ensure!( + end >= previous && end <= num_bytes, + "End offsets must be monotonically increasing within {num_bytes} bytes, \ + got {end} after {previous}" + ); + slot.write((data_start + end).as_()); + previous = end; + } + vortex_ensure!( + end_offsets.next().is_none(), + "End offset count exceeds the validity length {count}" + ); + vortex_ensure!( + previous == num_bytes, + "Final end offset {previous} does not match the value byte count {num_bytes}" + ); + + // SAFETY: the loop initialized the first `count` spare slots. + unsafe { self.offsets.set_len(offsets_len + count) }; + Ok(()) + } + + /// Copies each valid row's bytes from `value` into the byte storage and records the offsets. + /// See [`append_valid_slices`](Self::append_valid_slices); the validity bits are the caller's + /// job so that a failure here can be unwound. + fn gather_valid_slices<'a, F>( + &mut self, + num_bytes: usize, + validity: &Mask, + mut value: F, + ) -> VortexResult<()> + where + F: FnMut(usize) -> &'a [u8], + usize: AsPrimitive, + { + let count = validity.len(); + let data_start = self.data.len(); + let offsets_len = self.offsets.len(); + self.check_offset_limit(data_start, num_bytes)?; + self.offsets.reserve(count); + self.data.reserve(num_bytes); + + // Disjoint field borrows: the offsets spare capacity stays valid while the byte buffer + // grows, since the two are separate allocations. + let Self { offsets, data, .. } = self; + let spare = &mut offsets.spare_capacity_mut()[..count]; + + match validity.bit_buffer() { + AllOr::All => { + for (row, slot) in spare.iter_mut().enumerate() { + data.extend_from_slice(value(row)); + slot.write(data.len().as_()); + } + } + AllOr::None => { + spare.fill(MaybeUninit::new(data_start.as_())); + } + AllOr::Some(bits) => { + let mut row = 0; + bits.for_each_set_index(|index| { + // Null rows between the previous valid row and this one repeat its end offset. + spare[row..index].fill(MaybeUninit::new(data.len().as_())); + data.extend_from_slice(value(index)); + spare[index].write(data.len().as_()); + row = index + 1; + }); + spare[row..].fill(MaybeUninit::new(data.len().as_())); + } } + + // A caller whose slices overrun `num_bytes` only grows the byte buffer past the reservation, + // so the overrun is caught here rather than per value. The offsets are still uncommitted, so + // rejecting it now leaves them untouched. + vortex_ensure!( + data.len() == data_start + num_bytes, + "Value slices total {} bytes, expected {num_bytes}", + data.len() - data_start + ); + + // SAFETY: every branch above initialized all `count` spare slots. + unsafe { self.offsets.set_len(offsets_len + count) }; + Ok(()) + } + + /// Checks that an offset past `num_bytes` more bytes is representable as an `O`. + fn check_offset_limit(&self, data_start: usize, num_bytes: usize) -> VortexResult<()> { + let Some(limit) = data_start.checked_add(num_bytes) else { + vortex_bail!("Byte offset overflow: {data_start} + {num_bytes}"); + }; + vortex_ensure!( + u64::try_from(limit).is_ok_and(|limit| limit <= O::max_value_as_u64()), + "Byte offset {limit} does not fit in {}", + std::any::type_name::() + ); + Ok(()) } } +impl ArrayBuilder for VarBinBuilder { + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn len(&self) -> usize { + self.validity.len() + } + + fn append_zeros(&mut self, n: usize) { + self.offsets.push_n(self.last_offset(), n); + self.validity.append_n(true, n); + } + + unsafe fn append_nulls_unchecked(&mut self, n: usize) { + self.push_nulls(n); + } + + fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> { + self.append_scalar_repeated(scalar, 1) + } + + fn reserve_exact(&mut self, additional: usize) { + self.offsets.reserve(additional); + self.validity.reserve(additional); + } + + fn finish(&mut self) -> ArrayRef { + self.finish_into_varbin().into_array() + } + + fn finish_into_canonical(&mut self, ctx: &mut ExecutionCtx) -> Canonical { + self.finish() + .execute::(ctx) + .vortex_expect("varbin builder should canonicalize") + } +} + +/// Recovers a concrete [`VarBinBuilder`] from a `&mut dyn ArrayBuilder`. +/// +/// Evaluates to `Some($body)` with `$typed` bound to the `&mut VarBinBuilder` when the builder +/// is one, and `None` otherwise, so an encoding can fall through to its other output paths: +/// +/// ```ignore +/// if let Some(result) = match_each_varbin_builder!(builder, |builder| { +/// append_my_encoding(array, builder, ctx) +/// }) { +/// return result; +/// } +/// ``` +#[macro_export] +macro_rules! match_each_varbin_builder { + ($builder:expr, | $typed:ident | $body:expr) => { + $crate::__match_varbin_builder_widths!($builder, |$typed| $body, [u32, u64, i32, i64]) + }; +} + +/// Expands `$body` once per listed offset width, guarded by a downcast. See +/// [`match_each_varbin_builder!`]. +#[doc(hidden)] +#[macro_export] +macro_rules! __match_varbin_builder_widths { + ($builder:expr, | $typed:ident | $body:expr, [$($width:ty),+ $(,)?]) => {{ + let __varbin_builder: &mut dyn $crate::builders::ArrayBuilder = $builder; + $crate::__match_varbin_builder_arms!(__varbin_builder, |$typed| $body, [$($width),+]) + }}; +} + +/// The `if`/`else if` chain behind [`__match_varbin_builder_widths!`], one arm per width. +#[doc(hidden)] +#[macro_export] +macro_rules! __match_varbin_builder_arms { + ($builder:expr, | $typed:ident | $body:expr, []) => { + None + }; + ($builder:expr, | $typed:ident | $body:expr, [$head:ty $(, $tail:ty)*]) => { + if $builder + .as_any() + .is::<$crate::builders::VarBinBuilder<$head>>() + { + let $typed = match $builder + .as_any_mut() + .downcast_mut::<$crate::builders::VarBinBuilder<$head>>() + { + Some(typed) => typed, + None => unreachable!("builder type checked above"), + }; + Some($body) + } else { + $crate::__match_varbin_builder_arms!($builder, |$typed| $body, [$($tail),*]) + } + }; +} + +/// Running totals of `lengths`, wrapping so a corrupt lengths child is rejected rather than +/// panicking; [`VarBinBuilder::extend_offsets`] catches the resulting non-monotonic offsets. +#[inline] +fn prefix_sums>(lengths: &[P]) -> impl Iterator { + lengths.iter().scan(0usize, |end, length| { + *end = end.wrapping_add(length.as_()); + Some(*end) + }) +} + #[cfg(test)] mod tests { + use std::mem::MaybeUninit; + + use rstest::rstest; use vortex_error::VortexResult; + use vortex_mask::Mask; + use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; - use crate::arrays::varbin::VarBinArrayExt; + use crate::arrays::ChunkedArray; + use crate::arrays::ConstantArray; + use crate::arrays::VarBinArray; + use crate::arrays::VarBinViewArray; + use crate::arrays::varbin::VarBinArraySlotsExt; use crate::arrays::varbin::builder::VarBinBuilder; + use crate::assert_arrays_eq; + use crate::builders::ArrayBuilder; use crate::dtype::DType; use crate::dtype::Nullability::Nullable; use crate::expr::stats::Precision; @@ -144,11 +754,15 @@ mod tests { #[test] fn test_builder() { - let mut builder = VarBinBuilder::::with_capacity(0); + let mut builder = VarBinBuilder::::with_capacity_in( + DType::Utf8(Nullable), + 0, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); builder.append(Some(b"hello")); builder.append(None); builder.append(Some(b"world")); - let array = builder.finish(DType::Utf8(Nullable)); + let array = builder.finish_into_varbin(); assert_eq!(array.len(), 3); assert_eq!(array.dtype().nullability(), Nullable); @@ -166,13 +780,293 @@ mod tests { ); } + #[rstest] + #[case(false)] + #[case(true)] + fn test_append_varbin_to_builder(#[case] large_offsets: bool) -> VortexResult<()> { + let source = VarBinArray::from_iter( + [ + Some("prefix"), + Some("hello"), + None, + Some("world"), + Some("suffix"), + ], + DType::Utf8(Nullable), + ) + .into_array() + .slice(1..4)?; + let mut ctx = array_session().create_execution_ctx(); + + let actual = with_offsets(large_offsets, source.dtype().clone(), |builder| { + source.append_to_builder(builder, &mut ctx) + })?; + + assert_arrays_eq!(actual, source, &mut ctx); + Ok(()) + } + + #[test] + fn append_n_values_offset_overflow_returns_error() { + let mut builder = VarBinBuilder::::new_in( + DType::Utf8(Nullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + + // The limit is checked before anything is reserved, so the huge byte total is free. + let result = builder.append_n_values(b"hello", i32::MAX as usize / 5 + 1); + + assert!(result.is_err()); + assert_eq!(builder.offsets.len(), 1); + assert!(builder.data.is_empty()); + assert_eq!(builder.validity.len(), 0); + } + + #[test] + fn append_values_rejects_a_short_offset_count() { + let mut builder = VarBinBuilder::::new_in( + DType::Utf8(Nullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + + let result = builder.append_values(b"ab", [1usize, 2].into_iter(), &Mask::new_true(3)); + + assert!(result.is_err()); + assert_eq!(builder.offsets.len(), 1); + assert!(builder.data.is_empty()); + } + + #[test] + fn append_values_rejects_non_monotonic_offsets() { + let mut builder = VarBinBuilder::::new_in( + DType::Utf8(Nullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + + let result = builder.append_values(b"ab", [2usize, 1].into_iter(), &Mask::new_true(2)); + + assert!(result.is_err()); + assert_eq!(builder.offsets.len(), 1); + } + + #[test] + fn append_decoded_writes_into_the_builder_storage() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = VarBinBuilder::::new_in( + DType::Utf8(Nullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + + // SAFETY: the closure initializes exactly the 6 bytes it reports. + unsafe { + builder.append_decoded( + 6, + 4, + &[3usize, 0, 3], + &Mask::from_iter([true, false, true]), + &mut |spare: &mut [MaybeUninit]| { + for (slot, byte) in spare.iter_mut().zip(b"foobar") { + slot.write(*byte); + } + Ok(6) + }, + )?; + } + + let expected = + VarBinViewArray::from_iter([Some("foo"), None, Some("bar")], DType::Utf8(Nullable)); + assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx); + Ok(()) + } + + /// The offset width has to be checked before `decode` runs: an overflowing `i32` builder + /// should not pay for a full decompression first. + #[test] + fn append_decoded_rejects_an_offset_overflow_without_decoding() { + let mut builder = VarBinBuilder::::new_in( + DType::Utf8(Nullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + let num_bytes = i32::MAX as usize + 1; + let mut decoded = false; + + // SAFETY: the closure is never called, and reports only what it initializes if it were. + let result = unsafe { + builder.append_decoded(num_bytes, 0, &[num_bytes], &Mask::new_true(1), &mut |_| { + decoded = true; + Ok(0) + }) + }; + + assert!(result.is_err()); + assert!(!decoded, "decode ran before the offset width was checked"); + assert_eq!(builder.offsets.len(), 1); + assert_eq!(builder.validity.len(), 0); + } + + /// `append_scalar_repeated` reports an offset overflow rather than panicking, so a constant + /// column wider than the offset type errors like the bulk appends do. + #[test] + fn append_scalar_repeated_rejects_an_offset_overflow() { + let mut builder = VarBinBuilder::::new_in( + DType::Utf8(Nullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + + let result = builder + .append_scalar_repeated(&Scalar::utf8("hello", Nullable), i32::MAX as usize / 5 + 1); + + assert!(result.is_err()); + assert_eq!(builder.offsets.len(), 1); + assert!(builder.data.is_empty()); + assert_eq!(builder.validity.len(), 0); + } + + #[test] + fn append_decoded_rejects_a_short_decode() { + let mut builder = VarBinBuilder::::new_in( + DType::Utf8(Nullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + + // SAFETY: the closure initializes the 3 bytes it reports (none, and it reports 3 — but the + // length mismatch is rejected before anything is published). + let result = unsafe { + builder.append_decoded(6, 0, &[3usize, 3], &Mask::new_true(2), &mut |spare| { + spare[..3].fill(MaybeUninit::new(b'x')); + Ok(3) + }) + }; + + assert!(result.is_err()); + assert_eq!(builder.offsets.len(), 1); + assert_eq!(builder.validity.len(), 0); + } + + /// Slices that overrun the declared byte count are rejected, and the builder is left exactly + /// as it was — the overrun bytes are copied before the total can be checked, so the unwind has + /// to put them back. + #[rstest] + #[case::all_valid(Mask::new_true(2))] + #[case::some_valid(Mask::from_iter([true, false, true]))] + fn append_valid_slices_rejects_a_byte_count_mismatch(#[case] validity: Mask) { + let mut builder = VarBinBuilder::::new_in( + DType::Utf8(Nullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + let values = [b"foo".as_slice(), b"quux".as_slice()]; + + let mut next = 0; + let result = builder.append_valid_slices(6, &validity, |_| { + next += 1; + values[next - 1] + }); + + assert!(result.is_err()); + assert_eq!(builder.offsets.len(), 1); + assert!(builder.data.is_empty()); + assert_eq!(builder.validity.len(), 0); + } + + #[test] + #[should_panic(expected = "The offset count must be one more than the validity length")] + fn finish_rejects_mismatched_validity() { + let mut builder = VarBinBuilder::::new_in( + DType::Utf8(Nullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + builder.validity.append_true(); + drop(builder.finish_into_varbin()); + } + + #[rstest] + #[case(false)] + #[case(true)] + fn test_array_builder_methods(#[case] large_offsets: bool) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let result = with_offsets(large_offsets, DType::Utf8(Nullable), |builder| { + builder.reserve_exact(3); + builder.append_zero(); + builder.append_scalar(&Scalar::utf8("hello", Nullable))?; + builder.append_null(); + assert_eq!(builder.len(), 3); + Ok(()) + })?; + + assert_eq!( + result.validity()?.execute_mask(3, &mut ctx)?, + Mask::from_iter([true, true, false]) + ); + Ok(()) + } + + #[rstest] + #[case(false)] + #[case(true)] + fn test_append_varbinview_validity_to_builder(#[case] large_offsets: bool) -> VortexResult<()> { + let long = "a value that does not fit inline"; + let all_null = VarBinViewArray::from_iter([None::<&str>, None], DType::Utf8(Nullable)); + let mixed = + VarBinViewArray::from_iter([Some("hello"), None, Some(long)], DType::Utf8(Nullable)); + let expected = VarBinViewArray::from_iter( + [None, None, Some("hello"), None, Some(long)], + DType::Utf8(Nullable), + ); + let mut ctx = array_session().create_execution_ctx(); + + let actual = with_offsets(large_offsets, expected.dtype().clone(), |builder| { + all_null + .clone() + .into_array() + .append_to_builder(builder, &mut ctx)?; + mixed + .clone() + .into_array() + .append_to_builder(builder, &mut ctx) + })?; + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + /// `match_each_varbin_builder!` covers every offset width a `VarBinBuilder` can be built with, + /// so both canonical appends must reach their specialization for all of them. A macro that + /// only matched the signed pair would send an unsigned builder down a downcast that assumes + /// `VarBinViewBuilder` and panic. + #[rstest] + #[case::u32(VarBinBuilder::::new_in(DType::Utf8(Nullable), vortex_buffer::BufferAllocatorRef::static_ref()))] + #[case::u64(VarBinBuilder::::new_in(DType::Utf8(Nullable), vortex_buffer::BufferAllocatorRef::static_ref()))] + #[case::i32(VarBinBuilder::::new_in(DType::Utf8(Nullable), vortex_buffer::BufferAllocatorRef::static_ref()))] + #[case::i64(VarBinBuilder::::new_in(DType::Utf8(Nullable), vortex_buffer::BufferAllocatorRef::static_ref()))] + fn append_to_every_offset_width(#[case] mut builder: impl ArrayBuilder) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let long = "a string that is far too long to be inlined in a view"; + let values = [Some("hello"), None, Some(long), Some("")]; + + let view = VarBinViewArray::from_iter(values, DType::Utf8(Nullable)).into_array(); + let varbin = VarBinArray::from_iter(values, DType::Utf8(Nullable)).into_array(); + let constant = ConstantArray::new(Scalar::from("hello").into_nullable(), 2).into_array(); + + for chunk in [&view, &varbin, &constant] { + chunk.append_to_builder(&mut builder, &mut ctx)?; + } + + let expected = ChunkedArray::try_new(vec![view, varbin, constant], DType::Utf8(Nullable))?; + assert_arrays_eq!(builder.finish(), expected, &mut ctx); + Ok(()) + } + #[test] fn offsets_have_is_sorted_stat() -> VortexResult<()> { - let mut builder = VarBinBuilder::::with_capacity(0); + let mut builder = VarBinBuilder::::with_capacity_in( + DType::Utf8(Nullable), + 0, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); builder.append_value(b"aaa"); - builder.append_null(); + builder.push_null(); builder.append_value(b"bbb"); - let array = builder.finish(DType::Utf8(Nullable)); + let array = builder.finish_into_varbin(); let is_sorted = array .offsets() @@ -184,8 +1078,11 @@ mod tests { #[test] fn empty_builder_offsets_have_is_sorted_stat() -> VortexResult<()> { - let builder = VarBinBuilder::::new(); - let array = builder.finish(DType::Utf8(Nullable)); + let mut builder = VarBinBuilder::::new_in( + DType::Utf8(Nullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + let array = builder.finish_into_varbin(); let is_sorted = array .offsets() @@ -194,4 +1091,29 @@ mod tests { assert_eq!(is_sorted, Precision::Exact(true)); Ok(()) } + + /// Runs `f` against an `i32` or `i64` builder and returns the finished array. + fn with_offsets( + large_offsets: bool, + dtype: DType, + f: impl FnOnce(&mut dyn ArrayBuilder) -> VortexResult<()>, + ) -> VortexResult { + if large_offsets { + let mut builder = VarBinBuilder::::with_capacity_in( + dtype, + 8, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + f(&mut builder)?; + Ok(builder.finish_into_varbin()) + } else { + let mut builder = VarBinBuilder::::with_capacity_in( + dtype, + 8, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + f(&mut builder)?; + Ok(builder.finish_into_varbin()) + } + } } diff --git a/vortex-array/src/arrays/varbin/compute/cast.rs b/vortex-array/src/arrays/varbin/compute/cast.rs index b0c4325c774..cc05ec5ac24 100644 --- a/vortex-array/src/arrays/varbin/compute/cast.rs +++ b/vortex-array/src/arrays/varbin/compute/cast.rs @@ -9,7 +9,7 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::VarBin; use crate::arrays::VarBinArray; -use crate::arrays::varbin::VarBinArrayExt; +use crate::arrays::varbin::VarBinArraySlotsExt; use crate::dtype::DType; use crate::scalar_fn::fns::cast::CastKernel; use crate::scalar_fn::fns::cast::CastReduce; diff --git a/vortex-array/src/arrays/varbin/compute/compare.rs b/vortex-array/src/arrays/varbin/compute/compare.rs index f6fff972c9b..464cc2fcd63 100644 --- a/vortex-array/src/arrays/varbin/compute/compare.rs +++ b/vortex-array/src/arrays/varbin/compute/compare.rs @@ -15,7 +15,7 @@ use crate::array::ArrayView; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::VarBin; -use crate::arrays::varbin::VarBinArrayExt; +use crate::arrays::varbin::VarBinArraySlotsExt; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::match_each_integer_ptype; @@ -153,6 +153,7 @@ fn collect_lane_bits( /// /// Offsets at null positions are not validated, so an out-of-bounds or inverted range is /// possible there; such lanes answer `false`, and validity masks them out of the result anyway. +#[allow(clippy::inline_always)] #[inline(always)] fn value_eq(bytes: &[u8], start: usize, end: usize, constant: &[u8]) -> bool { // A lane can only match when its length equals the constant's, so lanes of a different @@ -164,6 +165,7 @@ fn value_eq(bytes: &[u8], start: usize, end: usize, constant: &[u8]) -> bool { /// Order `bytes[start..end]` against `constant`, treating the unvalidated garbage ranges that /// can appear at null positions as empty; validity masks those lanes out of the result anyway. +#[allow(clippy::inline_always)] #[inline(always)] fn value_cmp(bytes: &[u8], start: usize, end: usize, constant: &[u8]) -> Ordering { bytes.get(start..end).unwrap_or_default().cmp(constant) @@ -301,11 +303,15 @@ mod tests { #[test] fn varbin_i64_offsets_compare_constant() { let mut ctx = array_session().create_execution_ctx(); - let mut builder = VarBinBuilder::::with_capacity(3); + let mut builder = VarBinBuilder::::with_capacity_in( + DType::Utf8(Nullability::NonNullable), + 3, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); builder.append_value(b"abc"); builder.append_value(b"xyz"); builder.append_value(b"abc"); - let array = builder.finish(DType::Utf8(Nullability::NonNullable)); + let array = builder.finish_into_varbin(); let result = array .into_array() @@ -322,11 +328,15 @@ mod tests { #[test] fn varbin_i64_offsets_compare_constant_binary() { let mut ctx = array_session().create_execution_ctx(); - let mut builder = VarBinBuilder::::with_capacity(3); + let mut builder = VarBinBuilder::::with_capacity_in( + DType::Binary(Nullability::NonNullable), + 3, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); builder.append_value(b"abc"); builder.append_value(b"xyz"); builder.append_value(b"abc"); - let array = builder.finish(DType::Binary(Nullability::NonNullable)); + let array = builder.finish_into_varbin(); let result = array .into_array() diff --git a/vortex-array/src/arrays/varbin/compute/filter.rs b/vortex-array/src/arrays/varbin/compute/filter.rs index 400beb9aa0a..2e780e58f26 100644 --- a/vortex-array/src/arrays/varbin/compute/filter.rs +++ b/vortex-array/src/arrays/varbin/compute/filter.rs @@ -3,10 +3,10 @@ use itertools::Itertools; use num_traits::AsPrimitive; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; -use vortex_error::vortex_panic; use vortex_mask::AllOr; use vortex_mask::Mask; use vortex_mask::MaskIter; @@ -20,10 +20,13 @@ use crate::arrays::VarBin; use crate::arrays::VarBinArray; use crate::arrays::filter::FilterKernel; use crate::arrays::varbin::VarBinArrayExt; +use crate::arrays::varbin::VarBinArraySlotsExt; use crate::arrays::varbin::builder::VarBinBuilder; use crate::dtype::DType; use crate::dtype::IntegerPType; +use crate::dtype::OffsetBuilderPType; use crate::match_each_integer_ptype; +use crate::match_smallest_list_offset_type; impl FilterKernel for VarBin { fn filter( @@ -62,47 +65,52 @@ fn filter_select_var_bin_by_slice( ) -> VortexResult { let offsets = values.offsets().clone().execute::(ctx)?; match_each_integer_ptype!(offsets.ptype(), |O| { - filter_select_var_bin_by_slice_primitive_offset( - values.dtype().clone(), - offsets.as_slice::(), - values.bytes().as_slice(), - mask_slices, - values - .varbin_validity() - .execute_mask(values.as_ref().len(), ctx) - .vortex_expect("Failed to compute validity mask"), - selection_count, - ) + match_smallest_list_offset_type!(values.bytes().len(), |B| { + filter_select_var_bin_by_slice_primitive_offset::( + values.dtype().clone(), + offsets.as_slice::(), + values.bytes().as_slice(), + mask_slices, + values + .varbin_validity() + .execute_mask(values.as_ref().len(), ctx) + .vortex_expect("Failed to compute validity mask"), + selection_count, + ctx.allocator(), + ) + }) }) } -fn filter_select_var_bin_by_slice_primitive_offset( +fn filter_select_var_bin_by_slice_primitive_offset( dtype: DType, offsets: &[O], data: &[u8], mask_slices: &[(usize, usize)], logical_validity: Mask, selection_count: usize, + allocator: &BufferAllocatorRef, ) -> VortexResult where O: IntegerPType, - usize: AsPrimitive, + B: OffsetBuilderPType, + usize: AsPrimitive, { - let mut builder = VarBinBuilder::::with_capacity(selection_count); + let mut builder = VarBinBuilder::::with_capacity_in(dtype, selection_count, allocator); match logical_validity.bit_buffer() { AllOr::All => { - mask_slices.iter().for_each(|(start, end)| { - update_non_nullable_slice(data, offsets, &mut builder, *start, *end) - }); + for &(start, end) in mask_slices { + update_non_nullable_slice(data, offsets, &mut builder, start, end)?; + } } AllOr::None => { - builder.append_n_nulls(selection_count); + builder.push_nulls(selection_count); } AllOr::Some(validity) => { for (start, end) in mask_slices.iter().copied() { let null_sl = validity.slice(start..end); if null_sl.true_count() == null_sl.len() { - update_non_nullable_slice(data, offsets, &mut builder, start, end) + update_non_nullable_slice(data, offsets, &mut builder, start, end)?; } else { for (idx, valid) in null_sl.iter().enumerate() { if valid { @@ -120,40 +128,40 @@ where })?; builder.append_value(&data[s..e]) } else { - builder.append_null() + builder.push_null() } } } } } } - Ok(builder.finish(dtype)) + Ok(builder.finish_into_varbin()) } -fn update_non_nullable_slice( +fn update_non_nullable_slice( data: &[u8], offsets: &[O], - builder: &mut VarBinBuilder, + builder: &mut VarBinBuilder, start: usize, end: usize, -) where +) -> VortexResult<()> +where O: IntegerPType, - usize: AsPrimitive, + B: OffsetBuilderPType, + usize: AsPrimitive, { - let new_data = { - let offset_start = offsets[start].to_usize().unwrap_or_else(|| { - vortex_panic!("Failed to convert offset to usize: {}", offsets[start]) - }); - let offset_end = offsets[end].to_usize().unwrap_or_else(|| { - vortex_panic!("Failed to convert offset to usize: {}", offsets[end]) - }); - &data[offset_start..offset_end] - }; + let offset_start = offsets[start] + .to_usize() + .ok_or_else(|| vortex_err!("Failed to convert offset to usize: {}", offsets[start]))?; + let offset_end = offsets[end] + .to_usize() + .ok_or_else(|| vortex_err!("Failed to convert offset to usize: {}", offsets[end]))?; + let new_data = &data[offset_start..offset_end]; let new_offsets = offsets[start..end + 1] .iter() .map(|o| *o - offsets[start]) .dropping(1); - builder.append_values(new_data, new_offsets, end - start) + builder.append_values(new_data, new_offsets, &Mask::new_true(end - start)) } fn filter_select_var_bin_by_index( @@ -164,27 +172,31 @@ fn filter_select_var_bin_by_index( ) -> VortexResult { let offsets = values.offsets().clone().execute::(ctx)?; match_each_integer_ptype!(offsets.ptype(), |O| { - filter_select_var_bin_by_index_primitive_offset( - values.dtype().clone(), - offsets.as_slice::(), - values.bytes().as_slice(), - mask_indices, - values - .varbin_validity() - .execute_mask(values.as_ref().len(), ctx) - .vortex_expect("Failed to compute validity mask"), - selection_count, - ) + match_smallest_list_offset_type!(values.bytes().len(), |B| { + filter_select_var_bin_by_index_primitive_offset::( + values.dtype().clone(), + offsets.as_slice::(), + values.bytes().as_slice(), + mask_indices, + values + .varbin_validity() + .execute_mask(values.as_ref().len(), ctx) + .vortex_expect("Failed to compute validity mask"), + selection_count, + ctx.allocator(), + ) + }) }) } -fn filter_select_var_bin_by_index_primitive_offset( +fn filter_select_var_bin_by_index_primitive_offset( dtype: DType, offsets: &[O], data: &[u8], mask_indices: &[usize], mask: Mask, selection_count: usize, + allocator: &BufferAllocatorRef, ) -> VortexResult { let value_at = |idx: usize| -> VortexResult<&[u8]> { let start = offsets[idx] @@ -196,7 +208,7 @@ fn filter_select_var_bin_by_index_primitive_offset( Ok(&data[start..end]) }; - let mut builder = VarBinBuilder::::with_capacity(selection_count); + let mut builder = VarBinBuilder::::with_capacity_in(dtype, selection_count, allocator); match mask.bit_buffer() { AllOr::All => { for idx in mask_indices.iter().copied() { @@ -204,19 +216,19 @@ fn filter_select_var_bin_by_index_primitive_offset( } } AllOr::None => { - builder.append_n_nulls(selection_count); + builder.push_nulls(selection_count); } AllOr::Some(validity) => { for idx in mask_indices.iter().copied() { if validity.value(idx) { builder.append_value(value_at(idx)?) } else { - builder.append_null() + builder.push_null() } } } } - Ok(builder.finish(dtype)) + Ok(builder.finish_into_varbin()) } #[cfg(test)] diff --git a/vortex-array/src/arrays/varbin/compute/mask.rs b/vortex-array/src/arrays/varbin/compute/mask.rs index 19236cee47b..0d91b6fe3c3 100644 --- a/vortex-array/src/arrays/varbin/compute/mask.rs +++ b/vortex-array/src/arrays/varbin/compute/mask.rs @@ -8,7 +8,7 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::VarBin; use crate::arrays::VarBinArray; -use crate::arrays::varbin::VarBinArrayExt; +use crate::arrays::varbin::VarBinArraySlotsExt; use crate::scalar_fn::fns::mask::MaskReduce; use crate::validity::Validity; diff --git a/vortex-array/src/arrays/varbin/compute/mod.rs b/vortex-array/src/arrays/varbin/compute/mod.rs index 480c07a3031..ea1a387008f 100644 --- a/vortex-array/src/arrays/varbin/compute/mod.rs +++ b/vortex-array/src/arrays/varbin/compute/mod.rs @@ -10,6 +10,8 @@ mod filter; mod mask; mod take; +pub use take::take_varbin; + #[cfg(test)] mod tests { use rstest::rstest; diff --git a/vortex-array/src/arrays/varbin/compute/slice.rs b/vortex-array/src/arrays/varbin/compute/slice.rs index b3ab91cc578..04ed9e52b6e 100644 --- a/vortex-array/src/arrays/varbin/compute/slice.rs +++ b/vortex-array/src/arrays/varbin/compute/slice.rs @@ -11,7 +11,7 @@ use crate::array::ArrayView; use crate::arrays::VarBin; use crate::arrays::VarBinArray; use crate::arrays::slice::SliceReduce; -use crate::arrays::varbin::VarBinArrayExt; +use crate::arrays::varbin::VarBinArraySlotsExt; impl SliceReduce for VarBin { fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 6a6454d0603..59689d9116c 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -1,27 +1,47 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::iter; +use std::ptr; +use std::sync::Arc; + +use itertools::Itertools as _; +use num_traits::AsPrimitive; use vortex_buffer::BitBufferMut; +use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_buffer::ByteBufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_error::vortex_panic; +use vortex_mask::AllOr; use vortex_mask::Mask; use crate::ArrayRef; +use crate::Columnar; use crate::IntoArray; use crate::array::ArrayView; +use crate::arrays::PiecewiseSequence; use crate::arrays::PrimitiveArray; use crate::arrays::VarBin; use crate::arrays::VarBinArray; +use crate::arrays::VarBinViewArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::piecewise_sequence::constant_unsigned_usize; +use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::primitive::PrimitiveArrayExt; use crate::arrays::varbin::VarBinArrayExt; +use crate::arrays::varbin::VarBinArraySlotsExt; +use crate::arrays::varbinview::BinaryView; +use crate::arrays::varbinview::build_views::MAX_BUFFER_LEN; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::dtype::PType; +use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; +use crate::match_each_integer_ptype; use crate::match_each_unsigned_integer_ptype; use crate::validity::Validity; @@ -37,83 +57,290 @@ fn taken_offset_ptype(offsets_ptype: PType) -> PType { } } +// "take" can have offsets which don't fit in 32 bits. "signed" is needed for +// Arrow interoperability. +enum Offsets { + U32 { + offsets: BufferMut, + signed: bool, + }, + U64 { + offsets: BufferMut, + signed: bool, + }, +} + +impl Offsets { + fn with_capacity(offset_ptype: PType, capacity: usize) -> Self { + let signed = offset_ptype.is_signed_int(); + match offset_ptype { + PType::U64 | PType::I64 => Self::U64 { + offsets: BufferMut::with_capacity(capacity), + signed, + }, + _ => Self::U32 { + offsets: BufferMut::with_capacity(capacity), + signed, + }, + } + } + + fn len(&self) -> usize { + match self { + Self::U32 { offsets, .. } => offsets.len(), + Self::U64 { offsets, .. } => offsets.len(), + } + } + + fn push(&mut self, offset: usize) { + let (offsets, signed) = match self { + Self::U64 { offsets, .. } => { + offsets.push(offset as u64); + return; + } + Self::U32 { offsets, signed } => (offsets, *signed), + }; + let max_narrow = if signed { + i32::MAX as usize + } else { + u32::MAX as usize + }; + if offset <= max_narrow { + offsets.push(u32::try_from(offset).vortex_expect("offset fits u32")); + return; + } + + // Copying the buffer may look scary but in fact it isn't. + // If we don't copy-and-widen, we need to know target width + // beforehand. This is approach taken in ListArray, where first pass + // iterates over offsets and determines target width, and second offset + // copies the offset to the buffer. Both have similar performance + // Note this happens only if 32-bit offset overflows which means array + // must hold >4GB of data which is quite rare. + let mut wide = BufferMut::::with_capacity(offsets.capacity() + 1); + wide.extend(offsets.iter().map(|&o| u64::from(o))); + wide.push(offset as u64); + *self = Self::U64 { + offsets: wide, + signed, + }; + } + + fn into_array(self) -> ArrayRef { + let (offsets, out_offset_ptype) = match self { + Self::U32 { offsets, signed } => ( + PrimitiveArray::new(offsets.freeze(), Validity::NonNullable), + if signed { PType::I32 } else { PType::U32 }, + ), + Self::U64 { offsets, signed } => ( + PrimitiveArray::new(offsets.freeze(), Validity::NonNullable), + if signed { PType::I64 } else { PType::U64 }, + ), + }; + offsets.reinterpret_cast(out_offset_ptype).into_array() + } +} + impl TakeExecute for VarBin { fn take( array: ArrayView<'_, VarBin>, indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - // TODO(joe): Be lazy with execute let offsets = array.offsets().clone().execute::(ctx)?; - let data = array.bytes(); - let indices = indices.clone().execute::(ctx)?; + let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); + let last_offset = match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { + offsets.as_slice::().last().map_or(0usize, |&o| o.as_()) + }); + + // VarBinView can't hold this buffer, so we can't canonicalize and + // take() (take panics). Convert to VarBin + if last_offset > MAX_BUFFER_LEN { + return Ok(Some(take_varbin(array, indices, ctx)?.into_array())); + } + + let data = array.bytes().clone(); let dtype = array .dtype() .clone() .union_nullability(indices.dtype().nullability()); - let array_validity = array - .varbin_validity() - .execute_mask(array.as_ref().len(), ctx)?; - let indices_validity = indices + let validity = array.validity()?.take(indices)?; + + let indices = indices.clone().execute::(ctx)?; + let indices_mask = indices .as_ref() .validity()? .execute_mask(indices.as_ref().len(), ctx)?; - // Offsets and indices are non-negative; read them through their unsigned reinterpretations - // so we only monomorphize over the 4 unsigned widths each (4x4 instead of 8x8). On take, - // offsets get widened to either 32- or 64-bit (to avoid overflow); the built output offsets - // are reinterpreted back to `out_offset_ptype` to preserve the result's offset signedness. - let out_offset_ptype = taken_offset_ptype(offsets.ptype()); - let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); - let indices = indices.reinterpret_cast(indices.ptype().to_unsigned()); - - let array = match_each_unsigned_integer_ptype!(indices.ptype(), |I| { - match offsets.ptype() { - PType::U8 => take::( - dtype, - offsets.as_slice::(), + let views = match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { + match_each_integer_ptype!(indices.ptype(), |I| { + take_views( + offsets.as_slice::(), data.as_slice(), indices.as_slice::(), - array_validity, - indices_validity, - out_offset_ptype, - ), - PType::U16 => take::( - dtype, - offsets.as_slice::(), - data.as_slice(), - indices.as_slice::(), - array_validity, - indices_validity, - out_offset_ptype, - ), - PType::U32 => take::( - dtype, - offsets.as_slice::(), - data.as_slice(), - indices.as_slice::(), - array_validity, - indices_validity, - out_offset_ptype, - ), - PType::U64 => take::( - dtype, - offsets.as_slice::(), - data.as_slice(), - indices.as_slice::(), - array_validity, - indices_validity, - out_offset_ptype, - ), - _ => unreachable!("invalid PType for offsets"), - } + &indices_mask, + ) + }) }); - Ok(Some(array?.into_array())) + // SAFETY: every view references buffer 0 which is inside shared data buffer + unsafe { + Ok(Some( + VarBinViewArray::new_unchecked(views, Arc::from([data]), dtype, validity) + .into_array(), + )) + } + } +} + +fn take_views>( + offsets: &[O], + data: &[u8], + indices: &[I], + mask: &Mask, +) -> Buffer { + let build = |idx: usize| -> BinaryView { + let start: usize = offsets[idx].as_(); + let stop: usize = offsets[idx + 1].as_(); + let value = &data[start..stop]; + let len = stop - start; + + // Caller guarantees every offset is <= MAX_BUFFER_LEN + let start: u32 = start.as_(); + if len > BinaryView::MAX_INLINED_SIZE { + let mut prefix = [0u8; 4]; + prefix.copy_from_slice(&value[..4]); + let len: u32 = len.as_(); + BinaryView::new_ref(len, prefix, 0, start) + } else { + BinaryView::make_view(value, 0, start) + } + }; + + match mask.bit_buffer() { + AllOr::All => Buffer::from_trusted_len_iter(indices.iter().map(|i| build(i.as_()))), + AllOr::None => { + Buffer::from_trusted_len_iter(iter::repeat_n(BinaryView::default(), indices.len())) + } + AllOr::Some(buffer) => { + Buffer::from_trusted_len_iter(buffer.iter().zip(indices.iter()).map(|(valid, i)| { + if valid { + build(i.as_()) + } else { + BinaryView::default() + } + })) + } + } +} + +/// Take from a VarBin. Referenced bytes are copied +pub fn take_varbin( + array: ArrayView<'_, VarBin>, + indices: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)? + { + return Ok(taken); + } + + let offsets = array.offsets().clone().execute::(ctx)?; + let data = array.bytes(); + let indices = indices.clone().execute::(ctx)?; + let dtype = array + .dtype() + .clone() + .union_nullability(indices.dtype().nullability()); + let array_validity = array + .varbin_validity() + .execute_mask(array.as_ref().len(), ctx)?; + let indices_validity = indices + .as_ref() + .validity()? + .execute_mask(indices.as_ref().len(), ctx)?; + + // Offsets and indices are non-negative; read them through their unsigned reinterpretations + // so we only monomorphize over the 4 unsigned widths each (4x4 instead of 8x8). On take, + // offsets get widened to either 32- or 64-bit (to avoid overflow); the built output offsets + // are reinterpreted back to `out_offset_ptype` to preserve the result's offset signedness. + let out_offset_ptype = taken_offset_ptype(offsets.ptype()); + let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); + let indices = indices.reinterpret_cast(indices.ptype().to_unsigned()); + + match_each_unsigned_integer_ptype!(indices.ptype(), |I| { + match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { + take::( + dtype, + offsets.as_slice::(), + data.as_slice(), + indices.as_slice::(), + array_validity, + indices_validity, + out_offset_ptype, + ) + }) + }) +} + +fn take_contiguous_ranges( + array: ArrayView<'_, VarBin>, + indices: ArrayView<'_, PiecewiseSequence>, + indices_ref: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { + return Ok(None); + }; + let offsets = array.offsets().clone().execute::(ctx)?; + let out_offset_ptype = taken_offset_ptype(offsets.ptype()); + let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); + let bytes = array.bytes(); + let data = bytes.as_slice(); + let dtype = array.dtype().clone(); + let output_len = indices_ref.len(); + + let result = match lengths { + Columnar::Constant(lengths) => { + let length = constant_unsigned_usize(&lengths); + gather_slices_constant_dispatch( + &starts, + length, + &offsets, + data, + output_len, + out_offset_ptype, + )? + } + Columnar::Canonical(lengths) => { + let lengths = lengths.into_primitive(); + gather_slices_dispatch( + &starts, + &lengths, + &offsets, + data, + output_len, + out_offset_ptype, + )? + } + }; + + let validity = array.validity()?.take(indices_ref)?; + + // SAFETY: output offsets are built from valid input offsets, start at zero, are monotonically + // non-decreasing, and the copied data buffer has exactly the referenced byte length. + unsafe { + Ok(Some(VarBinArray::new_unchecked( + result.offsets, + result.data.freeze(), + dtype, + validity, + ))) } } -fn take( +fn take( dtype: DType, offsets: &[Offset], data: &[u8], @@ -123,7 +350,7 @@ fn take( out_offset_ptype: PType, ) -> VortexResult { if !validity_mask.all_true() || !indices_validity_mask.all_true() { - return Ok(take_nullable::( + return Ok(take_nullable::( dtype, offsets, data, @@ -134,9 +361,9 @@ fn take( )); } - let mut new_offsets = BufferMut::::with_capacity(indices.len() + 1); - new_offsets.push(NewOffset::zero()); - let mut current_offset = NewOffset::zero(); + let mut new_offsets = Offsets::with_capacity(out_offset_ptype, indices.len() + 1); + new_offsets.push(0); + let mut current_offset = 0usize; for &idx in indices { let idx = idx @@ -145,11 +372,13 @@ fn take( let start = offsets[idx]; let stop = offsets[idx + 1]; - current_offset += NewOffset::from(stop - start).vortex_expect("offset type overflow"); + current_offset += (stop - start) + .to_usize() + .vortex_expect("Failed to cast offset to usize"); new_offsets.push(current_offset); } - let mut new_data = ByteBufferMut::with_capacity(current_offset.as_()); + let mut new_data = ByteBufferMut::with_capacity(current_offset); for idx in indices { let idx = idx @@ -165,11 +394,7 @@ fn take( } let array_validity = Validity::from(dtype.nullability()); - - // Built unsigned; reinterpret back to the signedness-preserving result offset type. - let new_offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) - .reinterpret_cast(out_offset_ptype) - .into_array(); + let new_offsets = new_offsets.into_array(); // Safety: // All variants of VarBinArray are satisfied here. @@ -183,7 +408,359 @@ fn take( } } -fn take_nullable( +struct GatheredPiecewiseVarBin { + offsets: ArrayRef, + data: ByteBufferMut, +} + +fn gather_slices_constant_dispatch( + starts: &PrimitiveArray, + length: usize, + offsets: &PrimitiveArray, + data: &[u8], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + gather_slices_constant_start_dispatch::( + starts, + length, + offsets, + data, + output_len, + out_offset_ptype, + ) + }) +} + +fn gather_slices_constant_start_dispatch( + starts: &PrimitiveArray, + length: usize, + offsets: &PrimitiveArray, + data: &[u8], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult +where + S: UnsignedPType, +{ + match offsets.ptype() { + PType::U8 => gather_slices_constant_length::( + offsets.as_slice::(), + data, + starts.as_slice::(), + length, + output_len, + out_offset_ptype, + ), + PType::U16 => gather_slices_constant_length::( + offsets.as_slice::(), + data, + starts.as_slice::(), + length, + output_len, + out_offset_ptype, + ), + PType::U32 => gather_slices_constant_length::( + offsets.as_slice::(), + data, + starts.as_slice::(), + length, + output_len, + out_offset_ptype, + ), + PType::U64 => gather_slices_constant_length::( + offsets.as_slice::(), + data, + starts.as_slice::(), + length, + output_len, + out_offset_ptype, + ), + _ => unreachable!("offsets were reinterpreted to an unsigned integer ptype"), + } +} + +fn gather_slices_dispatch( + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + offsets: &PrimitiveArray, + data: &[u8], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + gather_slices_start_dispatch::( + starts, + lengths, + offsets, + data, + output_len, + out_offset_ptype, + ) + }) +} + +fn gather_slices_start_dispatch( + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + offsets: &PrimitiveArray, + data: &[u8], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult +where + S: UnsignedPType, +{ + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + gather_slices_start_length_dispatch::( + starts, + lengths, + offsets, + data, + output_len, + out_offset_ptype, + ) + }) +} + +fn gather_slices_start_length_dispatch( + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + offsets: &PrimitiveArray, + data: &[u8], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, +{ + match offsets.ptype() { + PType::U8 => gather_slices::( + offsets.as_slice::(), + data, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + out_offset_ptype, + ), + PType::U16 => gather_slices::( + offsets.as_slice::(), + data, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + out_offset_ptype, + ), + PType::U32 => gather_slices::( + offsets.as_slice::(), + data, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + out_offset_ptype, + ), + PType::U64 => gather_slices::( + offsets.as_slice::(), + data, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + out_offset_ptype, + ), + _ => unreachable!("offsets were reinterpreted to an unsigned integer ptype"), + } +} + +fn gather_slices_constant_length( + offsets: &[Offset], + data: &[u8], + starts: &[S], + length: usize, + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult +where + S: UnsignedPType, + Offset: IntegerPType, +{ + let computed_len = starts + .len() + .checked_mul(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); + + let mut new_offsets = Offsets::with_capacity(out_offset_ptype, output_len + 1); + new_offsets.push(0); + let mut output_bytes = 0usize; + + for start in starts { + let start = start.as_(); + if length == 0 { + continue; + } + + let offset_range = &offsets[start..][..=length]; + let byte_start = offset_range[0].as_(); + let byte_end = offset_range[length].as_(); + vortex_ensure!( + byte_start <= byte_end && byte_end <= data.len(), + "VarBin offsets range {byte_start}..{byte_end} exceeds data length {}", + data.len() + ); + + for &offset in &offset_range[1..] { + let offset = offset.as_(); + let relative = offset.checked_sub(byte_start).ok_or_else(|| { + vortex_err!("VarBin offsets are not monotonic at offset {offset}") + })?; + let output_offset = output_bytes.checked_add(relative).ok_or_else(|| { + vortex_err!("PiecewiseSequence VarBin output byte length overflow") + })?; + new_offsets.push(output_offset); + } + + output_bytes = output_bytes + .checked_add(byte_end - byte_start) + .ok_or_else(|| vortex_err!("PiecewiseSequence VarBin output byte length overflow"))?; + } + + let mut new_data = ByteBufferMut::with_capacity(output_bytes); + let spare = &mut new_data.spare_capacity_mut()[..output_bytes]; + let mut cursor = 0usize; + for start in starts { + let start = start.as_(); + if length == 0 { + continue; + } + + let offset_range = &offsets[start..][..=length]; + let byte_start = offset_range[0].as_(); + let byte_end = offset_range[length].as_(); + let src = &data[byte_start..byte_end]; + // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. + unsafe { + ptr::copy_nonoverlapping( + src.as_ptr(), + spare[cursor..][..src.len()].as_mut_ptr().cast::(), + src.len(), + ); + } + cursor += src.len(); + } + // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity. + unsafe { new_data.set_len(cursor) }; + vortex_ensure!( + new_data.len() == output_bytes, + "PiecewiseSequenceArray gathered byte length {} does not match declared byte length {output_bytes}", + new_data.len() + ); + + let offsets = new_offsets.into_array(); + Ok(GatheredPiecewiseVarBin { + offsets, + data: new_data, + }) +} + +fn gather_slices( + offsets: &[Offset], + data: &[u8], + starts: &[S], + lengths: &[L], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, + Offset: IntegerPType, +{ + let mut new_offsets = Offsets::with_capacity(out_offset_ptype, output_len + 1); + new_offsets.push(0); + let mut output_bytes = 0usize; + + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = start.as_(); + let length = length.as_(); + if length == 0 { + continue; + } + + let offset_range = &offsets[start..][..=length]; + let byte_start = offset_range[0].as_(); + let byte_end = offset_range[length].as_(); + vortex_ensure!( + byte_start <= byte_end && byte_end <= data.len(), + "VarBin offsets range {byte_start}..{byte_end} exceeds data length {}", + data.len() + ); + + for &offset in &offset_range[1..] { + let offset = offset.as_(); + let relative = offset.checked_sub(byte_start).ok_or_else(|| { + vortex_err!("VarBin offsets are not monotonic at offset {offset}") + })?; + let output_offset = output_bytes.checked_add(relative).ok_or_else(|| { + vortex_err!("PiecewiseSequence VarBin output byte length overflow") + })?; + new_offsets.push(output_offset); + } + + output_bytes = output_bytes + .checked_add(byte_end - byte_start) + .ok_or_else(|| vortex_err!("PiecewiseSequence VarBin output byte length overflow"))?; + } + vortex_ensure!( + new_offsets.len() == output_len + 1, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + new_offsets.len() - 1 + ); + + let mut new_data = ByteBufferMut::with_capacity(output_bytes); + let spare = &mut new_data.spare_capacity_mut()[..output_bytes]; + let mut cursor = 0usize; + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = start.as_(); + let length = length.as_(); + if length == 0 { + continue; + } + + let offset_range = &offsets[start..][..=length]; + let byte_start = offset_range[0].as_(); + let byte_end = offset_range[length].as_(); + let src = &data[byte_start..byte_end]; + // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. + unsafe { + ptr::copy_nonoverlapping( + src.as_ptr(), + spare[cursor..][..src.len()].as_mut_ptr().cast::(), + src.len(), + ); + } + cursor += src.len(); + } + // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity. + unsafe { new_data.set_len(cursor) }; + vortex_ensure!( + new_data.len() == output_bytes, + "PiecewiseSequenceArray gathered byte length {} does not match declared byte length {output_bytes}", + new_data.len() + ); + + let offsets = new_offsets.into_array(); + Ok(GatheredPiecewiseVarBin { + offsets, + data: new_data, + }) +} + +fn take_nullable( dtype: DType, offsets: &[Offset], data: &[u8], @@ -192,9 +769,9 @@ fn take_nullable VarBinArray { - let mut new_offsets = BufferMut::::with_capacity(indices.len() + 1); - new_offsets.push(NewOffset::zero()); - let mut current_offset = NewOffset::zero(); + let mut new_offsets = Offsets::with_capacity(out_offset_ptype, indices.len() + 1); + new_offsets.push(0); + let mut current_offset = 0usize; let mut validity_buffer = BitBufferMut::with_capacity(indices.len()); @@ -215,7 +792,9 @@ fn take_nullable VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let big = u32::MAX as usize + 5; + let mut sink = super::Offsets::with_capacity(PType::U32, 3); + sink.push(0); + sink.push(7); + sink.push(big); + + let array = sink.into_array().execute::(&mut ctx)?; + assert_eq!(array.as_slice::(), &[0, 7, big as u64]); + Ok(()) + } + + #[test] + fn test_take_signed() { + let mut ctx = array_session().create_execution_ctx(); + let array = VarBinArray::new( + buffer![0i32, 5, 10].into_array(), + ByteBuffer::copy_from(b"helloworld"), + DType::Utf8(Nullability::NonNullable), + Validity::NonNullable, + ); + + let taken = array.take(buffer![1u32, 0].into_array()).unwrap(); + + let expected = VarBinViewArray::from_iter( + [Some("world"), Some("hello")], + DType::Utf8(Nullability::NonNullable), + ); + assert_arrays_eq!(expected, taken, &mut ctx); + } } diff --git a/vortex-array/src/arrays/varbin/mod.rs b/vortex-array/src/arrays/varbin/mod.rs index 4bca67fad81..e7bdca23e70 100644 --- a/vortex-array/src/arrays/varbin/mod.rs +++ b/vortex-array/src/arrays/varbin/mod.rs @@ -3,11 +3,15 @@ mod array; pub use array::VarBinArrayExt; +pub use array::VarBinArraySlotsExt; pub use array::VarBinData; pub use array::VarBinDataParts; +pub use array::VarBinSlots; +pub use array::VarBinSlotsView; pub use vtable::VarBinArray; pub(crate) mod compute; +pub use compute::take_varbin; mod vtable; pub use vtable::VarBin; diff --git a/vortex-array/src/arrays/varbin/vtable/canonical.rs b/vortex-array/src/arrays/varbin/vtable/canonical.rs index fca9671e3b4..825d920f1c7 100644 --- a/vortex-array/src/arrays/varbin/vtable/canonical.rs +++ b/vortex-array/src/arrays/varbin/vtable/canonical.rs @@ -3,7 +3,6 @@ use std::sync::Arc; -use num_traits::AsPrimitive; use vortex_error::VortexResult; use crate::ExecutionCtx; @@ -12,47 +11,50 @@ use crate::arrays::PrimitiveArray; use crate::arrays::VarBin; use crate::arrays::VarBinViewArray; use crate::arrays::varbinview::build_views::MAX_BUFFER_LEN; -use crate::arrays::varbinview::build_views::build_views; -use crate::arrays::varbinview::build_views::offsets_to_lengths; +use crate::arrays::varbinview::build_views::build_views_from_offsets; use crate::match_each_integer_ptype; /// Converts a VarBinArray to its canonical form (VarBinViewArray). /// /// This is a shared helper used by both `canonicalize` and `execute`. +/// +/// The value bytes are handed over as they are — only the offsets are consumed, to derive the view +/// lengths — so this costs one view per row and no byte copy. pub(crate) fn varbin_to_canonical( array: ArrayView<'_, VarBin>, ctx: &mut ExecutionCtx, ) -> VortexResult { let parts = array.into_owned().into_data_parts(); - let offsets = parts.offsets.execute::(ctx)?; - - match_each_integer_ptype!(offsets.ptype(), |P| { - let offsets_slice = offsets.as_slice::

(); - let first: usize = offsets_slice[0].as_(); - let last: usize = offsets_slice[offsets_slice.len() - 1].as_(); - let bytes = parts.bytes.unwrap_host().slice(first..last).into_mut(); - - let lens = offsets_to_lengths(offsets_slice); - let (buffers, views) = build_views(0, MAX_BUFFER_LEN, bytes, lens.as_slice()); - - // SAFETY: views are correctly computed from valid offsets - Ok(unsafe { - VarBinViewArray::new_unchecked(views, Arc::from(buffers), parts.dtype, parts.validity) - }) + let (buffers, views) = match_each_integer_ptype!(offsets.ptype(), |P| { + build_views_from_offsets( + 0, + MAX_BUFFER_LEN, + parts.bytes.unwrap_host(), + offsets.as_slice::

(), + ) + }); + + // SAFETY: views are correctly computed from valid offsets + Ok(unsafe { + VarBinViewArray::new_unchecked(views, Arc::from(buffers), parts.dtype, parts.validity) }) } #[cfg(test)] mod tests { use rstest::rstest; + use vortex_error::VortexResult; + use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; + use crate::arrays::ChunkedArray; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::arrays::varbin::builder::VarBinBuilder; use crate::assert_arrays_eq; + use crate::builders::VarBinViewBuilder; use crate::dtype::DType; use crate::dtype::Nullability; @@ -60,14 +62,18 @@ mod tests { #[case(DType::Utf8(Nullability::Nullable))] #[case(DType::Binary(Nullability::Nullable))] fn test_canonical_varbin_sliced(#[case] dtype: DType) { - let mut varbin = VarBinBuilder::::with_capacity(10); - varbin.append_null(); - varbin.append_null(); + let mut varbin = VarBinBuilder::::with_capacity_in( + dtype.clone(), + 10, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + varbin.push_null(); + varbin.push_null(); // inlined value varbin.append_value("123456789012".as_bytes()); // non-inlinable value varbin.append_value("1234567890123".as_bytes()); - let varbin = varbin.finish(dtype.clone()); + let varbin = varbin.finish_into_varbin(); let varbin = varbin.slice(1..4).unwrap(); @@ -108,6 +114,90 @@ mod tests { assert_arrays_eq!(canonical, expected, &mut ctx); } + /// Appending a `VarBin` array to a `VarBinViewBuilder` builds views over its bytes directly + /// instead of canonicalizing first, so the views must be numbered against the buffers the + /// builder already holds. Interleaving `VarBin` appends with value appends (which stage an + /// in-progress buffer) and with a `VarBinView` append exercises that numbering. + #[rstest] + #[case(DType::Utf8(Nullability::Nullable))] + #[case(DType::Binary(Nullability::Nullable))] + fn append_varbin_to_varbinview_builder(#[case] dtype: DType) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let long = "a value long enough that its view has to reference a buffer"; + let longer = "another value long enough that its view has to reference a buffer"; + + // Two chunks, each with an inlined value, a buffer-referencing value and a null, so both + // pushed buffers are non-empty and the second must not reuse the first's index. + let first = VarBinArray::from_iter([Some("short"), None, Some(long)], dtype.clone()); + let second = VarBinArray::from_iter([Some(longer), Some("tiny"), None], dtype.clone()); + let view = VarBinViewArray::from_iter([Some(long), None], dtype.clone()); + + let mut builder = VarBinViewBuilder::with_capacity_in( + dtype.clone(), + 8, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); + first + .as_array() + .clone() + .append_to_builder(&mut builder, &mut ctx)?; + // Stages an in-progress buffer, which the next append has to account for. + builder.append_value(longer); + second + .as_array() + .clone() + .append_to_builder(&mut builder, &mut ctx)?; + view.clone() + .into_array() + .append_to_builder(&mut builder, &mut ctx)?; + + let expected = ChunkedArray::try_new( + vec![ + first.as_array().clone(), + VarBinViewArray::from_iter([Some(longer)], dtype.clone()).into_array(), + second.as_array().clone(), + view.into_array(), + ], + dtype, + )?; + assert_arrays_eq!(builder.finish_into_varbinview(), expected, &mut ctx); + Ok(()) + } + + /// A builder configured to compact must not be handed a raw buffer behind its back: the + /// inlined values leave the pushed buffer only partly referenced, and skipping compaction + /// would keep those bytes alive. The builder measures utilization from the value lengths and + /// drops the fully-inlined heap. + #[test] + fn append_varbin_to_a_compacting_builder_still_compacts() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DType::Utf8(Nullability::NonNullable); + // Every value inlines, so nothing references the value bytes at all. + let array = VarBinArray::from_iter_nonnull(["short", "tiny", "small"], dtype.clone()); + + let mut builder = VarBinViewBuilder::with_compaction_in( + dtype, + 4, + 1.0, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); + array + .as_array() + .clone() + .append_to_builder(&mut builder, &mut ctx)?; + let compacted = builder.finish_into_varbinview(); + + assert!( + compacted + .data_buffers() + .iter() + .all(|buffer| buffer.is_empty()), + "a fully-inlined append should not retain any value bytes" + ); + assert_arrays_eq!(compacted, array.as_array().clone(), &mut ctx); + Ok(()) + } + // Empty array: offsets has exactly one element; no elements to canonicalize. #[test] fn test_canonical_varbin_empty() { diff --git a/vortex-array/src/arrays/varbin/vtable/mod.rs b/vortex-array/src/arrays/varbin/vtable/mod.rs index 301fdebc084..397b0f8084b 100644 --- a/vortex-array/src/arrays/varbin/vtable/mod.rs +++ b/vortex-array/src/arrays/varbin/vtable/mod.rs @@ -20,18 +20,22 @@ use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::VTable; +use crate::arrays::PrimitiveArray; use crate::arrays::varbin::VarBinArrayExt; +use crate::arrays::varbin::VarBinArraySlotsExt; use crate::arrays::varbin::VarBinData; -use crate::arrays::varbin::array::NUM_SLOTS; -use crate::arrays::varbin::array::OFFSETS_SLOT; -use crate::arrays::varbin::array::SLOT_NAMES; +use crate::arrays::varbin::VarBinSlots; use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; +use crate::builders::VarBinViewBuilder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; +use crate::match_each_integer_ptype; +use crate::match_each_varbin_builder; use crate::serde::ArrayChildren; use crate::validity::Validity; -mod canonical; +pub(crate) mod canonical; mod kernel; mod operations; mod validity; @@ -91,11 +95,12 @@ impl VTable for VarBin { slots: &[Option], ) -> VortexResult<()> { vortex_ensure!( - slots.len() == NUM_SLOTS, - "VarBinArray expected {NUM_SLOTS} slots, found {}", + slots.len() == VarBinSlots::COUNT, + "VarBinArray expected {} slots, found {}", + VarBinSlots::COUNT, slots.len() ); - let offsets = slots[OFFSETS_SLOT] + let offsets = slots[VarBinSlots::OFFSETS] .as_ref() .vortex_expect("VarBinArray offsets slot"); vortex_ensure!( @@ -192,7 +197,7 @@ impl VTable for VarBin { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + VarBinSlots::NAMES[idx].to_string() } fn reduce_parent( @@ -203,6 +208,25 @@ impl VTable for VarBin { PARENT_RULES.evaluate(array, parent, child_idx) } + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if let Some(result) = + match_each_varbin_builder!(builder, |builder| builder.append_varbin(array, ctx)) + { + return result; + } + + // The two arms here are every builder a `Utf8`/`Binary` dtype has: all four + // `VarBinBuilder` widths above, and `VarBinViewBuilder` below. + let Some(builder) = builder.as_any_mut().downcast_mut::() else { + vortex_bail!("append_to_builder for VarBin requires a variable-binary builder") + }; + append_to_varbinview(array, builder, ctx) + } + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { Ok(ExecutionResult::done( varbin_to_canonical(array.as_view(), ctx)?.into_array(), @@ -210,5 +234,33 @@ impl VTable for VarBin { } } +/// Hands the value bytes to `builder` as a data buffer with views built over them. +/// +/// Canonicalizing first would build the same views, then pay for them twice more: once to wrap +/// them in a `VarBinViewArray` the builder immediately unwraps, and once for +/// `append_varbinview_array` to rewrite every view so its buffer index is rebased onto the +/// builder's. Handing the heap and offsets to the builder instead makes the whole append one view +/// per row with no byte copy — the builder adopts the referenced range of the heap as it is. That +/// range is fully covered by the new views, so this stays valid for a compacting builder too. +fn append_to_varbinview( + array: ArrayView<'_, VarBin>, + builder: &mut VarBinViewBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let len = array.as_ref().len(); + let validity = array.varbin_validity().execute_mask(len, ctx)?; + + let parts = array.into_owned().into_data_parts(); + let offsets = parts.offsets.execute::(ctx)?; + match_each_integer_ptype!(offsets.ptype(), |P| { + builder.append_buffer_with_offsets( + parts.bytes.unwrap_host(), + offsets.as_slice::

(), + &validity, + ) + }); + Ok(()) +} + #[derive(Clone, Debug)] pub struct VarBin; diff --git a/vortex-array/src/arrays/varbinview/array.rs b/vortex-array/src/arrays/varbinview/array.rs index 7fa66eaea7a..cfccfa2a00a 100644 --- a/vortex-array/src/arrays/varbinview/array.rs +++ b/vortex-array/src/arrays/varbinview/array.rs @@ -9,6 +9,7 @@ use std::sync::Arc; use smallvec::smallvec; use vortex_buffer::Alignment; use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -16,14 +17,20 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_error::vortex_panic; +use vortex_mask::AllOr; +use vortex_mask::Mask; +use crate::ArrayRef; use crate::ArraySlots; +use crate::ExecutionCtx; +#[cfg(debug_assertions)] use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayParts; use crate::array::TypedArrayRef; use crate::array::child_to_validity; use crate::array::validity_to_child; +use crate::array_slots; use crate::arrays::VarBinView; use crate::arrays::varbinview::BinaryView; use crate::buffer::BufferHandle; @@ -31,13 +38,16 @@ use crate::builders::ArrayBuilder; use crate::builders::VarBinViewBuilder; use crate::dtype::DType; use crate::dtype::Nullability; +#[cfg(debug_assertions)] use crate::legacy_session; use crate::validity::Validity; -/// The validity bitmap indicating which elements are non-null. -pub(super) const VALIDITY_SLOT: usize = 0; -pub(super) const NUM_SLOTS: usize = 1; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"]; +#[array_slots(VarBinView)] +pub struct VarBinViewSlots { + /// The validity bitmap indicating which elements are non-null. + #[slot(0)] + pub validity: Option, +} /// A variable-length binary view array that stores strings and binary data efficiently. /// @@ -117,6 +127,26 @@ pub struct VarBinViewDataParts { pub validity: Validity, } +// Walk invalid (null) runs of "mask" as [start, end) ranges +fn for_each_invalid_range(mask: &Mask, len: usize, mut f: impl FnMut(usize, usize)) { + match mask.bit_buffer() { + AllOr::All => {} + AllOr::None => f(0, len), + AllOr::Some(buffer) => { + let mut prev = 0; + for (start, end) in buffer.set_slices() { + if start > prev { + f(prev, start); + } + prev = end; + } + if prev < len { + f(prev, len); + } + } + } +} + impl VarBinViewData { fn dtype_parts(dtype: &DType) -> VortexResult<(bool, Nullability)> { match dtype { @@ -142,8 +172,9 @@ impl VarBinViewData { buffers: Arc<[ByteBuffer]>, dtype: DType, validity: Validity, + ctx: &mut ExecutionCtx, ) -> Self { - Self::try_new(views, buffers, dtype, validity) + Self::try_new(views, buffers, dtype, validity, ctx) .vortex_expect("VarBinViewArray construction failed") } @@ -176,8 +207,9 @@ impl VarBinViewData { buffers: Arc<[ByteBuffer]>, dtype: DType, validity: Validity, + ctx: &mut ExecutionCtx, ) -> VortexResult { - Self::validate(&views, &buffers, &dtype, &validity)?; + let views = Self::validate_and_fix(views, &buffers, &dtype, &validity, ctx)?; // SAFETY: validate ensures all invariants are met. Ok(unsafe { Self::new_unchecked(views, buffers, dtype, validity) }) @@ -252,8 +284,15 @@ impl VarBinViewData { validity: Validity, ) -> Self { #[cfg(debug_assertions)] - Self::validate(&views, &buffers, &dtype, &validity) - .vortex_expect("[Debug Assertion]: Invalid `VarBinViewArray` parameters"); + #[expect(clippy::disallowed_methods)] + Self::validate( + &views, + &buffers, + &dtype, + &validity, + &mut legacy_session().create_execution_ctx(), + ) + .vortex_expect("[Debug Assertion]: Invalid `VarBinViewArray` parameters"); let handles: Vec = buffers .iter() @@ -284,98 +323,25 @@ impl VarBinViewData { /// Validates the components that would be used to create a `VarBinViewArray`. /// - /// This function checks all the invariants required by `VarBinViewArray::new_unchecked`. + /// This function checks all the invariants required by `VarBinViewArray::new_unchecked`, + /// validating only the views at non-null slots. pub fn validate( views: &Buffer, buffers: &Arc<[ByteBuffer]>, dtype: &DType, validity: &Validity, + ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - vortex_ensure!( - validity.nullability() == dtype.nullability(), - InvalidArgument: "validity {:?} incompatible with nullability {:?}", - validity, - dtype.nullability() - ); - - match dtype { - DType::Utf8(_) => Self::validate_views(views, buffers, validity, |string| { - simdutf8::basic::from_utf8(string).is_ok() - })?, - DType::Binary(_) => Self::validate_views(views, buffers, validity, |_| true)?, - _ => vortex_bail!(InvalidArgument: "invalid DType {dtype} for `VarBinViewArray`"), - } - - Ok(()) - } - - #[allow(clippy::disallowed_methods)] - fn validate_views( - views: &Buffer, - buffers: &Arc<[ByteBuffer]>, - validity: &Validity, - validator: F, - ) -> VortexResult<()> - where - F: Fn(&[u8]) -> bool, - { - let validate_view = |idx: usize, view: &BinaryView| -> VortexResult<()> { - if view.is_inlined() { - // Validate the inline bytestring - let bytes = &view.as_inlined().data[..view.len() as usize]; - vortex_ensure!( - validator(bytes), - InvalidArgument: "view at index {idx}: inlined bytes failed utf-8 validation" - ); - } else { - // Validate the view pointer - let view = view.as_view(); - let buf_index = view.buffer_index as usize; - let start_offset = view.offset as usize; - let end_offset = start_offset.saturating_add(view.size as usize); - - let buf = buffers.get(buf_index).ok_or_else(|| - vortex_err!(InvalidArgument: "view at index {idx} references invalid buffer: {buf_index} out of bounds for VarBinViewData with {} buffers", - buffers.len()))?; - - vortex_ensure!( - start_offset < buf.len(), - InvalidArgument: "start offset {start_offset} out of bounds for buffer {buf_index} with size {}", - buf.len(), - ); - - vortex_ensure!( - end_offset <= buf.len(), - InvalidArgument: "end offset {end_offset} out of bounds for buffer {buf_index} with size {}", - buf.len(), - ); - - // Make sure the prefix data matches the buffer data. - let bytes = &buf[start_offset..end_offset]; - vortex_ensure!( - view.prefix == bytes[..4], - InvalidArgument: "VarBinView prefix does not match full string" - ); - - // Validate the full string - vortex_ensure!( - validator(bytes), - InvalidArgument: "view at index {idx}: outlined bytes fails utf-8 validation" - ); - } - Ok(()) - }; - + let check_utf8 = Self::check_nullability(dtype, validity)?; match validity { // Array-backed validity is the only variant that needs an execution context: execute it // into a mask once and zip it with the views, validating only the valid (non-null) // entries. Validity::Array(_) => { - let mut ctx = legacy_session().create_execution_ctx(); - let mask = validity.execute_mask(views.len(), &mut ctx)?; + let mask = validity.execute_mask(views.len(), ctx)?; for ((idx, view), valid) in views.iter().enumerate().zip(mask.iter()) { if valid { - validate_view(idx, view)?; + Self::validate_view(idx, view, buffers, check_utf8)?; } } } @@ -384,11 +350,145 @@ impl VarBinViewData { // No nulls: validate every view. Validity::NonNullable | Validity::AllValid => { for (idx, view) in views.iter().enumerate() { - validate_view(idx, view)?; + Self::validate_view(idx, view, buffers, check_utf8)?; + } + } + } + Ok(()) + } + + /// Validates components like validate() and replaces views at null slots to empty views + pub(crate) fn validate_and_fix( + views: Buffer, + buffers: &Arc<[ByteBuffer]>, + dtype: &DType, + validity: &Validity, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let check_utf8 = Self::check_nullability(dtype, validity)?; + let empty = BinaryView::empty_view(); + let len = views.len(); + + match validity { + Validity::Array(_) => { + let mask = validity.execute_mask(len, ctx)?; + match views.try_into_mut() { + Ok(mut views) => { + let slice = views.as_mut_slice(); + for (idx, valid) in mask.iter().enumerate() { + if valid { + Self::validate_view(idx, &slice[idx], buffers, check_utf8)?; + } else { + slice[idx] = empty; + } + } + Ok(views.freeze()) + } + Err(views) => { + let mut needs_replace = false; + for ((idx, view), valid) in views.iter().enumerate().zip(mask.iter()) { + if valid { + Self::validate_view(idx, view, buffers, check_utf8)?; + } else if *view != empty { + needs_replace = true; + } + } + if !needs_replace { + return Ok(views); + } + let mut views = views.into_mut(); + let slice = views.as_mut_slice(); + for_each_invalid_range(&mask, len, |start, end| { + slice[start..end].fill(empty) + }); + Ok(views.freeze()) + } + } + } + // Every entry is null, so there is nothing to validate: replace all views with empty. + Validity::AllInvalid => match views.try_into_mut() { + Ok(mut views) => { + views.as_mut_slice().fill(empty); + Ok(views.freeze()) + } + Err(views) if views.iter().all(|view| *view == empty) => Ok(views), + Err(views) => { + let mut views = views.into_mut(); + views.as_mut_slice().fill(empty); + Ok(views.freeze()) + } + }, + // No nulls: validate every view, nothing to replace. + Validity::NonNullable | Validity::AllValid => { + for (idx, view) in views.iter().enumerate() { + Self::validate_view(idx, view, buffers, check_utf8)?; } + Ok(views) } } + } + + fn check_nullability(dtype: &DType, validity: &Validity) -> VortexResult { + let (is_utf8, nullability) = Self::dtype_parts(dtype)?; + vortex_ensure!( + validity.nullability() == nullability, + InvalidArgument: "validity {:?} incompatible with nullability {:?}", + validity, + nullability + ); + Ok(is_utf8) + } + + fn validate_view( + idx: usize, + view: &BinaryView, + buffers: &Arc<[ByteBuffer]>, + check_utf8: bool, + ) -> VortexResult<()> { + let valid_utf8 = |bytes: &[u8]| !check_utf8 || simdutf8::basic::from_utf8(bytes).is_ok(); + if view.is_inlined() { + // Validate the inline bytestring + let bytes = &view.as_inlined().data[..view.len() as usize]; + vortex_ensure!( + valid_utf8(bytes), + InvalidArgument: "view at index {idx}: inlined bytes failed utf-8 validation" + ); + } else { + // Validate the view pointer + let view = view.as_view(); + let buf_index = view.buffer_index as usize; + let start_offset = view.offset as usize; + let end_offset = start_offset.saturating_add(view.size as usize); + + let buf = buffers.get(buf_index).ok_or_else(|| + vortex_err!(InvalidArgument: "view at index {idx} references invalid buffer: {buf_index} out of bounds for VarBinViewData with {} buffers", + buffers.len()))?; + + vortex_ensure!( + start_offset < buf.len(), + InvalidArgument: "start offset {start_offset} out of bounds for buffer {buf_index} with size {}", + buf.len(), + ); + + vortex_ensure!( + end_offset <= buf.len(), + InvalidArgument: "end offset {end_offset} out of bounds for buffer {buf_index} with size {}", + buf.len(), + ); + + // Make sure the prefix data matches the buffer data. + let bytes = &buf[start_offset..end_offset]; + vortex_ensure!( + view.prefix == bytes[..4], + InvalidArgument: "VarBinView prefix does not match full string" + ); + // Validate the full string + vortex_ensure!( + valid_utf8(bytes), + InvalidArgument: "view at index {idx}: outlined bytes fails utf-8 validation" + ); + } Ok(()) } @@ -476,7 +576,11 @@ impl VarBinViewData { dtype: DType, ) -> Self { let iter = iter.into_iter(); - let mut builder = VarBinViewBuilder::with_capacity(dtype, iter.size_hint().0); + let mut builder = VarBinViewBuilder::with_capacity_in( + dtype, + iter.size_hint().0, + BufferAllocatorRef::statically_allocated(), + ); for item in iter { match item { @@ -490,9 +594,10 @@ impl VarBinViewData { pub fn from_iter_str, I: IntoIterator>(iter: I) -> Self { let iter = iter.into_iter(); - let mut builder = VarBinViewBuilder::with_capacity( + let mut builder = VarBinViewBuilder::with_capacity_in( DType::Utf8(Nullability::NonNullable), iter.size_hint().0, + BufferAllocatorRef::statically_allocated(), ); for item in iter { @@ -506,9 +611,10 @@ impl VarBinViewData { iter: I, ) -> Self { let iter = iter.into_iter(); - let mut builder = VarBinViewBuilder::with_capacity( + let mut builder = VarBinViewBuilder::with_capacity_in( DType::Utf8(Nullability::Nullable), iter.size_hint().0, + BufferAllocatorRef::statically_allocated(), ); for item in iter { @@ -523,9 +629,10 @@ impl VarBinViewData { pub fn from_iter_bin, I: IntoIterator>(iter: I) -> Self { let iter = iter.into_iter(); - let mut builder = VarBinViewBuilder::with_capacity( + let mut builder = VarBinViewBuilder::with_capacity_in( DType::Binary(Nullability::NonNullable), iter.size_hint().0, + BufferAllocatorRef::statically_allocated(), ); for item in iter { @@ -539,9 +646,10 @@ impl VarBinViewData { iter: I, ) -> Self { let iter = iter.into_iter(); - let mut builder = VarBinViewBuilder::with_capacity( + let mut builder = VarBinViewBuilder::with_capacity_in( DType::Binary(Nullability::Nullable), iter.size_hint().0, + BufferAllocatorRef::statically_allocated(), ); for item in iter { @@ -566,7 +674,7 @@ pub trait VarBinViewArrayExt: TypedArrayRef { fn varbinview_validity(&self) -> Validity { child_to_validity( - self.as_ref().slots()[VALIDITY_SLOT].as_ref(), + self.as_ref().slots()[VarBinViewSlots::VALIDITY].as_ref(), self.dtype_parts().1, ) } @@ -594,7 +702,11 @@ impl Array { dtype: DType, ) -> Self { let iter = iter.into_iter(); - let mut builder = VarBinViewBuilder::with_capacity(dtype, iter.size_hint().0); + let mut builder = VarBinViewBuilder::with_capacity_in( + dtype, + iter.size_hint().0, + BufferAllocatorRef::statically_allocated(), + ); for value in iter { match value { Some(value) => builder.append_value(value), @@ -606,9 +718,10 @@ impl Array { pub fn from_iter_str, I: IntoIterator>(iter: I) -> Self { let iter = iter.into_iter(); - let mut builder = VarBinViewBuilder::with_capacity( + let mut builder = VarBinViewBuilder::with_capacity_in( DType::Utf8(Nullability::NonNullable), iter.size_hint().0, + BufferAllocatorRef::statically_allocated(), ); for value in iter { builder.append_value(value.as_ref()); @@ -620,9 +733,10 @@ impl Array { iter: I, ) -> Self { let iter = iter.into_iter(); - let mut builder = VarBinViewBuilder::with_capacity( + let mut builder = VarBinViewBuilder::with_capacity_in( DType::Utf8(Nullability::Nullable), iter.size_hint().0, + BufferAllocatorRef::statically_allocated(), ); for value in iter { match value { @@ -635,9 +749,10 @@ impl Array { pub fn from_iter_bin, I: IntoIterator>(iter: I) -> Self { let iter = iter.into_iter(); - let mut builder = VarBinViewBuilder::with_capacity( + let mut builder = VarBinViewBuilder::with_capacity_in( DType::Binary(Nullability::NonNullable), iter.size_hint().0, + BufferAllocatorRef::statically_allocated(), ); for value in iter { builder.append_value(value.as_ref()); @@ -649,9 +764,10 @@ impl Array { iter: I, ) -> Self { let iter = iter.into_iter(); - let mut builder = VarBinViewBuilder::with_capacity( + let mut builder = VarBinViewBuilder::with_capacity_in( DType::Binary(Nullability::Nullable), iter.size_hint().0, + BufferAllocatorRef::statically_allocated(), ); for value in iter { match value { @@ -668,8 +784,9 @@ impl Array { buffers: Arc<[ByteBuffer]>, dtype: DType, validity: Validity, + ctx: &mut ExecutionCtx, ) -> VortexResult { - let data = VarBinViewData::try_new(views, buffers, dtype.clone(), validity.clone())?; + let data = VarBinViewData::try_new(views, buffers, dtype.clone(), validity.clone(), ctx)?; let slots = VarBinViewData::make_slots(&validity, data.len()); Ok(Self::from_prevalidated_data(dtype, data, slots)) } diff --git a/vortex-array/src/arrays/varbinview/build_views.rs b/vortex-array/src/arrays/varbinview/build_views.rs index 4c35aab721a..8a6a5a0530d 100644 --- a/vortex-array/src/arrays/varbinview/build_views.rs +++ b/vortex-array/src/arrays/varbinview/build_views.rs @@ -6,7 +6,6 @@ use num_traits::AsPrimitive; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; -use vortex_buffer::ByteBufferMut; pub use crate::arrays::varbinview::BinaryView; use crate::dtype::NativePType; @@ -26,25 +25,107 @@ pub const MAX_BUFFER_LEN: usize = i32::MAX as usize; /// Split a large buffer of input `bytes` holding string data into `VarBinView` buffers and views. /// +/// The values must be laid end-to-end in `bytes`, one per entry of `lens`, describing the whole +/// buffer exactly. The returned buffers are zero-copy slices of `bytes`, numbered sequentially +/// from `start_buf_index`. +/// /// `max_buffer_len` must not exceed [`MAX_BUFFER_LEN`], since every view offset is stored in a /// `u32` and offsets are bounded by `max_buffer_len`. +/// +/// # Panics +/// +/// Panics if the lengths do not describe `bytes` exactly, or if a single value exceeds +/// `max_buffer_len`. pub fn build_views>( start_buf_index: u32, max_buffer_len: usize, - bytes: ByteBufferMut, + bytes: ByteBuffer, lens: &[P], ) -> (Vec, Buffer) { + let mut views = BufferMut::with_capacity(lens.len()); + let buffers = extend_views( + &mut views, + start_buf_index, + max_buffer_len, + &bytes, + lens.len(), + |i| lens[i].as_(), + ); + (buffers, views.freeze()) +} + +/// [`build_views`] for values described by an offsets buffer instead of lengths. +/// +/// `offsets` are absolute positions into `bytes` — the layout a `VarBinArray` stores — so there +/// is one more offset than there are values, and the values need not start at the beginning of +/// `bytes`: only `offsets[0]..offsets[last]` is referenced, and the returned buffers are zero-copy +/// slices of that range. +/// +/// # Panics +/// +/// Panics if `offsets` is empty, not monotonically non-decreasing within `bytes`, or if a single +/// value exceeds `max_buffer_len`. +pub fn build_views_from_offsets>( + start_buf_index: u32, + max_buffer_len: usize, + bytes: ByteBuffer, + offsets: &[P], +) -> (Vec, Buffer) { + assert!(!offsets.is_empty(), "offsets must hold at least one entry"); + let first: usize = offsets[0].as_(); + let last: usize = offsets[offsets.len() - 1].as_(); + let bytes = bytes.slice(first..last); + + let count = offsets.len() - 1; + let mut views = BufferMut::with_capacity(count); + // Wrapping keeps corrupt non-monotonic offsets from panicking on the subtraction itself; the + // wrapped length then fails the in-bounds slicing (or `max_buffer_len`) checks in the loop. + let buffers = extend_views( + &mut views, + start_buf_index, + max_buffer_len, + &bytes, + count, + |i| { + AsPrimitive::::as_(offsets[i + 1]) + .wrapping_sub(AsPrimitive::::as_(offsets[i])) + }, + ); + (buffers, views.freeze()) +} + +/// Appends one view per value straight into `views`, splitting `bytes` into buffers. +/// +/// This is the core behind [`build_views`]: it writes into an existing views buffer so that a +/// [`VarBinViewBuilder`](crate::builders::VarBinViewBuilder) can build views directly into its +/// storage without an intermediate allocation. `len_at(i)` is the byte length of value `i`, and +/// the `count` lengths must describe `bytes` exactly. The returned buffers are zero-copy slices +/// of `bytes`, numbered sequentially from `start_buf_index`. +pub(crate) fn extend_views( + views: &mut BufferMut, + start_buf_index: u32, + max_buffer_len: usize, + bytes: &ByteBuffer, + count: usize, + len_at: impl Fn(usize) -> usize, +) -> Vec { assert!( max_buffer_len <= MAX_BUFFER_LEN, "max_buffer_len cannot exceed MAX_BUFFER_LEN, offsets must fit in u32" ); if bytes.len() <= max_buffer_len { - // Common case: the whole decoded heap fits within a single buffer, so no rollover can occur - // (`bytes.len()` is the total decoded size and therefore an upper bound on every offset). - build_views_single_buffer(start_buf_index, bytes, lens) + // Common case: the whole decoded heap fits within a single buffer, so no rollover can + // occur (`bytes.len()` is the total decoded size and therefore an upper bound on every + // offset). + extend_views_single_buffer(views, start_buf_index, bytes, count, len_at); + if bytes.is_empty() { + Vec::new() + } else { + vec![bytes.clone()] + } } else { - build_views_rolling(start_buf_index, max_buffer_len, bytes, lens) + extend_views_rolling(views, start_buf_index, max_buffer_len, bytes, count, len_at) } } @@ -54,12 +135,15 @@ pub fn build_views>( /// reference views inline, avoiding the out-of-line `BinaryView::make_view` call for the common /// long-string case. Every offset is bounded by `bytes.len()`, which the caller has guaranteed is /// at most [`MAX_BUFFER_LEN`], so the `usize -> u32` conversions cannot truncate. -fn build_views_single_buffer>( - start_buf_index: u32, - bytes: ByteBufferMut, - lens: &[P], -) -> (Vec, Buffer) { - let mut views = BufferMut::::with_capacity(lens.len()); +fn extend_views_single_buffer( + views: &mut BufferMut, + buf_index: u32, + bytes: &ByteBuffer, + count: usize, + len_at: impl Fn(usize) -> usize, +) { + views.reserve(count); + let base = views.len(); let data = bytes.as_slice(); let mut offset = 0usize; @@ -68,72 +152,77 @@ fn build_views_single_buffer>( // loop-invariant, so it reloads and rewrites the output cursor through the stack each // iteration. Writing into the spare slice keeps the cursor in a register and the length is // set once after the loop. - let spare = views.spare_capacity_mut(); - for (slot, &len) in spare.iter_mut().zip(lens) { - let len = len.as_(); + let spare = &mut views.spare_capacity_mut()[..count]; + for (i, slot) in spare.iter_mut().enumerate() { + let len = len_at(i); let value = &data[offset..offset + len]; let view = if len > BinaryView::MAX_INLINED_SIZE { let mut prefix = [0u8; 4]; prefix.copy_from_slice(&value[..4]); - BinaryView::new_ref(len.as_(), prefix, start_buf_index, offset.as_()) + BinaryView::new_ref(len.as_(), prefix, buf_index, offset.as_()) } else { - BinaryView::make_view(value, start_buf_index, offset.as_()) + BinaryView::make_view(value, buf_index, offset.as_()) }; slot.write(view); offset += len; } - // SAFETY: the loop initialized exactly `lens.len()` contiguous views (`spare` has at least - // `lens.len()` slots, and `zip` stops at the shorter operand). - unsafe { views.set_len(lens.len()) }; - - let buffers = if bytes.is_empty() { - Vec::new() - } else { - vec![bytes.freeze()] - }; - (buffers, views.freeze()) + assert_eq!( + offset, + data.len(), + "value lengths must describe the byte heap exactly" + ); + // SAFETY: the loop initialized exactly `count` contiguous views (`spare` has at least + // `count` slots). + unsafe { views.set_len(base + count) }; } /// Build views when the heap exceeds `max_buffer_len` and must be split across multiple buffers. /// /// The buffer is rolled over every `max_buffer_len` bytes so that no view offset overflows the -/// `u32` offset field. -fn build_views_rolling>( +/// `u32` offset field. Each output buffer is a zero-copy slice of `bytes`. +fn extend_views_rolling( + views: &mut BufferMut, start_buf_index: u32, max_buffer_len: usize, - mut bytes: ByteBufferMut, - lens: &[P], -) -> (Vec, Buffer) { - let mut views = BufferMut::::with_capacity(lens.len()); + bytes: &ByteBuffer, + count: usize, + len_at: impl Fn(usize) -> usize, +) -> Vec { + views.reserve(count); let mut buffers = Vec::new(); let mut buf_index = start_buf_index; - let mut offset = 0; - for &len in lens { - let len = len.as_(); + let data = bytes.as_slice(); + // The absolute start of the current segment, and the offset of the next value within it. + let mut segment_start = 0usize; + let mut offset = 0usize; + for i in 0..count { + let len = len_at(i); assert!(len <= max_buffer_len, "values cannot exceed max_buffer_len"); - if (offset + len) > max_buffer_len { + if offset + len > max_buffer_len { // Roll the buffer every 2GiB, to avoid overflowing VarBinView offset field - let rest = bytes.split_off(offset); - - buffers.push(bytes.freeze()); + buffers.push(bytes.slice(segment_start..segment_start + offset)); buf_index += 1; + segment_start += offset; offset = 0; - - bytes = rest; } - let view = BinaryView::make_view(&bytes[offset..][..len], buf_index, offset.as_()); - // SAFETY: we reserved the right capacity beforehand - unsafe { views.push_unchecked(view) }; + let start = segment_start + offset; + let view = BinaryView::make_view(&data[start..start + len], buf_index, offset.as_()); + views.push(view); offset += len; } + assert_eq!( + segment_start + offset, + data.len(), + "value lengths must describe the byte heap exactly" + ); - if !bytes.is_empty() { - buffers.push(bytes.freeze()); + if segment_start < data.len() { + buffers.push(bytes.slice(segment_start..data.len())); } - (buffers, views.freeze()) + buffers } #[cfg(test)] @@ -145,17 +234,18 @@ mod tests { use crate::arrays::varbinview::BinaryView; use crate::arrays::varbinview::build_views::MAX_BUFFER_LEN; use crate::arrays::varbinview::build_views::build_views; + use crate::arrays::varbinview::build_views::build_views_from_offsets; /// Concatenate `values` into a single byte heap and return it alongside the per-element lengths, /// matching the `(bytes, lens)` inputs that `build_views` consumes. - fn flatten(values: &[&[u8]]) -> (ByteBufferMut, Vec) { + fn flatten(values: &[&[u8]]) -> (ByteBuffer, Vec) { let mut bytes = ByteBufferMut::empty(); let mut lens = Vec::with_capacity(values.len()); for v in values { bytes.extend_from_slice(v); lens.push(u32::try_from(v.len()).unwrap()); } - (bytes, lens) + (bytes.freeze(), lens) } /// Reconstruct the logical value behind each view by dereferencing it through the output @@ -206,7 +296,7 @@ mod tests { assert!(buffers.is_empty(), "empty heap must not allocate a buffer"); } else { assert_eq!(buffers.len(), 1, "whole heap must stay in one buffer"); - // The fast path freezes the input heap unchanged. + // The fast path adopts the input heap unchanged. let concatenated: Vec = values.concat(); assert_eq!(buffers[0].as_slice(), concatenated.as_slice()); } @@ -220,6 +310,37 @@ mod tests { assert_eq!(reconstruct(&buffers, &views, start_buf_index), expected); } + /// The output buffers must be zero-copy slices of the input heap, on both paths — a copy here + /// silently doubles the memory cost of every decode that feeds views. + #[test] + fn output_buffers_are_zero_copy() { + let values: &[&[u8]] = &[ + b"first long reference value", + b"tiny", + b"second long reference value!!", + b"third looooong reference value", + ]; + let (bytes, lens) = flatten(values); + let base = bytes.as_ptr(); + + // Fast path: the single output buffer is the input buffer. + let (buffers, _views) = build_views(0, bytes.len() + 1, bytes.clone(), &lens); + assert_eq!(buffers.len(), 1); + assert_eq!(buffers[0].as_ptr(), base, "fast path must not copy"); + + // Rolling path: every output buffer points into the input allocation. + let longest = values.iter().map(|v| v.len()).max().unwrap(); + let (buffers, _views) = build_views(0, longest, bytes, &lens); + assert!(buffers.len() > 1); + let mut expected_ptr = base; + for buffer in &buffers { + assert_eq!(buffer.as_ptr(), expected_ptr, "rolling path must not copy"); + // SAFETY: the buffers partition the input heap, so the next one starts where + // this one ends, still within (or one past) the original allocation. + expected_ptr = unsafe { expected_ptr.add(buffer.len()) }; + } + } + /// Offsets and sizes are written into the `u32` `Ref` fields via `as_` truncation, so we must /// confirm they stay correct once the running offset grows well past the 16-bit range (i.e. is /// not narrowed to a smaller width). A ~9 MiB heap pushes offsets above 2^23 while remaining far @@ -230,7 +351,7 @@ mod tests { const N: usize = 9000; const LEN: usize = 1000; // The final offset is (N - 1) * LEN, which must exceed 2^23 to be a meaningful check. - const _: () = assert!((N - 1) * LEN > (1 << 23)); + const { assert!((N - 1) * LEN > (1 << 23)) }; let values: Vec> = (0..N) .map(|i| { @@ -314,7 +435,7 @@ mod tests { #[test] fn fast_path_empty_input() { let lens: Vec = Vec::new(); - let (buffers, views) = build_views(0, 1024, ByteBufferMut::empty(), &lens); + let (buffers, views) = build_views(0, 1024, ByteBuffer::empty(), &lens); assert!(buffers.is_empty()); assert!(views.is_empty()); } @@ -336,6 +457,41 @@ mod tests { assert_eq!(views.as_slice(), &expected); } + /// The offsets-driven variant must agree with the lengths-driven one, reference only the + /// `offsets[0]..offsets[last]` range, and stay zero-copy — it exists so a `VarBinArray` heap + /// can feed views without materializing a lengths buffer or copying its bytes. + #[test] + fn from_offsets_matches_lengths_and_is_zero_copy() { + // A heap with a prefix and suffix outside the offsets range, as a sliced VarBin has. + let heap = ByteBuffer::copy_from(b"..a long value that is referenced!tiny..".as_slice()); + let offsets: Vec = vec![2, 34, 38]; + + let (buffers, views) = build_views_from_offsets(5, MAX_BUFFER_LEN, heap.clone(), &offsets); + + assert_eq!(buffers.len(), 1); + // Zero-copy: the buffer points at offset 2 of the original allocation. + // SAFETY: offset 2 is in bounds of the 40-byte heap. + assert_eq!(buffers[0].as_ptr(), unsafe { heap.as_ptr().add(2) }); + assert_eq!(buffers[0].len(), 36); + + assert_eq!( + reconstruct(&buffers, &views, 5), + vec![ + b"a long value that is referenced!".to_vec(), + b"tiny".to_vec() + ] + ); + } + + /// Lengths that do not cover the heap exactly are a caller bug and must be rejected rather + /// than silently emitting views over a partially-covered buffer. + #[test] + #[should_panic(expected = "value lengths must describe the byte heap exactly")] + fn short_lengths_panic() { + let (bytes, _) = flatten(&[b"a long value that is referenced", b"tiny"]); + build_views(0, MAX_BUFFER_LEN, bytes, &[31u32]); + } + // TODO(someone): ideally CI would run this in release mode as well, since debug builds make the // ~2.25 GiB allocation and fill loop substantially slower. /// Slow regression for the single-buffer fast-path guard. The fast path is only valid when the @@ -380,7 +536,7 @@ mod tests { } let lens = vec![u32::try_from(STRING_LEN).unwrap(); N]; - let (buffers, views) = build_views(0, MAX_BUFFER_LEN, bytes, &lens); + let (buffers, views) = build_views(0, MAX_BUFFER_LEN, bytes.freeze(), &lens); assert_eq!(views.len(), N); assert!( @@ -419,7 +575,7 @@ mod tests { // In real code, this would all fit in one buffer, but to unit test the splitting logic // we split buffers at length 26, which should result in two buffers for the output array. let raw_data = - ByteBufferMut::copy_from("aaaaaaaaaaaaabbbbbbbbbbbbbcccccccccccccddddddddddddd"); + ByteBuffer::copy_from("aaaaaaaaaaaaabbbbbbbbbbbbbcccccccccccccddddddddddddd"); let lens = vec![13u8; 4]; let (buffers, views) = build_views(0, 26, raw_data, &lens); @@ -446,11 +602,6 @@ mod tests { #[test] #[should_panic(expected = "max_buffer_len cannot exceed MAX_BUFFER_LEN")] fn test_max_buffer_len_too_large_panics() { - build_views( - 0, - MAX_BUFFER_LEN + 1, - ByteBufferMut::copy_from("abc"), - &[3u32], - ); + build_views(0, MAX_BUFFER_LEN + 1, ByteBuffer::copy_from("abc"), &[3u32]); } } diff --git a/vortex-array/src/arrays/varbinview/compact.rs b/vortex-array/src/arrays/varbinview/compact.rs index 2aba21ae448..38549b9cf85 100644 --- a/vortex-array/src/arrays/varbinview/compact.rs +++ b/vortex-array/src/arrays/varbinview/compact.rs @@ -65,6 +65,7 @@ impl VarBinViewArray { /// Iterates over all valid, non-inlined views, calling the provided /// closure for each one. + #[allow(clippy::inline_always)] #[inline(always)] fn iter_valid_views(&self, ctx: &mut ExecutionCtx, mut f: F) -> VortexResult<()> where @@ -141,10 +142,11 @@ impl VarBinViewArray { buffer_utilization_threshold: f64, // [0, 1] ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut builder = VarBinViewBuilder::with_compaction( + let mut builder = VarBinViewBuilder::with_compaction_in( self.dtype().clone(), self.len(), buffer_utilization_threshold, + ctx.allocator().clone(), ); builder.append_varbinview_array(self, ctx)?; Ok(builder.finish_into_varbinview()) @@ -159,7 +161,7 @@ pub(crate) struct BufferUtilization { } impl BufferUtilization { - fn zero(len: u32) -> Self { + pub(crate) fn zero(len: u32) -> Self { BufferUtilization { len, used: 0u32, @@ -168,7 +170,7 @@ impl BufferUtilization { } } - fn add(&mut self, offset: u32, size: u32) { + pub(crate) fn add(&mut self, offset: u32, size: u32) { self.used += size; self.min_offset = self.min_offset.min(offset); self.max_offset_end = self.max_offset_end.max(offset + size); diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index ce10abda3d0..c10b16cd419 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -2,25 +2,36 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::iter; +use std::ptr; use std::sync::Arc; +use itertools::Itertools as _; use num_traits::AsPrimitive; use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_mask::AllOr; use vortex_mask::Mask; use crate::ArrayRef; +use crate::Columnar; use crate::IntoArray; use crate::array::ArrayView; +use crate::arrays::PiecewiseSequence; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::piecewise_sequence::constant_unsigned_usize; +use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::varbinview::BinaryView; use crate::buffer::BufferHandle; +use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_integer_ptype; +use crate::match_each_unsigned_integer_ptype; impl TakeExecute for VarBinView { /// Take involves creating a new array that references the old array, just with the given set of views. @@ -29,6 +40,12 @@ impl TakeExecute for VarBinView { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)? + { + return Ok(Some(taken)); + } + let validity = array.validity()?.take(indices)?; let indices = indices.clone().execute::(ctx)?; @@ -57,6 +74,60 @@ impl TakeExecute for VarBinView { } } +fn take_contiguous_ranges( + array: ArrayView<'_, VarBinView>, + indices: ArrayView<'_, PiecewiseSequence>, + indices_ref: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { + return Ok(None); + }; + let source = array.views(); + let output_len = indices_ref.len(); + let views = match lengths { + Columnar::Constant(lengths) => { + let length = constant_unsigned_usize(&lengths); + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + gather_view_slices_constant_length( + source, + starts.as_slice::(), + length, + output_len, + )? + }) + } + Columnar::Canonical(lengths) => { + let lengths = lengths.into_primitive(); + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + gather_view_slices( + source, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + )? + }) + }) + } + }; + let validity = array.validity()?.take(indices_ref)?; + + // SAFETY: ranges were validated against the source views, and copied views still reference the + // same backing data buffers. + unsafe { + Ok(Some( + VarBinViewArray::new_handle_unchecked( + BufferHandle::new_host(views.into_byte_buffer()), + Arc::clone(array.data_buffers()), + array.dtype().clone(), + validity, + ) + .into_array(), + )) + } +} + fn take_views>( views_ref: &[BinaryView], indices: &[I], @@ -85,6 +156,91 @@ fn take_views>( } } +fn gather_view_slices_constant_length( + source: &[BinaryView], + starts: &[S], + length: usize, + output_len: usize, +) -> VortexResult> +where + S: UnsignedPType, +{ + let computed_len = starts + .len() + .checked_mul(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); + + let mut views = BufferMut::::with_capacity(output_len); + let spare = &mut views.spare_capacity_mut()[..output_len]; + let mut cursor = 0usize; + for &start in starts { + let start = start.as_(); + let src = &source[start..][..length]; + // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. + unsafe { + ptr::copy_nonoverlapping( + src.as_ptr(), + spare[cursor..][..src.len()] + .as_mut_ptr() + .cast::(), + src.len(), + ); + } + cursor += src.len(); + } + // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity. + unsafe { views.set_len(cursor) }; + vortex_ensure!( + views.len() == output_len, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + views.len() + ); + Ok(views.freeze()) +} + +fn gather_view_slices( + source: &[BinaryView], + starts: &[S], + lengths: &[L], + output_len: usize, +) -> VortexResult> +where + S: UnsignedPType, + L: UnsignedPType, +{ + let mut views = BufferMut::::with_capacity(output_len); + let spare = &mut views.spare_capacity_mut()[..output_len]; + let mut cursor = 0usize; + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = start.as_(); + let length = length.as_(); + let src = &source[start..][..length]; + // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. + unsafe { + ptr::copy_nonoverlapping( + src.as_ptr(), + spare[cursor..][..src.len()] + .as_mut_ptr() + .cast::(), + src.len(), + ); + } + cursor += src.len(); + } + // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity. + unsafe { views.set_len(cursor) }; + vortex_ensure!( + views.len() == output_len, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + views.len() + ); + Ok(views.freeze()) +} + #[cfg(test)] mod tests { use rstest::rstest; diff --git a/vortex-array/src/arrays/varbinview/compute/zip.rs b/vortex-array/src/arrays/varbinview/compute/zip.rs index 8966a327f9d..2711a7e4fc6 100644 --- a/vortex-array/src/arrays/varbinview/compute/zip.rs +++ b/vortex-array/src/arrays/varbinview/compute/zip.rs @@ -52,7 +52,7 @@ impl ZipKernel for VarBinView { buffers.extend_from_iter(if_false.data_buffers().iter().map(|b| b.as_host().clone())); let mut views_builder = BufferMut::::with_capacity(len); - let mut validity_builder = LazyBitBufferBuilder::new(len); + let mut validity_builder = LazyBitBufferBuilder::new(len, ctx.allocator().clone()); let true_validity = if_true.varbinview_validity().execute_mask(len, ctx)?; let false_validity = if_false.varbinview_validity().execute_mask(len, ctx)?; diff --git a/vortex-array/src/arrays/varbinview/mod.rs b/vortex-array/src/arrays/varbinview/mod.rs index 7d08123a971..1c6160b13eb 100644 --- a/vortex-array/src/arrays/varbinview/mod.rs +++ b/vortex-array/src/arrays/varbinview/mod.rs @@ -3,8 +3,10 @@ mod array; pub use array::VarBinViewArrayExt; +pub use array::VarBinViewArraySlotsExt; pub use array::VarBinViewData; pub use array::VarBinViewDataParts; +pub use array::VarBinViewSlots; pub use vtable::VarBinViewArray; pub(crate) mod compact; diff --git a/vortex-array/src/arrays/varbinview/tests.rs b/vortex-array/src/arrays/varbinview/tests.rs index 981d4f602a1..80238403536 100644 --- a/vortex-array/src/arrays/varbinview/tests.rs +++ b/vortex-array/src/arrays/varbinview/tests.rs @@ -1,11 +1,30 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; + +use vortex_buffer::BitBuffer; +use vortex_buffer::Buffer; +use vortex_buffer::ByteBuffer; +use vortex_buffer::ByteBufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::registry::ReadContext; + +use crate::ArrayContext; +use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; +use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; use crate::arrays::varbinview::BinaryView; +use crate::arrays::varbinview::VarBinViewData; use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::serde::SerializeOptions; +use crate::serde::SerializedArray; +use crate::validity::Validity; #[test] pub fn varbin_view() { @@ -49,3 +68,76 @@ pub fn binary_view_size_and_alignment() { assert_eq!(size_of::(), 16); assert_eq!(align_of::(), 16); } + +#[test] +pub fn validate_replaces_null_views() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let views = Buffer::::copy_from(vec![ + BinaryView::new_inlined(b"ololo"), + BinaryView::new_ref(13, *b"AAAA", 0xDEAD_BEEF, 0xF000_0000), + ]); + let buffers: Arc<[ByteBuffer]> = Arc::new([]); + let dtype = DType::Utf8(Nullability::Nullable); + let buffer = BitBuffer::from_iter([true, false]); + let validity = Validity::from_bit_buffer(buffer, Nullability::Nullable); + + let replaced = + VarBinViewData::validate_and_fix(views.clone(), &buffers, &dtype, &validity, &mut ctx)?; + assert_eq!(replaced[0], views[0]); + assert_eq!(replaced[1], BinaryView::empty_view()); + + let replaced = + VarBinViewData::validate_and_fix(views, &buffers, &dtype, &Validity::AllInvalid, &mut ctx)?; + assert!( + replaced + .iter() + .all(|view| *view == BinaryView::empty_view()) + ); + Ok(()) +} + +#[test] +pub fn deserialize_null_views() -> VortexResult<()> { + let views = Buffer::::copy_from(vec![ + BinaryView::new_ref(14, *b"hell", 0, 0), + BinaryView::new_ref(13, *b"AAAA", 0xDEAD_BEEF, 0xF000_0000), + ]); + let buffers = Arc::new([ByteBuffer::from(b"hello world ololo".to_vec())]); + let dtype = DType::Utf8(Nullability::Nullable); + let buffer = BitBuffer::from_iter([true, false]); + let validity = Validity::from_bit_buffer(buffer, Nullability::Nullable); + let session = array_session(); + let array = VarBinViewArray::try_new( + views.clone(), + buffers, + dtype.clone(), + validity, + &mut session.create_execution_ctx(), + )?; + + let array_ctx = ArrayContext::empty(); + let serialized = + array + .clone() + .into_array() + .serialize(&array_ctx, &session, &SerializeOptions::default())?; + + let mut concat = ByteBufferMut::empty(); + for buf in serialized { + concat.extend_from_slice(buf.as_ref()); + } + let parts = SerializedArray::try_from(concat.freeze())?; + let decoded = parts.decode( + &dtype, + array.len(), + &ReadContext::new(array_ctx.to_ids()), + &session, + )?; + + let decoded = decoded + .as_opt::() + .ok_or_else(|| vortex_err!("expected VarBinView"))?; + assert_eq!(decoded.views()[0], views[0]); + assert_eq!(decoded.views()[1], BinaryView::empty_view()); + Ok(()) +} diff --git a/vortex-array/src/arrays/varbinview/view.rs b/vortex-array/src/arrays/varbinview/view.rs index 1b07baa4576..94392529715 100644 --- a/vortex-array/src/arrays/varbinview/view.rs +++ b/vortex-array/src/arrays/varbinview/view.rs @@ -99,6 +99,8 @@ impl Ref { impl BinaryView { /// Maximum size of an inlined binary value. + /// + /// cbindgen:ignore pub const MAX_INLINED_SIZE: usize = 12; /// Create a view from a value, block and offset @@ -244,6 +246,26 @@ impl BinaryView { unsafe { &mut self._ref } } + /// Returns the bytes of this value, reading out of `buffers` when it is not inlined. + /// + /// Unlike [`VarBinViewData::bytes_at`](crate::arrays::varbinview::VarBinViewData::bytes_at) + /// this borrows + /// from slices the caller has already resolved instead of cloning a buffer handle, so it is + /// safe to call once per row in a loop. + /// + /// # Panics + /// + /// Panics if the view references a buffer or range not covered by `buffers`. + #[inline] + pub fn bytes<'a>(&'a self, buffers: &[&'a [u8]]) -> &'a [u8] { + if self.is_inlined() { + self.as_inlined().value() + } else { + let view = self.as_view(); + &buffers[view.buffer_index as usize][view.as_range()] + } + } + /// Returns the binary view as u128 representation. pub fn as_u128(&self) -> u128 { // SAFETY: binary view always safe to read as u128 LE bytes diff --git a/vortex-array/src/arrays/varbinview/vtable/mod.rs b/vortex-array/src/arrays/varbinview/vtable/mod.rs index f730282ddb4..0d2963cd6b2 100644 --- a/vortex-array/src/arrays/varbinview/vtable/mod.rs +++ b/vortex-array/src/arrays/varbinview/vtable/mod.rs @@ -19,14 +19,14 @@ use crate::ArrayRef; use crate::EqMode; use crate::ExecutionCtx; use crate::ExecutionResult; +use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::VTable; use crate::arrays::varbinview::BinaryView; use crate::arrays::varbinview::VarBinViewData; -use crate::arrays::varbinview::array::NUM_SLOTS; -use crate::arrays::varbinview::array::SLOT_NAMES; +use crate::arrays::varbinview::array::VarBinViewSlots; use crate::arrays::varbinview::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; @@ -34,6 +34,7 @@ use crate::builders::VarBinViewBuilder; use crate::dtype::DType; use crate::hash::ArrayEq; use crate::hash::ArrayHash; +use crate::match_each_varbin_builder; use crate::serde::ArrayChildren; use crate::validity::Validity; mod kernel; @@ -93,8 +94,9 @@ impl VTable for VarBinView { slots: &[Option], ) -> VortexResult<()> { vortex_ensure!( - slots.len() == NUM_SLOTS, - "VarBinViewArray expected {NUM_SLOTS} slots, found {}", + slots.len() == VarBinViewSlots::COUNT, + "VarBinViewArray expected {} slots, found {}", + VarBinViewSlots::COUNT, slots.len() ); vortex_ensure!( @@ -167,7 +169,7 @@ impl VTable for VarBinView { buffers: &[BufferHandle], children: &dyn ArrayChildren, - _session: &VortexSession, + session: &VortexSession, ) -> VortexResult> { if !metadata.is_empty() { vortex_bail!( @@ -223,13 +225,14 @@ impl VTable for VarBinView { Arc::from(data_buffers), dtype.clone(), validity.clone(), + &mut session.create_execution_ctx(), )?; let slots = VarBinViewData::make_slots(&validity, len); Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + VarBinViewSlots::NAMES[idx].to_string() } fn reduce_parent( @@ -245,10 +248,15 @@ impl VTable for VarBinView { builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - let Some(builder) = builder.as_any_mut().downcast_mut::() else { - vortex_bail!("append_to_builder for VarBinView requires a VarBinViewBuilder"); - }; - builder.append_varbinview_array(&array.into_owned(), ctx) + if let Some(builder) = builder.as_any_mut().downcast_mut::() { + return builder.append_varbinview_array(&array.into_owned(), ctx); + } + if let Some(result) = + match_each_varbin_builder!(builder, |builder| builder.append_varbinview(array, ctx)) + { + return result; + } + vortex_bail!("append_to_builder for VarBinView requires a variable-binary builder") } fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { diff --git a/vortex-array/src/arrays/variant/compute/filter.rs b/vortex-array/src/arrays/variant/compute/filter.rs index ade0df74d62..cd6178f8d12 100644 --- a/vortex-array/src/arrays/variant/compute/filter.rs +++ b/vortex-array/src/arrays/variant/compute/filter.rs @@ -10,7 +10,7 @@ use crate::array::ArrayView; use crate::arrays::Variant; use crate::arrays::VariantArray; use crate::arrays::filter::FilterReduce; -use crate::arrays::variant::VariantArrayExt; +use crate::arrays::variant::VariantArraySlotsExt; impl FilterReduce for Variant { fn filter(array: ArrayView<'_, Variant>, mask: &Mask) -> VortexResult> { diff --git a/vortex-array/src/arrays/variant/compute/slice.rs b/vortex-array/src/arrays/variant/compute/slice.rs index 6dc8c5132d4..675e98e62d7 100644 --- a/vortex-array/src/arrays/variant/compute/slice.rs +++ b/vortex-array/src/arrays/variant/compute/slice.rs @@ -11,7 +11,7 @@ use crate::array::ArrayView; use crate::arrays::Variant; use crate::arrays::VariantArray; use crate::arrays::slice::SliceReduce; -use crate::arrays::variant::VariantArrayExt; +use crate::arrays::variant::VariantArraySlotsExt; impl SliceReduce for Variant { fn slice(array: ArrayView<'_, Variant>, range: Range) -> VortexResult> { diff --git a/vortex-array/src/arrays/variant/compute/take.rs b/vortex-array/src/arrays/variant/compute/take.rs index 8676eeebba4..d009538d5b4 100644 --- a/vortex-array/src/arrays/variant/compute/take.rs +++ b/vortex-array/src/arrays/variant/compute/take.rs @@ -9,7 +9,7 @@ use crate::array::ArrayView; use crate::arrays::Variant; use crate::arrays::VariantArray; use crate::arrays::dict::TakeReduce; -use crate::arrays::variant::VariantArrayExt; +use crate::arrays::variant::VariantArraySlotsExt; impl TakeReduce for Variant { fn take(array: ArrayView<'_, Variant>, indices: &ArrayRef) -> VortexResult> { diff --git a/vortex-array/src/arrays/variant/mod.rs b/vortex-array/src/arrays/variant/mod.rs index 1a5157e9aa1..b5610320719 100644 --- a/vortex-array/src/arrays/variant/mod.rs +++ b/vortex-array/src/arrays/variant/mod.rs @@ -5,7 +5,6 @@ mod vtable; pub(crate) mod compute; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -20,15 +19,10 @@ use crate::ArrayRef; use crate::array::Array; use crate::array::ArrayParts; use crate::array::EmptyArrayData; -use crate::array::TypedArrayRef; +use crate::array_slots; use crate::dtype::DType; -pub(super) const CORE_STORAGE_SLOT: usize = 0; -pub(super) const SHREDDED_SLOT: usize = 1; -pub(super) const NUM_SLOTS: usize = 2; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["core_storage", "shredded"]; - -/// Accessors for canonical variant storage. +/// Slots for canonical variant storage. /// /// A canonical variant array keeps the full variant value for every row in `core_storage` and may /// carry a row-aligned, storage-agnostic `shredded` typed tree for selected paths. @@ -37,28 +31,23 @@ pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["core_storage", "shredded"]; /// chunked, constant, or otherwise encoded. Callers must use normal array operations instead of /// assuming a particular slot layout. The shredded child may have any dtype; its dtype is recorded /// during serialization and validated by normal child deserialization. -pub trait VariantArrayExt: TypedArrayRef { - /// Returns the logical variant storage that preserves the full value for every row. - fn core_storage(&self) -> &ArrayRef { - self.as_ref().slots()[CORE_STORAGE_SLOT] - .as_ref() - .vortex_expect("validated variant core_storage slot") - } - - /// Returns the optional row-aligned typed shredded tree for selected variant paths. - /// This functions returns `Some` only if the array was canonicalized and the shredded data +#[array_slots(Variant)] +pub struct VariantSlots { + /// The logical variant storage that preserves the full value for every row. + #[slot(0)] + pub core_storage: ArrayRef, + /// The optional row-aligned typed shredded tree for selected variant paths. + /// This slot is `Some` only if the array was canonicalized and the shredded data /// was pulled out of the underlying variant storage. - fn shredded(&self) -> Option<&ArrayRef> { - self.as_ref().slots()[SHREDDED_SLOT].as_ref() - } + #[slot(1)] + pub shredded: Option, } -impl> VariantArrayExt for T {} impl Array { /// Creates a new `VariantArray` with logical variant core storage and optional shredded storage. /// /// `core_storage` must have `DType::Variant`, but it may use any Variant-typed physical - /// encoding. See [`VariantArrayExt`] for the higher-level storage contract. + /// encoding. See [`VariantSlots`] for the higher-level storage contract. /// /// `shredded`, when present, must be row-aligned with `core_storage` and stores typed values for /// selected variant paths. @@ -71,8 +60,13 @@ impl Array { let len = core_storage.len(); let stats = core_storage.statistics().to_owned(); Ok(Array::try_from_parts( - ArrayParts::new(Variant, dtype, len, EmptyArrayData) - .with_slots(vec![Some(core_storage), shredded].into()), + ArrayParts::new(Variant, dtype, len, EmptyArrayData).with_slots( + VariantSlots { + core_storage, + shredded, + } + .into_slots(), + ), )? .with_stats_set(stats)) } @@ -96,7 +90,7 @@ mod tests { use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::VariantArray; - use crate::arrays::variant::VariantArrayExt; + use crate::arrays::variant::VariantArraySlotsExt; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; @@ -116,18 +110,17 @@ mod tests { } fn row_storage(values: impl IntoIterator) -> VortexResult { - let chunks = values - .into_iter() - .map(|value| { + Ok(ChunkedArray::try_new( + values.into_iter().map(|value| { ConstantArray::new( Scalar::variant(Scalar::primitive(value, Nullability::NonNullable)), 1, ) .into_array() - }) - .collect(); - - Ok(ChunkedArray::try_new(chunks, DType::Variant(Nullability::NonNullable))?.into_array()) + }), + DType::Variant(Nullability::NonNullable), + )? + .into_array()) } fn variant_with_shredded( diff --git a/vortex-array/src/arrays/variant/vtable/kernel.rs b/vortex-array/src/arrays/variant/vtable/kernel.rs index f289748eded..22d52ecd6a7 100644 --- a/vortex-array/src/arrays/variant/vtable/kernel.rs +++ b/vortex-array/src/arrays/variant/vtable/kernel.rs @@ -18,9 +18,8 @@ use crate::arrays::Variant; use crate::arrays::VariantArray; use crate::arrays::scalar_fn::ExactScalarFn; use crate::arrays::scalar_fn::ScalarFnArrayView; -use crate::arrays::scalar_fn::ScalarFnFactoryExt; use crate::arrays::struct_::StructArrayExt; -use crate::arrays::variant::VariantArrayExt; +use crate::arrays::variant::VariantArraySlotsExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; @@ -60,12 +59,7 @@ impl ExecuteParentKernel for VariantGetKernel { // perfectly represented by the canonical shredded tree. let core_validity = array.core_storage().validity()?; let make_fallback = |ctx: &mut ExecutionCtx| { - execute_fallback_variant_get( - array.len(), - parent.options.clone(), - array.core_storage().clone(), - ctx, - ) + execute_fallback_variant_get(parent.options.clone(), array.core_storage().clone(), ctx) }; // Canonical shredded storage is a logical typed tree. We can only walk object @@ -95,7 +89,6 @@ impl ExecuteParentKernel for VariantGetKernel { .is_some_and(|dtype| !dtype.is_variant()) { return execute_fallback_variant_get( - array.len(), VariantGetOptions::new(VariantPath::root(), parent.options.dtype().cloned()), typed, ctx, @@ -202,12 +195,11 @@ fn merge_typed_as_variant( } fn execute_fallback_variant_get( - len: usize, options: VariantGetOptions, core_storage: ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult { - VariantGet - .try_new_array(len, options, [core_storage])? + VariantGet::try_new(core_storage, options)? + .into_array() .execute::(ctx) } diff --git a/vortex-array/src/arrays/variant/vtable/mod.rs b/vortex-array/src/arrays/variant/vtable/mod.rs index 4d84511ae6a..12064cc626a 100644 --- a/vortex-array/src/arrays/variant/vtable/mod.rs +++ b/vortex-array/src/arrays/variant/vtable/mod.rs @@ -25,10 +25,7 @@ use crate::array::ArrayView; use crate::array::EmptyArrayData; use crate::array::VTable; use crate::array::with_empty_buffers; -use crate::arrays::variant::CORE_STORAGE_SLOT; -use crate::arrays::variant::NUM_SLOTS; -use crate::arrays::variant::SHREDDED_SLOT; -use crate::arrays::variant::SLOT_NAMES; +use crate::arrays::variant::VariantSlots; use crate::arrays::variant::compute::rules::RULES; use crate::buffer::BufferHandle; use crate::dtype::DType; @@ -76,15 +73,16 @@ impl VTable for Variant { slots: &[Option], ) -> VortexResult<()> { vortex_ensure!( - slots.len() == NUM_SLOTS, - "VariantArray expects {NUM_SLOTS} slots, got {}", + slots.len() == VariantSlots::COUNT, + "VariantArray expects {} slots, got {}", + VariantSlots::COUNT, slots.len() ); vortex_ensure!( - slots[CORE_STORAGE_SLOT].is_some(), + slots[VariantSlots::CORE_STORAGE].is_some(), "VariantArray core_storage slot must be present" ); - let core_storage = slots[CORE_STORAGE_SLOT] + let core_storage = slots[VariantSlots::CORE_STORAGE] .as_ref() .vortex_expect("validated core_storage slot presence"); vortex_ensure!( @@ -103,7 +101,7 @@ impl VTable for Variant { core_storage.len(), len ); - if let Some(shredded) = slots[SHREDDED_SLOT].as_ref() { + if let Some(shredded) = slots[VariantSlots::SHREDDED].as_ref() { vortex_ensure!( shredded.len() == len, "VariantArray shredded length {} does not match outer length {}", @@ -138,7 +136,7 @@ impl VTable for Variant { array: ArrayView<'_, Self>, _session: &VortexSession, ) -> VortexResult>> { - let shredded_dtype = array.slots()[SHREDDED_SLOT] + let shredded_dtype = array.slots()[VariantSlots::SHREDDED] .as_ref() .map(|shredded| shredded.dtype().try_into()) .transpose()?; @@ -180,13 +178,18 @@ impl VTable for Variant { .map(|dtype| children.get(1, &dtype, len)) .transpose()?; Ok( - ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData) - .with_slots(vec![Some(core_storage), shredded].into()), + ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData).with_slots( + VariantSlots { + core_storage, + shredded, + } + .into_slots(), + ), ) } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - match SLOT_NAMES.get(idx) { + match VariantSlots::NAMES.get(idx) { Some(name) => (*name).to_string(), None => vortex_panic!("VariantArray slot_name index {idx} out of bounds"), } diff --git a/vortex-array/src/arrays/variant/vtable/operations.rs b/vortex-array/src/arrays/variant/vtable/operations.rs index f01f0ca6c8b..c7f7d36fd97 100644 --- a/vortex-array/src/arrays/variant/vtable/operations.rs +++ b/vortex-array/src/arrays/variant/vtable/operations.rs @@ -8,7 +8,7 @@ use crate::ExecutionCtx; use crate::array::ArrayView; use crate::array::OperationsVTable; use crate::arrays::Variant; -use crate::arrays::variant::VariantArrayExt; +use crate::arrays::variant::VariantArraySlotsExt; use crate::scalar::Scalar; impl OperationsVTable for Variant { diff --git a/vortex-array/src/arrays/variant/vtable/validity.rs b/vortex-array/src/arrays/variant/vtable/validity.rs index 47287b0442a..598b2461135 100644 --- a/vortex-array/src/arrays/variant/vtable/validity.rs +++ b/vortex-array/src/arrays/variant/vtable/validity.rs @@ -6,7 +6,7 @@ use vortex_error::VortexResult; use crate::array::ArrayView; use crate::array::ValidityVTable; use crate::arrays::Variant; -use crate::arrays::variant::VariantArrayExt; +use crate::arrays::variant::VariantArraySlotsExt; use crate::validity::Validity; impl ValidityVTable for Variant { diff --git a/vortex-array/src/builders/bool.rs b/vortex-array/src/builders/bool.rs index 73d2089e221..c9f222fde7f 100644 --- a/vortex-array/src/builders/bool.rs +++ b/vortex-array/src/builders/bool.rs @@ -5,9 +5,9 @@ use std::any::Any; use std::mem; use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; @@ -29,14 +29,32 @@ pub struct BoolBuilder { } impl BoolBuilder { + /// Creates a builder with the default capacity. + #[deprecated(note = "use `new_in` with an explicit allocator")] pub fn new(nullability: Nullability) -> Self { - Self::with_capacity(nullability, DEFAULT_BUILDER_CAPACITY) + Self::new_in(nullability, BufferAllocatorRef::static_ref()) } + /// Creates a builder with the default capacity using `allocator`. + pub fn new_in(nullability: Nullability, allocator: &BufferAllocatorRef) -> Self { + Self::with_capacity_in(nullability, DEFAULT_BUILDER_CAPACITY, allocator) + } + + /// Creates a builder with the given capacity. + #[deprecated(note = "use `with_capacity_in` with an explicit allocator")] pub fn with_capacity(nullability: Nullability, capacity: usize) -> Self { + Self::with_capacity_in(nullability, capacity, BufferAllocatorRef::static_ref()) + } + + /// Creates a builder with the given capacity and allocator. + pub fn with_capacity_in( + nullability: Nullability, + capacity: usize, + allocator: &BufferAllocatorRef, + ) -> Self { Self { - inner: BitBufferMut::with_capacity(capacity), - nulls: LazyBitBufferBuilder::new(capacity), + inner: BitBufferMut::with_capacity_in(capacity, allocator.clone()), + nulls: LazyBitBufferBuilder::new(capacity, allocator.clone()), dtype: DType::Bool(nullability), } } @@ -62,8 +80,10 @@ impl BoolBuilder { "Null count and value count should match when calling BoolBuilder::finish." ); + let allocator = self.inner.allocator().clone(); + let inner = mem::replace(&mut self.inner, BitBufferMut::empty_in(allocator)).freeze(); BoolArray::new( - mem::take(&mut self.inner).freeze(), + inner, self.nulls.finish_with_nullability(self.dtype.nullability()), ) } @@ -127,10 +147,6 @@ impl ArrayBuilder for BoolBuilder { self.nulls.reserve_exact(additional); } - unsafe fn set_validity_unchecked(&mut self, validity: Mask) { - self.nulls = LazyBitBufferBuilder::from_validity_mask(validity); - } - fn finish(&mut self) -> ArrayRef { self.finish_into_bool().into_array() } @@ -145,6 +161,7 @@ mod tests { use rand::RngExt; use rand::SeedableRng; use rand::prelude::StdRng; + use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use crate::ArrayRef; @@ -157,7 +174,7 @@ mod tests { use crate::builders::ArrayBuilder; use crate::builders::BoolBuilder; use crate::builders::bool::BoolArray; - use crate::builders::builder_with_capacity; + use crate::builders::builder_with_capacity_in; use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar::Scalar; @@ -186,7 +203,11 @@ mod tests { let chunk = make_opt_bool_chunks(len, chunk_count); let mut ctx = array_session().create_execution_ctx(); - let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count); + let mut builder = builder_with_capacity_in( + chunk.dtype(), + len * chunk_count, + BufferAllocatorRef::static_ref(), + ); chunk .clone() .append_to_builder(builder.as_mut(), &mut ctx)?; @@ -206,7 +227,11 @@ mod tests { #[test] fn test_append_scalar() { let mut ctx = array_session().create_execution_ctx(); - let mut builder = BoolBuilder::with_capacity(Nullability::Nullable, 10); + let mut builder = BoolBuilder::with_capacity_in( + Nullability::Nullable, + 10, + BufferAllocatorRef::static_ref(), + ); // Test appending true value. let true_scalar = Scalar::bool(true, Nullability::Nullable); @@ -225,7 +250,11 @@ mod tests { assert_arrays_eq!(&array, &expected, &mut ctx); // Test wrong dtype error. - let mut builder = BoolBuilder::with_capacity(Nullability::NonNullable, 10); + let mut builder = BoolBuilder::with_capacity_in( + Nullability::NonNullable, + 10, + BufferAllocatorRef::static_ref(), + ); let wrong_scalar = Scalar::from(42i32); assert!(builder.append_scalar(&wrong_scalar).is_err()); } diff --git a/vortex-array/src/builders/child.rs b/vortex-array/src/builders/child.rs new file mode 100644 index 00000000000..a9f5b46fc84 --- /dev/null +++ b/vortex-array/src/builders/child.rs @@ -0,0 +1,371 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::BufferAllocatorRef; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::ChunkedArray; +use crate::builders::ArrayBuilder; +use crate::builders::builder_with_capacity_in; +use crate::dtype::DType; +use crate::scalar::Scalar; + +/// Accumulates the child of a nested [`ArrayBuilder`] without canonicalizing appended arrays. +/// +/// Nested builders receive values from two sources: individual [`Scalar`]s, which have to be +/// materialized into a canonical builder, and whole arrays, whose encoding the builder has no +/// reason to decode. A `ChildBuilder` keeps the latter as chunks and materializes only the former, +/// stitching everything back together into a [`ChunkedArray`] on [`finish`](Self::finish) when +/// more than one chunk accumulated. +/// +/// This keeps canonical arrays canonical only at the top level, which is all that [`Canonical`] +/// promises: the fields of a `StructArray`, the elements of a list, and the storage of an +/// extension array may all stay compressed. +/// +/// [`Canonical`]: crate::Canonical +pub struct ChildBuilder { + /// The [`DType`] shared by every chunk and by the scalar builder. + dtype: DType, + + /// Completed chunks, in logical order. Never contains an empty chunk. + chunks: Vec, + + /// The summed length of `chunks`. + chunks_len: usize, + + /// Builder holding the scalars appended after the last chunk. + pending: Box, +} + +impl ChildBuilder { + /// Creates a child builder with the provided allocator and capacity. + pub fn with_capacity(dtype: &DType, capacity: usize, allocator: &BufferAllocatorRef) -> Self { + Self { + dtype: dtype.clone(), + chunks: Vec::new(), + chunks_len: 0, + pending: builder_with_capacity_in(dtype, capacity, allocator), + } + } + + /// The number of values appended so far. + pub fn len(&self) -> usize { + self.chunks_len + self.pending.len() + } + + /// Appends every value of `array` to the child as a chunk of its own, keeping its encoding. + /// + /// However short the array, it becomes a chunk: the caller had a whole array to hand, and + /// deciding on its behalf that its values are cheaper copied than referenced would be guessing + /// at a boundary only the caller can see. Callers that would rather have the values copied + /// should append them as scalars. + /// + /// An appended [`ChunkedArray`] stays one chunk rather than giving up its own. Unpacking it + /// would spill a chunk per row into this child for a caller expressing a repeated run as + /// chunks - a tiled fixed-size list, say - and the child's chunk list is what every later + /// append and the final [`finish`](Self::finish) walk. Nesting costs one level, not one per + /// appended array. + /// + /// Nothing is decoded here, so `_ctx` goes unused; it stays in the signature so that the + /// nested builders forwarding their [`ExecutionCtx`] here do not have to explain why they + /// don't. + pub fn append_array(&mut self, array: &ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult<()> { + vortex_ensure!( + array.dtype() == &self.dtype, + "Cannot append an array of dtype {} to a child builder of dtype {}", + array.dtype(), + self.dtype, + ); + + if array.is_empty() { + return Ok(()); + } + + self.flush_pending(); + self.chunks_len += array.len(); + + self.chunks.push(array.clone()); + + Ok(()) + } + + /// Appends a single [`Scalar`] to the child. + pub fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> { + self.pending.append_scalar(scalar) + } + + /// Appends `n` "zero" values to the child. + /// + /// See [`ArrayBuilder::append_zeros`]. + pub fn append_zeros(&mut self, n: usize) { + self.pending.append_zeros(n) + } + + /// Appends `n` null values to the child. + /// + /// See [`ArrayBuilder::append_nulls`]. + pub fn append_nulls(&mut self, n: usize) { + self.pending.append_nulls(n) + } + + /// Appends `n` default values to the child. + /// + /// See [`ArrayBuilder::append_defaults`]. + pub fn append_defaults(&mut self, n: usize) { + self.pending.append_defaults(n) + } + + /// Allocates space for `additional` more values in the scalar builder. + pub fn reserve_exact(&mut self, additional: usize) { + self.pending.reserve_exact(additional) + } + + /// Finishes the child, combining the accumulated chunks into a [`ChunkedArray`] when there is + /// more than one of them. + pub fn finish(&mut self) -> ArrayRef { + if self.chunks.is_empty() { + return self.pending.finish(); + } + + self.flush_pending(); + self.chunks_len = 0; + + let mut chunks = std::mem::take(&mut self.chunks); + if chunks.len() == 1 { + return chunks.remove(0); + } + + unsafe { ChunkedArray::new_unchecked(chunks, self.dtype.clone()) }.into_array() + } + + /// Moves whatever the scalar builder holds into `chunks`, keeping the chunks in logical order. + fn flush_pending(&mut self) { + if self.pending.is_empty() { + return; + } + self.chunks_len += self.pending.len(); + let pending = self.pending.finish(); + self.chunks.push(pending); + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_buffer::BufferAllocatorRef; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use super::ChildBuilder; + use crate::ArrayRef; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::Chunked; + use crate::arrays::ChunkedArray; + use crate::arrays::Constant; + use crate::arrays::ConstantArray; + use crate::arrays::Primitive; + use crate::arrays::PrimitiveArray; + use crate::arrays::chunked::ChunkedArrayExt; + use crate::assert_arrays_eq; + use crate::dtype::DType; + use crate::dtype::Nullability::NonNullable; + use crate::dtype::Nullability::Nullable; + use crate::dtype::PType::I32; + use crate::scalar::Scalar; + + /// An arbitrary array length. `ChildBuilder` treats no length specially, so the tests only + /// need a length long enough to tell chunks apart. + const CHUNK_LEN: usize = 64; + + /// A non-canonical array of `len` values, all equal to `value`. + fn constant(value: i32, len: usize) -> ArrayRef { + ConstantArray::new(value, len).into_array() + } + + /// A non-canonical *nullable* array of `len` non-null values, all equal to `value`. + fn nullable_constant(value: i32, len: usize) -> ArrayRef { + ConstantArray::new(Scalar::primitive(value, Nullable), len).into_array() + } + + #[test] + fn test_appended_arrays_are_kept_as_chunks() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = + ChildBuilder::with_capacity(&DType::from(I32), 0, BufferAllocatorRef::static_ref()); + + builder.append_array(&constant(1, CHUNK_LEN), &mut ctx)?; + builder.append_array(&constant(2, CHUNK_LEN), &mut ctx)?; + assert_eq!(builder.len(), 2 * CHUNK_LEN); + + let child = builder.finish(); + let chunked = child.as_::(); + assert_eq!(chunked.nchunks(), 2); + // The chunks were never decoded. + assert!(chunked.iter_chunks().all(|c| c.is::())); + + Ok(()) + } + + /// However short the appended array, its encoding survives: a single-value array is a chunk + /// too. A caller that wants the values copied appends them as scalars instead. + #[test] + fn test_short_arrays_are_kept_as_chunks_too() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = + ChildBuilder::with_capacity(&DType::from(I32), 0, BufferAllocatorRef::static_ref()); + + builder.append_array(&constant(1, 1), &mut ctx)?; + builder.append_array(&constant(2, 1), &mut ctx)?; + assert_eq!(builder.len(), 2); + + let child = builder.finish(); + let chunked = child.as_::(); + assert_eq!(chunked.nchunks(), 2); + assert!(chunked.iter_chunks().all(|c| c.is::())); + + Ok(()) + } + + #[test] + fn test_scalars_interleaved_with_chunks_keep_their_order() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = + ChildBuilder::with_capacity(&DType::from(I32), 0, BufferAllocatorRef::static_ref()); + + builder.append_scalar(&1i32.into())?; + builder.append_array(&constant(2, CHUNK_LEN), &mut ctx)?; + builder.append_scalar(&3i32.into())?; + + let child = builder.finish(); + assert_eq!(child.len(), CHUNK_LEN + 2); + + let expected = ChunkedArray::try_new( + vec![ + buffer![1i32].into_array(), + constant(2, CHUNK_LEN), + buffer![3i32].into_array(), + ], + DType::from(I32), + )? + .into_array(); + assert_arrays_eq!(&child, &expected, &mut ctx); + + Ok(()) + } + + #[test] + fn test_single_chunk_is_not_wrapped() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = + ChildBuilder::with_capacity(&DType::from(I32), 0, BufferAllocatorRef::static_ref()); + + builder.append_array(&constant(7, CHUNK_LEN), &mut ctx)?; + + let child = builder.finish(); + assert!(child.is::()); + + Ok(()) + } + + #[test] + fn test_empty_arrays_never_become_chunks() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = + ChildBuilder::with_capacity(&DType::from(I32), 0, BufferAllocatorRef::static_ref()); + let empty = constant(1, CHUNK_LEN).slice(0..0)?; + + builder.append_array(&empty, &mut ctx)?; + builder.append_array(&constant(1, CHUNK_LEN), &mut ctx)?; + builder.append_array(&empty, &mut ctx)?; + builder.append_array(&constant(2, CHUNK_LEN), &mut ctx)?; + builder.append_array(&empty, &mut ctx)?; + + assert_eq!(builder.len(), 2 * CHUNK_LEN); + assert_eq!(builder.finish().as_::().nchunks(), 2); + + Ok(()) + } + + /// An empty child must still finish as an empty array rather than as an empty [`ChunkedArray`]. + #[test] + fn test_empty_child_finishes_without_chunks() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = + ChildBuilder::with_capacity(&DType::from(I32), 0, BufferAllocatorRef::static_ref()); + + builder.append_array(&constant(1, CHUNK_LEN).slice(0..0)?, &mut ctx)?; + + let child = builder.finish(); + assert!(child.is_empty()); + assert!(child.is::()); + + Ok(()) + } + + /// The dtype check has to run before the empty check, so that a mismatched array is rejected + /// whether or not it would have become a chunk. + #[rstest] + #[case::empty(0)] + #[case::non_empty(CHUNK_LEN)] + fn test_appending_a_mismatched_dtype_is_rejected(#[case] len: usize) { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = + ChildBuilder::with_capacity(&DType::from(I32), 0, BufferAllocatorRef::static_ref()); + + let wrong_dtype = ConstantArray::new(1i64, len).into_array(); + assert!(builder.append_array(&wrong_dtype, &mut ctx).is_err()); + } + + /// Everything the scalar builder can produce has to be flushed ahead of the next chunk. + #[test] + fn test_zeros_and_nulls_around_chunks_keep_their_order() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DType::Primitive(I32, Nullable); + let mut builder = ChildBuilder::with_capacity(&dtype, 0, BufferAllocatorRef::static_ref()); + + builder.append_array(&nullable_constant(1, CHUNK_LEN), &mut ctx)?; + builder.append_nulls(2); + builder.append_array(&nullable_constant(2, CHUNK_LEN), &mut ctx)?; + builder.append_zeros(1); + + let child = builder.finish(); + assert_eq!(child.len(), 2 * CHUNK_LEN + 3); + assert_eq!(child.as_::().nchunks(), 4); + + let expected = PrimitiveArray::from_option_iter( + std::iter::repeat_n(Some(1i32), CHUNK_LEN) + .chain([None, None]) + .chain(std::iter::repeat_n(Some(2i32), CHUNK_LEN)) + .chain([Some(0)]), + ) + .into_array(); + assert_arrays_eq!(&child, &expected, &mut ctx); + + Ok(()) + } + + #[test] + fn test_finish_resets_the_builder() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = + ChildBuilder::with_capacity(&DType::from(I32), 0, BufferAllocatorRef::static_ref()); + + builder.append_array(&constant(1, CHUNK_LEN), &mut ctx)?; + builder.append_scalar(&2i32.into())?; + assert_eq!(builder.finish().len(), CHUNK_LEN + 1); + + assert_eq!(builder.len(), 0); + builder.append_scalar(&3i32.into())?; + + let expected = PrimitiveArray::new(buffer![3i32], NonNullable.into()).into_array(); + assert_arrays_eq!(&builder.finish(), &expected, &mut ctx); + + Ok(()) + } +} diff --git a/vortex-array/src/builders/decimal.rs b/vortex-array/src/builders/decimal.rs index c08dfbc4d87..d65dc41ad4a 100644 --- a/vortex-array/src/builders/decimal.rs +++ b/vortex-array/src/builders/decimal.rs @@ -3,13 +3,13 @@ use std::any::Any; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_error::vortex_panic; -use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; @@ -95,22 +95,51 @@ macro_rules! delegate_fn { impl DecimalBuilder { /// Creates a new `DecimalBuilder` with a capacity of [`DEFAULT_BUILDER_CAPACITY`]. + #[deprecated(note = "use `new_in` with an explicit allocator")] pub fn new(decimal: DecimalDType, nullability: Nullability) -> Self { - Self::with_capacity::(DEFAULT_BUILDER_CAPACITY, decimal, nullability) + Self::new_in::(decimal, nullability, BufferAllocatorRef::static_ref()) + } + + /// Creates a new `DecimalBuilder` with the default capacity using `allocator`. + pub fn new_in( + decimal: DecimalDType, + nullability: Nullability, + allocator: &BufferAllocatorRef, + ) -> Self { + Self::with_capacity_in::(DEFAULT_BUILDER_CAPACITY, decimal, nullability, allocator) } /// Creates a new `DecimalBuilder` with the given `capacity`. + #[deprecated(note = "use `with_capacity_in` with an explicit allocator")] pub fn with_capacity( capacity: usize, decimal: DecimalDType, nullability: Nullability, + ) -> Self { + Self::with_capacity_in::( + capacity, + decimal, + nullability, + BufferAllocatorRef::static_ref(), + ) + } + + /// Creates a new `DecimalBuilder` with `capacity` using `allocator`. + pub fn with_capacity_in( + capacity: usize, + decimal: DecimalDType, + nullability: Nullability, + allocator: &BufferAllocatorRef, ) -> Self { Self { dtype: DType::Decimal(decimal, nullability), values: match_each_decimal_value_type!(T::DECIMAL_TYPE, |D| { - DecimalBuffer::from(BufferMut::::with_capacity(capacity)) + DecimalBuffer::from(BufferMut::::with_capacity_in( + capacity, + allocator.clone(), + )) }), - nulls: LazyBitBufferBuilder::new(capacity), + nulls: LazyBitBufferBuilder::new(capacity, allocator.clone()), } } @@ -154,7 +183,7 @@ impl DecimalBuilder { let decimal_dtype = *self.decimal_dtype(); - delegate_fn!(std::mem::take(&mut self.values), |T, values| { + delegate_fn!(self.values.take(), |T, values| { DecimalArray::new::(values.freeze(), decimal_dtype, validity) }) } @@ -219,10 +248,6 @@ impl ArrayBuilder for DecimalBuilder { self.nulls.reserve_exact(additional); } - unsafe fn set_validity_unchecked(&mut self, validity: Mask) { - self.nulls = LazyBitBufferBuilder::from_validity_mask(validity); - } - fn finish(&mut self) -> ArrayRef { self.finish_into_decimal().into_array() } @@ -233,6 +258,13 @@ impl ArrayBuilder for DecimalBuilder { } impl DecimalBuffer { + fn take(&mut self) -> Self { + delegate_fn!(self, |T, buffer| { + let allocator = buffer.allocator(); + DecimalBuffer::from(std::mem::replace(buffer, allocator.with_capacity(0))) + }) + } + fn push(&mut self, value: V) { delegate_fn!(self, |T, buffer| { buffer.push( @@ -296,14 +328,10 @@ impl_from_buffer!(i64, I64); impl_from_buffer!(i128, I128); impl_from_buffer!(i256, I256); -impl Default for DecimalBuffer { - fn default() -> Self { - Self::I8(BufferMut::::empty()) - } -} - #[cfg(test)] mod tests { + use vortex_buffer::BufferAllocatorRef; + use crate::VortexSessionExecute; use crate::array_session; use crate::assert_arrays_eq; @@ -316,13 +344,21 @@ mod tests { fn test_mixed_extend() { let values = 42i8; - let mut i8s = DecimalBuilder::new::(DecimalDType::new(2, 1), false.into()); + let mut i8s = DecimalBuilder::new_in::( + DecimalDType::new(2, 1), + false.into(), + BufferAllocatorRef::static_ref(), + ); for v in 0..values { i8s.append_value(v); } let i8s = i8s.finish(); - let mut i128s = DecimalBuilder::new::(DecimalDType::new(2, 1), false.into()); + let mut i128s = DecimalBuilder::new_in::( + DecimalDType::new(2, 1), + false.into(), + BufferAllocatorRef::static_ref(), + ); i8s.append_to_builder(&mut i128s, &mut array_session().create_execution_ctx()) .unwrap(); let i128s = i128s.finish(); @@ -344,7 +380,11 @@ mod tests { use crate::scalar::Scalar; // Simply test that the builder accepts its own finish output via scalar. - let mut builder = DecimalBuilder::new::(DecimalDType::new(10, 2), true.into()); + let mut builder = DecimalBuilder::new_in::( + DecimalDType::new(10, 2), + true.into(), + BufferAllocatorRef::static_ref(), + ); builder.append_value(1234i64); builder.append_value(5678i64); builder.append_null(); @@ -357,7 +397,11 @@ mod tests { assert_arrays_eq!(&array, &expected, &mut ctx); // Test by taking a scalar from the array and appending it to a new builder. - let mut builder2 = DecimalBuilder::new::(DecimalDType::new(10, 2), true.into()); + let mut builder2 = DecimalBuilder::new_in::( + DecimalDType::new(10, 2), + true.into(), + BufferAllocatorRef::static_ref(), + ); for i in 0..array.len() { let scalar = array .execute_scalar(i, &mut array_session().create_execution_ctx()) @@ -369,7 +413,11 @@ mod tests { assert_arrays_eq!(&array2, &array, &mut ctx); // Test wrong dtype error. - let mut builder = DecimalBuilder::new::(DecimalDType::new(10, 2), false.into()); + let mut builder = DecimalBuilder::new_in::( + DecimalDType::new(10, 2), + false.into(), + BufferAllocatorRef::static_ref(), + ); let wrong_scalar = Scalar::from(true); assert!(builder.append_scalar(&wrong_scalar).is_err()); } diff --git a/vortex-array/src/builders/dict/bytes.rs b/vortex-array/src/builders/dict/bytes.rs index 33b2e5225f7..4e822967df1 100644 --- a/vortex-array/src/builders/dict/bytes.rs +++ b/vortex-array/src/builders/dict/bytes.rs @@ -8,7 +8,9 @@ use std::sync::Arc; use num_traits::AsPrimitive; use vortex_array::ExecutionCtx; +use vortex_buffer::Alignment; use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; @@ -31,7 +33,7 @@ use crate::arrays::PrimitiveArray; use crate::arrays::VarBin; use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; -use crate::arrays::varbin::VarBinArrayExt; +use crate::arrays::varbin::VarBinArraySlotsExt; use crate::arrays::varbinview::build_views::BinaryView; use crate::dtype::DType; use crate::dtype::PType; @@ -50,29 +52,58 @@ pub struct BytesDictBuilder { dtype: DType, max_dict_bytes: usize, max_dict_len: usize, + allocator: BufferAllocatorRef, } -pub fn bytes_dict_builder(dtype: DType, constraints: &DictConstraints) -> Box { +pub fn bytes_dict_builder( + dtype: DType, + constraints: &DictConstraints, + allocator: BufferAllocatorRef, +) -> Box { match constraints.max_len as u64 { - max if max <= u8::MAX as u64 => Box::new(BytesDictBuilder::::new(dtype, constraints)), - max if max <= u16::MAX as u64 => Box::new(BytesDictBuilder::::new(dtype, constraints)), - max if max <= u32::MAX as u64 => Box::new(BytesDictBuilder::::new(dtype, constraints)), - _ => Box::new(BytesDictBuilder::::new(dtype, constraints)), + max if max <= u8::MAX as u64 => Box::new(BytesDictBuilder::::new_in( + dtype, + constraints, + allocator, + )), + max if max <= u16::MAX as u64 => Box::new(BytesDictBuilder::::new_in( + dtype, + constraints, + allocator, + )), + max if max <= u32::MAX as u64 => Box::new(BytesDictBuilder::::new_in( + dtype, + constraints, + allocator, + )), + _ => Box::new(BytesDictBuilder::::new_in( + dtype, + constraints, + allocator, + )), } } impl BytesDictBuilder { - pub fn new(dtype: DType, constraints: &DictConstraints) -> Self { + pub fn new_in( + dtype: DType, + constraints: &DictConstraints, + allocator: BufferAllocatorRef, + ) -> Self { Self { lookup: Some(HashTable::new()), - views: BufferMut::::empty(), + views: BufferMut::::empty_aligned_in( + Alignment::of::(), + allocator.clone(), + ), null_code: OnceCell::new(), - values: BufferMut::empty(), - values_nulls: BitBufferMut::empty(), + values: BufferMut::empty_aligned_in(Alignment::of::(), allocator.clone()), + values_nulls: BitBufferMut::empty_in(allocator.clone()), hasher: DefaultHashBuilder::default(), dtype, - max_dict_bytes: constraints.max_bytes, + max_dict_bytes: constraints.max_bytes.min(u32::MAX as usize), max_dict_len: constraints.max_len, + allocator, } } @@ -172,7 +203,7 @@ impl BytesDictBuilder { F: FnMut(usize) -> &'a [u8], { let mut local_lookup = self.lookup.take().vortex_expect("Must have a lookup dict"); - let mut codes: BufferMut = BufferMut::with_capacity(len); + let mut codes = BufferMut::::with_capacity_in(len, self.allocator.clone()); match validity_mask.bit_buffer() { AllOr::All => { @@ -286,9 +317,25 @@ impl DictEncoder for BytesDictBuilder { } fn reset(&mut self) -> ArrayRef { - let views = mem::take(&mut self.views).freeze(); - let buffer = mem::take(&mut self.values).freeze(); - let value_nulls = mem::take(&mut self.values_nulls).freeze(); + if let Some(lookup) = self.lookup.as_mut() { + lookup.clear(); + } + self.null_code = OnceCell::new(); + let views = mem::replace( + &mut self.views, + BufferMut::empty_aligned_in(Alignment::of::(), self.allocator.clone()), + ) + .freeze(); + let buffer = mem::replace( + &mut self.values, + BufferMut::empty_aligned_in(Alignment::of::(), self.allocator.clone()), + ) + .freeze(); + let value_nulls = mem::replace( + &mut self.values_nulls, + BitBufferMut::empty_in(self.allocator.clone()), + ) + .freeze(); // SAFETY: we build the views explicitly and the bytes should be checked before feeding // to the encoder. @@ -314,10 +361,12 @@ mod test { use std::sync::LazyLock; use vortex_buffer::Buffer; + use vortex_buffer::BufferAllocatorRef; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; use vortex_session::VortexSession; + use super::BytesDictBuilder; use crate::IntoArray; use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; @@ -325,8 +374,11 @@ mod test { use crate::arrays::VarBinViewArray; use crate::arrays::dict::DictArraySlotsExt; use crate::arrays::varbinview::BinaryView; + use crate::assert_arrays_eq; use crate::buffer::BufferHandle; + use crate::builders::dict::UNCONSTRAINED; use crate::builders::dict::dict_encode; + use crate::builders::dict::dict_encoder_in; use crate::dtype::DType; use crate::dtype::Nullability; use crate::validity::Validity; @@ -430,4 +482,38 @@ mod test { assert_eq!(codes.as_slice::(), &[0, 0, 1, 1, 0, 1, 0, 1]); Ok(()) } + + #[test] + fn reset_clears_dict() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let first = VarBinViewArray::from_iter_str(["one", "two"]).into_array(); + let mut encoder = dict_encoder_in(&first, &UNCONSTRAINED, ctx.allocator().clone()); + + assert_arrays_eq!( + encoder.encode(&first, &mut ctx)?, + PrimitiveArray::from_iter([0u64, 1]), + &mut ctx + ); + assert_arrays_eq!(encoder.reset(), first, &mut ctx); + + let second = VarBinViewArray::from_iter_str(["one", "three"]).into_array(); + assert_arrays_eq!( + encoder.encode(&second, &mut ctx)?, + PrimitiveArray::from_iter([0u64, 1]), + &mut ctx + ); + assert_arrays_eq!(encoder.reset(), second, &mut ctx); + + Ok(()) + } + + #[test] + fn max_dict_bytes_cannot_exceed_the_view_offset_range() { + let builder = BytesDictBuilder::::new_in( + DType::Utf8(Nullability::NonNullable), + &UNCONSTRAINED, + BufferAllocatorRef::statically_allocated(), + ); + assert_eq!(builder.max_dict_bytes, u32::MAX as usize); + } } diff --git a/vortex-array/src/builders/dict/mod.rs b/vortex-array/src/builders/dict/mod.rs index 94834f83cc4..988bf3f7295 100644 --- a/vortex-array/src/builders/dict/mod.rs +++ b/vortex-array/src/builders/dict/mod.rs @@ -3,6 +3,7 @@ use bytes::bytes_dict_builder; use primitive::primitive_dict_builder; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_panic; @@ -44,15 +45,30 @@ pub trait DictEncoder: Send { fn codes_ptype(&self) -> PType; } +/// Create a dictionary encoder using the default allocator. +#[deprecated(note = "use `dict_encoder_in` with an explicit allocator")] pub fn dict_encoder(array: &ArrayRef, constraints: &DictConstraints) -> Box { + dict_encoder_in( + array, + constraints, + BufferAllocatorRef::statically_allocated(), + ) +} + +/// Create a dictionary encoder using `allocator`. +pub fn dict_encoder_in( + array: &ArrayRef, + constraints: &DictConstraints, + allocator: BufferAllocatorRef, +) -> Box { let dict_builder: Box = if let Some(pa) = array.as_opt::() { match_each_native_ptype!(pa.ptype(), |P| { - primitive_dict_builder::

(pa.dtype().nullability(), constraints) + primitive_dict_builder::

(pa.dtype().nullability(), constraints, allocator) }) } else if let Some(vbv) = array.as_opt::() { - bytes_dict_builder(vbv.dtype().clone(), constraints) + bytes_dict_builder(vbv.dtype().clone(), constraints, allocator) } else if let Some(vb) = array.as_opt::() { - bytes_dict_builder(vb.dtype().clone(), constraints) + bytes_dict_builder(vb.dtype().clone(), constraints, allocator) } else { vortex_panic!("Can only encode primitive or varbin/view arrays") }; @@ -67,7 +83,7 @@ pub fn dict_encode_with_constraints( constraints: &DictConstraints, ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut encoder = dict_encoder(array, constraints); + let mut encoder = dict_encoder_in(array, constraints, ctx.allocator().clone()); let codes = encoder.encode(array, ctx)?.narrow(ctx)?; // SAFETY: The encoding process will produce a value set of codes and values // All values in the dictionary are guaranteed to be referenced by at least one code diff --git a/vortex-array/src/builders/dict/primitive.rs b/vortex-array/src/builders/dict/primitive.rs index dea2932367b..ba2e86c6ee3 100644 --- a/vortex-array/src/builders/dict/primitive.rs +++ b/vortex-array/src/builders/dict/primitive.rs @@ -6,7 +6,9 @@ use std::hash::Hash; use std::mem; use rustc_hash::FxBuildHasher; +use vortex_buffer::Alignment; use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -31,6 +33,7 @@ use crate::validity::Validity; pub fn primitive_dict_builder( nullability: Nullability, constraints: &DictConstraints, + allocator: BufferAllocatorRef, ) -> Box where NativeValue: Hash + Eq, @@ -44,20 +47,25 @@ where width => vortex_panic!("invalid bit_width: {width}"), }); match max_possible_len { - max if max <= u8::MAX as u64 => { - Box::new(PrimitiveDictBuilder::::new(nullability, constraints)) - } - max if max <= u16::MAX as u64 => Box::new(PrimitiveDictBuilder::::new( + max if max <= u8::MAX as u64 => Box::new(PrimitiveDictBuilder::::new_in( nullability, constraints, + allocator, )), - max if max <= u32::MAX as u64 => Box::new(PrimitiveDictBuilder::::new( + max if max <= u16::MAX as u64 => Box::new(PrimitiveDictBuilder::::new_in( nullability, constraints, + allocator, )), - _ => Box::new(PrimitiveDictBuilder::::new( + max if max <= u32::MAX as u64 => Box::new(PrimitiveDictBuilder::::new_in( nullability, constraints, + allocator, + )), + _ => Box::new(PrimitiveDictBuilder::::new_in( + nullability, + constraints, + allocator, )), } } @@ -68,15 +76,20 @@ where NativeValue: Hash + Eq, Code: UnsignedPType, { - pub fn new(nullability: Nullability, constraints: &DictConstraints) -> Self { + pub fn new_in( + nullability: Nullability, + constraints: &DictConstraints, + allocator: BufferAllocatorRef, + ) -> Self { let max_dict_len = constraints .max_len .min(constraints.max_bytes / T::PTYPE.byte_width()); Self { lookup: HashMap::with_hasher(FxBuildHasher), null_code: OnceCell::new(), - values: BufferMut::::empty(), - values_nulls: BitBufferMut::empty(), + values: BufferMut::::empty_aligned_in(Alignment::of::(), allocator.clone()), + values_nulls: BitBufferMut::empty_in(allocator.clone()), + allocator, nullability, max_dict_len, } @@ -135,6 +148,7 @@ pub struct PrimitiveDictBuilder { values_nulls: BitBufferMut, nullability: Nullability, max_dict_len: usize, + allocator: BufferAllocatorRef, } impl DictEncoder for PrimitiveDictBuilder @@ -144,7 +158,7 @@ where Code: UnsignedPType, { fn encode(&mut self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - let mut codes = BufferMut::::with_capacity(array.len()); + let mut codes = BufferMut::::with_capacity_in(array.len(), self.allocator.clone()); let prim = array.clone().execute::(ctx)?; match prim.validity()?.execute_mask(array.len(), ctx)? { @@ -183,11 +197,18 @@ where } fn reset(&mut self) -> ArrayRef { - PrimitiveArray::new( - self.values.clone(), - Validity::from_bit_buffer(mem::take(&mut self.values_nulls).freeze(), self.nullability), + self.lookup.clear(); + self.null_code = OnceCell::new(); + let values = mem::replace( + &mut self.values, + BufferMut::empty_aligned_in(Alignment::of::(), self.allocator.clone()), + ); + let nulls = mem::replace( + &mut self.values_nulls, + BitBufferMut::empty_in(self.allocator.clone()), ) - .into_array() + .freeze(); + PrimitiveArray::new(values, Validity::from_bit_buffer(nulls, self.nullability)).into_array() } fn codes_ptype(&self) -> PType { @@ -206,7 +227,9 @@ mod test { use crate::VortexSessionExecute; use crate::arrays::dict::DictArraySlotsExt; use crate::assert_arrays_eq; + use crate::builders::dict::UNCONSTRAINED; use crate::builders::dict::dict_encode; + use crate::builders::dict::dict_encoder_in; use crate::builders::dict::primitive::PrimitiveArray; static SESSION: LazyLock = LazyLock::new(crate::array_session); @@ -247,4 +270,26 @@ mod test { PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(); assert_arrays_eq!(dict.values(), expected_values, &mut ctx); } + + #[test] + fn reset_clears_dict() { + let mut ctx = SESSION.create_execution_ctx(); + let first = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(); + let mut encoder = dict_encoder_in(&first, &UNCONSTRAINED, ctx.allocator().clone()); + + assert_arrays_eq!( + encoder.encode(&first, &mut ctx).unwrap(), + buffer![0u32, 1, 2].into_array(), + &mut ctx + ); + assert_arrays_eq!(encoder.reset(), first, &mut ctx); + + let second = PrimitiveArray::from_option_iter([Some(4i32), Some(5)]).into_array(); + assert_arrays_eq!( + encoder.encode(&second, &mut ctx).unwrap(), + buffer![0u32, 1].into_array(), + &mut ctx + ); + assert_arrays_eq!(encoder.reset(), second, &mut ctx); + } } diff --git a/vortex-array/src/builders/extension.rs b/vortex-array/src/builders/extension.rs index aa8e1b76aa2..544b119a0a6 100644 --- a/vortex-array/src/builders/extension.rs +++ b/vortex-array/src/builders/extension.rs @@ -3,9 +3,9 @@ use std::any::Any; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; @@ -13,8 +13,8 @@ use crate::IntoArray; use crate::arrays::ExtensionArray; use crate::arrays::extension::ExtensionArrayExt; use crate::builders::ArrayBuilder; +use crate::builders::ChildBuilder; use crate::builders::DEFAULT_BUILDER_CAPACITY; -use crate::builders::builder_with_capacity; use crate::canonical::Canonical; use crate::dtype::DType; use crate::dtype::extension::ExtDTypeRef; @@ -24,19 +24,35 @@ use crate::scalar::Scalar; /// The builder for building a [`ExtensionArray`]. pub struct ExtensionBuilder { dtype: DType, - storage: Box, + storage: ChildBuilder, } impl ExtensionBuilder { /// Creates a new `ExtensionBuilder` with a capacity of [`DEFAULT_BUILDER_CAPACITY`]. + #[deprecated(note = "use `new_in` with an explicit allocator")] pub fn new(ext_dtype: ExtDTypeRef) -> Self { - Self::with_capacity(ext_dtype, DEFAULT_BUILDER_CAPACITY) + Self::new_in(ext_dtype, BufferAllocatorRef::static_ref()) + } + + /// Creates a new `ExtensionBuilder` with the default capacity using `allocator`. + pub fn new_in(ext_dtype: ExtDTypeRef, allocator: &BufferAllocatorRef) -> Self { + Self::with_capacity_in(ext_dtype, DEFAULT_BUILDER_CAPACITY, allocator) } /// Creates a new `ExtensionBuilder` with the given `capacity`. + #[deprecated(note = "use `with_capacity_in` with an explicit allocator")] pub fn with_capacity(ext_dtype: ExtDTypeRef, capacity: usize) -> Self { + Self::with_capacity_in(ext_dtype, capacity, BufferAllocatorRef::static_ref()) + } + + /// Creates a new `ExtensionBuilder` with `capacity` using `allocator`. + pub fn with_capacity_in( + ext_dtype: ExtDTypeRef, + capacity: usize, + allocator: &BufferAllocatorRef, + ) -> Self { Self { - storage: builder_with_capacity(ext_dtype.storage_dtype(), capacity), + storage: ChildBuilder::with_capacity(ext_dtype.storage_dtype(), capacity, allocator), dtype: DType::Extension(ext_dtype), } } @@ -53,9 +69,7 @@ impl ExtensionBuilder { array: &ExtensionArray, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - array - .storage_array() - .append_to_builder(self.storage.as_mut(), ctx) + self.storage.append_array(array.storage_array(), ctx) } /// Finishes the builder directly into a [`ExtensionArray`]. @@ -116,10 +130,6 @@ impl ArrayBuilder for ExtensionBuilder { self.storage.reserve_exact(capacity) } - unsafe fn set_validity_unchecked(&mut self, validity: Mask) { - unsafe { self.storage.set_validity_unchecked(validity) }; - } - fn finish(&mut self) -> ArrayRef { self.finish_into_extension().into_array() } @@ -147,7 +157,8 @@ mod tests { let mut ctx = array_session().create_execution_ctx(); let ext_dtype = Date::new(TimeUnit::Days, Nullability::Nullable).erased(); - let mut builder = ExtensionBuilder::new(ext_dtype.clone()); + let mut builder = + ExtensionBuilder::new_in(ext_dtype.clone(), BufferAllocatorRef::static_ref()); // Test appending a valid extension value. let storage1 = Scalar::from(Some(42i32)); @@ -177,7 +188,7 @@ mod tests { assert_eq!(array.len(), 3); // Test wrong dtype error. - let mut builder = ExtensionBuilder::new(ext_dtype); + let mut builder = ExtensionBuilder::new_in(ext_dtype, BufferAllocatorRef::static_ref()); let wrong_scalar = Scalar::from(true); assert!(builder.append_scalar(&wrong_scalar).is_err()); } diff --git a/vortex-array/src/builders/fixed_size_list.rs b/vortex-array/src/builders/fixed_size_list.rs index 9150d37b50b..e9a66b62d83 100644 --- a/vortex-array/src/builders/fixed_size_list.rs +++ b/vortex-array/src/builders/fixed_size_list.rs @@ -4,22 +4,23 @@ use std::any::Any; use std::sync::Arc; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; -use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; +use crate::arrays::ChunkedArray; use crate::arrays::FixedSizeListArray; -use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::builders::ArrayBuilder; +use crate::builders::ChildBuilder; use crate::builders::DEFAULT_BUILDER_CAPACITY; -use crate::builders::LazyBitBufferBuilder; -use crate::builders::builder_with_capacity; +use crate::builders::ValidityBuilder; use crate::canonical::Canonical; use crate::dtype::DType; use crate::dtype::Nullability; @@ -34,37 +35,73 @@ pub struct FixedSizeListBuilder { /// The builder for the underlying elements of the [`FixedSizeListArray`]. /// /// This builder will have a capacity equal to the `list_size * capacity`. - elements_builder: Box, + elements_builder: ChildBuilder, /// The null map builder of the [`FixedSizeListArray`]. /// /// We also use this type to store the length of the final output array. - nulls: LazyBitBufferBuilder, + nulls: ValidityBuilder, } impl FixedSizeListBuilder { /// Creates a new `FixedSizeListBuilder` with a capacity of [`DEFAULT_BUILDER_CAPACITY`]. + #[deprecated(note = "use `new_in` with an explicit allocator")] pub fn new(element_dtype: Arc, list_size: u32, nullability: Nullability) -> Self { - Self::with_capacity( + Self::new_in( + element_dtype, + list_size, + nullability, + BufferAllocatorRef::static_ref(), + ) + } + + /// Creates a new `FixedSizeListBuilder` with the default capacity using `allocator`. + pub fn new_in( + element_dtype: Arc, + list_size: u32, + nullability: Nullability, + allocator: &BufferAllocatorRef, + ) -> Self { + Self::with_capacity_in( element_dtype, list_size, nullability, DEFAULT_BUILDER_CAPACITY, + allocator, ) } /// Creates a new `FixedSizeListBuilder` with the given `capacity`. + #[deprecated(note = "use `with_capacity_in` with an explicit allocator")] pub fn with_capacity( element_dtype: Arc, list_size: u32, nullability: Nullability, capacity: usize, + ) -> Self { + Self::with_capacity_in( + element_dtype, + list_size, + nullability, + capacity, + BufferAllocatorRef::static_ref(), + ) + } + + /// Creates a new `FixedSizeListBuilder` with `capacity` using `allocator`. + pub fn with_capacity_in( + element_dtype: Arc, + list_size: u32, + nullability: Nullability, + capacity: usize, + allocator: &BufferAllocatorRef, ) -> Self { let elements_capacity = capacity * list_size as usize; - let elements_builder = builder_with_capacity(&element_dtype, elements_capacity); + let elements_builder = + ChildBuilder::with_capacity(&element_dtype, elements_capacity, allocator); let fsl_dtype = DType::FixedSizeList(element_dtype, list_size, nullability); - let nulls = LazyBitBufferBuilder::new(capacity); + let nulls = ValidityBuilder::new(capacity, allocator); Self { dtype: fsl_dtype, @@ -98,12 +135,59 @@ impl FixedSizeListBuilder { self.list_size() ); - array.append_to_builder(self.elements_builder.as_mut(), ctx)?; + self.elements_builder.append_array(array, ctx)?; self.nulls.append_non_null(); Ok(()) } + /// Appends `array` as `n` identical non-null lists. + /// + /// A fixed-size list array holds its elements back to back, so `n` identical lists are the + /// array's elements tiled `n` times - there is no layout that lets the rows share one range of + /// elements the way a list view's can. The tiling costs nothing to build even so: the elements + /// go in as a [`ChunkedArray`] of `n` clones of the same array, so the tile's values are stored + /// once however many rows reference them, and the child holds the whole run as one chunk. + /// + /// A caller with a run of appends to make should hand over the same `array` each time rather + /// than rebuild it, which is the whole reason this takes an array instead of a scalar. + pub fn append_array_as_repeated_list( + &mut self, + array: &ArrayRef, + n: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + vortex_ensure!( + array.dtype() == self.element_dtype(), + "Array dtype {:?} does not match list element dtype {:?}", + array.dtype(), + self.element_dtype() + ); + vortex_ensure!( + array.len() == self.list_size() as usize, + "Array length {} does not match fixed list size {}", + array.len(), + self.list_size() + ); + + if n == 0 { + return Ok(()); + } + + // SAFETY: every chunk is `array` itself, so they share its dtype and none is empty. + let tiled = unsafe { + ChunkedArray::new_unchecked( + std::iter::repeat_n(array.clone(), n).collect::>(), + self.element_dtype().clone(), + ) + }; + self.elements_builder + .append_array(&tiled.into_array(), ctx)?; + self.nulls.append_n_non_nulls(n); + + Ok(()) + } + /// Appends the values of a canonical [`FixedSizeListArray`] to the builder, recursing into the /// elements builder. pub(crate) fn append_fixed_size_list_array( @@ -115,11 +199,8 @@ impl FixedSizeListBuilder { return Ok(()); } - array - .elements() - .append_to_builder(self.elements_builder.as_mut(), ctx)?; - self.nulls - .append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?); + self.elements_builder.append_array(array.elements(), ctx)?; + self.nulls.append_validity(array.validity()?, array.len()); Ok(()) } @@ -263,10 +344,6 @@ impl ArrayBuilder for FixedSizeListBuilder { self.nulls.reserve_exact(additional); } - unsafe fn set_validity_unchecked(&mut self, validity: Mask) { - self.nulls = LazyBitBufferBuilder::from_validity_mask(validity); - } - fn finish(&mut self) -> ArrayRef { self.finish_into_fixed_size_list().into_array() } @@ -280,6 +357,7 @@ impl ArrayBuilder for FixedSizeListBuilder { mod tests { use std::sync::Arc; + use vortex_buffer::BufferAllocatorRef; use vortex_buffer::buffer; use vortex_error::VortexExpect; @@ -289,6 +367,7 @@ mod tests { use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; + use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::builders::ArrayBuilder; use crate::builders::fixed_size_list::FixedSizeListArray; use crate::dtype::DType; @@ -300,8 +379,13 @@ mod tests { #[test] fn test_empty() { - let mut builder = - FixedSizeListBuilder::with_capacity(Arc::new(I32.into()), 3, NonNullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity_in( + Arc::new(I32.into()), + 3, + NonNullable, + 0, + BufferAllocatorRef::static_ref(), + ); let fsl = builder.finish(); assert_eq!(fsl.len(), 0); @@ -310,8 +394,13 @@ mod tests { #[test] fn test_values() { let dtype: Arc = Arc::new(I32.into()); - let mut builder = - FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 3, NonNullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity_in( + Arc::clone(&dtype), + 3, + NonNullable, + 0, + BufferAllocatorRef::static_ref(), + ); builder .append_value( @@ -347,8 +436,13 @@ mod tests { #[test] fn test_degenerate_size_zero_non_nullable() { let dtype: Arc = Arc::new(I32.into()); - let mut builder = - FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 0, NonNullable, 10000000); + let mut builder = FixedSizeListBuilder::with_capacity_in( + Arc::clone(&dtype), + 0, + NonNullable, + 10000000, + BufferAllocatorRef::static_ref(), + ); // Append multiple "empty" lists. for _ in 0..100 { @@ -373,8 +467,13 @@ mod tests { fn test_degenerate_size_zero_nullable() { // Use nullable elements since we'll be appending nulls let dtype: Arc = Arc::new(DType::Primitive(I32, Nullable)); - let mut builder = - FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 0, Nullable, 10000000); + let mut builder = FixedSizeListBuilder::with_capacity_in( + Arc::clone(&dtype), + 0, + Nullable, + 10000000, + BufferAllocatorRef::static_ref(), + ); // Mix of null and non-null empty lists. for i in 0..100 { @@ -402,8 +501,13 @@ mod tests { fn test_capacity_growth() { let dtype: Arc = Arc::new(I32.into()); // Start with capacity 0. - let mut builder = - FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 2, NonNullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity_in( + Arc::clone(&dtype), + 2, + NonNullable, + 0, + BufferAllocatorRef::static_ref(), + ); // Add more items than initial capacity. for i in 0..5 { @@ -431,7 +535,13 @@ mod tests { fn test_large_size_zero_capacity_empty_result() { let dtype: Arc = Arc::new(I32.into()); // Large list size but zero capacity and no appends. - let mut builder = FixedSizeListBuilder::with_capacity(dtype, 100000000, NonNullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity_in( + dtype, + 100000000, + NonNullable, + 0, + BufferAllocatorRef::static_ref(), + ); let fsl = builder.finish(); assert_eq!(fsl.len(), 0); @@ -446,7 +556,13 @@ mod tests { fn test_nullable_lists_non_nullable_elements() { let mut ctx = array_session().create_execution_ctx(); let dtype: Arc = Arc::new(DType::Primitive(I32, NonNullable)); - let mut builder = FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 2, Nullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity_in( + Arc::clone(&dtype), + 2, + Nullable, + 0, + BufferAllocatorRef::static_ref(), + ); builder .append_value( @@ -497,8 +613,13 @@ mod tests { #[test] fn test_non_nullable_lists_nullable_elements() { let dtype: Arc = Arc::new(DType::Primitive(I32, Nullable)); - let mut builder = - FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 3, NonNullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity_in( + Arc::clone(&dtype), + 3, + NonNullable, + 0, + BufferAllocatorRef::static_ref(), + ); builder .append_value( @@ -541,7 +662,13 @@ mod tests { #[test] fn test_append_zeros() { let dtype: Arc = Arc::new(I32.into()); - let mut builder = FixedSizeListBuilder::with_capacity(dtype, 3, NonNullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity_in( + dtype, + 3, + NonNullable, + 0, + BufferAllocatorRef::static_ref(), + ); builder.append_zeros(5); @@ -568,7 +695,13 @@ mod tests { let mut ctx = array_session().create_execution_ctx(); // Elements must be nullable if we're going to append null lists let dtype: Arc = Arc::new(DType::Primitive(I32, Nullable)); - let mut builder = FixedSizeListBuilder::with_capacity(dtype, 2, Nullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity_in( + dtype, + 2, + Nullable, + 0, + BufferAllocatorRef::static_ref(), + ); assert_eq!(builder.dtype().nullability(), Nullable); builder.append_nulls(3); @@ -597,7 +730,13 @@ mod tests { let mut ctx = array_session().create_execution_ctx(); // Elements must be nullable if we're going to append null lists let dtype: Arc = Arc::new(DType::Primitive(I32, Nullable)); - let mut builder = FixedSizeListBuilder::with_capacity(dtype, 2, Nullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity_in( + dtype, + 2, + Nullable, + 0, + BufferAllocatorRef::static_ref(), + ); assert_eq!(builder.dtype().nullability(), Nullable); builder @@ -624,7 +763,13 @@ mod tests { #[test] fn test_append_zeros_degenerate() { let dtype: Arc = Arc::new(I32.into()); - let mut builder = FixedSizeListBuilder::with_capacity(dtype, 0, NonNullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity_in( + dtype, + 0, + NonNullable, + 0, + BufferAllocatorRef::static_ref(), + ); assert_eq!(builder.len(), 0); builder.append_zeros(1000); @@ -642,8 +787,13 @@ mod tests { #[test] fn test_invalid_size_error() { let dtype: Arc = Arc::new(I32.into()); - let mut builder = - FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 3, NonNullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity_in( + Arc::clone(&dtype), + 3, + NonNullable, + 0, + BufferAllocatorRef::static_ref(), + ); // Try to append a list with wrong size. let result = builder.append_value( @@ -677,7 +827,13 @@ mod tests { 3, ); - let mut builder = FixedSizeListBuilder::with_capacity(dtype, 2, Nullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity_in( + dtype, + 2, + Nullable, + 0, + BufferAllocatorRef::static_ref(), + ); let source_array = source.into_array(); source_array @@ -758,7 +914,13 @@ mod tests { 2, ); - let mut builder = FixedSizeListBuilder::with_capacity(dtype, 0, Nullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity_in( + dtype, + 0, + Nullable, + 0, + BufferAllocatorRef::static_ref(), + ); source1 .into_array() @@ -827,8 +989,13 @@ mod tests { 0, ); - let mut builder = - FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 3, NonNullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity_in( + Arc::clone(&dtype), + 3, + NonNullable, + 0, + BufferAllocatorRef::static_ref(), + ); // Add some initial data. builder @@ -857,7 +1024,13 @@ mod tests { let mut ctx = array_session().create_execution_ctx(); // Use nullable elements since we'll be appending nulls let dtype: Arc = Arc::new(DType::Primitive(I32, Nullable)); - let mut builder = FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 2, Nullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity_in( + Arc::clone(&dtype), + 2, + Nullable, + 0, + BufferAllocatorRef::static_ref(), + ); // Mix of operations. builder @@ -944,7 +1117,13 @@ mod tests { fn test_append_scalar() { let mut ctx = array_session().create_execution_ctx(); let dtype: Arc = Arc::new(I32.into()); - let mut builder = FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 2, Nullable, 10); + let mut builder = FixedSizeListBuilder::with_capacity_in( + Arc::clone(&dtype), + 2, + Nullable, + 10, + BufferAllocatorRef::static_ref(), + ); // Test appending a valid fixed-size list. let list_scalar1 = @@ -1004,7 +1183,13 @@ mod tests { ); // Test wrong dtype error. - let mut builder = FixedSizeListBuilder::with_capacity(dtype, 2, NonNullable, 10); + let mut builder = FixedSizeListBuilder::with_capacity_in( + dtype, + 2, + NonNullable, + 10, + BufferAllocatorRef::static_ref(), + ); let wrong_scalar = Scalar::from(42i32); assert!(builder.append_scalar(&wrong_scalar).is_err()); } @@ -1013,8 +1198,13 @@ mod tests { fn test_append_array_as_list() { let dtype: Arc = Arc::new(I32.into()); let mut ctx = array_session().create_execution_ctx(); - let mut builder = - FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 3, NonNullable, 10); + let mut builder = FixedSizeListBuilder::with_capacity_in( + Arc::clone(&dtype), + 3, + NonNullable, + 10, + BufferAllocatorRef::static_ref(), + ); // Append a primitive array as a single list entry. let arr1 = buffer![1i32, 2, 3].into_array(); @@ -1064,8 +1254,13 @@ mod tests { ); // Test dtype mismatch error. - let mut builder = - FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 3, NonNullable, 10); + let mut builder = FixedSizeListBuilder::with_capacity_in( + Arc::clone(&dtype), + 3, + NonNullable, + 10, + BufferAllocatorRef::static_ref(), + ); let wrong_dtype_arr = buffer![1i64, 2, 3].into_array(); assert!( builder @@ -1074,7 +1269,13 @@ mod tests { ); // Test length mismatch error. - let mut builder = FixedSizeListBuilder::with_capacity(dtype, 3, NonNullable, 10); + let mut builder = FixedSizeListBuilder::with_capacity_in( + dtype, + 3, + NonNullable, + 10, + BufferAllocatorRef::static_ref(), + ); let wrong_len_arr = buffer![1i32, 2].into_array(); assert!( builder diff --git a/vortex-array/src/builders/lazy_null_builder.rs b/vortex-array/src/builders/lazy_null_builder.rs index 24abe9ba17d..d46d731cf25 100644 --- a/vortex-array/src/builders/lazy_null_builder.rs +++ b/vortex-array/src/builders/lazy_null_builder.rs @@ -3,6 +3,7 @@ use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexExpect; use vortex_error::vortex_panic; use vortex_mask::Mask; @@ -18,52 +19,18 @@ pub struct LazyBitBufferBuilder { inner: Option, len: usize, capacity: usize, + allocator: BufferAllocatorRef, } impl LazyBitBufferBuilder { /// Creates a new empty builder. /// `capacity` is the number of bits in the null buffer. - pub fn new(capacity: usize) -> Self { + pub fn new(capacity: usize, allocator: BufferAllocatorRef) -> Self { Self { inner: None, len: 0, capacity, - } - } - - /// Creates a builder pre-populated from a validity mask, taking ownership of the mask's buffer - /// instead of copying it where possible. - /// - /// This is the counterpart to [`append_validity_mask`](Self::append_validity_mask) for callers - /// that want to *replace* the builder's contents with the mask rather than extend them: because - /// we own the mask, we can move its buffer in instead of copying it. - pub fn from_validity_mask(validity_mask: Mask) -> Self { - match validity_mask { - // An unmaterialized builder already represents `len` non-null values, so an all-valid - // mask stays lazy. - Mask::AllTrue(len) => Self { - inner: None, - len, - capacity: len, - }, - Mask::AllFalse(len) => Self::from_buffer(BitBufferMut::new_unset(len)), - // Take ownership of the underlying buffer; `into_bit_buffer` and `try_into_mut` only - // copy when the buffer is shared, otherwise this is a move. - values @ Mask::Values(_) => Self::from_buffer( - values - .into_bit_buffer() - .try_into_mut() - .unwrap_or_else(|buffer| BitBufferMut::copy_from(&buffer)), - ), - } - } - - /// Creates a builder backed by an already-materialized buffer. - fn from_buffer(inner: BitBufferMut) -> Self { - Self { - inner: Some(inner), - len: 0, - capacity: 0, + allocator, } } @@ -184,7 +151,8 @@ impl LazyBitBufferBuilder { #[inline(never)] fn materialize(&mut self) { if self.inner.is_none() { - let mut bit_mut = BitBufferMut::with_capacity(self.len.max(self.capacity)); + let mut bit_mut = + BitBufferMut::with_capacity_in(self.len.max(self.capacity), self.allocator.clone()); bit_mut.append_n(true, self.len); self.inner = Some(bit_mut); } diff --git a/vortex-array/src/builders/list.rs b/vortex-array/src/builders/list.rs index 130f37534cf..5e97f122cba 100644 --- a/vortex-array/src/builders/list.rs +++ b/vortex-array/src/builders/list.rs @@ -5,12 +5,12 @@ use std::any::Any; use std::sync::Arc; use num_traits::AsPrimitive; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; -use vortex_mask::Mask; use crate::ArrayRef; use crate::Canonical; @@ -22,47 +22,60 @@ use crate::arrays::ListArray; use crate::arrays::ListView; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; -use crate::arrays::list::ListArrayExt; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::list::ListArraySlotsExt; +use crate::arrays::listview::ListViewArraySlotsExt; +use crate::arrays::listview::ListViewRebuildMode; use crate::builders::ArrayBuilder; +use crate::builders::ChildBuilder; use crate::builders::DEFAULT_BUILDER_CAPACITY; -use crate::builders::LazyBitBufferBuilder; use crate::builders::PrimitiveBuilder; -use crate::builders::builder_with_capacity; +use crate::builders::ValidityBuilder; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::dtype::Nullability; use crate::dtype::Nullability::NonNullable; +use crate::dtype::OffsetBuilderPType; use crate::match_each_integer_ptype; use crate::scalar::ListScalar; use crate::scalar::Scalar; -/// The builder for building a [`ListArray`], parametrized by the [`IntegerPType`] of the `offsets` -/// builder. -pub struct ListBuilder { +/// The builder for building a [`ListArray`], parametrized by the [`OffsetBuilderPType`] of the +/// `offsets` builder. +pub struct ListBuilder { /// The [`DType`] of the [`ListArray`]. This **must** be a [`DType::List`]. dtype: DType, /// The builder for the underlying elements of the [`ListArray`]. - elements_builder: Box, + elements_builder: ChildBuilder, /// The builder for the `offsets` into the `elements` array. offsets_builder: PrimitiveBuilder, /// The null map builder of the [`ListArray`]. - nulls: LazyBitBufferBuilder, + nulls: ValidityBuilder, } -impl ListBuilder { +impl ListBuilder { /// Creates a new `ListBuilder` with a capacity of [`DEFAULT_BUILDER_CAPACITY`]. + #[deprecated(note = "use `new_in` with an explicit allocator")] pub fn new(value_dtype: Arc, nullability: Nullability) -> Self { - Self::with_capacity( + Self::new_in(value_dtype, nullability, BufferAllocatorRef::static_ref()) + } + + /// Creates a new `ListBuilder` with the default capacity using `allocator`. + pub fn new_in( + value_dtype: Arc, + nullability: Nullability, + allocator: &BufferAllocatorRef, + ) -> Self { + Self::with_capacity_in( value_dtype, nullability, // We arbitrarily choose 2 times the number of list scalars for the capacity of the // elements builder since we cannot know this ahead of time. DEFAULT_BUILDER_CAPACITY * 2, DEFAULT_BUILDER_CAPACITY, + allocator, ) } @@ -73,14 +86,34 @@ impl ListBuilder { /// # Notes /// /// The number of offsets is one more than the length (# of list scalars) in the array. + #[deprecated(note = "use `with_capacity_in` with an explicit allocator")] pub fn with_capacity( value_dtype: Arc, nullability: Nullability, elements_capacity: usize, capacity: usize, ) -> Self { - let elements_builder = builder_with_capacity(value_dtype.as_ref(), elements_capacity); - let mut offsets_builder = PrimitiveBuilder::::with_capacity(NonNullable, capacity + 1); + Self::with_capacity_in( + value_dtype, + nullability, + elements_capacity, + capacity, + BufferAllocatorRef::static_ref(), + ) + } + + /// Creates a list builder with the given capacities using `allocator`. + pub fn with_capacity_in( + value_dtype: Arc, + nullability: Nullability, + elements_capacity: usize, + capacity: usize, + allocator: &BufferAllocatorRef, + ) -> Self { + let elements_builder = + ChildBuilder::with_capacity(value_dtype.as_ref(), elements_capacity, allocator); + let mut offsets_builder = + PrimitiveBuilder::::with_capacity_in(NonNullable, capacity + 1, allocator); // The first offset is always 0 and represents an empty list. offsets_builder.append_zero(); @@ -88,7 +121,7 @@ impl ListBuilder { Self { elements_builder, offsets_builder, - nulls: LazyBitBufferBuilder::new(capacity), + nulls: ValidityBuilder::new(capacity, allocator), dtype: DType::List(value_dtype, nullability), } } @@ -111,8 +144,7 @@ impl ListBuilder { self.element_dtype() ); - self.elements_builder.reserve_exact(array.len()); - array.append_to_builder(self.elements_builder.as_mut(), ctx)?; + self.elements_builder.append_array(array, ctx)?; self.nulls.append_non_null(); self.offsets_builder.append_value( O::from_usize(self.elements_builder.len()) @@ -174,10 +206,109 @@ impl ListBuilder { element_dtype } + + /// Appends the values of a [`List`]-encoded `array` to this builder. + /// + /// List encodings dispatch here through + /// [`match_each_list_builder!`](crate::match_each_list_builder) because the concrete list + /// builders are generic over their offset integer type, which cannot be named through a + /// `dyn ArrayBuilder`. + pub fn append_list_array( + &mut self, + array: ArrayView<'_, List>, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if array.is_empty() { + return Ok(()); + } + + self.nulls.append_validity(array.validity()?, array.len()); + + let num_lists = array.len(); + let offsets = array.offsets().clone().execute::(ctx)?; + match_each_integer_ptype!(offsets.ptype(), |OffsetType| { + let offsets = offsets.as_slice::(); + let first: usize = offsets[0].as_(); + let last: usize = offsets[num_lists].as_(); + + // Lists in a `ListArray` are contiguous, so the referenced elements can be appended + // in bulk and the offsets rebased onto this builder's elements. + let elements_base = self.elements_builder.len(); + if last > first { + self.elements_builder + .append_array(&array.elements().slice(first..last)?, ctx)?; + } + + self.offsets_builder.reserve_exact(num_lists); + let mut offsets_range = self.offsets_builder.uninit_range(num_lists); + for i in 0..num_lists { + let end: usize = offsets[i + 1].as_(); + offsets_range.set_value( + i, + O::from_usize(end - first + elements_base) + .vortex_expect("Failed to convert offset"), + ); + } + // SAFETY: We have initialized all `num_lists` values, and since the `offsets` array is + // non-nullable, we are done. + unsafe { offsets_range.finish() }; + }); + Ok(()) + } + + /// Appends the values of a [`ListView`]-encoded `array` to this builder. + /// + /// See [`append_list_array`](Self::append_list_array); this is the same hook for the canonical + /// [`ListViewArray`] encoding. + /// + /// A `ListArray`'s offsets can only describe contiguous, in-order lists, so views laid out any + /// other way (overlapping, out of order, or with interior gaps) are flattened first. + pub fn append_listview_array( + &mut self, + array: ArrayView<'_, ListView>, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if array.is_empty() { + return Ok(()); + } + + self.nulls.append_validity(array.validity()?, array.len()); + + // Flatten the views into the only layout `ListArray` offsets can express. This is a cheap + // clone when they already are laid out that way, and the flattened result keeps the + // original validity, so the null map appended above still describes it. + let array = array + .into_owned() + .rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx)?; + debug_assert!(array.is_zero_copy_to_list()); + + // Note that `ListViewArray` has `n` offsets and sizes, not `n+1` offsets like `ListArray`. + let elements = array.elements(); + let offsets = array.offsets().clone().execute::(ctx)?; + let sizes = array.sizes().clone().execute::(ctx)?; + + match_each_integer_ptype!(offsets.ptype(), |OffsetType| { + match_each_integer_ptype!(sizes.ptype(), |SizeType| { + extend_from_listview( + self, + elements, + offsets.as_slice::(), + sizes.as_slice::(), + ctx, + )? + }) + }); + Ok(()) + } } -/// Appends `ListViewArray`-layout lists (`n` offsets and sizes) into a [`ListBuilder`], converting -/// into the `ListArray` (`n + 1` offsets) layout. +/// Appends the lists of a zero-copy-to-list [`ListViewArray`] (`n` offsets and sizes) into a +/// [`ListBuilder`], converting into the `ListArray` (`n + 1` offsets) layout. +/// +/// The caller must have made `new_offsets` and `new_sizes` zero-copyable to a `ListArray`, so the +/// lists they describe are contiguous and in order — which is the only layout `ListArray` offsets +/// can express. That lets the referenced elements be appended in bulk, with the offsets rebased +/// onto this builder's elements, instead of appending a slice per list. fn extend_from_listview( builder: &mut ListBuilder, new_elements: &ArrayRef, @@ -186,37 +317,33 @@ fn extend_from_listview( ctx: &mut ExecutionCtx, ) -> VortexResult<()> where - O: IntegerPType, + O: OffsetBuilderPType, OffsetType: IntegerPType, SizeType: IntegerPType, { let num_lists = new_offsets.len(); debug_assert_eq!(num_lists, new_sizes.len()); - let total_elements: usize = new_sizes.iter().map(|size| size.as_()).sum(); - builder.elements_builder.reserve_exact(total_elements); + // Leading and trailing unreferenced elements are allowed even in a zero-copy-to-list layout, + // so the referenced range is bounded by the first list's start and the last list's end. + let first: usize = new_offsets[0].as_(); + let last: usize = new_offsets[num_lists - 1].as_() + new_sizes[num_lists - 1].as_(); + + let elements_base = builder.elements_builder.len(); + if last > first { + builder + .elements_builder + .append_array(&new_elements.slice(first..last)?, ctx)?; + } - let mut curr_offset = builder.elements_builder.len(); builder.offsets_builder.reserve_exact(num_lists); let mut offsets_range = builder.offsets_builder.uninit_range(num_lists); - - // We need to append each list individually, converting from `ListViewArray` format to - // the `ListArray` format that `ListBuilder` expects. - for i in 0..new_offsets.len() { - let offset: usize = new_offsets[i].as_(); - let size: usize = new_sizes[i].as_(); - - if size > 0 { - let list_elements = new_elements - .slice(offset..offset + size) - .vortex_expect("list builder slice"); - list_elements.append_to_builder(builder.elements_builder.as_mut(), ctx)?; - curr_offset += size; - } - - let new_offset = O::from_usize(curr_offset).vortex_expect("Failed to convert offset"); - - offsets_range.set_value(i, new_offset); + for i in 0..num_lists { + let end: usize = new_offsets[i].as_() + new_sizes[i].as_(); + offsets_range.set_value( + i, + O::from_usize(end - first + elements_base).vortex_expect("Failed to convert offset"), + ); } // SAFETY: We have initialized all `num_lists` values, and since the `offsets` array is @@ -225,7 +352,7 @@ where Ok(()) } -impl ArrayBuilder for ListBuilder { +impl ArrayBuilder for ListBuilder { fn as_any(&self) -> &dyn Any { self } @@ -281,10 +408,6 @@ impl ArrayBuilder for ListBuilder { self.nulls.reserve_exact(additional); } - unsafe fn set_validity_unchecked(&mut self, validity: Mask) { - self.nulls = LazyBitBufferBuilder::from_validity_mask(validity); - } - fn finish(&mut self) -> ArrayRef { self.finish_into_list().into_array() } @@ -296,84 +419,6 @@ impl ArrayBuilder for ListBuilder { .vortex_expect("list builder should canonicalize to listview"); Canonical::List(listview) } - - fn append_list_array( - &mut self, - array: ArrayView<'_, List>, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - if array.is_empty() { - return Ok(()); - } - - self.nulls - .append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?); - - let num_lists = array.len(); - let offsets = array.offsets().clone().execute::(ctx)?; - match_each_integer_ptype!(offsets.ptype(), |OffsetType| { - let offsets = offsets.as_slice::(); - let first: usize = offsets[0].as_(); - let last: usize = offsets[num_lists].as_(); - - // Lists in a `ListArray` are contiguous, so the referenced elements can be appended - // in bulk and the offsets rebased onto this builder's elements. - let elements_base = self.elements_builder.len(); - if last > first { - self.elements_builder.reserve_exact(last - first); - array - .elements() - .slice(first..last)? - .append_to_builder(self.elements_builder.as_mut(), ctx)?; - } - - self.offsets_builder.reserve_exact(num_lists); - let mut offsets_range = self.offsets_builder.uninit_range(num_lists); - for i in 0..num_lists { - let end: usize = offsets[i + 1].as_(); - offsets_range.set_value( - i, - O::from_usize(end - first + elements_base) - .vortex_expect("Failed to convert offset"), - ); - } - // SAFETY: We have initialized all `num_lists` values, and since the `offsets` array is - // non-nullable, we are done. - unsafe { offsets_range.finish() }; - }); - Ok(()) - } - - fn append_listview_array( - &mut self, - array: ArrayView<'_, ListView>, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - if array.is_empty() { - return Ok(()); - } - - self.nulls - .append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?); - - // Note that `ListViewArray` has `n` offsets and sizes, not `n+1` offsets like `ListArray`. - let elements = array.elements(); - let offsets = array.offsets().clone().execute::(ctx)?; - let sizes = array.sizes().clone().execute::(ctx)?; - - match_each_integer_ptype!(offsets.ptype(), |OffsetType| { - match_each_integer_ptype!(sizes.ptype(), |SizeType| { - extend_from_listview( - self, - elements, - offsets.as_slice::(), - sizes.as_slice::(), - ctx, - )? - }) - }); - Ok(()) - } } #[cfg(test)] @@ -382,6 +427,7 @@ mod tests { use Nullability::NonNullable; use Nullability::Nullable; + use vortex_buffer::BufferAllocatorRef; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -391,17 +437,18 @@ mod tests { use crate::arrays::ChunkedArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; - use crate::arrays::list::ListArrayExt; + use crate::arrays::list::ListArraySlotsExt; use crate::arrays::listview::ListViewArrayExt; + use crate::arrays::listview::ListViewArraySlotsExt; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; use crate::builders::ListViewBuilder; - use crate::builders::builder_with_capacity; + use crate::builders::builder_with_capacity_in; use crate::builders::list::ListArray; use crate::builders::list::ListBuilder; use crate::dtype::DType; - use crate::dtype::IntegerPType; use crate::dtype::Nullability; + use crate::dtype::OffsetBuilderPType; use crate::dtype::PType::I32; use crate::executor::VortexSessionExecute; use crate::scalar::Scalar; @@ -409,8 +456,13 @@ mod tests { #[test] fn test_empty() { - let mut builder = - ListBuilder::::with_capacity(Arc::new(I32.into()), NonNullable, 0, 0); + let mut builder = ListBuilder::::with_capacity_in( + Arc::new(I32.into()), + NonNullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); let list = builder.finish(); assert_eq!(list.len(), 0); @@ -419,7 +471,13 @@ mod tests { #[test] fn test_values() { let dtype: Arc = Arc::new(I32.into()); - let mut builder = ListBuilder::::with_capacity(Arc::clone(&dtype), NonNullable, 0, 0); + let mut builder = ListBuilder::::with_capacity_in( + Arc::clone(&dtype), + NonNullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); builder .append_value( @@ -456,7 +514,13 @@ mod tests { #[test] fn test_append_empty_list() { let dtype: Arc = Arc::new(I32.into()); - let mut builder = ListBuilder::::with_capacity(Arc::clone(&dtype), NonNullable, 0, 0); + let mut builder = ListBuilder::::with_capacity_in( + Arc::clone(&dtype), + NonNullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); assert!( builder @@ -468,7 +532,13 @@ mod tests { #[test] fn test_nullable_values() { let dtype: Arc = Arc::new(I32.into()); - let mut builder = ListBuilder::::with_capacity(Arc::clone(&dtype), Nullable, 0, 0); + let mut builder = ListBuilder::::with_capacity_in( + Arc::clone(&dtype), + Nullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); builder .append_value( @@ -507,7 +577,7 @@ mod tests { assert_eq!(list_array.list_elements_at(2).unwrap().len(), 3); } - fn test_extend_builder_gen() { + fn test_extend_builder_gen() { let list = ListArray::from_iter_opt_slow::( [Some(vec![0, 1, 2]), None, Some(vec![4, 5])], Arc::new(I32.into()), @@ -518,7 +588,13 @@ mod tests { let mut ctx = array_session().create_execution_ctx(); - let mut builder = ListBuilder::::with_capacity(Arc::new(I32.into()), Nullable, 18, 9); + let mut builder = ListBuilder::::with_capacity_in( + Arc::new(I32.into()), + Nullable, + 18, + 9, + BufferAllocatorRef::static_ref(), + ); list.append_to_builder(&mut builder, &mut ctx).unwrap(); list.append_to_builder(&mut builder, &mut ctx).unwrap(); list.slice(0..0) @@ -590,35 +666,64 @@ mod tests { // `builder_with_capacity` produces a `ListViewBuilder` for `DType::List`; appending the // `List`-encoded array must dispatch into it instead of bailing. - let mut listview_builder = builder_with_capacity(list.dtype(), list.len()); + let mut listview_builder = + builder_with_capacity_in(list.dtype(), list.len(), BufferAllocatorRef::static_ref()); list.append_to_builder(listview_builder.as_mut(), &mut ctx)?; assert_arrays_eq!(listview_builder.finish(), list, &mut ctx); // A `ListViewBuilder` with non-`u64` (including signed) offset and size types must work // for both source encodings. - let mut lv_u32_u8 = ListViewBuilder::::with_capacity(elem_dtype(), Nullable, 8, 4); - list.append_to_builder(&mut lv_u32_u8, &mut ctx)?; - assert_arrays_eq!(lv_u32_u8.finish(), list, &mut ctx); + let mut lv_u64_u32 = ListViewBuilder::::with_capacity_in( + elem_dtype(), + Nullable, + 8, + 4, + BufferAllocatorRef::static_ref(), + ); + list.append_to_builder(&mut lv_u64_u32, &mut ctx)?; + assert_arrays_eq!(lv_u64_u32.finish(), list, &mut ctx); - let mut lv_i32_i16 = - ListViewBuilder::::with_capacity(elem_dtype(), Nullable, 8, 4); - list.append_to_builder(&mut lv_i32_i16, &mut ctx)?; - assert_arrays_eq!(lv_i32_i16.finish(), list, &mut ctx); + let mut lv_i64_i32 = ListViewBuilder::::with_capacity_in( + elem_dtype(), + Nullable, + 8, + 4, + BufferAllocatorRef::static_ref(), + ); + list.append_to_builder(&mut lv_i64_i32, &mut ctx)?; + assert_arrays_eq!(lv_i64_i32.finish(), list, &mut ctx); - let mut lv_u16_u16 = - ListViewBuilder::::with_capacity(elem_dtype(), Nullable, 8, 4); - listview.append_to_builder(&mut lv_u16_u16, &mut ctx)?; - assert_arrays_eq!(lv_u16_u16.finish(), list, &mut ctx); + let mut lv_u32_u32 = ListViewBuilder::::with_capacity_in( + elem_dtype(), + Nullable, + 8, + 4, + BufferAllocatorRef::static_ref(), + ); + listview.append_to_builder(&mut lv_u32_u32, &mut ctx)?; + assert_arrays_eq!(lv_u32_u32.finish(), list, &mut ctx); // Both source encodings appended into `ListBuilder`s with non-`u64` (including signed) // offset types. - let mut list_builder = ListBuilder::::with_capacity(elem_dtype(), Nullable, 8, 4); + let mut list_builder = ListBuilder::::with_capacity_in( + elem_dtype(), + Nullable, + 8, + 4, + BufferAllocatorRef::static_ref(), + ); list.append_to_builder(&mut list_builder, &mut ctx)?; assert_arrays_eq!(list_builder.finish(), list, &mut ctx); - let mut list_builder_i16 = ListBuilder::::with_capacity(elem_dtype(), Nullable, 8, 4); - listview.append_to_builder(&mut list_builder_i16, &mut ctx)?; - assert_arrays_eq!(list_builder_i16.finish(), list, &mut ctx); + let mut list_builder_i32 = ListBuilder::::with_capacity_in( + elem_dtype(), + Nullable, + 8, + 4, + BufferAllocatorRef::static_ref(), + ); + listview.append_to_builder(&mut list_builder_i32, &mut ctx)?; + assert_arrays_eq!(list_builder_i32.finish(), list, &mut ctx); Ok(()) } @@ -639,13 +744,25 @@ mod tests { )?; // Appending twice checks growth from a non-empty builder and offset rebasing. - let mut builder = ListBuilder::::with_capacity(Arc::clone(&dtype), Nullable, 0, 0); + let mut builder = ListBuilder::::with_capacity_in( + Arc::clone(&dtype), + Nullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); builder.append_list_array(source.as_view(), &mut ctx)?; builder.append_list_array(source.as_view(), &mut ctx)?; assert_arrays_eq!(builder.finish(), expected, &mut ctx); let source_listview = source.into_array().execute::(&mut ctx)?; - let mut builder = ListBuilder::::with_capacity(dtype, Nullable, 0, 0); + let mut builder = ListBuilder::::with_capacity_in( + dtype, + Nullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); builder.append_listview_array(source_listview.as_view(), &mut ctx)?; builder.append_listview_array(source_listview.as_view(), &mut ctx)?; assert_arrays_eq!(builder.finish(), expected, &mut ctx); @@ -653,15 +770,73 @@ mod tests { Ok(()) } + /// A `ListArray`'s offsets can only describe contiguous, in-order lists, so an overlapping + /// source has to be flattened before its elements can be appended in bulk. A sliced source, + /// meanwhile, keeps the layout it has and is appended from wherever its first list starts. + #[test] + fn test_append_listview_array_flattens_overlaps_and_skips_leading_elements() -> VortexResult<()> + { + let mut ctx = array_session().create_execution_ctx(); + let dtype: Arc = Arc::new(I32.into()); + + // Overlapping source, so not zero-copyable to a list: + // - List 0: [10, 20] + // - List 1: null (size is intentionally non-zero in the source metadata) + // - List 2: [10], sharing the elements list 0 already referenced + let overlapping = unsafe { + ListViewArray::new_unchecked( + buffer![10i32, 20, 30].into_array(), + buffer![0u32, 1, 0].into_array(), + buffer![2u8, 2, 1].into_array(), + Validity::from_iter([true, false, true]), + ) + }; + assert!(!overlapping.is_zero_copy_to_list()); + + // Zero-copyable source sliced past its first list, so its elements start at offset 2. + let sliced = unsafe { + ListViewArray::new_unchecked( + buffer![40i32, 50, 60, 70].into_array(), + buffer![0u32, 2].into_array(), + buffer![2u32, 2].into_array(), + Validity::AllValid, + ) + .with_zero_copy_to_list(true) + } + .into_array() + .slice(1..2)? + .execute::(&mut ctx)?; + + let mut builder = ListBuilder::::with_capacity_in( + dtype, + Nullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); + builder.append_listview_array(overlapping.as_view(), &mut ctx)?; + builder.append_listview_array(sliced.as_view(), &mut ctx)?; + + let list = builder.finish_into_list(); + assert_arrays_eq!( + list.elements(), + PrimitiveArray::from_iter([10i32, 20, 10, 60, 70]), + &mut ctx + ); + assert_arrays_eq!( + list.offsets(), + PrimitiveArray::from_iter([0u32, 2, 2, 3, 5]), + &mut ctx + ); + + Ok(()) + } + #[test] fn test_extend_builder() { - test_extend_builder_gen::(); - test_extend_builder_gen::(); test_extend_builder_gen::(); test_extend_builder_gen::(); - test_extend_builder_gen::(); - test_extend_builder_gen::(); test_extend_builder_gen::(); test_extend_builder_gen::(); } @@ -719,7 +894,13 @@ mod tests { #[test] fn test_append_scalar() { let dtype: Arc = Arc::new(I32.into()); - let mut builder = ListBuilder::::with_capacity(Arc::clone(&dtype), Nullable, 20, 10); + let mut builder = ListBuilder::::with_capacity_in( + Arc::clone(&dtype), + Nullable, + 20, + 10, + BufferAllocatorRef::static_ref(), + ); // Test appending a valid list. let list_scalar1 = @@ -790,7 +971,13 @@ mod tests { ); // Test wrong dtype error. - let mut builder = ListBuilder::::with_capacity(dtype, NonNullable, 20, 10); + let mut builder = ListBuilder::::with_capacity_in( + dtype, + NonNullable, + 20, + 10, + BufferAllocatorRef::static_ref(), + ); let wrong_scalar = Scalar::from(42i32); assert!(builder.append_scalar(&wrong_scalar).is_err()); } @@ -799,8 +986,13 @@ mod tests { fn test_append_array_as_list() { let dtype: Arc = Arc::new(I32.into()); let mut ctx = array_session().create_execution_ctx(); - let mut builder = - ListBuilder::::with_capacity(Arc::clone(&dtype), NonNullable, 20, 10); + let mut builder = ListBuilder::::with_capacity_in( + Arc::clone(&dtype), + NonNullable, + 20, + 10, + BufferAllocatorRef::static_ref(), + ); // Append a primitive array as a single list entry. let arr1 = buffer![1i32, 2, 3].into_array(); @@ -849,7 +1041,13 @@ mod tests { ); // Test dtype mismatch error. - let mut builder = ListBuilder::::with_capacity(dtype, NonNullable, 20, 10); + let mut builder = ListBuilder::::with_capacity_in( + dtype, + NonNullable, + 20, + 10, + BufferAllocatorRef::static_ref(), + ); let wrong_dtype_arr = buffer![1i64, 2, 3].into_array(); assert!( builder diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index 4e6523bd7d3..a0c326b7f37 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -12,11 +12,12 @@ use std::sync::Arc; +use num_traits::ToPrimitive; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; -use vortex_mask::Mask; use crate::ArrayRef; use crate::Canonical; @@ -27,38 +28,38 @@ use crate::arrays::List; use crate::arrays::ListView; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; -use crate::arrays::list::ListArrayExt; +use crate::arrays::list::ListArraySlotsExt; use crate::arrays::listview::ListViewArrayExt; -use crate::arrays::listview::ListViewRebuildMode; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::builders::ArrayBuilder; +use crate::builders::ChildBuilder; use crate::builders::DEFAULT_BUILDER_CAPACITY; use crate::builders::PrimitiveBuilder; use crate::builders::UninitRange; -use crate::builders::builder_with_capacity; -use crate::builders::lazy_null_builder::LazyBitBufferBuilder; -use crate::builtins::ArrayBuiltins; +use crate::builders::ValidityBuilder; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::dtype::Nullability; +use crate::dtype::OffsetBuilderPType; use crate::match_each_integer_ptype; use crate::scalar::ListScalar; use crate::scalar::Scalar; -/// A builder for creating [`ListViewArray`] instances, parameterized by the [`IntegerPType`] of -/// the `offsets` and the `sizes` builders. +/// A builder for creating [`ListViewArray`] instances, parameterized by the [`OffsetBuilderPType`] +/// of the `offsets` and the `sizes` builders. /// /// This builder tracks both offsets and sizes using potentially different integer types for memory /// efficiency. For example, you might use `u64` for offsets but only `u8` for sizes if your lists /// are small. /// -/// Any combination of [`IntegerPType`] types are valid, as long as the type of `sizes` can fit into -/// the type of `offsets`. -pub struct ListViewBuilder { +/// Any combination of [`OffsetBuilderPType`] types is valid, as long as the type of `sizes` can fit +/// into the type of `offsets`. +pub struct ListViewBuilder { /// The [`DType`] of the [`ListViewArray`]. This **must** be a [`DType::List`]. dtype: DType, /// The builder for the underlying elements of the [`ListArray`](crate::arrays::ListArray). - elements_builder: Box, + elements_builder: ChildBuilder, /// The builder for the `offsets` into the `elements` array. offsets_builder: PrimitiveBuilder, @@ -67,19 +68,38 @@ pub struct ListViewBuilder { sizes_builder: PrimitiveBuilder, /// The null map builder of the [`ListViewArray`]. - nulls: LazyBitBufferBuilder, + nulls: ValidityBuilder, + + /// Whether the appends so far leave the result zero-copyable to a [`ListArray`]. + /// + /// Only [`append_listview_array`](Self::append_listview_array) can clear this; every other + /// append writes its lists back to back. + /// + /// [`ListArray`]: crate::arrays::ListArray + zero_copy_to_list: bool, } -impl ListViewBuilder { +impl ListViewBuilder { /// Creates a new `ListViewBuilder` with a capacity of [`DEFAULT_BUILDER_CAPACITY`]. + #[deprecated(note = "use `new_in` with an explicit allocator")] pub fn new(element_dtype: Arc, nullability: Nullability) -> Self { - Self::with_capacity( + Self::new_in(element_dtype, nullability, BufferAllocatorRef::static_ref()) + } + + /// Creates a new `ListViewBuilder` with the default capacity using `allocator`. + pub fn new_in( + element_dtype: Arc, + nullability: Nullability, + allocator: &BufferAllocatorRef, + ) -> Self { + Self::with_capacity_in( element_dtype, nullability, // We arbitrarily choose 2 times the number of list scalars for the capacity of the // elements builder since we cannot know this ahead of time. DEFAULT_BUILDER_CAPACITY * 2, DEFAULT_BUILDER_CAPACITY, + allocator, ) } @@ -90,20 +110,43 @@ impl ListViewBuilder { /// # Panics /// /// Panics if the size type `S` cannot fit within the offset type `O`. + #[deprecated(note = "use `with_capacity_in` with an explicit allocator")] pub fn with_capacity( element_dtype: Arc, nullability: Nullability, elements_capacity: usize, capacity: usize, ) -> Self { - let elements_builder = builder_with_capacity(&element_dtype, elements_capacity); + Self::with_capacity_in( + element_dtype, + nullability, + elements_capacity, + capacity, + BufferAllocatorRef::static_ref(), + ) + } + + /// Creates a list-view builder with the given capacities using `allocator`. + /// + /// # Panics + /// + /// Panics if the size type `S` cannot fit within the offset type `O`. + pub fn with_capacity_in( + element_dtype: Arc, + nullability: Nullability, + elements_capacity: usize, + capacity: usize, + allocator: &BufferAllocatorRef, + ) -> Self { + let elements_builder = + ChildBuilder::with_capacity(&element_dtype, elements_capacity, allocator); let offsets_builder = - PrimitiveBuilder::::with_capacity(Nullability::NonNullable, capacity); + PrimitiveBuilder::::with_capacity_in(Nullability::NonNullable, capacity, allocator); let sizes_builder = - PrimitiveBuilder::::with_capacity(Nullability::NonNullable, capacity); + PrimitiveBuilder::::with_capacity_in(Nullability::NonNullable, capacity, allocator); - let nulls = LazyBitBufferBuilder::new(capacity); + let nulls = ValidityBuilder::new(capacity, allocator); Self { dtype: DType::List(element_dtype, nullability), @@ -111,6 +154,7 @@ impl ListViewBuilder { offsets_builder, sizes_builder, nulls, + zero_copy_to_list: true, } } @@ -142,8 +186,7 @@ impl ListViewBuilder { "appending this list would cause an offset overflow" ); - self.elements_builder.reserve_exact(num_elements); - array.append_to_builder(self.elements_builder.as_mut(), ctx)?; + self.elements_builder.append_array(array, ctx)?; self.nulls.append_non_null(); self.offsets_builder.append_value( @@ -196,6 +239,58 @@ impl ListViewBuilder { Ok(()) } + /// Appends `array` as `n` identical non-null lists, storing its elements once. + /// + /// A `ListViewArray` can point many views at one range of elements, so a repeated list costs + /// its elements once however many rows it covers. The elements go in as one appended array, so + /// a caller that hands over the same `array` on every call - a sparse array filling the gaps + /// between its patches, say - stores those elements once for the whole result. + /// + /// The views share their elements, so the result is no longer zero-copyable to a + /// [`ListArray`](crate::arrays::ListArray) unless it covers a single row or empty lists. + pub fn append_array_as_repeated_list( + &mut self, + array: &ArrayRef, + n: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + vortex_ensure!( + array.dtype() == self.element_dtype(), + "Array dtype {:?} does not match list element dtype {:?}", + array.dtype(), + self.element_dtype() + ); + + if n == 0 { + return Ok(()); + } + + let curr_offset = self.elements_builder.len(); + let num_elements = array.len(); + + // We must assert this even in release mode to ensure that the safety comment in + // `finish_into_listview` is correct. + assert!( + ((curr_offset + num_elements) as u64) < O::max_value_as_u64(), + "appending this list would cause an offset overflow" + ); + + self.elements_builder.append_array(array, ctx)?; + + let offset = + O::from_usize(curr_offset).vortex_expect("Failed to convert from usize to `O`"); + let size = S::from_usize(num_elements).vortex_expect("Failed to convert from usize to `S`"); + self.offsets_builder.append_n_values(offset, n); + self.sizes_builder.append_n_values(size, n); + self.nulls.append_n_non_nulls(n); + + if n > 1 && num_elements > 0 { + self.zero_copy_to_list = false; + } + + Ok(()) + } + /// Finishes the builder directly into a [`ListViewArray`]. pub fn finish_into_listview(&mut self) -> ListViewArray { debug_assert_eq!(self.offsets_builder.len(), self.sizes_builder.len()); @@ -206,6 +301,8 @@ impl ListViewBuilder { let sizes = self.sizes_builder.finish(); let validity = self.nulls.finish_with_nullability(self.dtype.nullability()); + let zero_copy_to_list = std::mem::replace(&mut self.zero_copy_to_list, true); + // SAFETY: // - Both the offsets and the sizes are non-nullable. // - The offsets, sizes, and validity have the same length since we always appended the same @@ -213,12 +310,13 @@ impl ListViewBuilder { // - We checked on construction that the sizes type fits into the offsets. // - In every method that adds values to this builder (`append_value`, `append_scalar`, // `append_list_array`, and `append_listview_array`), we checked that `offset + size` - // does not overflow. - // - We constructed everything in a way that builds the `ListViewArray` similar to the shape - // of a `ListArray`, so we know the resulting array is zero-copyable to a `ListArray`. + // does not overflow. `append_listview_array` rebases the offsets it was handed onto + // exactly the elements it appended, so the source's bound carries over. + // - Every append writes its lists back to back, so the result is zero-copyable to a + // `ListArray` unless `zero_copy_to_list` recorded an appended layout we left alone. unsafe { ListViewArray::new_unchecked(elements, offsets, sizes, validity) - .with_zero_copy_to_list(true) + .with_zero_copy_to_list(zero_copy_to_list) } } @@ -231,9 +329,139 @@ impl ListViewBuilder { element_dtype } + + /// Appends the values of a [`List`]-encoded `array` to this builder. + /// + /// List encodings dispatch here through + /// [`match_each_list_builder!`](crate::match_each_list_builder) because the concrete list + /// builders are generic over their offset/size integer types, which cannot be named through a + /// `dyn ArrayBuilder`. + pub fn append_list_array( + &mut self, + array: ArrayView<'_, List>, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if array.is_empty() { + return Ok(()); + } + + self.nulls.append_validity(array.validity()?, array.len()); + + let offsets = array.offsets().clone().execute::(ctx)?; + match_each_integer_ptype!(offsets.ptype(), |OffsetType| { + extend_from_list( + self, + array.elements(), + offsets.as_slice::(), + ctx, + )? + }); + Ok(()) + } + + /// Appends the values of a [`ListView`]-encoded `array` to this builder. + /// + /// See [`append_list_array`](Self::append_list_array); this is the same hook for the canonical + /// [`ListViewArray`] encoding. + /// + /// The views keep the layout they arrived in, so overlapping sources keep sharing their + /// elements and the finished array reports [`is_zero_copy_to_list`] as `false`. Callers that + /// need an exact layout should [`rebuild`](ListViewArray::rebuild) it. + /// + /// [`is_zero_copy_to_list`]: crate::arrays::listview::ListViewData::is_zero_copy_to_list + pub fn append_listview_array( + &mut self, + array: ArrayView<'_, ListView>, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if array.is_empty() { + return Ok(()); + } + + let len = array.len(); + + // Materialize the metadata once and do the trimming and the rebase by hand. Going through + // `rebuild(ListViewRebuildMode::TrimElements)` would subtract the window start from every + // offset with a compute kernel, only for the rebase below to add this builder's elements + // base straight back on - two passes, one of them through the compute stack, for one + // addition per offset. Casting the sizes to the builder's type is another kernel for what + // is a copy. + let offsets = array.offsets().clone().execute::(ctx)?; + let sizes = array.sizes().clone().execute::(ctx)?; + + // The window of `elements` that the views actually reference. Everything outside it is + // unreachable and must not be appended. An exact source covers its window back to back and + // in order, so the first and last view bound it; any other layout has to be searched. + let (start, end) = if array.is_zero_copy_to_list() { + let last = len - 1; + ( + metadata_at(&offsets, 0), + metadata_at(&offsets, last) + metadata_at(&sizes, last), + ) + } else { + array.into_owned().referenced_element_bounds(ctx)? + }; + + // An exact source references every element it carries, back to back, so it lands flush + // against the elements already in the builder. Any other layout does not. + self.zero_copy_to_list &= array.is_zero_copy_to_list(); + + self.nulls.append_validity(array.validity()?, len); + + // Bulk append the referenced elements; the offsets are rebased onto them below. + let elements_base = self.elements_builder.len(); + + // We must assert this even in release mode to ensure that the safety comment in + // `finish_into_listview` is correct. + assert!( + ((elements_base + (end - start)) as u64) < O::max_value_as_u64(), + "appending this list would cause an offset overflow" + ); + + if end > start { + self.elements_builder + .append_array(&array.elements().slice(start..end)?, ctx)?; + } + + // Every view lies inside `start..end`, so rebasing it onto `elements_base` keeps it inside + // the elements this builder now holds - which is what lets `finish_into_listview` build the + // array unchecked. + assert_eq!( + self.elements_builder.len(), + elements_base + (end - start), + "appending the referenced elements did not extend the child by the window's length" + ); + + self.offsets_builder.reserve_exact(len); + let offsets_range = self.offsets_builder.uninit_range(len); + match_each_integer_ptype!(offsets.ptype(), |A| { + extend_rebased_offsets::( + offsets_range, + offsets.as_slice::(), + start, + elements_base, + ); + }); + + self.sizes_builder.reserve_exact(len); + let mut sizes_range = self.sizes_builder.uninit_range(len); + if sizes.ptype() == S::PTYPE { + // The sizes already have the builder's type, so there is nothing to convert. + sizes_range.copy_from_slice(0, sizes.as_slice::()); + // SAFETY: `copy_from_slice` initialized all `len` values, and the sizes builder is + // non-nullable. + unsafe { sizes_range.finish() }; + } else { + match_each_integer_ptype!(sizes.ptype(), |A| { + extend_converted_sizes::(sizes_range, sizes.as_slice::()); + }); + } + + Ok(()) + } } -impl ArrayBuilder for ListViewBuilder { +impl ArrayBuilder for ListViewBuilder { fn as_any(&self) -> &dyn std::any::Any { self } @@ -307,10 +535,6 @@ impl ArrayBuilder for ListViewBuilder { self.nulls.reserve_exact(capacity); } - unsafe fn set_validity_unchecked(&mut self, validity: Mask) { - self.nulls = LazyBitBufferBuilder::from_validity_mask(validity); - } - fn finish(&mut self) -> ArrayRef { self.finish_into_listview().into_array() } @@ -318,88 +542,6 @@ impl ArrayBuilder for ListViewBuilder { fn finish_into_canonical(&mut self, _ctx: &mut ExecutionCtx) -> Canonical { Canonical::List(self.finish_into_listview()) } - - fn append_list_array( - &mut self, - array: ArrayView<'_, List>, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - if array.is_empty() { - return Ok(()); - } - - self.nulls - .append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?); - - let offsets = array.offsets().clone().execute::(ctx)?; - match_each_integer_ptype!(offsets.ptype(), |OffsetType| { - extend_from_list( - self, - array.elements(), - offsets.as_slice::(), - ctx, - )? - }); - Ok(()) - } - - fn append_listview_array( - &mut self, - array: ArrayView<'_, ListView>, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - if array.is_empty() { - return Ok(()); - } - - // Normalize to an exact zero-copy-to-list layout and then bulk append. This avoids the - // very expensive scalar_at-per-list path for overlapping / out-of-order list views. - let listview = array - .into_owned() - .rebuild(ListViewRebuildMode::MakeExact, ctx)?; - debug_assert!(listview.is_zero_copy_to_list()); - - self.nulls - .append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?); - - // Bulk append the new elements (which should have no gaps or overlaps). - let old_elements_len = self.elements_builder.len(); - self.elements_builder - .reserve_exact(listview.elements().len()); - listview - .elements() - .append_to_builder(self.elements_builder.as_mut(), ctx)?; - let new_elements_len = self.elements_builder.len(); - - // Reserve enough space for the new views. - let extend_length = listview.len(); - self.sizes_builder.reserve_exact(extend_length); - self.offsets_builder.reserve_exact(extend_length); - - // The incoming sizes might have a different type than the builder, so we need to cast. - let cast_sizes = listview - .sizes() - .clone() - .cast(self.sizes_builder.dtype().clone())?; - cast_sizes.append_to_builder(&mut self.sizes_builder, ctx)?; - - // Now we need to adjust all of the offsets by adding the current number of elements in the - // builder. - let uninit_range = self.offsets_builder.uninit_range(extend_length); - - // This should be cheap because we didn't compress after rebuilding. - let new_offsets = listview.offsets().clone().execute::(ctx)?; - - match_each_integer_ptype!(new_offsets.ptype(), |A| { - adjust_and_extend_offsets::( - uninit_range, - new_offsets, - old_elements_len, - new_elements_len, - ); - }); - Ok(()) - } } /// Appends `ListArray`-layout lists (`n + 1` cumulative offsets) into a [`ListViewBuilder`]. @@ -414,8 +556,8 @@ fn extend_from_list( ctx: &mut ExecutionCtx, ) -> VortexResult<()> where - O: IntegerPType, - S: IntegerPType, + O: OffsetBuilderPType, + S: OffsetBuilderPType, OffsetType: IntegerPType, { let num_lists = offsets.len() - 1; @@ -432,10 +574,9 @@ where ); if last > first { - builder.elements_builder.reserve_exact(last - first); - elements - .slice(first..last)? - .append_to_builder(builder.elements_builder.as_mut(), ctx)?; + builder + .elements_builder + .append_array(&elements.slice(first..last)?, ctx)?; } builder.offsets_builder.reserve_exact(num_lists); @@ -462,50 +603,74 @@ where Ok(()) } -/// Given new offsets, adds them to the `UninitRange` after adding the `old_elements_len` to each -/// offset. -fn adjust_and_extend_offsets( - mut uninit_range: UninitRange, - new_offsets: PrimitiveArray, - old_elements_len: usize, - new_elements_len: usize, -) { - let new_offsets_slice = new_offsets.as_slice::(); - let old_elements_len = O::from_usize(old_elements_len) - .vortex_expect("the old elements length did not fit into the offset type (impossible)"); - let new_elements_len = O::from_usize(new_elements_len) - .vortex_expect("the current elements length did not fit into the offset type (impossible)"); - - for i in 0..uninit_range.len() { - let new_offset = O::from_usize( - new_offsets_slice[i] - .to_usize() - .vortex_expect("Offsets must always fit in usize"), - ) - .vortex_expect("New offset somehow did not fit into the builder's offset type"); +/// Reads one non-nullable integer value of list view metadata as a `usize`. +fn metadata_at(metadata: &PrimitiveArray, index: usize) -> usize { + match_each_integer_ptype!(metadata.ptype(), |A| { + metadata.as_slice::()[index] + .to_usize() + .vortex_expect("list view metadata must fit in a usize") + }) +} - // We have to check this even in release mode to ensure the final `new_unchecked` - // construction in `finish_into_listview` is valid. - let adjusted_new_offset = new_offset + old_elements_len; - assert!( - adjusted_new_offset <= new_elements_len, - "[{i}/{}]: {new_offset} + {old_elements_len} \ - = {adjusted_new_offset} <= {new_elements_len} failed", - uninit_range.len() +/// Writes `offsets` into `range`, moved off the source's element window and onto the elements the +/// builder already holds. +/// +/// `window_start` is the first element the source's views reference, so every offset is at least +/// `window_start` and the subtraction cannot underflow. +fn extend_rebased_offsets( + mut range: UninitRange, + offsets: &[A], + window_start: usize, + elements_base: usize, +) { + debug_assert_eq!(range.len(), offsets.len()); + + for (i, &offset) in offsets.iter().enumerate() { + let offset = offset + .to_usize() + .vortex_expect("offsets must always fit in usize"); + debug_assert!( + offset >= window_start, + "offset {offset} precedes the referenced window at {window_start}" ); - - uninit_range.set_value(i, adjusted_new_offset); + let rebased = O::from_usize(offset - window_start + elements_base) + .vortex_expect("rebased offset did not fit into the builder's offset type"); + range.set_value(i, rebased); } // SAFETY: We have set all the values in the range, and since `offsets` are non-nullable, we are // done. - unsafe { uninit_range.finish() }; + unsafe { range.finish() }; +} + +/// Writes `sizes` into `range`, converting them to the builder's size type. +/// +/// Sizes that already have the builder's type are copied in bulk by the caller instead. +fn extend_converted_sizes( + mut range: UninitRange, + sizes: &[A], +) { + debug_assert_eq!(range.len(), sizes.len()); + + for (i, &size) in sizes.iter().enumerate() { + let size = S::from_usize( + size.to_usize() + .vortex_expect("sizes must always fit in usize"), + ) + .vortex_expect("size did not fit into the builder's size type"); + range.set_value(i, size); + } + + // SAFETY: We have set all the values in the range, and since `sizes` are non-nullable, we are + // done. + unsafe { range.finish() }; } #[cfg(test)] mod tests { use std::sync::Arc; + use vortex_buffer::BufferAllocatorRef; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -514,9 +679,11 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; + use crate::arrays::ConstantArray; use crate::arrays::ListArray; use crate::arrays::ListViewArray; use crate::arrays::listview::ListViewArrayExt; + use crate::arrays::listview::ListViewArraySlotsExt; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; use crate::builders::listview::PrimitiveArray; @@ -529,8 +696,13 @@ mod tests { #[test] fn test_empty() { - let mut builder = - ListViewBuilder::::with_capacity(Arc::new(I32.into()), NonNullable, 0, 0); + let mut builder = ListViewBuilder::::with_capacity_in( + Arc::new(I32.into()), + NonNullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); let listview = builder.finish(); assert_eq!(listview.len(), 0); @@ -540,8 +712,13 @@ mod tests { fn test_basic_append_and_nulls() { let mut ctx = array_session().create_execution_ctx(); let dtype: Arc = Arc::new(I32.into()); - let mut builder = - ListViewBuilder::::with_capacity(Arc::clone(&dtype), Nullable, 0, 0); + let mut builder = ListViewBuilder::::with_capacity_in( + Arc::clone(&dtype), + Nullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); // Append a regular list. builder @@ -603,10 +780,15 @@ mod tests { #[test] fn test_different_offset_size_types() { let mut ctx = array_session().create_execution_ctx(); - // Test u32 offsets with u8 sizes. + // Test u64 offsets with u32 sizes. let dtype: Arc = Arc::new(I32.into()); - let mut builder = - ListViewBuilder::::with_capacity(Arc::clone(&dtype), NonNullable, 0, 0); + let mut builder = ListViewBuilder::::with_capacity_in( + Arc::clone(&dtype), + NonNullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); builder .append_value( @@ -647,10 +829,15 @@ mod tests { &mut ctx ); - // Test u64 offsets with u16 sizes. + // Test i64 offsets with i32 sizes. let dtype2: Arc = Arc::new(I32.into()); - let mut builder2 = - ListViewBuilder::::with_capacity(Arc::clone(&dtype2), NonNullable, 0, 0); + let mut builder2 = ListViewBuilder::::with_capacity_in( + Arc::clone(&dtype2), + NonNullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); for i in 0..5 { builder2 @@ -677,8 +864,13 @@ mod tests { fn test_builder_trait_methods() { let mut ctx = array_session().create_execution_ctx(); let dtype: Arc = Arc::new(I32.into()); - let mut builder = - ListViewBuilder::::with_capacity(Arc::clone(&dtype), Nullable, 0, 0); + let mut builder = ListViewBuilder::::with_capacity_in( + Arc::clone(&dtype), + Nullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); // Test append_zeros (creates empty lists). builder.append_zeros(2); @@ -738,8 +930,13 @@ mod tests { ) .unwrap(); - let mut builder = - ListViewBuilder::::with_capacity(Arc::clone(&dtype), Nullable, 0, 0); + let mut builder = ListViewBuilder::::with_capacity_in( + Arc::clone(&dtype), + Nullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); // Add initial data. builder @@ -815,8 +1012,13 @@ mod tests { (0..100).map(|i| (i % 10 != 0).then(|| vec![i])).collect(); let source = ListArray::from_iter_opt_slow::(lists.clone(), Arc::clone(&dtype))?; - let mut builder = - ListViewBuilder::::with_capacity(Arc::clone(&dtype), Nullable, 0, 0); + let mut builder = ListViewBuilder::::with_capacity_in( + Arc::clone(&dtype), + Nullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); builder.append_list_array(source.as_view(), &mut ctx)?; // Append a second time to check growth from a non-empty builder and offset rebasing. builder.append_list_array(source.as_view(), &mut ctx)?; @@ -833,6 +1035,90 @@ mod tests { Ok(()) } + /// A constant list array points every view at a single copy of the value; flattening it in the + /// builder would materialize a copy per row. + #[test] + fn test_constant_list_append_keeps_one_copy_of_the_value() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let element_dtype: Arc = Arc::new(I32.into()); + + const ROWS: usize = 10_000; + let fill = Scalar::list( + Arc::clone(&element_dtype), + vec![1i32.into(), 2i32.into(), 3i32.into()], + NonNullable, + ); + let constant = ConstantArray::new(fill, ROWS).into_array(); + + let mut builder = ListViewBuilder::::with_capacity_in( + element_dtype, + NonNullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); + constant.append_to_builder(&mut builder, &mut ctx)?; + let listview = builder.finish_into_listview(); + + assert_eq!(listview.len(), ROWS); + assert_eq!( + listview.elements().len(), + 3, + "the fill value should be stored once, not once per row", + ); + assert!(!listview.is_zero_copy_to_list()); + assert_arrays_eq!(&listview.into_array(), &constant, &mut ctx); + + Ok(()) + } + + /// Only the elements the views reference land in the builder: appending a slice out of the + /// middle of a list array leaves the elements on either side of it behind. + #[test] + fn test_append_listview_array_trims_unreferenced_elements() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype: Arc = Arc::new(I32.into()); + + // Five lists of two elements each, of which we append only the middle three. + let source = ListArray::from_iter_slow::( + (0..5).map(|i| vec![2 * i, 2 * i + 1]), + Arc::clone(&dtype), + )? + .into_array() + .execute::(&mut ctx)?; + let middle = source.slice(1..4)?.execute::(&mut ctx)?; + + let mut builder = ListViewBuilder::::with_capacity_in( + Arc::clone(&dtype), + NonNullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); + builder.append_listview_array(middle.as_view(), &mut ctx)?; + // A second append has to rebase onto the elements already in the builder. + builder.append_listview_array(middle.as_view(), &mut ctx)?; + let listview = builder.finish_into_listview(); + + assert_eq!( + listview.elements().len(), + 12, + "only the six referenced elements of each append should have landed", + ); + assert!( + listview.is_zero_copy_to_list(), + "trimming an exact source keeps the result exact", + ); + + let expected = ListArray::from_iter_slow::( + (1..4).chain(1..4).map(|i| vec![2 * i, 2 * i + 1]), + dtype, + )?; + assert_arrays_eq!(listview, expected, &mut ctx); + + Ok(()) + } + #[test] fn test_extend_from_array_overlapping_listview() { let mut ctx = array_session().create_execution_ctx(); @@ -852,15 +1138,21 @@ mod tests { }; assert!(!source.is_zero_copy_to_list()); - let mut builder = - ListViewBuilder::::with_capacity(Arc::clone(&dtype), Nullable, 0, 0); + let mut builder = ListViewBuilder::::with_capacity_in( + Arc::clone(&dtype), + Nullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); builder .append_listview_array(source.as_view(), &mut ctx) .unwrap(); let listview = builder.finish_into_listview(); assert_eq!(listview.len(), 3); - assert!(listview.is_zero_copy_to_list()); + // The builder kept the source's overlapping layout. + assert!(!listview.is_zero_copy_to_list()); assert_arrays_eq!( listview.list_elements_at(0).unwrap(), @@ -874,7 +1166,8 @@ mod tests { .execute_is_valid(1, &mut ctx) .unwrap() ); - assert_eq!(listview.list_elements_at(1).unwrap().len(), 0); + // List 1 is null, so the builder no longer rewrites its size to zero. + assert_eq!(listview.size_at(1), source.size_at(1)); assert_arrays_eq!( listview.list_elements_at(2).unwrap(), PrimitiveArray::from_iter([10i32]), @@ -885,8 +1178,13 @@ mod tests { #[test] fn test_error_append_null_to_non_nullable() { let dtype: Arc = Arc::new(I32.into()); - let mut builder = - ListViewBuilder::::with_capacity(Arc::clone(&dtype), NonNullable, 0, 0); + let mut builder = ListViewBuilder::::with_capacity_in( + Arc::clone(&dtype), + NonNullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); // Create a null list with nullable type (since Scalar::null requires nullable type). let null_scalar = Scalar::null(DType::List(dtype, Nullable)); @@ -907,8 +1205,13 @@ mod tests { fn test_append_array_as_list() { let dtype: Arc = Arc::new(I32.into()); let mut ctx = array_session().create_execution_ctx(); - let mut builder = - ListViewBuilder::::with_capacity(Arc::clone(&dtype), NonNullable, 20, 10); + let mut builder = ListViewBuilder::::with_capacity_in( + Arc::clone(&dtype), + NonNullable, + 20, + 10, + BufferAllocatorRef::static_ref(), + ); // Append a primitive array as a single list entry. let arr1 = buffer![1i32, 2, 3].into_array(); @@ -964,7 +1267,13 @@ mod tests { ); // Test dtype mismatch error. - let mut builder = ListViewBuilder::::with_capacity(dtype, NonNullable, 20, 10); + let mut builder = ListViewBuilder::::with_capacity_in( + dtype, + NonNullable, + 20, + 10, + BufferAllocatorRef::static_ref(), + ); let wrong_dtype_arr = buffer![1i64, 2, 3].into_array(); assert!( builder diff --git a/vortex-array/src/builders/map.rs b/vortex-array/src/builders/map.rs new file mode 100644 index 00000000000..4169000d96f --- /dev/null +++ b/vortex-array/src/builders/map.rs @@ -0,0 +1,184 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::any::Any; +use std::sync::Arc; + +use vortex_buffer::BufferAllocatorRef; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::ArrayRef; +use crate::Canonical; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::ListView; +use crate::arrays::Map; +use crate::arrays::MapArray; +use crate::arrays::map::MapArraySlotsExt; +use crate::builders::ArrayBuilder; +use crate::builders::DEFAULT_BUILDER_CAPACITY; +use crate::builders::ListViewBuilder; +use crate::dtype::DType; +use crate::dtype::MapDType; +use crate::dtype::Nullability; +use crate::dtype::OffsetBuilderPType; +use crate::scalar::MapScalar; +use crate::scalar::Scalar; + +/// A builder for canonical [`MapArray`] values. +/// +/// The builder owns a [`ListViewBuilder`] whose elements are non-nullable `{key, value}` structs. +/// It preserves the map dtype's `keys_sorted` assertion while delegating offsets, sizes, and outer +/// validity to that list-view builder. +pub struct MapBuilder { + dtype: DType, + map_dtype: MapDType, + entries_builder: ListViewBuilder, +} + +impl MapBuilder { + /// Creates a map builder with the default capacity. + #[deprecated(note = "use `new_in` with an explicit allocator")] + pub fn new(map_dtype: MapDType, nullability: Nullability) -> Self { + Self::new_in(map_dtype, nullability, BufferAllocatorRef::static_ref()) + } + + /// Creates a map builder with the default capacity using `allocator`. + pub fn new_in( + map_dtype: MapDType, + nullability: Nullability, + allocator: &BufferAllocatorRef, + ) -> Self { + Self::with_capacity_in(map_dtype, nullability, DEFAULT_BUILDER_CAPACITY, allocator) + } + + /// Creates a map builder with space for `capacity` map rows. + #[deprecated(note = "use `with_capacity_in` with an explicit allocator")] + pub fn with_capacity(map_dtype: MapDType, nullability: Nullability, capacity: usize) -> Self { + Self::with_capacity_in( + map_dtype, + nullability, + capacity, + BufferAllocatorRef::static_ref(), + ) + } + + /// Creates a map builder with space for `capacity` rows using `allocator`. + pub fn with_capacity_in( + map_dtype: MapDType, + nullability: Nullability, + capacity: usize, + allocator: &BufferAllocatorRef, + ) -> Self { + let entries_builder = ListViewBuilder::with_capacity_in( + Arc::new(map_dtype.entries_dtype()), + nullability, + capacity.saturating_mul(2), + capacity, + allocator, + ); + let dtype = DType::Map(map_dtype.clone(), nullability); + Self { + dtype, + map_dtype, + entries_builder, + } + } + + /// Appends one map scalar. + pub fn append_value(&mut self, value: MapScalar<'_>) -> VortexResult<()> { + vortex_ensure!( + value.dtype() == &self.dtype, + "MapBuilder expected map scalar with dtype {}, got {}", + self.dtype, + value.dtype() + ); + + if value.is_null() { + self.entries_builder.append_null(); + return Ok(()); + } + + let entry_dtype = self.map_dtype.entries_dtype(); + let entries = value + .entries() + .map(|(key, value)| Scalar::struct_(entry_dtype.clone(), vec![key, value])) + .collect(); + let entries = Scalar::list(Arc::new(entry_dtype), entries, self.dtype.nullability()); + self.entries_builder.append_value(entries.as_list()) + } + + /// Finishes the builder directly into a [`MapArray`]. + pub fn finish_into_map(&mut self) -> MapArray { + MapArray::new( + self.map_dtype.clone(), + self.entries_builder.finish_into_listview(), + ) + } + + /// Appends the values of a [`Map`]-encoded `array` to this builder. + pub fn append_map_array( + &mut self, + array: ArrayView<'_, Map>, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + vortex_ensure!( + array.dtype() == self.dtype(), + "MapBuilder expected map array with dtype {}, got {}", + self.dtype(), + array.dtype() + ); + self.entries_builder + .append_listview_array(array.entries().as_::(), ctx) + } +} + +impl ArrayBuilder for MapBuilder { + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn len(&self) -> usize { + self.entries_builder.len() + } + + fn append_zeros(&mut self, n: usize) { + self.entries_builder.append_zeros(n); + } + + unsafe fn append_nulls_unchecked(&mut self, n: usize) { + unsafe { self.entries_builder.append_nulls_unchecked(n) }; + } + + fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> { + vortex_ensure!( + scalar.dtype() == self.dtype(), + "MapBuilder expected scalar with dtype {}, got {}", + self.dtype(), + scalar.dtype() + ); + self.append_value(scalar.as_map()) + } + + fn reserve_exact(&mut self, additional: usize) { + self.entries_builder.reserve_exact(additional); + } + + fn finish(&mut self) -> ArrayRef { + self.finish_into_map().into_array() + } + + fn finish_into_canonical(&mut self, _ctx: &mut ExecutionCtx) -> Canonical { + Canonical::Map(self.finish_into_map()) + } +} diff --git a/vortex-array/src/builders/mod.rs b/vortex-array/src/builders/mod.rs index 561efc2711e..afebe74c225 100644 --- a/vortex-array/src/builders/mod.rs +++ b/vortex-array/src/builders/mod.rs @@ -6,15 +6,22 @@ //! Every logical type in Vortex has a canonical (uncompressed) in-memory encoding. This module //! provides pre-allocated builders to construct new canonical arrays. //! +//! Canonical form is not recursive, and neither are these builders: appending an array to a nested +//! builder keeps the child in the encoding it arrived in instead of decoding it. The fields of a +//! [`StructArray`](crate::arrays::StructArray), the elements of a list, and the storage of an +//! [`ExtensionArray`](crate::arrays::ExtensionArray) may therefore come back compressed, or as a +//! [`ChunkedArray`](crate::arrays::ChunkedArray) when several arrays were appended in turn. +//! //! ## Example: //! //! ``` -//! use vortex_array::builders::{builder_with_capacity, ArrayBuilder}; +//! use vortex_array::builders::{ArrayBuilder, builder_with_capacity_in}; //! use vortex_array::dtype::{DType, Nullability}; +//! use vortex_array::memory::BufferAllocatorRef; //! use vortex_array::{VortexSessionExecute, array_session}; //! //! // Create a new builder for string data. -//! let mut builder = builder_with_capacity(&DType::Utf8(Nullability::NonNullable), 4); +//! let mut builder = builder_with_capacity_in(&DType::Utf8(Nullability::NonNullable), 4, BufferAllocatorRef::static_ref()); //! //! builder.append_scalar(&"a".into()).unwrap(); //! builder.append_scalar(&"b".into()).unwrap(); @@ -34,47 +41,50 @@ use std::any::Any; use std::sync::Arc; use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; -use crate::array::ArrayView; -use crate::arrays::List; -use crate::arrays::ListView; use crate::canonical::Canonical; use crate::dtype::DType; use crate::match_each_decimal_value_type; use crate::match_each_native_ptype; -use crate::memory::HostAllocatorRef; +use crate::memory::BufferAllocatorRef; use crate::scalar::Scalar; mod lazy_null_builder; pub(crate) use lazy_null_builder::LazyBitBufferBuilder; mod bool; +mod child; mod decimal; pub mod dict; mod extension; mod fixed_size_list; mod list; mod listview; +mod map; mod null; mod primitive; mod struct_; +mod validity; mod varbinview; pub use bool::*; +pub(crate) use child::ChildBuilder; pub use decimal::*; pub use extension::*; pub use fixed_size_list::*; pub use list::*; pub use listview::*; +pub use map::*; pub use null::*; pub use primitive::*; pub use struct_::*; +pub(crate) use validity::ValidityBuilder; pub use varbinview::*; +pub use crate::arrays::varbin::builder::VarBinBuilder; + #[cfg(test)] mod tests; @@ -160,26 +170,11 @@ pub trait ArrayBuilder: Send { /// Allocate space for extra `additional` items fn reserve_exact(&mut self, additional: usize); - /// Override builders validity with the one provided. - /// - /// Note that this will have no effect on the final array if the array builder is non-nullable. - fn set_validity(&mut self, validity: Mask) { - if !self.dtype().is_nullable() { - return; - } - assert_eq!(self.len(), validity.len()); - unsafe { self.set_validity_unchecked(validity) } - } - - /// override validity with the one provided, without checking lengths - /// - /// # Safety - /// - /// Given validity must have an equal length to [`self.len()`](Self::len). - unsafe fn set_validity_unchecked(&mut self, validity: Mask); - /// Constructs an Array from the builder components. /// + /// The returned array is canonical at the top level only; its children keep whatever encoding + /// they were appended with. + /// /// # Panics /// /// This function may panic if the builder's methods are called with invalid arguments. If only @@ -195,53 +190,195 @@ pub trait ArrayBuilder: Send { /// then converts it to canonical form. Specific builders can override this with optimized /// implementations that avoid the intermediate [`ArrayRef`] creation. fn finish_into_canonical(&mut self, ctx: &mut ExecutionCtx) -> Canonical; +} - /// Appends the values of a [`List`]-encoded `array` to this builder. - /// - /// Only list-typed builders support this; the default implementation returns an error. List - /// encodings dispatch through this method because the concrete list builders are generic over - /// their offset/size integer types, which cannot be named through a `dyn ArrayBuilder`. - fn append_list_array( - &mut self, - array: ArrayView<'_, List>, - _ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - vortex_bail!( - "cannot append a List array of dtype {} to a {} builder", - array.dtype(), - self.dtype() - ) - } +/// Matches a `&mut dyn ArrayBuilder` against every concrete list builder type, i.e. every +/// [`ListBuilder`]`` and [`ListViewBuilder`]`` instantiation over the +/// [`OffsetBuilderPType`](crate::dtype::OffsetBuilderPType) offset/size types (`u32`, `u64`, `i32`, +/// `i64`). +/// +/// Binds the downcast builder as `$builder` and evaluates `$body` with it, yielding +/// `Some($body)`; yields `None` when the builder is not a list builder. List encodings dispatch +/// through this matcher because the concrete list builders are generic over their offset/size +/// integer types, which cannot be named through a `dyn ArrayBuilder`. The matcher is exhaustive +/// because `OffsetBuilderPType` is sealed, so no other instantiations can be constructed. +#[macro_export] +macro_rules! match_each_list_builder { + ($dyn_builder:expr, | $builder:ident | $body:expr) => {{ + let __dyn_builder: &mut dyn $crate::builders::ArrayBuilder = $dyn_builder; + match $crate::__match_each_list_builder!( + __dyn_builder, + $builder, + $body, + [u32, u64, i32, i64] + ) { + ::core::option::Option::Some(__result) => ::core::option::Option::Some(__result), + ::core::option::Option::None => $crate::__match_each_listview_builder!( + __dyn_builder, + $builder, + $body, + [u32, u64, i32, i64] + ), + } + }}; +} - /// Appends the values of a [`ListView`]-encoded `array` to this builder. - /// - /// See [`append_list_array`](Self::append_list_array); this is the same hook for the canonical - /// [`ListViewArray`](crate::arrays::ListViewArray) encoding. - fn append_listview_array( - &mut self, - array: ArrayView<'_, ListView>, - _ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - vortex_bail!( - "cannot append a ListView array of dtype {} to a {} builder", - array.dtype(), - self.dtype() - ) - } +/// Matches a `&mut dyn ArrayBuilder` against every concrete [`ListViewBuilder`]`` +/// instantiation over the [`OffsetBuilderPType`](crate::dtype::OffsetBuilderPType) offset/size +/// types (`u32`, `u64`, `i32`, `i64`), and only those. +/// +/// Binds the downcast builder as `$builder` and evaluates `$body` with it, yielding +/// `Some($body)`; yields `None` when the builder is not a list-view builder - including when it +/// is a [`ListBuilder`]. Callers reach for this instead of +/// [`match_each_list_builder!`](crate::match_each_list_builder) when the body needs methods only +/// a list-view builder has, such as +/// [`append_array_as_repeated_list`](ListViewBuilder::append_array_as_repeated_list). +#[macro_export] +macro_rules! match_each_listview_builder { + ($dyn_builder:expr, | $builder:ident | $body:expr) => {{ + let __dyn_builder: &mut dyn $crate::builders::ArrayBuilder = $dyn_builder; + $crate::__match_each_listview_builder!(__dyn_builder, $builder, $body, [u32, u64, i32, i64]) + }}; } -/// Construct a new canonical builder for the given [`DType`]. +#[doc(hidden)] +#[macro_export] +macro_rules! __match_each_list_builder { + ($target:ident, $builder:ident, $body:expr, []) => { + ::core::option::Option::None + }; + ($target:ident, $builder:ident, $body:expr, [$offset:ty $(, $rest:ty)*]) => { + if let ::core::option::Option::Some($builder) = + $crate::builders::ArrayBuilder::as_any_mut($target) + .downcast_mut::<$crate::builders::ListBuilder<$offset>>() + { + ::core::option::Option::Some($body) + } else { + $crate::__match_each_list_builder!($target, $builder, $body, [$($rest),*]) + } + }; +} + +#[doc(hidden)] +#[macro_export] +macro_rules! __match_each_listview_builder { + ($target:ident, $builder:ident, $body:expr, []) => { + ::core::option::Option::None + }; + ($target:ident, $builder:ident, $body:expr, [$offset:ty $(, $rest:ty)*]) => { + match $crate::__match_each_listview_builder_size!( + $target, + $builder, + $body, + $offset, + [u32, u64, i32, i64] + ) { + ::core::option::Option::Some(__result) => ::core::option::Option::Some(__result), + ::core::option::Option::None => $crate::__match_each_listview_builder!( + $target, + $builder, + $body, + [$($rest),*] + ), + } + }; +} + +#[doc(hidden)] +#[macro_export] +macro_rules! __match_each_listview_builder_size { + ($target:ident, $builder:ident, $body:expr, $offset:ty, []) => { + ::core::option::Option::None + }; + ($target:ident, $builder:ident, $body:expr, $offset:ty, [$size:ty $(, $rest:ty)*]) => { + if let ::core::option::Option::Some($builder) = + $crate::builders::ArrayBuilder::as_any_mut($target) + .downcast_mut::<$crate::builders::ListViewBuilder<$offset, $size>>() + { + ::core::option::Option::Some($body) + } else { + $crate::__match_each_listview_builder_size!( + $target, $builder, $body, $offset, [$($rest),*] + ) + } + }; +} + +/// Matches a `&mut dyn ArrayBuilder` against every concrete map builder type. /// +/// Binds the downcast builder as `$builder` and evaluates `$body` with it, yielding +/// `Some($body)`; yields `None` when the builder is not a map builder. +#[macro_export] +macro_rules! match_each_map_builder { + ($dyn_builder:expr, | $builder:ident | $body:expr) => {{ + let __dyn_builder: &mut dyn $crate::builders::ArrayBuilder = $dyn_builder; + $crate::__match_each_map_builder!(__dyn_builder, $builder, $body, [u32, u64, i32, i64]) + }}; +} + +#[doc(hidden)] +#[macro_export] +macro_rules! __match_each_map_builder { + ($target:ident, $builder:ident, $body:expr, []) => { + ::core::option::Option::None + }; + ($target:ident, $builder:ident, $body:expr, [$offset:ty $(, $rest:ty)*]) => { + match $crate::__match_each_map_builder_size!( + $target, + $builder, + $body, + $offset, + [u32, u64, i32, i64] + ) { + ::core::option::Option::Some(__result) => ::core::option::Option::Some(__result), + ::core::option::Option::None => $crate::__match_each_map_builder!( + $target, + $builder, + $body, + [$($rest),*] + ), + } + }; +} + +#[doc(hidden)] +#[macro_export] +macro_rules! __match_each_map_builder_size { + ($target:ident, $builder:ident, $body:expr, $offset:ty, []) => { + ::core::option::Option::None + }; + ($target:ident, $builder:ident, $body:expr, $offset:ty, [$size:ty $(, $rest:ty)*]) => { + if let ::core::option::Option::Some($builder) = + $crate::builders::ArrayBuilder::as_any_mut($target) + .downcast_mut::<$crate::builders::MapBuilder<$offset, $size>>() + { + ::core::option::Option::Some($body) + } else { + $crate::__match_each_map_builder_size!( + $target, $builder, $body, $offset, [$($rest),*] + ) + } + }; +} + +/// Construct a new canonical builder for the given [`DType`]. +#[deprecated(note = "use `builder_with_capacity_in` with an explicit allocator")] +pub fn builder_with_capacity(dtype: &DType, capacity: usize) -> Box { + builder_with_capacity_in(dtype, capacity, BufferAllocatorRef::static_ref()) +} + +/// Construct a new canonical builder using `allocator`. /// /// # Example /// /// ``` -/// use vortex_array::builders::{builder_with_capacity, ArrayBuilder}; +/// use vortex_array::builders::{ArrayBuilder, builder_with_capacity_in}; /// use vortex_array::dtype::{DType, Nullability}; +/// use vortex_array::memory::BufferAllocatorRef; /// use vortex_array::{VortexSessionExecute, array_session}; /// /// // Create a new builder for string data. -/// let mut builder = builder_with_capacity(&DType::Utf8(Nullability::NonNullable), 4); +/// let mut builder = builder_with_capacity_in(&DType::Utf8(Nullability::NonNullable), 4, BufferAllocatorRef::static_ref()); /// /// builder.append_scalar(&"a".into()).unwrap(); /// builder.append_scalar(&"b".into()).unwrap(); @@ -256,68 +393,80 @@ pub trait ArrayBuilder: Send { /// assert_eq!(strings.execute_scalar(2, &mut ctx).unwrap(), "c".into()); /// assert_eq!(strings.execute_scalar(3, &mut ctx).unwrap(), "d".into()); /// ``` -pub fn builder_with_capacity(dtype: &DType, capacity: usize) -> Box { +pub fn builder_with_capacity_in( + dtype: &DType, + capacity: usize, + allocator: &BufferAllocatorRef, +) -> Box { match dtype { DType::Null => Box::new(NullBuilder::new()), - DType::Bool(n) => Box::new(BoolBuilder::with_capacity(*n, capacity)), + DType::Bool(n) => Box::new(BoolBuilder::with_capacity_in(*n, capacity, allocator)), DType::Primitive(ptype, n) => { match_each_native_ptype!(ptype, |P| { - Box::new(PrimitiveBuilder::

::with_capacity(*n, capacity)) + Box::new(PrimitiveBuilder::

::with_capacity_in( + *n, capacity, allocator, + )) }) } DType::Decimal(decimal_type, n) => { match_each_decimal_value_type!( DecimalType::smallest_decimal_value_type(decimal_type), |D| { - Box::new(DecimalBuilder::with_capacity::( + Box::new(DecimalBuilder::with_capacity_in::( capacity, *decimal_type, *n, + allocator, )) } ) } - DType::Utf8(n) => Box::new(VarBinViewBuilder::with_capacity(DType::Utf8(*n), capacity)), - DType::Binary(n) => Box::new(VarBinViewBuilder::with_capacity( + DType::Utf8(n) => Box::new(VarBinViewBuilder::with_capacity_in( + DType::Utf8(*n), + capacity, + allocator.clone(), + )), + DType::Binary(n) => Box::new(VarBinViewBuilder::with_capacity_in( DType::Binary(*n), capacity, + allocator.clone(), )), - DType::List(dtype, n) => Box::new(ListViewBuilder::::with_capacity( + DType::List(dtype, n) => Box::new(ListViewBuilder::::with_capacity_in( Arc::clone(dtype), *n, 2 * capacity, // Arbitrarily choose 2 times the `offsets` capacity here. capacity, + allocator, + )), + DType::Map(map_dtype, nullability) => Box::new(MapBuilder::::with_capacity_in( + map_dtype.clone(), + *nullability, + capacity, + allocator, )), DType::FixedSizeList(elem_dtype, list_size, null) => { - Box::new(FixedSizeListBuilder::with_capacity( + Box::new(FixedSizeListBuilder::with_capacity_in( Arc::clone(elem_dtype), *list_size, *null, capacity, + allocator, )) } - DType::Struct(struct_dtype, n) => Box::new(StructBuilder::with_capacity( + DType::Struct(struct_dtype, n) => Box::new(StructBuilder::with_capacity_in( struct_dtype.clone(), *n, capacity, + allocator, )), DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), DType::Variant(_) => { unimplemented!() } - DType::Extension(ext_dtype) => { - Box::new(ExtensionBuilder::with_capacity(ext_dtype.clone(), capacity)) - } + DType::Extension(ext_dtype) => Box::new(ExtensionBuilder::with_capacity_in( + ext_dtype.clone(), + capacity, + allocator, + )), } } - -/// Construct a new canonical builder for the given [`DType`] using a host -/// [`crate::memory::HostAllocator`]. -pub fn builder_with_capacity_in( - allocator: HostAllocatorRef, - dtype: &DType, - capacity: usize, -) -> Box { - let _allocator = allocator; - builder_with_capacity(dtype, capacity) -} diff --git a/vortex-array/src/builders/null.rs b/vortex-array/src/builders/null.rs index 28eab14dd7a..541a27c2c89 100644 --- a/vortex-array/src/builders/null.rs +++ b/vortex-array/src/builders/null.rs @@ -5,7 +5,6 @@ use std::any::Any; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; @@ -72,8 +71,6 @@ impl ArrayBuilder for NullBuilder { fn reserve_exact(&mut self, _additional: usize) {} - unsafe fn set_validity_unchecked(&mut self, _validity: Mask) {} - fn finish(&mut self) -> ArrayRef { NullArray::new(self.length).into_array() } diff --git a/vortex-array/src/builders/primitive.rs b/vortex-array/src/builders/primitive.rs index 4ce1aafe1f7..647854176c5 100644 --- a/vortex-array/src/builders/primitive.rs +++ b/vortex-array/src/builders/primitive.rs @@ -4,6 +4,7 @@ use std::any::Any; use std::mem::MaybeUninit; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -32,15 +33,31 @@ pub struct PrimitiveBuilder { impl PrimitiveBuilder { /// Creates a new `PrimitiveBuilder` with a capacity of [`DEFAULT_BUILDER_CAPACITY`]. + #[deprecated(note = "use `new_in` with an explicit allocator")] pub fn new(nullability: Nullability) -> Self { - Self::with_capacity(nullability, DEFAULT_BUILDER_CAPACITY) + Self::new_in(nullability, BufferAllocatorRef::static_ref()) + } + + /// Creates a new `PrimitiveBuilder` with the default capacity using `allocator`. + pub fn new_in(nullability: Nullability, allocator: &BufferAllocatorRef) -> Self { + Self::with_capacity_in(nullability, DEFAULT_BUILDER_CAPACITY, allocator) } /// Creates a new `PrimitiveBuilder` with the given `capacity`. + #[deprecated(note = "use `with_capacity_in` with an explicit allocator")] pub fn with_capacity(nullability: Nullability, capacity: usize) -> Self { + Self::with_capacity_in(nullability, capacity, BufferAllocatorRef::static_ref()) + } + + /// Creates a new `PrimitiveBuilder` with `capacity` using `allocator`. + pub fn with_capacity_in( + nullability: Nullability, + capacity: usize, + allocator: &BufferAllocatorRef, + ) -> Self { Self { - values: BufferMut::with_capacity(capacity), - nulls: LazyBitBufferBuilder::new(capacity), + values: BufferMut::with_capacity_in(capacity, allocator.clone()), + nulls: LazyBitBufferBuilder::new(capacity, allocator.clone()), dtype: DType::Primitive(T::PTYPE, nullability), } } @@ -83,10 +100,11 @@ impl PrimitiveBuilder { /// use std::mem::MaybeUninit; /// use vortex_array::builders::{ArrayBuilder, PrimitiveBuilder}; /// use vortex_array::dtype::Nullability; + /// use vortex_array::memory::BufferAllocatorRef; /// /// // Create a new builder. /// let mut builder: PrimitiveBuilder = - /// PrimitiveBuilder::with_capacity(Nullability::NonNullable, 5); + /// PrimitiveBuilder::with_capacity_in(Nullability::NonNullable, 5, BufferAllocatorRef::static_ref()); /// /// // Populate the values. /// let mut uninit_range = builder.uninit_range(5); @@ -120,7 +138,9 @@ impl PrimitiveBuilder { .nulls .finish_with_nullability(self.dtype().nullability()); - PrimitiveArray::new(std::mem::take(&mut self.values).freeze(), validity) + let allocator = self.values.allocator().clone(); + let values = std::mem::replace(&mut self.values, allocator.with_capacity(0)).freeze(); + PrimitiveArray::new(values, validity) } /// Extends the primitive array with an iterator. @@ -201,10 +221,6 @@ impl ArrayBuilder for PrimitiveBuilder { self.nulls.reserve_exact(additional); } - unsafe fn set_validity_unchecked(&mut self, validity: Mask) { - self.nulls = LazyBitBufferBuilder::from_validity_mask(validity); - } - fn finish(&mut self) -> ArrayRef { self.finish_into_primitive().into_array() } @@ -375,7 +391,11 @@ mod tests { #[test] fn test_multiple_uninit_ranges_correct_offsets() { let mut ctx = array_session().create_execution_ctx(); - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 10); + let mut builder = PrimitiveBuilder::::with_capacity_in( + Nullability::NonNullable, + 10, + BufferAllocatorRef::static_ref(), + ); // First range. let mut range1 = builder.uninit_range(3); @@ -418,7 +438,11 @@ mod tests { /// This test ensures the new API works correctly. #[test] fn test_append_mask_on_uninit_range() { - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::Nullable, 5); + let mut builder = PrimitiveBuilder::::with_capacity_in( + Nullability::Nullable, + 5, + BufferAllocatorRef::static_ref(), + ); let mut range = builder.uninit_range(3); // Create a mask for 3 values. @@ -468,7 +492,11 @@ mod tests { expected = "Tried to append a mask to an `UninitRange` that was beyond the allowed range" )] fn test_append_mask_wrong_length_panics() { - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::Nullable, 10); + let mut builder = PrimitiveBuilder::::with_capacity_in( + Nullability::Nullable, + 10, + BufferAllocatorRef::static_ref(), + ); let mut range = builder.uninit_range(5); // Try to append a mask with wrong length (3 instead of 5). @@ -486,7 +514,11 @@ mod tests { #[test] fn test_copy_from_slice_with_offsets() { let mut ctx = array_session().create_execution_ctx(); - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 10); + let mut builder = PrimitiveBuilder::::with_capacity_in( + Nullability::NonNullable, + 10, + BufferAllocatorRef::static_ref(), + ); let mut range = builder.uninit_range(6); // Copy to different offsets. @@ -514,7 +546,11 @@ mod tests { /// modify individual bits with relative indexing. #[test] fn test_set_bit_relative_indexing() { - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::Nullable, 10); + let mut builder = PrimitiveBuilder::::with_capacity_in( + Nullability::Nullable, + 10, + BufferAllocatorRef::static_ref(), + ); // First add some values to the builder. builder.append_value(100); @@ -588,7 +624,10 @@ mod tests { #[test] #[should_panic(expected = "cannot create an uninit range of length 0")] fn test_zero_length_uninit_range_panics() { - let mut builder = PrimitiveBuilder::::new(Nullability::NonNullable); + let mut builder = PrimitiveBuilder::::new_in( + Nullability::NonNullable, + BufferAllocatorRef::static_ref(), + ); let _range = builder.uninit_range(0); } @@ -596,7 +635,11 @@ mod tests { #[test] #[should_panic(expected = "uninit_range of len 261 exceeds builder with length 0 and capacity")] fn test_uninit_range_exceeds_capacity_panics() { - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 5); + let mut builder = PrimitiveBuilder::::with_capacity_in( + Nullability::NonNullable, + 5, + BufferAllocatorRef::static_ref(), + ); let _range = builder.uninit_range(261); } @@ -607,7 +650,11 @@ mod tests { #[cfg(debug_assertions)] #[should_panic(expected = "tried to copy a slice into a `UninitRange` past its boundary")] fn test_copy_from_slice_out_of_bounds() { - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 10); + let mut builder = PrimitiveBuilder::::with_capacity_in( + Nullability::NonNullable, + 10, + BufferAllocatorRef::static_ref(), + ); let mut range = builder.uninit_range(3); // Try to copy 3 elements starting at offset 1 (would need 4 slots total). @@ -619,7 +666,11 @@ mod tests { /// This test demonstrates proper usage of the unsafe `finish` method. #[test] fn test_finish_unsafe_contract() { - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::Nullable, 5); + let mut builder = PrimitiveBuilder::::with_capacity_in( + Nullability::Nullable, + 5, + BufferAllocatorRef::static_ref(), + ); let mut range = builder.uninit_range(3); // Set validity mask. @@ -647,7 +698,11 @@ mod tests { use crate::dtype::DType; use crate::scalar::Scalar; - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::Nullable, 10); + let mut builder = PrimitiveBuilder::::with_capacity_in( + Nullability::Nullable, + 10, + BufferAllocatorRef::static_ref(), + ); // Test appending a valid primitive value. let scalar1 = Scalar::primitive(42i32, Nullability::Nullable); @@ -698,7 +753,11 @@ mod tests { ); // Test wrong dtype error. - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 10); + let mut builder = PrimitiveBuilder::::with_capacity_in( + Nullability::NonNullable, + 10, + BufferAllocatorRef::static_ref(), + ); let wrong_scalar = Scalar::from(true); assert!(builder.append_scalar(&wrong_scalar).is_err()); } diff --git a/vortex-array/src/builders/struct_.rs b/vortex-array/src/builders/struct_.rs index da08e262760..21f4361ee36 100644 --- a/vortex-array/src/builders/struct_.rs +++ b/vortex-array/src/builders/struct_.rs @@ -4,12 +4,12 @@ use std::any::Any; use itertools::Itertools; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; -use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; @@ -17,9 +17,9 @@ use crate::IntoArray; use crate::arrays::StructArray; use crate::arrays::struct_::StructArrayExt; use crate::builders::ArrayBuilder; +use crate::builders::ChildBuilder; use crate::builders::DEFAULT_BUILDER_CAPACITY; -use crate::builders::LazyBitBufferBuilder; -use crate::builders::builder_with_capacity; +use crate::builders::ValidityBuilder; use crate::canonical::Canonical; use crate::dtype::DType; use crate::dtype::Nullability; @@ -30,30 +30,61 @@ use crate::scalar::StructScalar; /// The builder for building a [`StructArray`]. pub struct StructBuilder { dtype: DType, - builders: Vec>, - nulls: LazyBitBufferBuilder, + builders: Vec, + nulls: ValidityBuilder, } impl StructBuilder { /// Creates a new `StructBuilder` with a capacity of [`DEFAULT_BUILDER_CAPACITY`]. + #[deprecated(note = "use `new_in` with an explicit allocator")] pub fn new(struct_dtype: StructFields, nullability: Nullability) -> Self { - Self::with_capacity(struct_dtype, nullability, DEFAULT_BUILDER_CAPACITY) + Self::new_in(struct_dtype, nullability, BufferAllocatorRef::static_ref()) + } + + /// Creates a new `StructBuilder` with the default capacity using `allocator`. + pub fn new_in( + struct_dtype: StructFields, + nullability: Nullability, + allocator: &BufferAllocatorRef, + ) -> Self { + Self::with_capacity_in( + struct_dtype, + nullability, + DEFAULT_BUILDER_CAPACITY, + allocator, + ) } /// Creates a new `StructBuilder` with the given `capacity`. + #[deprecated(note = "use `with_capacity_in` with an explicit allocator")] pub fn with_capacity( struct_dtype: StructFields, nullability: Nullability, capacity: usize, + ) -> Self { + Self::with_capacity_in( + struct_dtype, + nullability, + capacity, + BufferAllocatorRef::static_ref(), + ) + } + + /// Creates a new `StructBuilder` with `capacity` using `allocator`. + pub fn with_capacity_in( + struct_dtype: StructFields, + nullability: Nullability, + capacity: usize, + allocator: &BufferAllocatorRef, ) -> Self { let builders = struct_dtype .fields() - .map(|dt| builder_with_capacity(&dt, capacity)) + .map(|dt| ChildBuilder::with_capacity(&dt, capacity, allocator)) .collect(); Self { builders, - nulls: LazyBitBufferBuilder::new(capacity), + nulls: ValidityBuilder::new(capacity, allocator), dtype: DType::Struct(struct_dtype, nullability), } } @@ -131,11 +162,10 @@ impl StructBuilder { .iter_unmasked_fields() .zip_eq(self.builders.iter_mut()) { - field.append_to_builder(builder.as_mut(), ctx)?; + builder.append_array(field, ctx)?; } - self.nulls - .append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?); + self.nulls.append_validity(array.validity()?, array.len()); Ok(()) } } @@ -191,10 +221,6 @@ impl ArrayBuilder for StructBuilder { self.nulls.reserve_exact(capacity); } - unsafe fn set_validity_unchecked(&mut self, validity: Mask) { - self.nulls = LazyBitBufferBuilder::from_validity_mask(validity); - } - fn finish(&mut self) -> ArrayRef { self.finish_into_struct().into_array() } @@ -206,6 +232,8 @@ impl ArrayBuilder for StructBuilder { #[cfg(test)] mod tests { + use vortex_buffer::BufferAllocatorRef; + use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; @@ -226,7 +254,12 @@ mod tests { fn test_struct_builder() { let sdt = StructFields::new(["a", "b"].into(), vec![I32.into(), I32.into()]); let dtype = DType::Struct(sdt.clone(), Nullability::NonNullable); - let mut builder = StructBuilder::with_capacity(sdt, Nullability::NonNullable, 0); + let mut builder = StructBuilder::with_capacity_in( + sdt, + Nullability::NonNullable, + 0, + BufferAllocatorRef::static_ref(), + ); builder .append_value(Scalar::struct_(dtype.clone(), vec![1.into(), 2.into()]).as_struct()) @@ -241,7 +274,12 @@ mod tests { fn test_append_nullable_struct() { let sdt = StructFields::new(["a", "b"].into(), vec![I32.into(), I32.into()]); let dtype = DType::Struct(sdt.clone(), Nullability::Nullable); - let mut builder = StructBuilder::with_capacity(sdt, Nullability::Nullable, 0); + let mut builder = StructBuilder::with_capacity_in( + sdt, + Nullability::Nullable, + 0, + BufferAllocatorRef::static_ref(), + ); builder .append_value(Scalar::struct_(dtype.clone(), vec![1.into(), 2.into()]).as_struct()) @@ -277,7 +315,11 @@ mod tests { DType::Struct(fields, _) => fields.clone(), _ => panic!("Expected struct dtype"), }; - let mut builder = StructBuilder::new(struct_fields, Nullability::Nullable); + let mut builder = StructBuilder::new_in( + struct_fields, + Nullability::Nullable, + BufferAllocatorRef::static_ref(), + ); // Test appending a valid struct value. let struct_scalar1 = Scalar::struct_( @@ -332,7 +374,11 @@ mod tests { DType::Struct(fields, _) => fields.clone(), _ => panic!("Expected struct dtype"), }; - let mut builder = StructBuilder::new(struct_fields, Nullability::NonNullable); + let mut builder = StructBuilder::new_in( + struct_fields, + Nullability::NonNullable, + BufferAllocatorRef::static_ref(), + ); let wrong_scalar = Scalar::from(42i32); assert!(builder.append_scalar(&wrong_scalar).is_err()); } diff --git a/vortex-array/src/builders/tests.rs b/vortex-array/src/builders/tests.rs index 138e6941def..5b7977a0ef0 100644 --- a/vortex-array/src/builders/tests.rs +++ b/vortex-array/src/builders/tests.rs @@ -1,29 +1,171 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::alloc::Layout; +use std::ptr::NonNull; use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use allocator_api2::alloc::AllocError; +use allocator_api2::alloc::Allocator; +use allocator_api2::alloc::Global; use rstest::rstest; +use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; +use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_mask::Mask; +use crate::ArrayRef; use crate::Canonical; use crate::ExecutionCtx; +use crate::RecursiveCanonical; use crate::VortexSessionExecute; use crate::array::IntoArray; use crate::array_session; +use crate::arrays::Chunked; +use crate::arrays::ChunkedArray; +use crate::arrays::Constant; +use crate::arrays::ConstantArray; +use crate::arrays::Extension; +use crate::arrays::ExtensionArray; +use crate::arrays::FixedSizeList; +use crate::arrays::FixedSizeListArray; +use crate::arrays::List; +use crate::arrays::ListArray; +use crate::arrays::ListView; +use crate::arrays::ListViewArray; +use crate::arrays::Primitive; +use crate::arrays::PrimitiveArray; +use crate::arrays::Struct; +use crate::arrays::StructArray; +use crate::arrays::chunked::ChunkedArrayExt; +use crate::arrays::extension::ExtensionArraySlotsExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use crate::arrays::list::ListArraySlotsExt; +use crate::arrays::listview::ListViewArraySlotsExt; +use crate::arrays::struct_::StructArrayExt; +use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; -use crate::builders::builder_with_capacity; +use crate::builders::ListBuilder; +use crate::builders::builder_with_capacity_in; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; +use crate::dtype::extension::ExtDTypeRef; use crate::dtype::half::f16; use crate::extension::datetime::TimeUnit; use crate::extension::datetime::Timestamp; use crate::scalar::Scalar; +use crate::validity::Validity; + +#[derive(Debug)] +struct CountingAllocator { + allocations: Arc, +} + +// SAFETY: this forwards memory operations to Global and only counts allocations. +unsafe impl Allocator for CountingAllocator { + fn allocate(&self, layout: Layout) -> Result, AllocError> { + self.allocations.fetch_add(1, Ordering::Relaxed); + Global.allocate(layout) + } + + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + // SAFETY: ptr and layout came from Global. + unsafe { Global.deallocate(ptr, layout) } + } +} + +#[test] +#[allow(deprecated, unused_qualifications)] +fn deprecated_builder_constructors_keep_their_signatures() { + let _: fn(Nullability) -> super::BoolBuilder = super::BoolBuilder::new; + let _: fn(Nullability, usize) -> super::BoolBuilder = super::BoolBuilder::with_capacity; + let _: fn(Nullability) -> super::PrimitiveBuilder = super::PrimitiveBuilder::::new; + let _: fn(Nullability, usize) -> super::PrimitiveBuilder = + super::PrimitiveBuilder::::with_capacity; + let _: fn(DecimalDType, Nullability) -> super::DecimalBuilder = + super::DecimalBuilder::new::; + let _: fn(usize, DecimalDType, Nullability) -> super::DecimalBuilder = + super::DecimalBuilder::with_capacity::; + let _: fn(ExtDTypeRef) -> super::ExtensionBuilder = super::ExtensionBuilder::new; + let _: fn(ExtDTypeRef, usize) -> super::ExtensionBuilder = + super::ExtensionBuilder::with_capacity; + let _: fn(Arc, u32, Nullability) -> super::FixedSizeListBuilder = + super::FixedSizeListBuilder::new; + let _: fn(Arc, u32, Nullability, usize) -> super::FixedSizeListBuilder = + super::FixedSizeListBuilder::with_capacity; + let _: fn(Arc, Nullability) -> super::ListBuilder = super::ListBuilder::::new; + let _: fn(Arc, Nullability, usize, usize) -> super::ListBuilder = + super::ListBuilder::::with_capacity; + let _: fn(Arc, Nullability) -> super::ListViewBuilder = + super::ListViewBuilder::::new; + let _: fn(Arc, Nullability, usize, usize) -> super::ListViewBuilder = + super::ListViewBuilder::::with_capacity; + let _: fn(crate::dtype::MapDType, Nullability) -> super::MapBuilder = + super::MapBuilder::::new; + let _: fn(crate::dtype::MapDType, Nullability, usize) -> super::MapBuilder = + super::MapBuilder::::with_capacity; + let _: fn(StructFields, Nullability) -> super::StructBuilder = super::StructBuilder::new; + let _: fn(StructFields, Nullability, usize) -> super::StructBuilder = + super::StructBuilder::with_capacity; + let _: fn(DType) -> super::VarBinBuilder = super::VarBinBuilder::::new; + let _: fn(DType, usize) -> super::VarBinBuilder = + super::VarBinBuilder::::with_capacity; + let _: fn(DType, usize, usize) -> super::VarBinBuilder = + super::VarBinBuilder::::with_capacity_bytes; + let _: fn(DType, usize) -> super::VarBinViewBuilder = super::VarBinViewBuilder::with_capacity; + let _: fn(DType, usize) -> super::VarBinViewBuilder = + super::VarBinViewBuilder::with_buffer_deduplication; + let _: fn(DType, usize, f64) -> super::VarBinViewBuilder = + super::VarBinViewBuilder::with_compaction; + let _: fn( + DType, + usize, + super::CompletedBuffers, + super::BufferGrowthStrategy, + f64, + ) -> super::VarBinViewBuilder = super::VarBinViewBuilder::new; + let _: fn(&DType, usize) -> Box = super::builder_with_capacity; + let _: fn(&DType, usize, &BufferAllocatorRef) -> Box = + super::builder_with_capacity_in; +} + +#[test] +fn builder_reuses_its_allocator_after_finish() { + let allocations = Arc::new(AtomicUsize::new(0)); + let allocator = BufferAllocatorRef::new(CountingAllocator { + allocations: Arc::clone(&allocations), + }); + let dtype = DType::Struct( + StructFields::from_iter([ + ( + "number", + DType::Primitive(PType::I32, Nullability::Nullable), + ), + ("text", DType::Utf8(Nullability::Nullable)), + ]), + Nullability::Nullable, + ); + let mut builder = builder_with_capacity_in(&dtype, 1, &allocator); + + builder.append_null(); + drop(builder.finish()); + let first = allocations.load(Ordering::Relaxed); + + builder.append_null(); + drop(builder.finish()); + let second = allocations.load(Ordering::Relaxed); + + assert!( + second - first >= 5, + "all nested buffers must reuse the allocator" + ); +} /// Test that `append_zeros` produces the same result as manually appending `Scalar::default_value`. /// @@ -79,12 +221,14 @@ fn test_append_zeros_matches_default_value(#[case] dtype: DType) { let num_elements = 5; // Builder 1: Use append_zeros. - let mut builder_zeros = builder_with_capacity(&dtype, num_elements); + let mut builder_zeros = + builder_with_capacity_in(&dtype, num_elements, BufferAllocatorRef::static_ref()); builder_zeros.append_zeros(num_elements); let array_zeros = builder_zeros.finish(); // Builder 2: Manually append default values. - let mut builder_manual = builder_with_capacity(&dtype, num_elements); + let mut builder_manual = + builder_with_capacity_in(&dtype, num_elements, BufferAllocatorRef::static_ref()); let default_scalar = Scalar::zero_value(&dtype); for _ in 0..num_elements { builder_manual.append_scalar(&default_scalar).unwrap(); @@ -179,7 +323,7 @@ fn test_append_zeros_matches_default_value(#[case] dtype: DType) { )] #[should_panic(expected = "non-nullable")] fn test_append_nulls_panics_on_non_nullable(#[case] dtype: DType, #[case] count: usize) { - let mut builder = builder_with_capacity(&dtype, count); + let mut builder = builder_with_capacity_in(&dtype, count, BufferAllocatorRef::static_ref()); builder.append_nulls(count); } @@ -192,7 +336,7 @@ fn test_append_nulls_panics_on_non_nullable(#[case] dtype: DType, #[case] count: #[case::nullable_utf8(DType::Utf8(Nullability::Nullable), true)] #[case::non_nullable_utf8(DType::Utf8(Nullability::NonNullable), false)] fn test_append_defaults_behavior(#[case] dtype: DType, #[case] should_be_null: bool) { - let mut builder = builder_with_capacity(&dtype, 3); + let mut builder = builder_with_capacity_in(&dtype, 3, BufferAllocatorRef::static_ref()); builder.append_defaults(3); let array = builder.finish(); @@ -231,8 +375,8 @@ where F: FnMut(&mut dyn ArrayBuilder), { // Create two identical builders. - let mut builder1 = builder_with_capacity(dtype, 10); - let mut builder2 = builder_with_capacity(dtype, 10); + let mut builder1 = builder_with_capacity_in(dtype, 10, BufferAllocatorRef::static_ref()); + let mut builder2 = builder_with_capacity_in(dtype, 10, BufferAllocatorRef::static_ref()); // Fill both builders with the same data. fill_builder(builder1.as_mut()); @@ -526,7 +670,8 @@ fn test_to_canonical_f32() { ))] fn test_append_scalar_comprehensive(#[case] dtype: DType) { let num_elements = 3; - let mut builder = builder_with_capacity(&dtype, num_elements * 2); + let mut builder = + builder_with_capacity_in(&dtype, num_elements * 2, BufferAllocatorRef::static_ref()); // Create test scalars based on the dtype. let scalars = create_test_scalars_for_dtype(&dtype, num_elements); @@ -631,6 +776,9 @@ fn create_test_scalars_for_dtype(dtype: &DType, count: usize) -> Vec { .collect(); Scalar::fixed_size_list(Arc::clone(element_dtype), elements, *n) } + DType::Map(..) => { + panic!("map builders are not supported until MapArray exists") + } DType::Struct(fields, n) => { // Create struct scalars with field values. let field_values: Vec = fields @@ -700,7 +848,7 @@ fn assert_scalars_equal(actual: &Scalar, expected: &Scalar, dtype: &DType, index #[case::utf8(DType::Utf8(Nullability::Nullable))] #[case::binary(DType::Binary(Nullability::Nullable))] fn test_append_scalar_mixed_nulls(#[case] dtype: DType) { - let mut builder = builder_with_capacity(&dtype, 6); + let mut builder = builder_with_capacity_in(&dtype, 6, BufferAllocatorRef::static_ref()); // Create a pattern of valid, null, valid, null, valid. let test_scalars = create_test_scalars_for_dtype(&dtype, 3); @@ -778,7 +926,11 @@ fn test_append_scalar_mixed_nulls(#[case] dtype: DType) { #[test] fn test_append_scalar_wrong_dtype_rejection() { // Test bool builder rejecting i32 scalar. - let mut bool_builder = builder_with_capacity(&DType::Bool(Nullability::NonNullable), 1); + let mut bool_builder = builder_with_capacity_in( + &DType::Bool(Nullability::NonNullable), + 1, + BufferAllocatorRef::static_ref(), + ); let i32_scalar = Scalar::from(42i32); assert!( bool_builder.append_scalar(&i32_scalar).is_err(), @@ -786,8 +938,11 @@ fn test_append_scalar_wrong_dtype_rejection() { ); // Test i32 builder rejecting string scalar. - let mut i32_builder = - builder_with_capacity(&DType::Primitive(PType::I32, Nullability::NonNullable), 1); + let mut i32_builder = builder_with_capacity_in( + &DType::Primitive(PType::I32, Nullability::NonNullable), + 1, + BufferAllocatorRef::static_ref(), + ); let string_scalar = Scalar::utf8("test", Nullability::NonNullable); assert!( i32_builder.append_scalar(&string_scalar).is_err(), @@ -795,7 +950,11 @@ fn test_append_scalar_wrong_dtype_rejection() { ); // Test string builder rejecting binary scalar. - let mut string_builder = builder_with_capacity(&DType::Utf8(Nullability::NonNullable), 1); + let mut string_builder = builder_with_capacity_in( + &DType::Utf8(Nullability::NonNullable), + 1, + BufferAllocatorRef::static_ref(), + ); let binary_scalar = Scalar::binary(vec![0u8, 1, 2], Nullability::NonNullable); assert!( string_builder.append_scalar(&binary_scalar).is_err(), @@ -808,7 +967,7 @@ fn test_append_scalar_wrong_dtype_rejection() { #[test] fn test_append_scalar_repeated_same_instance() { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut builder = builder_with_capacity(&dtype, 5); + let mut builder = builder_with_capacity_in(&dtype, 5, BufferAllocatorRef::static_ref()); let scalar = Scalar::primitive(42i32, Nullability::NonNullable); @@ -834,66 +993,368 @@ fn test_append_scalar_repeated_same_instance() { } } -/// Test that `set_validity` correctly overrides a builder's validity across all mask variants. +/// Builders only promise a canonical *top level*, so a child array that is long enough to be worth +/// a chunk must come back out of the builder in the encoding it went in with. /// -/// `set_validity` moves the mask's buffer into the builder rather than copying it, so the -/// `sliced_offset` case is important: slicing a `Mask::Values` at a non-byte-aligned boundary -/// yields a buffer with a non-zero bit offset, which the move path must preserve. +/// Each case appends the same array twice, which additionally checks that the two chunks are +/// stitched back together into a [`ChunkedArray`] rather than being decoded and concatenated. #[rstest] -#[case::all_true(Mask::new_true(8), vec![true; 8])] -#[case::all_false(Mask::new_false(8), vec![false; 8])] -#[case::values( - Mask::from_iter([true, false, true, true, false, false, true, false]), - vec![true, false, true, true, false, false, true, false] +#[case::struct_field( + StructArray::try_from_iter([("a", constant_i32())]) + .vortex_expect("struct array") + .into_array(), + |array: &ArrayRef| array.as_::().unmasked_field(0).clone() )] -#[case::sliced_offset( - Mask::from_iter([ - false, false, false, // dropped by the slice - true, false, true, true, false, false, true, false, // kept: indices 3..11 - true, true, true, true, true, // dropped by the slice - ]) - .slice(3..11), - vec![true, false, true, true, false, false, true, false] +#[case::list_elements( + ListViewArray::new( + constant_i32(), + (0..CHUNK_LEN as u64).collect::>().into_array(), + Buffer::full(1u64, CHUNK_LEN).into_array(), + Validity::NonNullable, + ) + .into_array(), + |array: &ArrayRef| array.as_::().elements().clone() )] -fn test_set_validity_overrides_validity( - #[case] mask: Mask, - #[case] expected: Vec, +#[case::fixed_size_list_elements( + FixedSizeListArray::new( + constant_i32(), + 2, + Validity::NonNullable, + CHUNK_LEN / 2, + ) + .into_array(), + |array: &ArrayRef| array.as_::().elements().clone() +)] +#[case::extension_storage( + ExtensionArray::new( + Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased(), + ConstantArray::new( + Scalar::primitive(0i64, Nullability::NonNullable), + CHUNK_LEN, + ) + .into_array(), + ) + .into_array(), + |array: &ArrayRef| array.as_::().storage().clone() +)] +fn test_children_are_not_canonicalized( + #[case] array: ArrayRef, + #[case] child_of: fn(&ArrayRef) -> ArrayRef, ) -> VortexResult<()> { - let dtype = DType::Primitive(PType::I32, Nullability::Nullable); - let mut builder = builder_with_capacity(&dtype, mask.len()); - builder.append_zeros(mask.len()); + let mut ctx = array_session().create_execution_ctx(); - builder.set_validity(mask); + let mut builder = builder_with_capacity_in(array.dtype(), 0, BufferAllocatorRef::static_ref()); + array.append_to_builder(builder.as_mut(), &mut ctx)?; + array.append_to_builder(builder.as_mut(), &mut ctx)?; + let built = builder.finish(); - let validity = builder.finish().validity()?; + let child = child_of(&built); + let chunked = child.as_::(); + assert_eq!( + chunked.nchunks(), + 2, + "expected one chunk per appended array" + ); + assert!( + chunked.iter_chunks().all(|chunk| chunk.is::()), + "the constant-encoded child was decoded by the builder", + ); + + let expected = ChunkedArray::try_new(vec![array.clone(), array], built.dtype().clone())?; + assert_arrays_eq!(&built, &expected, &mut ctx); + + Ok(()) +} + +/// A child is chunked on the boundaries it is appended on, however small the appends. Appending +/// scalars instead is what asks the builder to copy the values into one canonical child. +#[test] +fn test_children_are_chunked_on_the_boundaries_they_are_appended_on() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - for (i, &valid) in expected.iter().enumerate() { - assert_eq!( - validity.execute_is_valid(i, &mut ctx)?, - valid, - "validity mismatch at index {i}" - ); + + let elements = ConstantArray::new(1i32, 2).into_array(); + let array = FixedSizeListArray::new(elements, 2, Validity::NonNullable, 1).into_array(); + + let mut builder = builder_with_capacity_in(array.dtype(), 0, BufferAllocatorRef::static_ref()); + for _ in 0..CHUNK_LEN { + array.append_to_builder(builder.as_mut(), &mut ctx)?; + } + let built = builder.finish(); + + assert_eq!(built.len(), CHUNK_LEN); + assert_eq!( + built + .as_::() + .elements() + .as_::() + .nchunks(), + CHUNK_LEN, + "one chunk per appended array", + ); + + // The same values appended as scalars land in a single canonical child. + let mut builder = builder_with_capacity_in(array.dtype(), 0, BufferAllocatorRef::static_ref()); + let scalar = array.execute_scalar(0, &mut ctx)?; + for _ in 0..CHUNK_LEN { + builder.append_scalar(&scalar)?; } + let built_from_scalars = builder.finish(); + + assert!( + built_from_scalars + .as_::() + .elements() + .is::() + ); + assert_arrays_eq!(&built_from_scalars, &built, &mut ctx); + Ok(()) } -/// Test that `set_validity` is a no-op on a non-nullable builder. +/// A builder that mixes appended arrays with scalar appends must keep the two in order. #[test] -fn test_set_validity_noop_when_non_nullable() -> VortexResult<()> { - let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut builder = builder_with_capacity(&dtype, 4); - builder.append_zeros(4); +fn test_struct_builder_interleaves_arrays_and_scalars() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let array = StructArray::try_from_iter([("a", constant_i32())])?.into_array(); + let scalar = Scalar::struct_( + array.dtype().clone(), + vec![Scalar::primitive(1i32, Nullability::NonNullable)], + ); - // Providing an all-false mask must not make the non-nullable array invalid. - builder.set_validity(Mask::new_false(4)); + let mut builder = builder_with_capacity_in(array.dtype(), 0, BufferAllocatorRef::static_ref()); + builder.append_scalar(&scalar)?; + array.append_to_builder(builder.as_mut(), &mut ctx)?; + builder.append_scalar(&scalar)?; + let built = builder.finish(); + + let scalar_array = StructArray::try_from_iter([( + "a", + PrimitiveArray::new(buffer![1i32], Validity::NonNullable), + )])? + .into_array(); + let expected = ChunkedArray::try_new( + vec![scalar_array.clone(), array, scalar_array], + built.dtype().clone(), + )?; + assert_arrays_eq!(&built, &expected, &mut ctx); - let validity = builder.finish().validity()?; + Ok(()) +} + +/// An arbitrary array length. Nested builders treat no length specially, so the tests only need a +/// length long enough to tell chunks apart. +const CHUNK_LEN: usize = 64; + +/// A nested builder's own validity is accumulated the same way its children are: an appended +/// array's validity is kept as it arrived rather than executed into a mask and copied bit by bit. +#[test] +fn test_appended_validity_is_not_materialized() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - for i in 0..4 { - assert!( - validity.execute_is_valid(i, &mut ctx)?, - "index {i} should remain valid" - ); + + let array = StructArray::try_from_iter_with_validity( + [("a", iota(CHUNK_LEN))], + Validity::from_iter((0..CHUNK_LEN).map(|i| i % 3 != 0)), + )? + .into_array(); + + let mut builder = builder_with_capacity_in(array.dtype(), 0, BufferAllocatorRef::static_ref()); + array.append_to_builder(builder.as_mut(), &mut ctx)?; + array.append_to_builder(builder.as_mut(), &mut ctx)?; + let built = builder.finish(); + + let Validity::Array(validity) = built.validity()? else { + panic!("expected array-backed validity"); + }; + assert!( + validity.is::(), + "the two appended validities should have been concatenated, not copied into one buffer", + ); + + let expected = ChunkedArray::try_new(vec![array.clone(), array], built.dtype().clone())?; + assert_arrays_eq!(&built, &expected, &mut ctx); + + Ok(()) +} + +/// A non-canonical array of [`CHUNK_LEN`] `i32` values. +fn constant_i32() -> ArrayRef { + ConstantArray::new(0i32, CHUNK_LEN).into_array() +} + +/// Two lists of [`CHUNK_LEN`] elements each, so that appending them chunks the elements. +fn two_lists_of_chunk_len() -> ListViewArray { + ListViewArray::new( + iota(2 * CHUNK_LEN), + u64s([0, CHUNK_LEN]), + u64s([CHUNK_LEN, CHUNK_LEN]), + Validity::NonNullable, + ) +} + +/// `0..n` as an `i32` array, so that the values of one chunk are distinguishable from the next. +fn iota(n: usize) -> ArrayRef { + (0..n) + .map(|i| i32::try_from(i).vortex_expect("iota value fits in an i32")) + .collect::>() + .into_array() +} + +/// A `u64` array of list offsets or sizes. +fn u64s(values: impl IntoIterator) -> ArrayRef { + values + .into_iter() + .map(|v| u64::try_from(v).vortex_expect("list offset fits in a u64")) + .collect::>() + .into_array() +} + +/// Once a list builder keeps its elements as chunks, `elements_builder.len()` is a running total +/// across those chunks — every offset appended afterwards has to be rebased onto it. +/// +/// The two cases cover the two bulk paths into the elements builder: appending a `ListViewArray` +/// rebases the view's own offsets, while appending a `ListArray` slices the elements first. +#[rstest] +#[case::from_listview(two_lists_of_chunk_len().into_array())] +#[case::from_list( + ListArray::new( + iota(2 * CHUNK_LEN), + u64s([0, CHUNK_LEN, 2 * CHUNK_LEN]), + Validity::NonNullable, + ) + .into_array() +)] +fn test_list_offsets_are_rebased_across_element_chunks( + #[case] lists: ArrayRef, +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let mut builder = builder_with_capacity_in( + &DType::List( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + Nullability::NonNullable, + ), + 0, + BufferAllocatorRef::static_ref(), + ); + lists.append_to_builder(builder.as_mut(), &mut ctx)?; + lists.append_to_builder(builder.as_mut(), &mut ctx)?; + let built = builder.finish(); + + assert!( + built.as_::().elements().is::(), + "the elements should have been kept as chunks", + ); + + let expected = ChunkedArray::try_new(vec![lists.clone(), lists], built.dtype().clone())?; + assert_arrays_eq!(&built, &expected, &mut ctx); + + Ok(()) +} + +/// `ListBuilder` computes each offset from the running element count as well, and reaches the +/// elements builder through `append_array_as_list` rather than a bulk append. +#[test] +fn test_list_builder_offsets_are_rebased_across_element_chunks() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let element_dtype = Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)); + + let mut builder = ListBuilder::::with_capacity_in( + element_dtype, + Nullability::NonNullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); + for value in 0..3i32 { + builder + .append_array_as_list(&ConstantArray::new(value, CHUNK_LEN).into_array(), &mut ctx)?; } + let built = builder.finish(); + + assert!(built.as_::().elements().is::()); + + let expected = ListArray::new( + (0..3i32) + .flat_map(|value| std::iter::repeat_n(value, CHUNK_LEN)) + .collect::>() + .into_array(), + u64s((0..=3).map(|i| i * CHUNK_LEN)), + Validity::NonNullable, + ); + assert_arrays_eq!(&built, &expected, &mut ctx); + + Ok(()) +} + +/// A nested builder's own validity buffer is independent of its chunked child, so nulls appended +/// alongside chunks must survive. +#[rstest] +#[case::fixed_size_list( + FixedSizeListArray::new( + iota(CHUNK_LEN), + 4, + Validity::from_iter((0..CHUNK_LEN / 4).map(|i| i % 3 != 0)), + CHUNK_LEN / 4, + ) + .into_array(), + |array: &ArrayRef| array.as_::().elements().clone() +)] +#[case::struct_( + StructArray::try_from_iter_with_validity( + [("a", iota(CHUNK_LEN))], + Validity::from_iter((0..CHUNK_LEN).map(|i| i % 3 != 0)), + ) + .vortex_expect("struct array") + .into_array(), + |array: &ArrayRef| array.as_::().unmasked_field(0).clone() +)] +fn test_validity_survives_chunked_children( + #[case] array: ArrayRef, + #[case] child_of: fn(&ArrayRef) -> ArrayRef, +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let mut builder = builder_with_capacity_in(array.dtype(), 0, BufferAllocatorRef::static_ref()); + array.append_to_builder(builder.as_mut(), &mut ctx)?; + builder.append_nulls(1); + array.append_to_builder(builder.as_mut(), &mut ctx)?; + let built = builder.finish(); + + assert!(child_of(&built).is::()); + + let mut null = builder_with_capacity_in(array.dtype(), 1, BufferAllocatorRef::static_ref()); + null.append_nulls(1); + let expected = ChunkedArray::try_new( + vec![array.clone(), null.finish(), array], + built.dtype().clone(), + )?; + assert_arrays_eq!(&built, &expected, &mut ctx); + + Ok(()) +} + +/// Consumers that genuinely need a fully-decoded tree ask for it, and must still get one. +#[test] +fn test_chunked_children_canonicalize_recursively() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let array = StructArray::try_from_iter([("a", constant_i32())])?.into_array(); + let mut builder = builder_with_capacity_in(array.dtype(), 0, BufferAllocatorRef::static_ref()); + array.append_to_builder(builder.as_mut(), &mut ctx)?; + array.append_to_builder(builder.as_mut(), &mut ctx)?; + let built = builder.finish(); + + let recursive = built.clone().execute::(&mut ctx)?.0; + assert!( + recursive + .clone() + .into_array() + .as_::() + .unmasked_field(0) + .is::() + ); + assert_arrays_eq!(&recursive.into_array(), &built, &mut ctx); + Ok(()) } diff --git a/vortex-array/src/builders/validity.rs b/vortex-array/src/builders/validity.rs new file mode 100644 index 00000000000..61049798fa4 --- /dev/null +++ b/vortex-array/src/builders/validity.rs @@ -0,0 +1,258 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::BufferAllocatorRef; +use vortex_error::VortexExpect; +use vortex_error::vortex_panic; + +use crate::builders::LazyBitBufferBuilder; +use crate::dtype::Nullability; +use crate::dtype::Nullability::NonNullable; +use crate::dtype::Nullability::Nullable; +use crate::validity::Validity; + +/// Accumulates the validity of a nested [`ArrayBuilder`](crate::builders::ArrayBuilder) without +/// materializing a null buffer for it. +/// +/// A nested builder learns about validity from two sources: one row at a time, as scalars are +/// appended, and a whole array's worth at a time, as arrays are. Only the former needs a null +/// buffer. An appended array already carries its validity in whatever form it was stored in — +/// [`Validity::AllValid`] and [`Validity::AllInvalid`] cost nothing at all, and an array-backed +/// validity is a bool array that is already built — so this builder keeps those as runs and +/// concatenates them at the end, exactly as [`Validity::concat`] does. +/// +/// Materializing them instead would mean executing every appended array's validity into a +/// [`Mask`](vortex_mask::Mask) and copying its bits, which for a builder assembling many chunks is +/// the dominant cost of tracking validity at all. +pub(crate) struct ValidityBuilder { + /// Completed runs, in logical order, with the number of values each covers. Never contains an + /// empty run. + runs: Vec<(Validity, usize)>, + + /// The summed length of `runs`. + runs_len: usize, + + /// Null buffer holding the bits appended since the last run. + pending: LazyBitBufferBuilder, +} + +impl ValidityBuilder { + /// Creates a validity builder with the provided allocator and capacity. + pub fn new(capacity: usize, allocator: &BufferAllocatorRef) -> Self { + Self { + runs: Vec::new(), + runs_len: 0, + pending: LazyBitBufferBuilder::new(capacity, allocator.clone()), + } + } + + /// The number of values whose validity has been recorded so far. + pub fn len(&self) -> usize { + self.runs_len + self.pending.len() + } + + /// Records one valid value. + pub fn append_non_null(&mut self) { + self.pending.append_non_null() + } + + /// Records `n` valid values. + pub fn append_n_non_nulls(&mut self, n: usize) { + self.pending.append_n_non_nulls(n) + } + + /// Records `n` null values. + pub fn append_n_nulls(&mut self, n: usize) { + self.pending.append_n_nulls(n) + } + + /// Records the validity of a whole appended array, covering `len` values, as a run of its own. + /// + /// However few values the run covers, it is kept as it arrived rather than executed into a + /// mask, so a builder's validity is split on exactly the boundaries its children are. + pub fn append_validity(&mut self, validity: Validity, len: usize) { + if len == 0 { + return; + } + + self.flush_pending(); + self.runs_len += len; + self.runs.push((validity, len)); + } + + /// Allocates space for `additional` more bits in the null buffer. + pub fn reserve_exact(&mut self, additional: usize) { + self.pending.reserve_exact(additional) + } + + /// Finishes the validity, concatenating the accumulated runs. + /// + /// # Panics + /// + /// Panics if a non-nullable builder recorded a null, matching + /// [`LazyBitBufferBuilder::finish_with_nullability`]. + pub fn finish_with_nullability(&mut self, nullability: Nullability) -> Validity { + if self.runs.is_empty() { + return self.pending.finish_with_nullability(nullability); + } + + self.flush_pending(); + self.runs_len = 0; + let runs = std::mem::take(&mut self.runs); + + // `Validity::concat` treats `NonNullable` and `AllValid` as different kinds and falls back + // to a bool array when both appear, which they do as soon as a non-nullable array is + // appended next to a scalar. Both mean "no nulls", so answer from the nullability instead. + if runs + .iter() + .all(|(validity, _)| validity.definitely_no_nulls()) + { + return nullability.into(); + } + + let validity = Validity::concat(runs).vortex_expect("runs is not empty"); + if nullability == NonNullable { + vortex_panic!("cannot finish a non-nullable builder holding {validity:?} validity"); + } + validity + } + + /// Moves whatever the null buffer holds into `runs`, keeping the runs in logical order. + fn flush_pending(&mut self) { + let len = self.pending.len(); + if len == 0 { + return; + } + // A run is only ever read back through `Validity::concat`, which takes the nullability + // from the runs as a whole, so an all-valid null buffer can stay lazy here. + let validity = self.pending.finish_with_nullability(Nullable); + self.runs_len += len; + self.runs.push((validity, len)); + } +} + +#[cfg(test)] +mod tests { + use vortex_buffer::BufferAllocatorRef; + use vortex_error::VortexResult; + use vortex_mask::Mask; + + use super::ValidityBuilder; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::Chunked; + use crate::dtype::Nullability::NonNullable; + use crate::dtype::Nullability::Nullable; + use crate::validity::Validity; + + /// An arbitrary run length. `ValidityBuilder` treats no length specially, so the tests only + /// need a length long enough to tell runs apart. + const RUN_LEN: usize = 64; + + /// A `Validity` backed by a bool array, which is the case a run avoids executing. + fn array_backed(len: usize) -> Validity { + Validity::from_mask(Mask::from_iter((0..len).map(|i| i % 2 == 0)), Nullable) + } + + #[test] + fn test_whole_array_validity_is_kept_as_a_run() { + let mut builder = ValidityBuilder::new(0, BufferAllocatorRef::static_ref()); + + builder.append_validity(array_backed(RUN_LEN), RUN_LEN); + builder.append_validity(array_backed(RUN_LEN), RUN_LEN); + assert_eq!(builder.len(), 2 * RUN_LEN); + + let Validity::Array(array) = builder.finish_with_nullability(Nullable) else { + panic!("expected array-backed validity"); + }; + assert!( + array.is::(), + "the runs should have been concatenated, not copied into one buffer", + ); + } + + /// Uniform runs collapse instead of becoming a bool array. + #[test] + fn test_all_valid_runs_stay_lazy() { + let mut builder = ValidityBuilder::new(0, BufferAllocatorRef::static_ref()); + + builder.append_validity(Validity::AllValid, RUN_LEN); + builder.append_validity(Validity::AllValid, RUN_LEN); + + assert!(matches!( + builder.finish_with_nullability(Nullable), + Validity::AllValid + )); + } + + /// A one-row validity earns a run too. Uniform runs still collapse, so a builder fed an array + /// at a time does not pay a bool array for validity it never had. + #[test] + fn test_short_validity_is_kept_as_a_run_too() { + let mut builder = ValidityBuilder::new(0, BufferAllocatorRef::static_ref()); + + for _ in 0..RUN_LEN { + builder.append_validity(Validity::AllInvalid, 1); + } + assert_eq!(builder.len(), RUN_LEN); + + assert!(matches!( + builder.finish_with_nullability(Nullable), + Validity::AllInvalid + )); + } + + /// Bits and runs interleave, and have to come back out in the order they went in. + #[test] + fn test_bits_and_runs_keep_their_order() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = ValidityBuilder::new(0, BufferAllocatorRef::static_ref()); + + builder.append_n_nulls(1); + builder.append_validity(Validity::AllValid, RUN_LEN); + builder.append_non_null(); + builder.append_validity(Validity::AllInvalid, RUN_LEN); + + let validity = builder.finish_with_nullability(Nullable); + let mask = validity.execute_mask(2 * RUN_LEN + 2, &mut ctx)?; + + let expected = Mask::from_iter( + [false] + .into_iter() + .chain(std::iter::repeat_n(true, RUN_LEN)) + .chain([true]) + .chain(std::iter::repeat_n(false, RUN_LEN)), + ); + assert_eq!(mask, expected); + + Ok(()) + } + + #[test] + fn test_non_nullable_finishes_non_nullable() { + let mut builder = ValidityBuilder::new(0, BufferAllocatorRef::static_ref()); + + builder.append_validity(Validity::NonNullable, RUN_LEN); + builder.append_n_non_nulls(1); + + assert!(matches!( + builder.finish_with_nullability(NonNullable), + Validity::NonNullable + )); + } + + #[test] + fn test_finish_resets_the_builder() { + let mut builder = ValidityBuilder::new(0, BufferAllocatorRef::static_ref()); + + builder.append_validity(Validity::AllInvalid, RUN_LEN); + assert_eq!(builder.finish_with_nullability(Nullable).maybe_len(), None); + + assert_eq!(builder.len(), 0); + builder.append_n_nulls(1); + assert!(matches!( + builder.finish_with_nullability(Nullable), + Validity::AllInvalid + )); + } +} diff --git a/vortex-array/src/builders/varbinview.rs b/vortex-array/src/builders/varbinview.rs index 1d530a384ef..e27336fe998 100644 --- a/vortex-array/src/builders/varbinview.rs +++ b/vortex-array/src/builders/varbinview.rs @@ -6,8 +6,10 @@ use std::ops::Range; use std::sync::Arc; use itertools::Itertools; +use num_traits::AsPrimitive; use vortex_buffer::Alignment; use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; @@ -15,6 +17,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; +use vortex_mask::AllOr; use vortex_mask::Mask; use vortex_utils::aliases::hash_map::Entry; use vortex_utils::aliases::hash_map::HashMap; @@ -25,11 +28,14 @@ use crate::IntoArray; use crate::arrays::VarBinViewArray; use crate::arrays::varbinview::VarBinViewArrayExt; use crate::arrays::varbinview::build_views::BinaryView; +use crate::arrays::varbinview::build_views::MAX_BUFFER_LEN; +use crate::arrays::varbinview::build_views::extend_views; use crate::arrays::varbinview::compact::BufferUtilization; use crate::builders::ArrayBuilder; use crate::builders::LazyBitBufferBuilder; use crate::canonical::Canonical; use crate::dtype::DType; +use crate::dtype::NativePType; use crate::scalar::Scalar; /// The builder for building a [`VarBinViewArray`]. @@ -41,56 +47,128 @@ pub struct VarBinViewBuilder { in_progress: Option, growth_strategy: BufferGrowthStrategy, compaction_threshold: f64, + allocator: BufferAllocatorRef, } impl VarBinViewBuilder { + /// Creates a builder with room for `capacity` values. + #[deprecated(note = "use `with_capacity_in` with an explicit allocator")] pub fn with_capacity(dtype: DType, capacity: usize) -> Self { - Self::new(dtype, capacity, Default::default(), Default::default(), 0.0) + Self::with_capacity_in(dtype, capacity, BufferAllocatorRef::statically_allocated()) } + /// Creates a builder with room for `capacity` values using `allocator`. + pub fn with_capacity_in(dtype: DType, capacity: usize, allocator: BufferAllocatorRef) -> Self { + Self::new_in( + dtype, + capacity, + Default::default(), + Default::default(), + 0.0, + allocator, + ) + } + + /// Creates a builder that deduplicates completed buffers. + #[deprecated(note = "use `with_buffer_deduplication_in` with an explicit allocator")] pub fn with_buffer_deduplication(dtype: DType, capacity: usize) -> Self { - Self::new( + Self::with_buffer_deduplication_in( + dtype, + capacity, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Creates a builder that deduplicates completed buffers using `allocator`. + pub fn with_buffer_deduplication_in( + dtype: DType, + capacity: usize, + allocator: BufferAllocatorRef, + ) -> Self { + Self::new_in( dtype, capacity, CompletedBuffers::Deduplicated(Default::default()), Default::default(), 0.0, + allocator, ) } + /// Creates a builder that compacts buffers below `compaction_threshold` utilization. + #[deprecated(note = "use `with_compaction_in` with an explicit allocator")] pub fn with_compaction(dtype: DType, capacity: usize, compaction_threshold: f64) -> Self { - Self::new( + Self::with_compaction_in( + dtype, + capacity, + compaction_threshold, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Creates a compacting builder using `allocator`. + pub fn with_compaction_in( + dtype: DType, + capacity: usize, + compaction_threshold: f64, + allocator: BufferAllocatorRef, + ) -> Self { + Self::new_in( dtype, capacity, Default::default(), Default::default(), compaction_threshold, + allocator, ) } + /// Creates a builder with explicit buffer policies. + #[deprecated(note = "use `new_in` with an explicit allocator")] pub fn new( dtype: DType, capacity: usize, completed: CompletedBuffers, growth_strategy: BufferGrowthStrategy, compaction_threshold: f64, + ) -> Self { + Self::new_in( + dtype, + capacity, + completed, + growth_strategy, + compaction_threshold, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Creates a builder with explicit buffer policies and `allocator`. + pub fn new_in( + dtype: DType, + capacity: usize, + completed: CompletedBuffers, + growth_strategy: BufferGrowthStrategy, + compaction_threshold: f64, + allocator: BufferAllocatorRef, ) -> Self { assert!( matches!(dtype, DType::Utf8(_) | DType::Binary(_)), "VarBinViewBuilder DType must be Utf8 or Binary." ); Self { - views_builder: BufferMut::with_capacity_preferred_aligned( + views_builder: BufferMut::with_capacity_preferred_aligned_in( capacity, Alignment::of::(), None, + allocator.clone(), ), - nulls: LazyBitBufferBuilder::new(capacity), + nulls: LazyBitBufferBuilder::new(capacity, allocator.clone()), completed, in_progress: None, dtype, growth_strategy, compaction_threshold, + allocator, } } @@ -148,10 +226,11 @@ impl VarBinViewBuilder { fn init_in_progress(&mut self, min_len: usize) { let next_buffer_size = self.growth_strategy.next_size() as usize; let to_reserve = next_buffer_size.max(min_len); - self.in_progress = Some(ByteBufferMut::with_capacity_preferred_aligned( + self.in_progress = Some(ByteBufferMut::with_capacity_preferred_aligned_in( to_reserve, Alignment::of::(), None, + self.allocator.clone(), )); } @@ -188,47 +267,465 @@ impl VarBinViewBuilder { self.completed.len() } - /// Returns true if a non-empty in-progress buffer is staged (and would - /// become a completed buffer on the next flush), false otherwise. - pub fn in_progress(&self) -> bool { - self.in_progress.is_some() + /// Whether this builder compacts the data buffers it is handed. The lengths-driven and + /// gather appends use this to gate their utilization measurement; the buffer-adopting escape + /// hatches ([`append_views_built_at`](Self::append_views_built_at) and + /// [`push_buffers`](Self::push_buffers)) always bypass it. + fn compacts_buffers(&self) -> bool { + self.compaction_threshold > 0.0 } - /// Pushes buffers and pre-adjusted views into the builder. + /// Adopts the buffers and views that `build` produces against the index its first buffer + /// will land at. /// - /// The provided `buffers` contain sections of data from a `VarBinViewArray`, and the - /// `views` are `BinaryView`s that have already been adjusted to reference the correct buffer - /// indices and offsets for this builder. All views must point to valid sections within the - /// provided buffers, and the validity length must match the view length. + /// The builder flushes its staged bytes, then hands `build` the index the next data buffer + /// will occupy; `build` returns data buffers — which land contiguously from that index — and + /// one view per entry of `validity`, already referencing them. This is the escape hatch for + /// an encoding that only discovers its views while walking its own byte format (e.g. + /// length-prefixed frames), where the lengths-driven appends cannot apply; keeping the + /// numbering inside this call is what makes the views come out right without a rebase pass. /// /// # Warning /// - /// This method does not check utilization of the given buffers. Callers must provide - /// buffers that are fully utilized by the given adjusted views. + /// This method does not check utilization of the returned buffers. `build` must return + /// buffers that are fully utilized by its views. /// /// # Panics /// - /// Panics if this builder deduplicates buffers and any of the given buffers already - /// exist in this builder. - pub fn push_buffer_and_adjusted_views( + /// Panics if `build` returns a different view count than `validity.len()`, or if this + /// builder deduplicates buffers and already holds one of the returned buffers. + pub fn append_views_built_at( &mut self, - buffers: &[ByteBuffer], - views: &Buffer, - validity_mask: Mask, - ) { + validity: &Mask, + build: impl FnOnce(u32) -> VortexResult<(Vec, Buffer)>, + ) -> VortexResult<()> { self.flush_in_progress(); - let expected_completed_len = self.completed.len() as usize + buffers.len(); - self.completed.extend_from_slice_unchecked(buffers); + let start_index = self.completed.len(); + let (buffers, views) = build(start_index)?; + assert_eq!( + views.len(), + validity.len(), + "Must build one view per validity entry" + ); + + let expected_completed_len = start_index as usize + buffers.len(); + self.completed.extend_from_slice_unchecked(&buffers); assert_eq!( self.completed.len() as usize, expected_completed_len, "Some buffers already exist", ); self.views_builder.extend_trusted(views.iter().copied()); - self.push_only_validity_mask(&validity_mask); + self.nulls.append_validity_mask(validity); + + debug_assert_eq!(self.nulls.len(), self.views_builder.len()); + Ok(()) + } + + /// Adopts `buffers` as completed data buffers, returning the index each landed at. + /// + /// Views appended afterwards (e.g. via [`append_views_gathered`](Self::append_views_gathered) + /// or [`append_views_scattered`](Self::append_views_scattered)) reference values through the + /// returned indices. A deduplicating builder returns the existing index for a buffer it + /// already holds, so repeated appends over shared storage — chunks gathered through one + /// dictionary, slices of one array — adopt it once. + /// + /// # Warning + /// + /// Buffers are taken as they are, without utilization measurement; like + /// [`append_views_built_at`](Self::append_views_built_at), this bypasses any compaction the + /// builder was configured for. Scattered appends deliberately have these semantics; + /// compacting paths that measure utilization use + /// [`push_buffers_compacted`](Self::push_buffers_compacted) instead. + fn push_buffers(&mut self, buffers: impl IntoIterator) -> Vec { + self.flush_in_progress(); + buffers + .into_iter() + .map(|buffer| self.completed.push(buffer)) + .collect() + } - debug_assert_eq!(self.nulls.len(), self.views_builder.len()) + /// The compacting counterpart of [`push_buffers`](Self::push_buffers): adopts each buffer + /// according to the [`CompactionStrategy`] its measured utilization earns — whole, sliced to + /// its referenced range, or not at all. Views into the buffers must then be rebased through + /// [`remap_view_compacted`](Self::remap_view_compacted), which routes views into unadopted + /// buffers through the builder's own storage. + fn push_buffers_compacted( + &mut self, + buffers: Vec, + utilizations: &[BufferUtilization], + ) -> Vec { + self.flush_in_progress(); + buffers + .into_iter() + .zip(utilizations) + .map(|(buffer, utilization)| { + match compaction_strategy(utilization, self.compaction_threshold) { + CompactionStrategy::KeepFull => CompactedSlot::Kept { + index: self.completed.push(buffer), + }, + CompactionStrategy::Slice { start, end } => CompactedSlot::Sliced { + index: self + .completed + .push(buffer.slice(start as usize..end as usize)), + shift: start, + }, + CompactionStrategy::Rewrite => CompactedSlot::Rewrite { source: buffer }, + } + }) + .collect() + } + + /// [`remap_view`] against [`CompactedSlot`]s: views into kept or sliced buffers are rebased + /// onto the index those landed at, and a view into a rewritten buffer has its bytes copied + /// into the builder's own storage. Inlined views pass through unchanged. + fn remap_view_compacted(&mut self, view: BinaryView, slots: &[CompactedSlot]) -> BinaryView { + if view.is_inlined() { + return view; + } + let view_ref = view.as_view(); + match &slots[view_ref.buffer_index as usize] { + CompactedSlot::Kept { index } => view_ref + .with_buffer_and_offset(*index, view_ref.offset) + .into(), + CompactedSlot::Sliced { index, shift } => view_ref + .with_buffer_and_offset(*index, view_ref.offset - shift) + .into(), + CompactedSlot::Rewrite { source } => { + let bytes = &source.as_slice()[view_ref.as_range()]; + let (buffer_idx, offset) = self.append_value_to_buffer(bytes); + BinaryView::make_view(bytes, buffer_idx, offset) + } + } + } + + /// Appends `validity.len()` values, gathering each valid row's view from `views` through the + /// index `view_at` returns for it. + /// + /// `buffers` — the data buffers the views reference, in the numbering the views use — are + /// adopted through the builder's buffer storage, and every gathered view is rebased onto + /// the indices they land at as it is written, so the whole append is one view per row with no + /// byte copy and no intermediate array. `view_at` is only called for valid rows; null rows get + /// an empty view. + /// + /// When the builder is configured to compact buffers, utilization is measured from the views + /// the gather actually references — duplicate indices count once — and each buffer is + /// adopted, sliced, or rewritten accordingly, so callers never need a + /// canonicalize-and-compact fallback. + /// + /// # Panics + /// + /// Panics if `view_at` returns an index out of bounds of `views`, or if a gathered view + /// references a buffer index outside `buffers`. + pub fn append_views_gathered( + &mut self, + buffers: impl IntoIterator, + views: &[BinaryView], + validity: &Mask, + view_at: impl Fn(usize) -> usize, + ) { + if self.compacts_buffers() { + return self.append_views_gathered_compacted( + buffers.into_iter().collect(), + views, + validity, + view_at, + ); + } + + let mapping = self.push_buffers(buffers); + + self.views_builder.reserve(validity.len()); + match validity { + Mask::AllTrue(len) => self + .views_builder + .extend_trusted((0..*len).map(|row| remap_view(views[view_at(row)], &mapping))), + Mask::AllFalse(len) => self.views_builder.push_n(BinaryView::empty_view(), *len), + Mask::Values(values) => { + let bits = values.bit_buffer(); + let start = self.views_builder.len(); + self.views_builder + .push_n(BinaryView::empty_view(), bits.len()); + let gathered = &mut self.views_builder[start..]; + bits.for_each_set_index(|row| { + gathered[row] = remap_view(views[view_at(row)], &mapping); + }); + } + } + + self.nulls.append_validity_mask(validity); + debug_assert_eq!(self.nulls.len(), self.views_builder.len()); + } + + /// The compacting arm of [`append_views_gathered`](Self::append_views_gathered): marks the + /// source views the gather references, measures each buffer's utilization from just those — + /// so entries gathered by many rows count once — and adopts the buffers through + /// [`push_buffers_compacted`](Self::push_buffers_compacted). Each referenced source view is + /// remapped once, so rows sharing an entry of a rewritten buffer share one copy of its bytes. + fn append_views_gathered_compacted( + &mut self, + buffers: Vec, + views: &[BinaryView], + validity: &Mask, + view_at: impl Fn(usize) -> usize, + ) { + let mut referenced = vec![false; views.len()]; + match validity { + Mask::AllTrue(len) => (0..*len).for_each(|row| referenced[view_at(row)] = true), + Mask::AllFalse(_) => {} + Mask::Values(values) => values + .bit_buffer() + .for_each_set_index(|row| referenced[view_at(row)] = true), + } + + let mut utilizations = unmeasured_utilizations(&buffers); + for (view, _) in views.iter().zip(&referenced).filter(|(_, used)| **used) { + measure_view(&mut utilizations, view); + } + let slots = self.push_buffers_compacted(buffers, &utilizations); + + let mut remapped = vec![BinaryView::empty_view(); views.len()]; + for (idx, view) in views.iter().enumerate() { + if referenced[idx] { + remapped[idx] = self.remap_view_compacted(*view, &slots); + } + } + + self.views_builder.reserve(validity.len()); + match validity { + Mask::AllTrue(len) => self + .views_builder + .extend_trusted((0..*len).map(|row| remapped[view_at(row)])), + Mask::AllFalse(len) => self.views_builder.push_n(BinaryView::empty_view(), *len), + Mask::Values(values) => { + let bits = values.bit_buffer(); + let start = self.views_builder.len(); + self.views_builder + .push_n(BinaryView::empty_view(), bits.len()); + let gathered = &mut self.views_builder[start..]; + bits.for_each_set_index(|row| gathered[row] = remapped[view_at(row)]); + } + } + + self.nulls.append_validity_mask(validity); + debug_assert_eq!(self.nulls.len(), self.views_builder.len()); + } + + /// Appends `len` rows that are `fill` everywhere except at the given patch rows. + /// + /// `buffers` — the data buffers `fill` and the patch views reference, in the numbering they + /// use — are adopted through the builder's buffer storage and the views are rebased onto + /// the indices they land at. `patches` yields `(row, view)` pairs with rows below `len`; the + /// fill rows cost one bulk view fill and each patch one view write, with no byte copy. + /// Validity is appended exactly as given — invalid rows keep whichever view they got. + /// + /// This append deliberately adopts the supplied buffers as-is, even when the builder is + /// configured for compaction. Callers should use it only when retaining those buffers is + /// acceptable. + /// + /// # Panics + /// + /// Panics if `validity` is not `len` long, a patch row is out of bounds, or a view references + /// a buffer index outside `buffers`. + pub fn append_views_scattered( + &mut self, + buffers: impl IntoIterator, + len: usize, + fill: BinaryView, + patches: impl Iterator, + validity: &Mask, + ) { + assert_eq!(validity.len(), len, "Must have one validity entry per row"); + let mapping = self.push_buffers(buffers); + + let start = self.views_builder.len(); + self.views_builder.push_n(remap_view(fill, &mapping), len); + let scattered = &mut self.views_builder[start..]; + for (row, view) in patches { + scattered[row] = remap_view(view, &mapping); + } + + self.nulls.append_validity_mask(validity); + debug_assert_eq!(self.nulls.len(), self.views_builder.len()); + } + + /// Appends values laid end-to-end in `bytes`, one per entry of `lengths`. + /// + /// The builder adopts `bytes` as a data buffer without copying it (splitting it only past the + /// `u32` view-offset limit) and builds the views directly into its own storage, so the whole + /// append costs one view per row. Every length is consumed, including those for null rows, so + /// the lengths must describe `bytes` exactly; bytes belonging to null rows are not retained + /// when compaction rewrites the buffer. + /// + /// When the builder is configured to compact buffers, the utilization is measured from the + /// lengths alone — only values too long to inline reference the buffer — and a heap below + /// the threshold is rewritten to just those values instead of adopted, so callers never need + /// a canonicalize-and-compact fallback. + /// + /// # Panics + /// + /// Panics if `lengths` and `validity` disagree in length, if the lengths do not describe + /// `bytes` exactly, or if this builder deduplicates buffers and already holds `bytes`. + pub fn append_buffer_with_lengths>( + &mut self, + bytes: ByteBuffer, + lengths: &[P], + validity: &Mask, + ) { + assert_eq!( + lengths.len(), + validity.len(), + "Must have one length per validity entry" + ); + self.append_buffer_views(&bytes, lengths.len(), validity, |i| lengths[i].as_()); + } + + /// [`append_buffer_with_lengths`](Self::append_buffer_with_lengths) for values described by + /// an offsets buffer instead of lengths. + /// + /// `offsets` are absolute positions into `bytes` — the layout a + /// [`VarBinArray`](crate::arrays::VarBinArray) stores — so there is one more offset than there + /// are values, and only the `offsets[0]..offsets[last]` range of `bytes` is adopted, again + /// without copying. + /// + /// # Panics + /// + /// Panics if `offsets` does not hold exactly one more entry than `validity`, or if the offsets + /// are not monotonically non-decreasing positions within `bytes`. + pub fn append_buffer_with_offsets>( + &mut self, + bytes: ByteBuffer, + offsets: &[P], + validity: &Mask, + ) { + assert_eq!( + offsets.len(), + validity.len() + 1, + "Must have one more offset than validity entries" + ); + let first: usize = offsets[0].as_(); + let last: usize = offsets[offsets.len() - 1].as_(); + let bytes = bytes.slice(first..last); + // Wrapping keeps corrupt non-monotonic offsets from panicking on the subtraction itself; + // the wrapped length then fails the in-bounds checks of the view-building loop. + self.append_buffer_views(&bytes, validity.len(), validity, |i| { + AsPrimitive::::as_(offsets[i + 1]) + .wrapping_sub(AsPrimitive::::as_(offsets[i])) + }); + } + + /// Shared tail of the bulk buffer appends: builds the views straight into the builder's views + /// storage, then adopts the buffer segments and the validity. + fn append_buffer_views( + &mut self, + bytes: &ByteBuffer, + count: usize, + validity: &Mask, + len_at: impl Fn(usize) -> usize, + ) { + self.flush_in_progress(); + + // A compacting builder measures utilization before adopting the buffer. Only values too + // long to inline reference the heap, so the measurement is one pass over the lengths and + // never touches the bytes. Heaps past the single-buffer limit are adopted as they are: + // they roll over into multiple buffers, and per-segment accounting is not worth the rare + // >2GiB case. + if self.compacts_buffers() && bytes.len() <= MAX_BUFFER_LEN { + let referenced: usize = match validity.bit_buffer() { + AllOr::All => (0..count) + .map(&len_at) + .filter(|len| *len > BinaryView::MAX_INLINED_SIZE) + .sum(), + AllOr::None => 0, + AllOr::Some(b) => { + let mut sum = 0; + b.for_each_set_index(|idx| { + let len = len_at(idx); + if len > BinaryView::MAX_INLINED_SIZE { + sum += len; + } + }); + sum + } + }; + #[expect(clippy::cast_precision_loss)] + if (referenced as f64) < self.compaction_threshold * (bytes.len() as f64) { + return self + .append_buffer_views_rewritten(bytes, count, validity, len_at, referenced); + } + } + + let start_index = self.completed.len(); + let segments = extend_views( + &mut self.views_builder, + start_index, + MAX_BUFFER_LEN, + bytes, + count, + len_at, + ); + + let expected_completed_len = start_index as usize + segments.len(); + self.completed.extend_from_slice_unchecked(&segments); + assert_eq!( + self.completed.len() as usize, + expected_completed_len, + "Some buffers already exist", + ); + + self.nulls.append_validity_mask(validity); + debug_assert_eq!(self.nulls.len(), self.views_builder.len()); + } + + /// The under-utilized arm of [`append_buffer_views`](Self::append_buffer_views): copies only + /// the values that actually reference the heap into a compact buffer, sized `referenced`, + /// instead of adopting the whole heap. A fully-inlined append pushes no buffer at all. + fn append_buffer_views_rewritten( + &mut self, + bytes: &ByteBuffer, + count: usize, + validity: &Mask, + len_at: impl Fn(usize) -> usize, + referenced: usize, + ) { + let buf_index = self.completed.len(); + let mut compact = ByteBufferMut::with_capacity_in(referenced, self.allocator.clone()); + self.views_builder.reserve(count); + + let data = bytes.as_slice(); + let mut offset = 0usize; + for (i, is_valid) in validity.iter().enumerate() { + let len = len_at(i); + let value = &data[offset..offset + len]; + let view = if !is_valid { + BinaryView::empty_view() + } else if len > BinaryView::MAX_INLINED_SIZE { + // In `u32` range: `referenced <= bytes.len() <= MAX_BUFFER_LEN` (checked by the + // caller), and `compact` never grows past `referenced`. + #[expect(clippy::cast_possible_truncation)] + let view = BinaryView::make_view(value, buf_index, compact.len() as u32); + compact.extend_from_slice(value); + view + } else { + BinaryView::make_view(value, buf_index, 0) + }; + self.views_builder.push(view); + offset += len; + } + assert_eq!( + offset, + data.len(), + "value lengths must describe the byte heap exactly" + ); + + if !compact.is_empty() { + let pushed_index = self.completed.push(compact.freeze()); + assert_eq!(pushed_index, buf_index, "Buffer already exists"); + } + + self.nulls.append_validity_mask(validity); + debug_assert_eq!(self.nulls.len(), self.views_builder.len()); } /// Finishes the builder directly into a [`VarBinViewArray`]. @@ -244,22 +741,18 @@ impl VarBinViewBuilder { let validity = self.nulls.finish_with_nullability(self.dtype.nullability()); + let views = std::mem::replace( + &mut self.views_builder, + BufferMut::empty_aligned_in(Alignment::of::(), self.allocator.clone()), + ) + .freeze(); + // SAFETY: the builder methods check safety at each step. unsafe { - VarBinViewArray::new_unchecked( - std::mem::take(&mut self.views_builder).freeze(), - buffers.finish(), - self.dtype.clone(), - validity, - ) + VarBinViewArray::new_unchecked(views, buffers.finish(), self.dtype.clone(), validity) } } - // Pushes a validity mask into the builder not affecting the views or buffers - fn push_only_validity_mask(&mut self, validity_mask: &Mask) { - self.nulls.append_validity_mask(validity_mask); - } - pub(crate) fn append_varbinview_array( &mut self, array: &VarBinViewArray, @@ -269,7 +762,7 @@ impl VarBinViewBuilder { let mask = array.varbinview_validity().execute_mask(array.len(), ctx)?; - self.push_only_validity_mask(&mask); + self.nulls.append_validity_mask(&mask); let view_adjustment = self.completed @@ -298,16 +791,14 @@ impl VarBinViewBuilder { .push_n(BinaryView::empty_view(), array.len()); } Mask::Values(v) => { - for (idx, (&view, is_valid)) in - array.views().iter().zip(v.bit_buffer().iter()).enumerate() - { - let new_view = if !is_valid { - BinaryView::empty_view() - } else { - self.push_view(view, &adjustment, array, idx) - }; - self.views_builder.push(new_view); - } + let views = array.views(); + let start = self.views_builder.len(); + self.views_builder + .push_n(BinaryView::empty_view(), array.len()); + v.bit_buffer().for_each_set_index(|idx| { + let new_view = self.push_view(views[idx], &adjustment, array, idx); + self.views_builder[start + idx] = new_view; + }); } }, } @@ -374,10 +865,6 @@ impl ArrayBuilder for VarBinViewBuilder { self.nulls.reserve_exact(additional); } - unsafe fn set_validity_unchecked(&mut self, validity: Mask) { - self.nulls = LazyBitBufferBuilder::from_validity_mask(validity); - } - fn finish(&mut self) -> ArrayRef { self.finish_into_varbinview().into_array() } @@ -388,7 +875,8 @@ impl ArrayBuilder for VarBinViewBuilder { } impl VarBinViewBuilder { - #[inline] + #[allow(clippy::inline_always)] + #[inline(always)] fn push_view( &mut self, view: BinaryView, @@ -408,6 +896,53 @@ impl VarBinViewBuilder { } } +/// Rebases a view built against a local buffer numbering onto the builder indices those buffers +/// landed at, i.e. `mapping[i]` is where the caller's buffer `i` went. Inlined views carry no +/// buffer reference and pass through unchanged. +#[inline] +fn remap_view(view: BinaryView, mapping: &[u32]) -> BinaryView { + if view.is_inlined() { + view + } else { + let view_ref = view.as_view(); + view_ref + .with_buffer_and_offset(mapping[view_ref.buffer_index as usize], view_ref.offset) + .into() + } +} + +/// Where a caller's buffer went under `VarBinViewBuilder::push_buffers_compacted`. +enum CompactedSlot { + /// Adopted whole at this index. + Kept { index: u32 }, + /// Adopted at this index as the slice starting `shift` bytes in, so view offsets shift down. + Sliced { index: u32, shift: u32 }, + /// Not adopted: each referenced view's bytes are copied out of the source buffer into the + /// builder's own storage. + Rewrite { source: ByteBuffer }, +} + +/// One unmeasured [`BufferUtilization`] per buffer, ready for [`measure_view`] passes. +fn unmeasured_utilizations(buffers: &[ByteBuffer]) -> Vec { + buffers + .iter() + .map(|buffer| { + // Views address at most `u32` offsets, so measuring an oversized buffer against the + // saturated length under-reports utilization, which can only compact harder — and its + // unaddressable tail is dead weight worth compacting anyway. + BufferUtilization::zero(u32::try_from(buffer.len()).unwrap_or(u32::MAX)) + }) + .collect() +} + +/// Counts `view`'s bytes against the buffer it references; inlined views reference none. +fn measure_view(utilizations: &mut [BufferUtilization], view: &BinaryView) { + if !view.is_inlined() { + let view_ref = view.as_view(); + utilizations[view_ref.buffer_index as usize].add(view_ref.offset, view_ref.size); + } +} + pub enum CompletedBuffers { Default(Vec), Deduplicated(DeduplicatedBuffers), @@ -429,12 +964,14 @@ impl CompletedBuffers { } } + /// Push a new block, returning the index it landed at (or, when deduplicating, the index of + /// the identical block already held). fn push(&mut self, block: ByteBuffer) -> u32 { match self { Self::Default(buffers) => { assert!(buffers.len() < u32::MAX as usize, "Too many blocks"); buffers.push(block); - self.len() + self.len() - 1 } Self::Deduplicated(buffers) => buffers.push(block), } @@ -864,7 +1401,10 @@ impl RewritingViewAdjustment { #[cfg(test)] mod tests { + use vortex_buffer::BufferAllocatorRef; + use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; + use vortex_mask::Mask; use crate::IntoArray; use crate::VortexSessionExecute; @@ -876,10 +1416,437 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; + /// A long-enough value that a view over it must reference a data buffer. + const LONG: &str = "a value that is far too long to inline"; + + /// The heap is adopted zero-copy as a data buffer, the views are built against it in place, + /// and the append composes with staged in-progress bytes on either side. + #[test] + fn test_append_buffer_with_lengths() { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = VarBinViewBuilder::with_capacity_in( + DType::Utf8(Nullability::Nullable), + 8, + BufferAllocatorRef::statically_allocated(), + ); + + // Stages an in-progress buffer the bulk append has to flush first. + builder.append_value(LONG); + + let heap = ByteBuffer::copy_from([LONG.as_bytes(), b"", b"tiny"].concat()); + let heap_ptr = heap.as_ptr(); + let lengths = [u32::try_from(LONG.len()).unwrap(), 0, 4]; + builder.append_buffer_with_lengths(heap, &lengths, &Mask::from_iter([true, false, true])); + + builder.append_value("tail"); + + let actual = builder.finish_into_varbinview(); + // The adopted heap sits after the flushed in-progress buffer, untouched. + assert_eq!(actual.data_buffers()[1].as_host().as_ptr(), heap_ptr); + + let expected = >::from_iter([ + Some(LONG), + Some(LONG), + None, + Some("tiny"), + Some("tail"), + ]); + assert_arrays_eq!(actual, expected, &mut ctx); + } + + /// Offsets need not start at zero: only the referenced range of the heap is adopted. + #[test] + fn test_append_buffer_with_offsets() { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = VarBinViewBuilder::with_capacity_in( + DType::Utf8(Nullability::Nullable), + 8, + BufferAllocatorRef::statically_allocated(), + ); + + let heap = ByteBuffer::copy_from(format!("..{LONG}tiny!!")); + let long_len = u32::try_from(LONG.len()).unwrap(); + let offsets = [2u32, 2 + long_len, 2 + long_len, 2 + long_len + 4]; + builder.append_buffer_with_offsets( + heap.clone(), + &offsets, + &Mask::from_iter([true, false, true]), + ); + + let actual = builder.finish_into_varbinview(); + // Zero-copy adoption of just the `offsets[0]..offsets[last]` range. + // SAFETY: offset 2 is in bounds of the heap. + assert_eq!(actual.data_buffers()[0].as_host().as_ptr(), unsafe { + heap.as_ptr().add(2) + }); + assert_eq!( + actual.data_buffers()[0].len(), + LONG.len() + 4, + "only the referenced range must be adopted" + ); + + let expected = + >::from_iter([Some(LONG), None, Some("tiny")]); + assert_arrays_eq!(actual, expected, &mut ctx); + } + + /// A compacting builder must measure the heap it is handed: values short enough to inline + /// never reference it, so an under-utilized heap is rewritten to just the referencing values. + #[test] + fn test_append_buffer_with_lengths_compacts_underutilized_heap() { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = VarBinViewBuilder::with_compaction_in( + DType::Utf8(Nullability::Nullable), + 4, + 1.0, + BufferAllocatorRef::statically_allocated(), + ); + + let heap = ByteBuffer::copy_from([b"short".as_slice(), LONG.as_bytes(), b"tiny"].concat()); + let lengths = [5u32, u32::try_from(LONG.len()).unwrap(), 4]; + builder.append_buffer_with_lengths(heap, &lengths, &Mask::new_true(3)); + + let actual = builder.finish_into_varbinview(); + assert_eq!(actual.data_buffers().len(), 1); + assert_eq!( + actual.data_buffers()[0].len(), + LONG.len(), + "the compact buffer must hold only the non-inlined values" + ); + + let expected = >::from_iter([ + Some("short"), + Some(LONG), + Some("tiny"), + ]); + assert_arrays_eq!(actual, expected, &mut ctx); + } + + /// Rewriting an under-utilized heap must consume null values' spans without retaining their + /// bytes or producing views that reference them. + #[test] + fn test_append_buffer_with_lengths_compaction_skips_null_bytes() { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = VarBinViewBuilder::with_compaction_in( + DType::Utf8(Nullability::Nullable), + 2, + 1.0, + BufferAllocatorRef::statically_allocated(), + ); + + let heap = ByteBuffer::copy_from([LONG.as_bytes(), LONG.as_bytes()].concat()); + let lengths = [u32::try_from(LONG.len()).unwrap(); 2]; + builder.append_buffer_with_lengths(heap, &lengths, &Mask::from_iter([false, true])); + + let actual = builder.finish_into_varbinview(); + assert_eq!(actual.data_buffers().len(), 1); + assert_eq!( + actual.data_buffers()[0].len(), + LONG.len(), + "the compact buffer must omit bytes belonging to null rows" + ); + + let expected = >::from_iter([None::<&str>, Some(LONG)]); + assert_arrays_eq!(actual, expected, &mut ctx); + } + + /// `push_buffers` returns where each buffer landed, and a deduplicating builder maps a + /// re-pushed buffer back to its existing index instead of holding it twice. + #[test] + fn test_push_buffers_deduplicates() { + let mut builder = VarBinViewBuilder::with_buffer_deduplication_in( + DType::Utf8(Nullability::Nullable), + 8, + BufferAllocatorRef::statically_allocated(), + ); + + let first = ByteBuffer::copy_from(LONG); + let second = ByteBuffer::copy_from("another value far too long to inline"); + + assert_eq!( + builder.push_buffers([first.clone(), second.clone()]), + [0, 1] + ); + assert_eq!(builder.push_buffers([second, first]), [1, 0]); + assert_eq!(builder.completed_block_count(), 2); + } + + /// Gathered views are rebased onto wherever the pushed buffers landed, and null rows never + /// resolve their index. + #[test] + fn test_append_views_gathered() { + let mut ctx = array_session().create_execution_ctx(); + let dictionary = >::from_iter([ + Some("tiny"), + Some(LONG), + Some("small"), + ]); + let buffers = dictionary + .data_buffers() + .iter() + .map(|buffer| buffer.as_host().clone()) + .collect::>(); + + let mut builder = VarBinViewBuilder::with_capacity_in( + DType::Utf8(Nullability::Nullable), + 8, + BufferAllocatorRef::statically_allocated(), + ); + // Stages an in-progress buffer that the gather has to flush ahead of its own buffers. + builder.append_value(LONG); + + let codes: [usize; 4] = [1, 0, usize::MAX, 2]; + let views = dictionary.views(); + builder.append_views_gathered( + buffers, + views, + &Mask::from_iter([true, true, false, true]), + // The null row's code is garbage; the builder must not look it up. + |row| codes[row], + ); + + let expected = >::from_iter([ + Some(LONG), + Some(LONG), + Some("tiny"), + None, + Some("small"), + ]); + assert_arrays_eq!(builder.finish_into_varbinview(), expected, &mut ctx); + } + + /// Scattered patches overwrite the fill view at their rows, and both are rebased onto the + /// adopted buffers. + #[test] + fn test_append_views_scattered() { + use crate::arrays::varbinview::build_views::BinaryView; + + let mut ctx = array_session().create_execution_ctx(); + let patch_values = + >::from_iter([Some(LONG), Some("tiny")]); + let mut buffers = patch_values + .data_buffers() + .iter() + .map(|buffer| buffer.as_host().clone()) + .collect::>(); + + let fill_bytes = ByteBuffer::copy_from("a fill value too long to inline"); + buffers.push(fill_bytes.clone()); + let fill = BinaryView::make_view( + fill_bytes.as_slice(), + u32::try_from(buffers.len() - 1).unwrap(), + 0, + ); + + let mut builder = VarBinViewBuilder::with_capacity_in( + DType::Utf8(Nullability::Nullable), + 8, + BufferAllocatorRef::statically_allocated(), + ); + let views = patch_values.views(); + builder.append_views_scattered( + buffers, + 5, + fill, + [(1usize, views[0]), (3usize, views[1])].into_iter(), + &Mask::from_iter([true, true, false, true, true]), + ); + + let expected = >::from_iter([ + Some("a fill value too long to inline"), + Some(LONG), + None, + Some("tiny"), + Some("a fill value too long to inline"), + ]); + assert_arrays_eq!(builder.finish_into_varbinview(), expected, &mut ctx); + } + + /// A compacting builder must measure the buffers a gather hands it: with the middle + /// dictionary value never referenced, adopting the heap whole would keep its bytes alive. + /// Duplicate codes must share one rewritten copy, and null rows must not resolve their code. + #[test] + fn test_append_views_gathered_compacts_unreferenced_values() { + const DEAD: &str = "a dead value nobody gathers, far too long to inline"; + const OTHER: &str = "another value far too long to inline"; + + let mut ctx = array_session().create_execution_ctx(); + let dictionary = + >::from_iter([Some(LONG), Some(DEAD), Some(OTHER)]); + assert_eq!(dictionary.data_buffers().len(), 1); + let buffers = dictionary + .data_buffers() + .iter() + .map(|buffer| buffer.as_host().clone()) + .collect::>(); + + let mut builder = VarBinViewBuilder::with_compaction_in( + DType::Utf8(Nullability::Nullable), + 8, + 1.0, + BufferAllocatorRef::statically_allocated(), + ); + let codes: [usize; 5] = [2, 0, usize::MAX, 2, 0]; + builder.append_views_gathered( + buffers, + dictionary.views(), + &Mask::from_iter([true, true, false, true, true]), + |row| codes[row], + ); + + let actual = builder.finish_into_varbinview(); + assert_eq!(actual.data_buffers().len(), 1); + assert_eq!( + actual.data_buffers()[0].len(), + LONG.len() + OTHER.len(), + "the rewritten buffer must hold each referenced value exactly once" + ); + + let expected = >::from_iter([ + Some(OTHER), + Some(LONG), + None, + Some(OTHER), + Some(LONG), + ]); + assert_arrays_eq!(actual, expected, &mut ctx); + } + + /// A gather that references only a contiguous tail of the heap must adopt just that slice, + /// zero-copy. + #[test] + fn test_append_views_gathered_slices_contiguous_range() { + const TAIL: &str = "the referenced tail value, too long to inline"; + + let mut ctx = array_session().create_execution_ctx(); + let dictionary = >::from_iter([Some(LONG), Some(TAIL)]); + let heap = dictionary.data_buffers()[0].as_host().clone(); + + let mut builder = VarBinViewBuilder::with_compaction_in( + DType::Utf8(Nullability::Nullable), + 8, + 1.0, + BufferAllocatorRef::statically_allocated(), + ); + builder.append_views_gathered( + [heap.clone()], + dictionary.views(), + &Mask::new_true(2), + |_| 1, + ); + + let actual = builder.finish_into_varbinview(); + assert_eq!(actual.data_buffers().len(), 1); + assert_eq!(actual.data_buffers()[0].len(), TAIL.len()); + // SAFETY: LONG.len() is in bounds of the heap. + assert_eq!(actual.data_buffers()[0].as_host().as_ptr(), unsafe { + heap.as_ptr().add(LONG.len()) + }); + + let expected = >::from_iter([Some(TAIL), Some(TAIL)]); + assert_arrays_eq!(actual, expected, &mut ctx); + } + + /// A buffer whose referenced share clears the threshold is adopted whole, zero-copy. + #[test] + fn test_append_views_gathered_keeps_utilized_buffer() { + const SHORTER: &str = "a shorter long value"; + + let mut ctx = array_session().create_execution_ctx(); + let dictionary = + >::from_iter([Some(LONG), Some(SHORTER)]); + let heap = dictionary.data_buffers()[0].as_host().clone(); + + let mut builder = VarBinViewBuilder::with_compaction_in( + DType::Utf8(Nullability::Nullable), + 8, + 0.5, + BufferAllocatorRef::statically_allocated(), + ); + builder.append_views_gathered( + [heap.clone()], + dictionary.views(), + &Mask::new_true(1), + |_| 0, + ); + + let actual = builder.finish_into_varbinview(); + assert_eq!(actual.data_buffers().len(), 1); + assert_eq!(actual.data_buffers()[0].as_host().as_ptr(), heap.as_ptr()); + assert_eq!(actual.data_buffers()[0].len(), heap.len()); + + let expected = >::from_iter([Some(LONG)]); + assert_arrays_eq!(actual, expected, &mut ctx); + } + + /// An all-null gather references nothing, so a compacting builder must not retain any of the + /// buffers — and must never resolve a code. + #[test] + fn test_append_views_gathered_all_null_drops_buffers() { + let mut ctx = array_session().create_execution_ctx(); + let dictionary = >::from_iter([Some(LONG)]); + let buffers = dictionary + .data_buffers() + .iter() + .map(|buffer| buffer.as_host().clone()) + .collect::>(); + + let mut builder = VarBinViewBuilder::with_compaction_in( + DType::Utf8(Nullability::Nullable), + 8, + 1.0, + BufferAllocatorRef::statically_allocated(), + ); + builder.append_views_gathered(buffers, dictionary.views(), &Mask::new_false(3), |_| { + usize::MAX + }); + + let actual = builder.finish_into_varbinview(); + assert!( + actual.data_buffers().is_empty(), + "an all-null gather must not retain any buffers" + ); + + let expected = >::from_iter([None::<&str>, None, None]); + assert_arrays_eq!(actual, expected, &mut ctx); + } + + /// A fully-inlined heap has zero utilization; a compacting builder must not retain it at all. + #[test] + fn test_append_buffer_with_lengths_drops_fully_inlined_heap() { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = VarBinViewBuilder::with_compaction_in( + DType::Utf8(Nullability::Nullable), + 4, + 1.0, + BufferAllocatorRef::statically_allocated(), + ); + + let heap = ByteBuffer::copy_from(b"shorttinysmall".as_slice()); + builder.append_buffer_with_lengths(heap, &[5u32, 4, 5], &Mask::new_true(3)); + + let actual = builder.finish_into_varbinview(); + assert!( + actual.data_buffers().is_empty(), + "a fully-inlined append must not retain any value bytes" + ); + + let expected = >::from_iter([ + Some("short"), + Some("tiny"), + Some("small"), + ]); + assert_arrays_eq!(actual, expected, &mut ctx); + } + #[test] fn test_utf8_builder() { let mut ctx = array_session().create_execution_ctx(); - let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); + let mut builder = VarBinViewBuilder::with_capacity_in( + DType::Utf8(Nullability::Nullable), + 10, + BufferAllocatorRef::statically_allocated(), + ); builder.append_value("Hello"); builder.append_null(); @@ -908,13 +1875,20 @@ mod tests { fn test_utf8_builder_with_extend() { let mut ctx = array_session().create_execution_ctx(); let array = { - let mut builder = - VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); + let mut builder = VarBinViewBuilder::with_capacity_in( + DType::Utf8(Nullability::Nullable), + 10, + BufferAllocatorRef::statically_allocated(), + ); builder.append_null(); builder.append_value("Hello2"); builder.finish() }; - let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); + let mut builder = VarBinViewBuilder::with_capacity_in( + DType::Utf8(Nullability::Nullable), + 10, + BufferAllocatorRef::statically_allocated(), + ); builder.append_value("Hello1"); array.append_to_builder(&mut builder, &mut ctx).unwrap(); @@ -936,16 +1910,22 @@ mod tests { #[test] fn test_buffer_deduplication() -> VortexResult<()> { let array = { - let mut builder = - VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); + let mut builder = VarBinViewBuilder::with_capacity_in( + DType::Utf8(Nullability::Nullable), + 10, + BufferAllocatorRef::statically_allocated(), + ); builder.append_value("This is a long string that should not be inlined"); builder.append_value("short string"); builder.finish_into_varbinview() }; assert_eq!(array.data_buffers().len(), 1); - let mut builder = - VarBinViewBuilder::with_buffer_deduplication(DType::Utf8(Nullability::Nullable), 10); + let mut builder = VarBinViewBuilder::with_buffer_deduplication_in( + DType::Utf8(Nullability::Nullable), + 10, + BufferAllocatorRef::statically_allocated(), + ); let mut ctx = array_session().create_execution_ctx(); @@ -961,8 +1941,11 @@ mod tests { assert_eq!(builder.completed_block_count(), 1); let array2 = { - let mut builder = - VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); + let mut builder = VarBinViewBuilder::with_capacity_in( + DType::Utf8(Nullability::Nullable), + 10, + BufferAllocatorRef::statically_allocated(), + ); builder.append_value("This is a long string that should not be inlined"); builder.finish_into_varbinview() }; @@ -986,8 +1969,11 @@ mod tests { use crate::scalar::Scalar; // Test with Utf8 builder. - let mut utf8_builder = - VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); + let mut utf8_builder = VarBinViewBuilder::with_capacity_in( + DType::Utf8(Nullability::Nullable), + 10, + BufferAllocatorRef::statically_allocated(), + ); // Test appending a valid utf8 value. let utf8_scalar1 = Scalar::utf8("hello", Nullability::Nullable); @@ -1007,8 +1993,11 @@ mod tests { assert_arrays_eq!(&array, &expected, &mut ctx); // Test with Binary builder. - let mut binary_builder = - VarBinViewBuilder::with_capacity(DType::Binary(Nullability::Nullable), 10); + let mut binary_builder = VarBinViewBuilder::with_capacity_in( + DType::Binary(Nullability::Nullable), + 10, + BufferAllocatorRef::statically_allocated(), + ); let binary_scalar = Scalar::binary(vec![1u8, 2, 3], Nullability::Nullable); binary_builder.append_scalar(&binary_scalar).unwrap(); @@ -1022,8 +2011,11 @@ mod tests { assert_arrays_eq!(&binary_array, &expected, &mut ctx); // Test wrong dtype error. - let mut builder = - VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::NonNullable), 10); + let mut builder = VarBinViewBuilder::with_capacity_in( + DType::Utf8(Nullability::NonNullable), + 10, + BufferAllocatorRef::statically_allocated(), + ); let wrong_scalar = Scalar::from(42i32); assert!(builder.append_scalar(&wrong_scalar).is_err()); } @@ -1056,12 +2048,13 @@ mod tests { use super::BufferGrowthStrategy; use super::VarBinViewBuilder; - let mut builder = VarBinViewBuilder::new( + let mut builder = VarBinViewBuilder::new_in( DType::Binary(Nullability::Nullable), 10, Default::default(), BufferGrowthStrategy::exponential(1024, 4096), 0.0, + BufferAllocatorRef::statically_allocated(), ); // Create a value larger than max_size diff --git a/vortex-array/src/builtins.rs b/vortex-array/src/builtins.rs index d3b27bb8753..eb26931e753 100644 --- a/vortex-array/src/builtins.rs +++ b/vortex-array/src/builtins.rs @@ -15,7 +15,6 @@ use crate::ArrayRef; use crate::IntoArray; use crate::arrays::ConstantArray; use crate::arrays::InterleaveArray; -use crate::arrays::scalar_fn::ScalarFnFactoryExt; use crate::dtype::DType; use crate::dtype::FieldName; use crate::expr::Expression; @@ -173,8 +172,7 @@ impl ArrayBuiltins for ArrayRef { if self.dtype() == &dtype { return Ok(self.clone()); } - Cast.try_new_array(self.len(), dtype, [self.clone()])? - .optimize() + Cast::new(self.clone(), dtype).into_array().optimize() } fn fill_null(&self, fill_value: impl Into) -> VortexResult { @@ -182,48 +180,38 @@ impl ArrayBuiltins for ArrayRef { if !self.dtype().is_nullable() { return self.cast(fill_value.dtype().clone()); } - FillNull - .try_new_array( - self.len(), - EmptyOptions, - [ - self.clone(), - ConstantArray::new(fill_value, self.len()).into_array(), - ], - )? - .optimize() + FillNull::try_new( + self.clone(), + ConstantArray::new(fill_value, self.len()).into_array(), + )? + .into_array() + .optimize() } fn get_item(&self, field_name: impl Into) -> VortexResult { - GetItem - .try_new_array(self.len(), field_name.into(), [self.clone()])? + GetItem::try_new(self.clone(), field_name)? + .into_array() .optimize() } fn is_null(&self) -> VortexResult { - IsNull - .try_new_array(self.len(), EmptyOptions, [self.clone()])? - .optimize() + IsNull::new(self.clone()).into_array().optimize() } fn is_not_null(&self) -> VortexResult { - IsNotNull - .try_new_array(self.len(), EmptyOptions, [self.clone()])? - .optimize() + IsNotNull::new(self.clone()).into_array().optimize() } fn mask(self, mask: ArrayRef) -> VortexResult { - Mask.try_new_array(self.len(), EmptyOptions, [self, mask])? - .optimize() + Mask::try_new(self, mask)?.into_array().optimize() } fn not(&self) -> VortexResult { - Not.try_new_array(self.len(), EmptyOptions, [self.clone()])? - .optimize() + Not::try_new(self.clone())?.into_array().optimize() } fn zip(&self, if_true: ArrayRef, if_false: ArrayRef) -> VortexResult { - Zip.try_new_array(self.len(), EmptyOptions, [if_true, if_false, self.clone()]) + Ok(Zip::try_new(if_true, if_false, self.clone())?.into_array()) } fn interleave( @@ -238,14 +226,14 @@ impl ArrayBuiltins for ArrayRef { } fn list_contains(&self, value: ArrayRef) -> VortexResult { - ListContains - .try_new_array(self.len(), EmptyOptions, [self.clone(), value])? + ListContains::try_new(self.clone(), value)? + .into_array() .optimize() } fn binary(&self, rhs: ArrayRef, op: Operator) -> VortexResult { - Binary - .try_new_array(self.len(), op, [self.clone(), rhs])? + Binary::try_new(self.clone(), rhs, op)? + .into_array() .optimize() } @@ -255,8 +243,8 @@ impl ArrayBuiltins for ArrayRef { upper: ArrayRef, options: BetweenOptions, ) -> VortexResult { - Between - .try_new_array(self.len(), options, [self, lower, upper])? + Between::try_new(self, lower, upper, options)? + .into_array() .optimize() } } diff --git a/vortex-array/src/canonical.rs b/vortex-array/src/canonical.rs index 71da52e1300..5330aa2f1ed 100644 --- a/vortex-array/src/canonical.rs +++ b/vortex-array/src/canonical.rs @@ -29,12 +29,16 @@ use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; use crate::arrays::ListView; use crate::arrays::ListViewArray; +use crate::arrays::Map; +use crate::arrays::MapArray; use crate::arrays::Null; use crate::arrays::NullArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::Struct; use crate::arrays::StructArray; +use crate::arrays::Union; +use crate::arrays::UnionArray; use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; use crate::arrays::Variant; @@ -45,10 +49,13 @@ use crate::arrays::extension::ExtensionArrayExt; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::arrays::listview::ListViewDataParts; use crate::arrays::listview::ListViewRebuildMode; +use crate::arrays::map::MapArrayExt; +use crate::arrays::map::MapArraySlotsExt; use crate::arrays::primitive::PrimitiveDataParts; use crate::arrays::struct_::StructDataParts; +use crate::arrays::union::UnionDataParts; use crate::arrays::varbinview::VarBinViewDataParts; -use crate::arrays::variant::VariantArrayExt; +use crate::arrays::variant::VariantArraySlotsExt; use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::Nullability; @@ -69,7 +76,7 @@ use crate::validity::Validity; /// /// Each `Canonical` variant has a corresponding [`DType`] variant, with the notable exception of /// [`Canonical::VarBinView`], which is the canonical encoding for both [`DType::Utf8`] and -/// [`DType::Binary`]. [`DType::Union`] does not yet have a public canonical array. +/// [`DType::Binary`]. /// /// # Laziness /// @@ -80,8 +87,10 @@ use crate::validity::Validity; /// /// # Arrow interoperability /// -/// All of the Vortex canonical encodings have an equivalent Arrow encoding that can be built -/// zero-copy, and the corresponding Arrow array types can also be built directly. +/// Most Vortex canonical encodings have an equivalent Arrow encoding that can be built zero-copy, +/// and the corresponding Arrow array types can also be built directly. Map array Arrow transport is +/// not implemented yet, and [`UnionArray`]'s independent top-level validity cannot be represented +/// directly by an Arrow union. /// /// The full list of canonical types and their equivalent Arrow array types are: /// @@ -91,6 +100,7 @@ use crate::validity::Validity; /// * `DecimalArray`: `arrow_array::Decimal128Array` and `arrow_array::Decimal256Array` /// * `VarBinViewArray`: `arrow_array::GenericByteViewArray` /// * `ListViewArray`: `arrow_array::ListViewArray` +/// * `MapArray`: Vortex `ListView>` storage /// * `FixedSizeListArray`: `arrow_array::FixedSizeListArray` /// * `StructArray`: `arrow_array::StructArray` /// @@ -126,8 +136,10 @@ pub enum Canonical { Decimal(DecimalArray), VarBinView(VarBinViewArray), List(ListViewArray), + Map(MapArray), FixedSizeList(FixedSizeListArray), Struct(StructArray), + Union(UnionArray), /// Canonical storage for extension dtypes, wrapping the canonical form of the storage dtype. Extension(ExtensionArray), /// Canonical storage for dynamic variant values, optionally with typed shredded paths. @@ -144,8 +156,10 @@ macro_rules! match_each_canonical { Canonical::Decimal($ident) => $eval, Canonical::VarBinView($ident) => $eval, Canonical::List($ident) => $eval, + Canonical::Map($ident) => $eval, Canonical::FixedSizeList($ident) => $eval, Canonical::Struct($ident) => $eval, + Canonical::Union($ident) => $eval, Canonical::Variant($ident) => $eval, Canonical::Extension($ident) => $eval, } @@ -209,6 +223,14 @@ impl Canonical { // An empty list view is trivially copyable to a list. .with_zero_copy_to_list(true) }), + DType::Map(map_dtype, nullability) => Canonical::Map(MapArray::new( + map_dtype.clone(), + Canonical::empty(&DType::List( + Arc::new(map_dtype.entries_dtype()), + *nullability, + )) + .into_listview(), + )), DType::FixedSizeList(elem_dtype, list_size, null) => Canonical::FixedSizeList(unsafe { FixedSizeListArray::new_unchecked( Canonical::empty(elem_dtype).into_array(), @@ -222,13 +244,15 @@ impl Canonical { struct_dtype .fields() .map(|f| Canonical::empty(&f).into_array()) - .collect::>(), + .collect::>(), struct_dtype.clone(), 0, Validity::from(n), ) }), - DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), + DType::Union(variants, nullability) => { + Canonical::Union(UnionArray::empty(variants.clone(), *nullability)) + } DType::Variant(_) => { vortex_panic!(InvalidArgument: "Canonical empty is not supported for Variant") } @@ -266,6 +290,14 @@ impl Canonical { Canonical::List(array) => Ok(Canonical::List( array.rebuild(ListViewRebuildMode::TrimElements, ctx)?, )), + Canonical::Map(array) => Ok(Canonical::Map(MapArray::new( + array.map_dtype().clone(), + array + .entries() + .as_::() + .into_owned() + .rebuild(ListViewRebuildMode::TrimElements, ctx)?, + ))), _ => Ok(self.clone()), } } @@ -369,6 +401,22 @@ impl Canonical { } } + pub fn as_map(&self) -> &MapArray { + if let Canonical::Map(a) = self { + a + } else { + vortex_panic!("Cannot get MapArray from {:?}", &self) + } + } + + pub fn into_map(self) -> MapArray { + if let Canonical::Map(a) = self { + a + } else { + vortex_panic!("Cannot unwrap MapArray from {:?}", &self) + } + } + pub fn as_fixed_size_list(&self) -> &FixedSizeListArray { if let Canonical::FixedSizeList(a) = self { a @@ -401,6 +449,24 @@ impl Canonical { } } + /// Return this canonical array as a sparse [`UnionArray`]. + pub fn as_union(&self) -> &UnionArray { + if let Canonical::Union(a) = self { + a + } else { + vortex_panic!("Cannot get UnionArray from {:?}", &self) + } + } + + /// Unwrap this canonical array as a sparse [`UnionArray`]. + pub fn into_union(self) -> UnionArray { + if let Canonical::Union(a) = self { + a + } else { + vortex_panic!("Cannot unwrap UnionArray from {:?}", &self) + } + } + pub fn as_extension(&self) -> &ExtensionArray { if let Canonical::Extension(a) = self { a @@ -457,6 +523,10 @@ pub trait ToCanonical { #[deprecated(note = "use `array.execute::(ctx)` instead")] fn to_listview(&self) -> ListViewArray; + /// Canonicalize into a [`MapArray`] if the target is [`Map`](DType::Map) typed. + #[deprecated(note = "use `array.execute::(ctx)` instead")] + fn to_map(&self) -> MapArray; + /// Canonicalize into a [`FixedSizeListArray`] if the target is [`List`](DType::FixedSizeList) /// typed. #[deprecated(note = "use `array.execute::(ctx)` instead")] @@ -512,6 +582,12 @@ impl ToCanonical for ArrayRef { result.into_listview() } + fn to_map(&self) -> MapArray { + #[expect(deprecated)] + let result = self.to_canonical().vortex_expect("to_canonical failed"); + result.into_map() + } + fn to_fixed_size_list(&self) -> FixedSizeListArray { #[expect(deprecated)] let result = self.to_canonical().vortex_expect("to_canonical failed"); @@ -631,6 +707,14 @@ impl Executable for CanonicalValidity { .with_zero_copy_to_list(zctl) }))) } + Canonical::Map(map) => { + let map_dtype = map.map_dtype().clone(); + let entries = map.entries().clone(); + Ok(CanonicalValidity(Canonical::Map(MapArray::new( + map_dtype, + entries.execute::(ctx)?.0.into_listview(), + )))) + } Canonical::FixedSizeList(fsl) => { let list_size = fsl.list_size(); let len = fsl.len(); @@ -652,6 +736,18 @@ impl Executable for CanonicalValidity { StructArray::new_unchecked(fields, struct_fields, len, validity.execute(ctx)?) }))) } + Canonical::Union(union) => { + let UnionDataParts { + variants, + type_ids, + children, + } = union.into_data_parts(); + let type_ids = type_ids.execute::(ctx)?.0.into_array(); + + Ok(CanonicalValidity(Canonical::Union(unsafe { + UnionArray::new_unchecked(type_ids, variants, children.iter().cloned()) + }))) + } Canonical::Extension(ext) => Ok(CanonicalValidity(Canonical::Extension( ExtensionArray::new( ext.ext_dtype().clone(), @@ -793,6 +889,17 @@ impl Executable for RecursiveCanonical { .with_zero_copy_to_list(zctl) }))) } + Canonical::Map(map) => { + let map_dtype = map.map_dtype().clone(); + let entries = map.entries().clone(); + Ok(RecursiveCanonical(Canonical::Map(MapArray::new( + map_dtype, + entries + .execute::(ctx)? + .0 + .into_listview(), + )))) + } Canonical::FixedSizeList(fsl) => { let list_size = fsl.list_size(); let len = fsl.len(); @@ -816,9 +923,9 @@ impl Executable for RecursiveCanonical { validity, } = st.into_data_parts(); let executed_fields = fields - .iter() - .map(|f| Ok(f.clone().execute::(ctx)?.0.into_array())) - .collect::>>()?; + .into_iter() + .map(|f| Ok(f.execute::(ctx)?.0.into_array())) + .collect::>>()?; Ok(RecursiveCanonical(Canonical::Struct(unsafe { StructArray::new_unchecked( @@ -829,6 +936,27 @@ impl Executable for RecursiveCanonical { ) }))) } + Canonical::Union(union) => { + let UnionDataParts { + variants, + type_ids, + children, + } = union.into_data_parts(); + let type_ids = type_ids.execute::(ctx)?.0.into_array(); + let children = children + .iter() + .cloned() + .map(|child| { + child + .execute::(ctx) + .map(|canonical| canonical.0.into_array()) + }) + .collect::>>()?; + + Ok(RecursiveCanonical(Canonical::Union(unsafe { + UnionArray::new_unchecked(type_ids, variants, children) + }))) + } Canonical::Extension(ext) => Ok(RecursiveCanonical(Canonical::Extension( ExtensionArray::new( ext.ext_dtype().clone(), @@ -979,6 +1107,18 @@ impl Executable for ListViewArray { } } +/// Execute the array to canonical form and unwrap as a [`MapArray`]. +/// +/// This will panic if the array's dtype is not map. +impl Executable for MapArray { + fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + match array.try_downcast::() { + Ok(map) => Ok(map), + Err(array) => Ok(Canonical::execute(array, ctx)?.into_map()), + } + } +} + /// Execute the array to canonical form and unwrap as a [`FixedSizeListArray`]. /// /// This will panic if the array's dtype is not fixed size list. @@ -1003,6 +1143,18 @@ impl Executable for StructArray { } } +/// Execute the array to canonical form and unwrap as a [`UnionArray`]. +/// +/// This will panic if the array's dtype is not union. +impl Executable for UnionArray { + fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + match array.try_downcast::() { + Ok(union_array) => Ok(union_array), + Err(array) => Ok(Canonical::execute(array, ctx)?.into_union()), + } + } +} + /// Execute the array to canonical form and unwrap as a [`VariantArray`]. /// /// This will panic if the array's dtype is not variant. @@ -1030,8 +1182,10 @@ pub enum CanonicalView<'a> { Decimal(ArrayView<'a, Decimal>), VarBinView(ArrayView<'a, VarBinView>), List(ArrayView<'a, ListView>), + Map(ArrayView<'a, Map>), FixedSizeList(ArrayView<'a, FixedSizeList>), Struct(ArrayView<'a, Struct>), + Union(ArrayView<'a, Union>), Extension(ArrayView<'a, Extension>), Variant(ArrayView<'a, Variant>), } @@ -1045,8 +1199,10 @@ impl From> for Canonical { CanonicalView::Decimal(a) => Canonical::Decimal(a.into_owned()), CanonicalView::VarBinView(a) => Canonical::VarBinView(a.into_owned()), CanonicalView::List(a) => Canonical::List(a.into_owned()), + CanonicalView::Map(a) => Canonical::Map(a.into_owned()), CanonicalView::FixedSizeList(a) => Canonical::FixedSizeList(a.into_owned()), CanonicalView::Struct(a) => Canonical::Struct(a.into_owned()), + CanonicalView::Union(a) => Canonical::Union(a.into_owned()), CanonicalView::Extension(a) => Canonical::Extension(a.into_owned()), CanonicalView::Variant(a) => Canonical::Variant(a.into_owned()), } @@ -1063,8 +1219,10 @@ impl CanonicalView<'_> { CanonicalView::Decimal(a) => a.array().clone(), CanonicalView::VarBinView(a) => a.array().clone(), CanonicalView::List(a) => a.array().clone(), + CanonicalView::Map(a) => a.array().clone(), CanonicalView::FixedSizeList(a) => a.array().clone(), CanonicalView::Struct(a) => a.array().clone(), + CanonicalView::Union(a) => a.array().clone(), CanonicalView::Extension(a) => a.array().clone(), CanonicalView::Variant(a) => a.array().clone(), } @@ -1083,7 +1241,9 @@ impl Matcher for AnyCanonical { || array.is::() || array.is::() || array.is::() + || array.is::() || array.is::() + || array.is::() || array.is::() || array.is::() || array.is::() @@ -1102,8 +1262,12 @@ impl Matcher for AnyCanonical { Some(CanonicalView::Decimal(a)) } else if let Some(a) = array.as_opt::() { Some(CanonicalView::Struct(a)) + } else if let Some(a) = array.as_opt::() { + Some(CanonicalView::Union(a)) } else if let Some(a) = array.as_opt::() { Some(CanonicalView::List(a)) + } else if let Some(a) = array.as_opt::() { + Some(CanonicalView::Map(a)) } else if let Some(a) = array.as_opt::() { Some(CanonicalView::FixedSizeList(a)) } else if let Some(a) = array.as_opt::() { @@ -1136,7 +1300,7 @@ mod test { use crate::arrays::Variant; use crate::arrays::VariantArray; use crate::arrays::struct_::StructArrayExt; - use crate::arrays::variant::VariantArrayExt; + use crate::arrays::variant::VariantArraySlotsExt; use crate::canonical::StructArray; use crate::dtype::Nullability; use crate::scalar::Scalar; diff --git a/vortex-array/src/compute/conformance/binary_numeric.rs b/vortex-array/src/compute/conformance/binary_numeric.rs index 1fe409bf648..1750ffc8e5d 100644 --- a/vortex-array/src/compute/conformance/binary_numeric.rs +++ b/vortex-array/src/compute/conformance/binary_numeric.rs @@ -25,6 +25,10 @@ use std::fmt::Debug; use itertools::Itertools; use num_traits::Bounded; +use num_traits::CheckedAdd; +use num_traits::CheckedDiv; +use num_traits::CheckedMul; +use num_traits::CheckedSub; use num_traits::Float; use num_traits::Num; use num_traits::Signed; @@ -40,11 +44,16 @@ use crate::RecursiveCanonical; use crate::arrays::ConstantArray; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::NativeDecimalType; use crate::dtype::NativePType; use crate::dtype::PType; +use crate::dtype::i256; +use crate::scalar::DecimalValue; use crate::scalar::NumericOperator; use crate::scalar::PrimitiveScalar; use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::numeric_op_result_decimal_dtype; fn to_vec_of_scalar(array: &ArrayRef, ctx: &mut ExecutionCtx) -> Vec { // Not fast, but obviously correct @@ -240,12 +249,220 @@ pub fn test_binary_numeric_array(array: &ArrayRef, ctx: &mut ExecutionCtx) { PType::F32 => test_binary_numeric_conformance::(array, ctx), PType::F64 => test_binary_numeric_conformance::(array, ctx), }, + DType::Decimal(decimal_dtype, _) => { + test_binary_numeric_conformance_decimal(array, *decimal_dtype, ctx) + } dtype => vortex_panic!( - "Binary numeric tests are only supported for primitive numeric types, got {dtype}", + "Binary numeric tests are only supported for primitive and decimal types, got {dtype}", ), } } +/// Tests binary numeric operations on a decimal array against the decimal scalar +/// implementation, using a set of representative constants: stored `0`, `1`, `-1`, the decimal +/// `1.0` (when it fits the precision), and the largest stored value for the precision. +fn test_binary_numeric_conformance_decimal( + array: &ArrayRef, + decimal_dtype: DecimalDType, + ctx: &mut ExecutionCtx, +) { + let precision = decimal_dtype.precision() as usize; + let prec_max = ::MAX_BY_PRECISION[precision]; + + let mut constants = vec![ + i256::from_i128(0), + i256::from_i128(1), + i256::from_i128(-1), + prec_max, + ]; + // The decimal value 1.0 (stored 10^s), when representable within the precision. + if decimal_dtype.scale() >= 0 + && let Some(one) = i256::from_i128(10).checked_pow(decimal_dtype.scale() as u32) + && one <= prec_max + { + constants.push(one); + } + + for constant in constants { + let value = DecimalValue::try_from_i256(constant, decimal_dtype) + .vortex_expect("conformance constants fit the precision"); + test_decimal_binary_numeric_with_scalar(array, value, decimal_dtype, ctx); + } +} + +fn test_decimal_binary_numeric_with_scalar( + array: &ArrayRef, + value: DecimalValue, + decimal_dtype: DecimalDType, + ctx: &mut ExecutionCtx, +) { + let canonicalized_array = array + .clone() + .execute::(ctx) + .vortex_expect("Must be able to canonicalise") + .into_array(); + let original_values = to_vec_of_scalar(&canonicalized_array, ctx); + + let scalar = Scalar::decimal(value, decimal_dtype, array.dtype().nullability()); + + let mut operators = vec![ + NumericOperator::Add, + NumericOperator::Sub, + NumericOperator::Mul, + ]; + if !value.is_zero() { + operators.push(NumericOperator::Div); + } + + for operator in operators { + for lhs_is_array in [true, false] { + test_decimal_binary_numeric_direction( + array, + &original_values, + &scalar, + decimal_dtype, + operator, + lhs_is_array, + ctx, + ); + } + } +} + +fn test_decimal_binary_numeric_direction( + array: &ArrayRef, + original_values: &[Scalar], + scalar: &Scalar, + decimal_dtype: DecimalDType, + operator: NumericOperator, + lhs_is_array: bool, + ctx: &mut ExecutionCtx, +) { + let result_decimal_dtype = numeric_op_result_decimal_dtype(decimal_dtype, operator) + .vortex_expect("decimal arithmetic must have a result dtype"); + let result_dtype = DType::Decimal(result_decimal_dtype, array.dtype().nullability()); + let expected_results = expected_decimal_results( + original_values, + scalar, + operator, + result_decimal_dtype, + &result_dtype, + lhs_is_array, + ); + + let constant = ConstantArray::new(scalar.clone(), array.len()).into_array(); + let (lhs, rhs) = if lhs_is_array { + (array.clone(), constant) + } else { + (constant, array.clone()) + }; + let result = lhs + .binary(rhs, operator.into()) + .vortex_expect("binary decimal op shouldn't fail") + .execute::(ctx) + .map(|c| c.0.into_array()); + + assert_decimal_results( + result, + &expected_results, + array, + scalar, + operator, + lhs_is_array, + ctx, + ); +} + +fn expected_decimal_results( + original_values: &[Scalar], + scalar: &Scalar, + operator: NumericOperator, + result_decimal_dtype: DecimalDType, + result_dtype: &DType, + lhs_is_array: bool, +) -> Vec> { + original_values + .iter() + .map(|value| { + let (lhs, rhs) = if lhs_is_array { + (value.as_decimal(), scalar.as_decimal()) + } else { + (scalar.as_decimal(), value.as_decimal()) + }; + let (Some(lhs), Some(rhs)) = (lhs.decimal_value(), rhs.decimal_value()) else { + return Some(Scalar::null(result_dtype.clone())); + }; + let lhs = lhs.as_i256(); + let rhs = rhs.as_i256(); + let value = match operator { + NumericOperator::Add => lhs.checked_add(&rhs), + NumericOperator::Sub => lhs.checked_sub(&rhs), + NumericOperator::Mul => lhs.checked_mul(&rhs), + NumericOperator::Div => { + let scale_power = result_decimal_dtype.scale(); + let factor = + i256::from_i128(10).checked_pow(scale_power.unsigned_abs() as u32)?; + if scale_power >= 0 { + lhs.checked_mul(&factor)?.checked_div(&rhs) + } else { + lhs.checked_div(&rhs.checked_mul(&factor)?) + } + } + }?; + let value = DecimalValue::try_from_i256(value, result_decimal_dtype).ok()?; + Some(Scalar::decimal( + value, + result_decimal_dtype, + result_dtype.nullability(), + )) + }) + .collect() +} + +fn assert_decimal_results( + result: vortex_error::VortexResult, + expected_results: &[Option], + array: &ArrayRef, + scalar: &Scalar, + operator: NumericOperator, + lhs_is_array: bool, + ctx: &mut ExecutionCtx, +) { + if expected_results.iter().any(Option::is_none) { + assert!( + result.is_err(), + "Decimal binary numeric operation should overflow for encoding {}: \ + {operator:?} {scalar} (lhs_is_array: {lhs_is_array})", + array.encoding_id(), + ); + return; + } + + let result = result.unwrap_or_else(|err| { + vortex_panic!( + "Decimal binary numeric operation unexpectedly failed for encoding {}: \ + {operator:?} {scalar} (lhs_is_array: {lhs_is_array}): {err}", + array.encoding_id(), + ) + }); + + let actual_values = to_vec_of_scalar(&result, ctx); + for (idx, (actual, expected)) in actual_values.iter().zip(expected_results).enumerate() { + let expected_value = expected + .as_ref() + .vortex_expect("non-overflowing decimal lane must have an expected value"); + assert_eq!( + actual, + expected_value, + "Decimal binary numeric operation failed for encoding {} at index {}: \ + ({array:?})[{idx}] {operator:?} {scalar} (lhs_is_array: {lhs_is_array}) \ + expected {expected_value:?}, got {actual:?}", + array.encoding_id(), + idx, + ); + } +} + /// Tests binary numeric operations with edge case scalar values. /// /// This function tests operations with scalar values: diff --git a/vortex-array/src/compute/conformance/cast.rs b/vortex-array/src/compute/conformance/cast.rs index abe23e6e2dc..6f9bd4aefd0 100644 --- a/vortex-array/src/compute/conformance/cast.rs +++ b/vortex-array/src/compute/conformance/cast.rs @@ -132,6 +132,11 @@ fn test_cast_from_null(array: &ArrayRef, ctx: &mut ExecutionCtx) { } fn test_cast_to_non_nullable(array: &ArrayRef, ctx: &mut ExecutionCtx) { + // DType::Null has no non-nullable form + if &DType::Null == array.dtype() { + return; + } + if array .invalid_count(ctx) .vortex_expect("invalid_count should succeed in conformance test") @@ -170,8 +175,8 @@ fn test_cast_to_non_nullable(array: &ArrayRef, ctx: &mut ExecutionCtx) { } } else { if &DType::Null == array.dtype() { - // DType::Null.as_nonnullable() (confusingly) returns DType:Null. Of course, a null - // array can be casted to DType::Null. + // DType::Null has no non-nullable form. A null array can only + // be cast to DType::Null return; } cast_and_execute(&array.clone(), array.dtype().as_nonnullable(), ctx) diff --git a/vortex-array/src/compute/conformance/consistency.rs b/vortex-array/src/compute/conformance/consistency.rs index 73f98cce4a7..cc1d826bdcf 100644 --- a/vortex-array/src/compute/conformance/consistency.rs +++ b/vortex-array/src/compute/conformance/consistency.rs @@ -1092,6 +1092,26 @@ fn test_slice_aggregate_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { } } +fn widened_primitive_dtype(dtype: &DType) -> Option { + let DType::Primitive(ptype, nullability) = dtype else { + return None; + }; + let widened = match ptype { + PType::U8 => PType::U16, + PType::U16 => PType::U32, + PType::U32 => PType::U64, + PType::U64 => return None, + PType::I8 => PType::I16, + PType::I16 => PType::I32, + PType::I32 => PType::I64, + PType::I64 => return None, + PType::F16 => PType::F32, + PType::F32 => PType::F64, + PType::F64 => return None, + }; + Some(DType::Primitive(widened, *nullability)) +} + /// Tests that cast operations preserve array properties when sliced. /// /// # Invariant @@ -1227,6 +1247,45 @@ fn test_cast_slice_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { opposite, )] } + DType::Map(map_dtype, nullability) => { + let opposite = match nullability { + Nullability::NonNullable => Nullability::Nullable, + Nullability::Nullable => Nullability::NonNullable, + }; + let key_dtype = map_dtype.key_dtype(); + let value_dtype = map_dtype.value_dtype(); + let mut targets = vec![DType::Map(map_dtype.clone(), opposite)]; + + if let Some(widened_key) = widened_primitive_dtype(&key_dtype) + && let Ok(dtype) = DType::map( + widened_key, + value_dtype.clone(), + map_dtype.keys_sorted(), + *nullability, + ) + { + targets.push(dtype); + } + + if let Some(widened_value) = widened_primitive_dtype(&value_dtype) + && let Ok(dtype) = DType::map( + key_dtype.clone(), + widened_value, + map_dtype.keys_sorted(), + *nullability, + ) + { + targets.push(dtype); + } + + if map_dtype.keys_sorted() + && let Ok(dtype) = DType::map(key_dtype, value_dtype, false, *nullability) + { + targets.push(dtype); + } + + targets + } DType::Struct(fields, nullability) => { let opposite = match nullability { Nullability::NonNullable => Nullability::Nullable, diff --git a/vortex-array/src/display/extractor.rs b/vortex-array/src/display/extractor.rs index bbfc25a3666..0998872529d 100644 --- a/vortex-array/src/display/extractor.rs +++ b/vortex-array/src/display/extractor.rs @@ -1,9 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::fmt; +pub use vortex_utils::tree::IndentedFormatter; +use vortex_utils::tree::TreeDisplayContext; +pub use vortex_utils::tree::TreeDisplayExtractor as TreeExtractor; use crate::ArrayRef; +use crate::arrays::Chunked; /// Context threaded through tree traversal for percentage calculations etc. pub struct TreeContext { @@ -25,73 +28,18 @@ impl TreeContext { pub fn parent_total_size(&self) -> Option { self.ancestor_sizes.last().cloned().flatten() } - - pub(crate) fn push(&mut self, size: Option) { - self.ancestor_sizes.push(size); - } - - pub(crate) fn pop(&mut self) { - self.ancestor_sizes.pop(); - } } -/// Wrapper providing access to a [`fmt::Formatter`] and the current indentation string. -pub struct IndentedFormatter<'a, 'b> { - inner: &'a mut fmt::Formatter<'b>, - indent: &'a str, -} - -impl<'a, 'b> IndentedFormatter<'a, 'b> { - pub(crate) fn new(f: &'a mut fmt::Formatter<'b>, indent: &'a str) -> Self { - Self { inner: f, indent } - } - - /// Access the indent string and underlying [`fmt::Formatter`] together. - pub fn parts(&mut self) -> (&str, &mut fmt::Formatter<'b>) { - (self.indent, self.inner) +impl TreeDisplayContext for TreeContext { + fn push_parent(&mut self, parent: &ArrayRef) { + self.ancestor_sizes.push(if parent.is::() { + None + } else { + Some(parent.nbytes()) + }); } - /// The current indentation string. - pub fn indent(&self) -> &str { - self.indent - } - - /// Access the underlying [`fmt::Formatter`]. - pub fn formatter(&mut self) -> &mut fmt::Formatter<'b> { - self.inner - } -} - -/// Trait for contributing display information to tree nodes. -/// -/// Each extractor represents one "dimension" of display (e.g., nbytes, stats, metadata, buffers). -/// Extractors are composable: you can combine any number of them via [`TreeDisplay::with`]. -/// -/// [`TreeDisplay::with`]: super::TreeDisplay::with -pub trait TreeExtractor: Send + Sync { - /// Write header annotations (space-prefixed) to the formatter. - fn write_header( - &self, - array: &ArrayRef, - ctx: &TreeContext, - f: &mut fmt::Formatter<'_>, - ) -> fmt::Result { - let _ = (array, ctx, f); - Ok(()) - } - - /// Write detail lines below the header. - /// - /// Content written through `f` is automatically indented. Use - /// [`f.formatter()`](IndentedFormatter::formatter) to access the underlying - /// [`fmt::Formatter`] for formatting flags. - fn write_details( - &self, - array: &ArrayRef, - ctx: &TreeContext, - f: &mut IndentedFormatter<'_, '_>, - ) -> fmt::Result { - let _ = (array, ctx, f); - Ok(()) + fn pop_parent(&mut self, _parent: &ArrayRef) { + self.ancestor_sizes.pop(); } } diff --git a/vortex-array/src/display/extractors/buffer.rs b/vortex-array/src/display/extractors/buffer.rs index 17861a654f7..2fb1c716f3d 100644 --- a/vortex-array/src/display/extractors/buffer.rs +++ b/vortex-array/src/display/extractors/buffer.rs @@ -17,7 +17,7 @@ pub struct BufferExtractor { pub show_percent: bool, } -impl TreeExtractor for BufferExtractor { +impl TreeExtractor for BufferExtractor { fn write_details( &self, array: &ArrayRef, diff --git a/vortex-array/src/display/extractors/encoding_summary.rs b/vortex-array/src/display/extractors/encoding_summary.rs index bc678643ecb..43c924190df 100644 --- a/vortex-array/src/display/extractors/encoding_summary.rs +++ b/vortex-array/src/display/extractors/encoding_summary.rs @@ -23,7 +23,7 @@ impl EncodingSummaryExtractor { } } -impl TreeExtractor for EncodingSummaryExtractor { +impl TreeExtractor for EncodingSummaryExtractor { fn write_header( &self, array: &ArrayRef, diff --git a/vortex-array/src/display/extractors/metadata.rs b/vortex-array/src/display/extractors/metadata.rs index a355bafe63c..d68e9c7a1d8 100644 --- a/vortex-array/src/display/extractors/metadata.rs +++ b/vortex-array/src/display/extractors/metadata.rs @@ -11,7 +11,7 @@ use crate::display::extractor::TreeExtractor; /// Extractor that adds a `metadata: ...` detail line. pub struct MetadataExtractor; -impl TreeExtractor for MetadataExtractor { +impl TreeExtractor for MetadataExtractor { fn write_details( &self, array: &ArrayRef, diff --git a/vortex-array/src/display/extractors/nbytes.rs b/vortex-array/src/display/extractors/nbytes.rs index 37e100279fc..27d850444db 100644 --- a/vortex-array/src/display/extractors/nbytes.rs +++ b/vortex-array/src/display/extractors/nbytes.rs @@ -13,7 +13,7 @@ use crate::display::extractor::TreeExtractor; /// Extractor that adds `nbytes=X (Y%)` to the header line. pub struct NbytesExtractor; -impl TreeExtractor for NbytesExtractor { +impl TreeExtractor for NbytesExtractor { fn write_header( &self, array: &ArrayRef, diff --git a/vortex-array/src/display/extractors/stats.rs b/vortex-array/src/display/extractors/stats.rs index e39dafa2340..7459daccc5c 100644 --- a/vortex-array/src/display/extractors/stats.rs +++ b/vortex-array/src/display/extractors/stats.rs @@ -116,7 +116,7 @@ impl fmt::Display for StatsDisplay<'_> { /// Extractor that adds stats annotations (e.g. `[nulls=3, min=5]`) to the header line. pub struct StatsExtractor; -impl TreeExtractor for StatsExtractor { +impl TreeExtractor for StatsExtractor { fn write_header( &self, array: &ArrayRef, diff --git a/vortex-array/src/display/mod.rs b/vortex-array/src/display/mod.rs index 1c12afcf8d2..780140af2ac 100644 --- a/vortex-array/src/display/mod.rs +++ b/vortex-array/src/display/mod.rs @@ -556,7 +556,7 @@ impl ArrayRef { metadata, stats, } => { - let extractors: [(bool, Box); 5] = [ + let extractors: [(bool, Box>); 5] = [ (true, Box::new(EncodingSummaryExtractor)), (*stats, Box::new(NbytesExtractor)), (*stats, Box::new(StatsExtractor)), diff --git a/vortex-array/src/display/tree_display.rs b/vortex-array/src/display/tree_display.rs index 5d0753ba354..76217da0f8d 100644 --- a/vortex-array/src/display/tree_display.rs +++ b/vortex-array/src/display/tree_display.rs @@ -3,8 +3,10 @@ use std::fmt; +use vortex_utils::tree::TreeDisplayAdapter; +use vortex_utils::tree::write_indented_tree; + use crate::ArrayRef; -use crate::arrays::Chunked; use crate::display::extractor::IndentedFormatter; use crate::display::extractor::TreeContext; use crate::display::extractor::TreeExtractor; @@ -37,7 +39,7 @@ use crate::display::extractors::StatsExtractor; /// ``` pub struct TreeDisplay { array: ArrayRef, - extractors: Vec>, + extractors: Vec>>, } impl TreeDisplay { @@ -64,57 +66,60 @@ impl TreeDisplay { } /// Add an extractor to the display pipeline. - pub fn with(mut self, extractor: E) -> Self { + pub fn with + 'static>(mut self, extractor: E) -> Self { self.extractors.push(Box::new(extractor)); self } /// Add a pre-boxed extractor to the display pipeline. - pub fn with_boxed(mut self, extractor: Box) -> Self { + pub fn with_boxed(mut self, extractor: Box>) -> Self { self.extractors.push(extractor); self } +} + +impl TreeDisplayAdapter for TreeDisplay { + type Context = TreeContext; + type Node = ArrayRef; - /// Recursively write a node and all its descendants directly to the formatter. fn write_node( &self, - name: &str, array: &ArrayRef, - ctx: &mut TreeContext, - indent: &str, + ctx: &TreeContext, f: &mut fmt::Formatter<'_>, ) -> fmt::Result { - // Header line: "{indent}{name}:{annotations...}\n" - write!(f, "{indent}{name}:")?; for extractor in &self.extractors { extractor.write_header(array, ctx, f)?; } - writeln!(f)?; + Ok(()) + } - // Detail lines - let child_indent = format!("{indent} "); - { - let mut indented = IndentedFormatter::new(f, &child_indent); - for extractor in &self.extractors { - extractor.write_details(array, ctx, &mut indented)?; - } + fn write_details( + &self, + array: &ArrayRef, + ctx: &TreeContext, + f: &mut IndentedFormatter<'_, '_>, + ) -> fmt::Result { + for extractor in &self.extractors { + extractor.write_details(array, ctx, f)?; } + Ok(()) + } - // Push context for children: chunked arrays reset the percentage root - let child_size = if array.is::() { - None - } else { - Some(array.nbytes()) - }; - ctx.push(child_size); - - // Recurse into children - for (child_name, child) in array.children_names().into_iter().zip(array.children()) { - self.write_node(&child_name, &child, ctx, &child_indent, f)?; + fn visit_children( + &self, + array: &ArrayRef, + visit: &mut dyn FnMut(&str, &ArrayRef, bool) -> fmt::Result, + ) -> fmt::Result { + let mut children = array + .children_names() + .into_iter() + .zip(array.children()) + .peekable(); + while let Some((child_name, child)) = children.next() { + let is_last = children.peek().is_none(); + visit(&child_name, &child, is_last)?; } - - ctx.pop(); - Ok(()) } } @@ -122,6 +127,6 @@ impl TreeDisplay { impl fmt::Display for TreeDisplay { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { let mut ctx = TreeContext::new(); - self.write_node("root", &self.array, &mut ctx, "", f) + write_indented_tree(self, "root", &self.array, &mut ctx, f) } } diff --git a/vortex-array/src/dtype/arbitrary/mod.rs b/vortex-array/src/dtype/arbitrary/mod.rs index 9b95074ff6f..ddfe1712ec3 100644 --- a/vortex-array/src/dtype/arbitrary/mod.rs +++ b/vortex-array/src/dtype/arbitrary/mod.rs @@ -36,7 +36,7 @@ impl<'a> Arbitrary<'a> for FieldName { fn random_dtype(u: &mut Unstructured<'_>, depth: u8) -> Result { const BASE_TYPE_COUNT: i32 = 5; - const CONTAINER_TYPE_COUNT: i32 = 3; + const CONTAINER_TYPE_COUNT: i32 = 4; let max_dtype_kind = if depth == 0 { BASE_TYPE_COUNT } else { @@ -59,12 +59,23 @@ fn random_dtype(u: &mut Unstructured<'_>, depth: u8) -> Result { u.choose_index(3)?.try_into().vortex_expect("impossible"), u.arbitrary()?, ), + 9 => random_map_dtype(u, depth - 1)?, // Null, // Extension(ExtDType, Nullability), _ => unreachable!("Number out of range"), }) } +fn random_map_dtype(u: &mut Unstructured<'_>, depth: u8) -> Result { + Ok(DType::map( + random_dtype(u, depth)?.as_nonnullable(), + random_dtype(u, depth)?, + false, + u.arbitrary()?, + ) + .vortex_expect("non-nullable generated map keys are always valid")) +} + impl<'a> Arbitrary<'a> for Nullability { fn arbitrary(u: &mut Unstructured<'a>) -> Result { Ok(if u.arbitrary()? { @@ -123,3 +134,23 @@ fn random_struct_dtype(u: &mut Unstructured<'_>, depth: u8) -> Result>>()?; Ok(StructFields::new(names, dtypes)) } + +#[cfg(test)] +mod tests { + use arbitrary::Unstructured; + + use super::random_map_dtype; + use crate::dtype::DType; + + #[test] + fn random_map_dtype_never_asserts_sorted_keys() { + for byte in 0..=u8::MAX { + let bytes = [byte; 128]; + let mut u = Unstructured::new(&bytes); + let Ok(DType::Map(map_dtype, _)) = random_map_dtype(&mut u, 0) else { + continue; + }; + assert!(!map_dtype.keys_sorted()); + } + } +} diff --git a/vortex-array/src/dtype/coercion.rs b/vortex-array/src/dtype/coercion.rs deleted file mode 100644 index 4d30a80c216..00000000000 --- a/vortex-array/src/dtype/coercion.rs +++ /dev/null @@ -1,807 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Utilities for performing type coercion. - -use std::sync::Arc; - -use crate::dtype::DType; -use crate::dtype::PType; -use crate::dtype::decimal::DecimalDType; - -impl PType { - /// Returns the least supertype (widest common type) of two primitive types, - /// or `None` if no lossless promotion exists. - pub fn least_supertype(self, other: PType) -> Option { - if self == other { - return Some(self); - } - - // Same family — pick the wider. - if self.is_unsigned_int() && other.is_unsigned_int() { - return Some(self.max_unsigned_ptype(other)); - } - if self.is_signed_int() && other.is_signed_int() { - return Some(self.max_signed_ptype(other)); - } - if self.is_float() && other.is_float() { - return if self.byte_width() >= other.byte_width() { - Some(self) - } else { - Some(other) - }; - } - - // Unsigned + Signed crossover — promote to signed one width-step wider. - if self.is_unsigned_int() && other.is_signed_int() { - return Self::unsigned_signed_supertype(self, other); - } - if self.is_signed_int() && other.is_unsigned_int() { - return Self::unsigned_signed_supertype(other, self); - } - - // Int + Float — pick the smallest float that losslessly represents the integer. - let (int, float) = if self.is_float() { - (other, self) - } else { - (self, other) - }; - Self::int_float_supertype(int, float) - } - - fn unsigned_signed_supertype(unsigned: PType, signed: PType) -> Option { - use PType::*; - match unsigned.byte_width().max(signed.byte_width()) { - 1 => Some(I16), - 2 => Some(I32), - 4 => Some(I64), - _ => None, // U64 + I64 — no lossless 128-bit integer type - } - } - - fn int_float_supertype(int: PType, float: PType) -> Option { - use PType::*; - let min_float = match int.byte_width() { - 1 => F16, // f16 has 11-bit mantissa, enough for 8-bit ints - 2 => F32, // f32 has 24-bit mantissa, enough for 16-bit ints - 4 => F64, // f64 has 53-bit mantissa, enough for 32-bit ints - _ => return None, // no standard float for 64-bit ints - }; - if float.byte_width() >= min_float.byte_width() { - Some(float) - } else { - Some(min_float) - } - } -} - -impl DType { - /// The core primitive — what type can hold both `self` and `other`? - /// Returns `None` if no common supertype exists. - pub fn least_supertype(&self, other: &DType) -> Option { - if let (DType::Union(lhs, lhs_null), DType::Union(rhs, rhs_null)) = (self, other) { - return (lhs == rhs).then(|| DType::Union(lhs.clone(), *lhs_null | *rhs_null)); - } - - let union_null = self.nullability() | other.nullability(); - - if let ( - DType::FixedSizeList(lhs_elem, lhs_size, _), - DType::FixedSizeList(rhs_elem, rhs_size, _), - ) = (self, other) - && lhs_size == rhs_size - { - let elem = lhs_elem.least_supertype(rhs_elem)?; - return Some(DType::FixedSizeList(Arc::new(elem), *lhs_size, union_null)); - } - - if let (DType::List(lhs_elem, _), DType::List(rhs_elem, _)) = (self, other) { - let elem = lhs_elem.least_supertype(rhs_elem)?; - return Some(DType::List(Arc::new(elem), union_null)); - } - - // Identity (ignoring nullability): return self with union nullability - if self.eq_ignore_nullability(other) { - return Some(self.with_nullability(union_null)); - } - - // Null + X: return X as nullable - if matches!(self, DType::Null) { - return Some(other.as_nullable()); - } - if matches!(other, DType::Null) { - return Some(self.as_nullable()); - } - - // Bool + numeric: return the numeric type (with union nullability) - if self.is_boolean() && other.is_numeric() { - return Some(other.with_nullability(union_null)); - } - if other.is_boolean() && self.is_numeric() { - return Some(self.with_nullability(union_null)); - } - - // Primitive + Primitive (different ptypes): delegate to PType::least_supertype - if let (DType::Primitive(lhs_p, _), DType::Primitive(rhs_p, _)) = (self, other) { - return lhs_p - .least_supertype(*rhs_p) - .map(|p| DType::Primitive(p, union_null)); - } - - // Decimal + Decimal: compute wider decimal - if let (DType::Decimal(lhs_d, _), DType::Decimal(rhs_d, _)) = (self, other) { - return decimal_least_supertype(*lhs_d, *rhs_d).map(|d| DType::Decimal(d, union_null)); - } - - // Decimal + integer Primitive: convert integer to Decimal, then widen - if let (DType::Decimal(dec, _), DType::Primitive(p, _)) = (self, other) - && p.is_int() - { - let int_dec = DecimalDType::new(integer_decimal_precision(*p), 0); - return decimal_least_supertype(*dec, int_dec).map(|d| DType::Decimal(d, union_null)); - } - if let (DType::Primitive(p, _), DType::Decimal(dec, _)) = (self, other) - && p.is_int() - { - let int_dec = DecimalDType::new(integer_decimal_precision(*p), 0); - return decimal_least_supertype(int_dec, *dec).map(|d| DType::Decimal(d, union_null)); - } - - // Extension + anything: delegate to vtable - if let DType::Extension(ext) = self { - return ext - .least_supertype(other) - .map(|dt| dt.with_nullability(union_null)); - } - if let DType::Extension(ext) = other { - return ext - .least_supertype(self) - .map(|dt| dt.with_nullability(union_null)); - } - - None - } - - /// Fold over a slice — what type can hold all of these? - pub fn least_supertype_of(types: &[DType]) -> Option { - types - .iter() - .skip(1) - .try_fold(types[0].clone(), |acc, t| acc.least_supertype(t)) - } - - /// Is there any implicit coercion path from `other` to `self`? - pub fn can_coerce_from(&self, other: &DType) -> bool { - if let ( - DType::FixedSizeList(target_elem, target_size, _), - DType::FixedSizeList(source_elem, source_size, _), - ) = (self, other) - { - return target_size == source_size - && (self.is_nullable() || !other.is_nullable()) - && target_elem.can_coerce_from(source_elem); - } - - if let (DType::List(target_elem, _), DType::List(source_elem, _)) = (self, other) { - return (self.is_nullable() || !other.is_nullable()) - && target_elem.can_coerce_from(source_elem); - } - - // Same type (ignoring nullability): check nullability compatibility - if self.eq_ignore_nullability(other) { - return self.is_nullable() || !other.is_nullable(); - } - - // Null → nullable target - if matches!(other, DType::Null) { - return self.is_nullable(); - } - - // Bool → numeric - if other.is_boolean() && self.is_numeric() { - return self.is_nullable() || !other.is_nullable(); - } - - // Primitive widening: true if least_supertype(source, target) == target - if let (DType::Primitive(..), DType::Primitive(..)) = (self, other) { - return other - .least_supertype(self) - .is_some_and(|st| st.eq_ignore_nullability(self)) - && (self.is_nullable() || !other.is_nullable()); - } - - // Decimal widening - if let (DType::Decimal(target, _), DType::Decimal(source, _)) = (self, other) { - let target_integral = target.precision() as i16 - target.scale() as i16; - let source_integral = source.precision() as i16 - source.scale() as i16; - return target_integral >= source_integral - && target.scale() >= source.scale() - && (self.is_nullable() || !other.is_nullable()); - } - - // Integer → Decimal - if let (DType::Decimal(dec, _), DType::Primitive(p, _)) = (self, other) - && p.is_int() - { - let needed = integer_decimal_precision(*p); - let integral_digits = dec.precision() as i16 - dec.scale() as i16; - return integral_digits >= needed as i16 - && (self.is_nullable() || !other.is_nullable()); - } - - // Extension: delegate to vtable - if let DType::Extension(ext) = self { - return ext.can_coerce_from(other); - } - - false - } - - /// Convenience — is there a path from `self` to `other`? - pub fn can_coerce_to(&self, other: &DType) -> bool { - other.can_coerce_from(self) - } - - /// Are all types in the slice mutually coercible to a common type? - pub fn are_coercible(types: &[DType]) -> bool { - DType::least_supertype_of(types).is_some() - } - - /// Can all types in the slice be coerced to a specific target? - pub fn all_coercible_to(types: &[DType], target: &DType) -> bool { - types.iter().all(|t| target.can_coerce_from(t)) - } - - /// Coerce a slice to a specific target — returns the vec of targets - /// if all are coercible, `None` if any are not. - pub fn coerce_all_to(types: &[DType], target: &DType) -> Option> { - types - .iter() - .all(|t| target.can_coerce_from(t)) - .then(|| vec![target.clone(); types.len()]) - } - - /// Coerce a slice to their mutual least supertype. - pub fn coerce_to_supertype(types: &[DType]) -> Option> { - let supertype = DType::least_supertype_of(types)?; - Some(vec![supertype; types.len()]) - } - - /// Is this a numeric type (primitive int/float or decimal)? - pub fn is_numeric(&self) -> bool { - matches!(self, DType::Primitive(..) | DType::Decimal(..)) - } - - /// Is this a temporal type (date, time, timestamp, duration)? - pub fn is_temporal(&self) -> bool { - match self { - DType::Extension(ext) => { - use crate::dtype::extension::Matcher; - use crate::extension::datetime::AnyTemporal; - AnyTemporal::matches(ext) - } - _ => false, - } - } -} - -/// Maps integer PType widths to the minimum decimal precision needed. -fn integer_decimal_precision(ptype: PType) -> u8 { - match ptype { - PType::U8 | PType::I8 => 3, - PType::U16 | PType::I16 => 5, - PType::U32 | PType::I32 => 10, - PType::U64 | PType::I64 => 19, - _ => 19, - } -} - -/// Compute the least supertype of two decimal types using SQL-standard rules. -fn decimal_least_supertype(a: DecimalDType, b: DecimalDType) -> Option { - let a_integral = a.precision() as i16 - a.scale() as i16; - let b_integral = b.precision() as i16 - b.scale() as i16; - let max_integral = a_integral.max(b_integral); - let max_scale = a.scale().max(b.scale()); - let precision = u8::try_from(max_integral + max_scale as i16).ok()?; - DecimalDType::try_new(precision, max_scale).ok() -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use crate::dtype::DType; - use crate::dtype::PType; - use crate::dtype::UnionVariants; - use crate::dtype::decimal::DecimalDType; - use crate::dtype::nullability::Nullability::NonNullable; - use crate::dtype::nullability::Nullability::Nullable; - - #[test] - fn is_numeric() { - assert!(DType::Primitive(PType::I32, NonNullable).is_numeric()); - assert!(DType::Primitive(PType::F64, NonNullable).is_numeric()); - assert!(DType::Decimal(DecimalDType::new(10, 2), NonNullable).is_numeric()); - assert!(!DType::Bool(NonNullable).is_numeric()); - assert!(!DType::Utf8(NonNullable).is_numeric()); - assert!(!DType::Null.is_numeric()); - } - - #[test] - fn least_supertype_identity() { - let i32_nn = DType::Primitive(PType::I32, NonNullable); - assert_eq!(i32_nn.least_supertype(&i32_nn).unwrap(), i32_nn); - } - - #[test] - fn least_supertype_nullability_union() { - let i32_nn = DType::Primitive(PType::I32, NonNullable); - let i32_n = DType::Primitive(PType::I32, Nullable); - assert_eq!(i32_nn.least_supertype(&i32_n).unwrap(), i32_n); - assert_eq!(i32_n.least_supertype(&i32_nn).unwrap(), i32_n); - } - - #[test] - fn least_supertype_null_absorption() { - let i32_nn = DType::Primitive(PType::I32, NonNullable); - assert_eq!( - DType::Null.least_supertype(&i32_nn).unwrap(), - DType::Primitive(PType::I32, Nullable) - ); - assert_eq!( - i32_nn.least_supertype(&DType::Null).unwrap(), - DType::Primitive(PType::I32, Nullable) - ); - } - - #[test] - fn least_supertype_union_requires_identical_variants() { - let nonnullable = DType::Union( - UnionVariants::new( - ["value"].into(), - vec![DType::Primitive(PType::I32, NonNullable)], - ) - .unwrap(), - NonNullable, - ); - let nullable = DType::Union( - UnionVariants::new( - ["value"].into(), - vec![DType::Primitive(PType::I32, Nullable)], - ) - .unwrap(), - NonNullable, - ); - - assert_eq!( - nonnullable.least_supertype(&nonnullable), - Some(nonnullable.clone()) - ); - assert!(nonnullable.least_supertype(&nullable).is_none()); - assert!(nullable.least_supertype(&nonnullable).is_none()); - } - - #[test] - fn least_supertype_null_makes_union_outer_nullable() { - let nonnullable = DType::Union( - UnionVariants::new( - ["value"].into(), - vec![DType::Primitive(PType::I32, NonNullable)], - ) - .unwrap(), - NonNullable, - ); - let nullable = nonnullable.as_nullable(); - - assert_eq!( - DType::Null.least_supertype(&nonnullable), - Some(nullable.clone()) - ); - assert_eq!(nonnullable.least_supertype(&DType::Null), Some(nullable)); - } - - #[test] - fn least_supertype_unsigned_widening() { - let u8_nn = DType::Primitive(PType::U8, NonNullable); - let u32_nn = DType::Primitive(PType::U32, NonNullable); - assert_eq!(u8_nn.least_supertype(&u32_nn).unwrap(), u32_nn); - } - - #[test] - fn least_supertype_signed_widening() { - let i16_nn = DType::Primitive(PType::I16, NonNullable); - let i64_nn = DType::Primitive(PType::I64, NonNullable); - assert_eq!(i16_nn.least_supertype(&i64_nn).unwrap(), i64_nn); - } - - #[test] - fn least_supertype_cross_family() { - let u8_nn = DType::Primitive(PType::U8, NonNullable); - let i8_nn = DType::Primitive(PType::I8, NonNullable); - assert_eq!( - u8_nn.least_supertype(&i8_nn).unwrap(), - DType::Primitive(PType::I16, NonNullable) - ); - } - - #[test] - fn least_supertype_u64_i64_none() { - let u64_nn = DType::Primitive(PType::U64, NonNullable); - let i64_nn = DType::Primitive(PType::I64, NonNullable); - assert!(u64_nn.least_supertype(&i64_nn).is_none()); - } - - #[test] - fn least_supertype_int_float_promotion() { - let u8_nn = DType::Primitive(PType::U8, NonNullable); - let f32_nn = DType::Primitive(PType::F32, NonNullable); - assert_eq!(u8_nn.least_supertype(&f32_nn).unwrap(), f32_nn); - } - - #[test] - fn least_supertype_i32_f32_to_f64() { - let i32_nn = DType::Primitive(PType::I32, NonNullable); - let f32_nn = DType::Primitive(PType::F32, NonNullable); - assert_eq!( - i32_nn.least_supertype(&f32_nn).unwrap(), - DType::Primitive(PType::F64, NonNullable) - ); - } - - #[test] - fn least_supertype_bool_numeric() { - let bool_nn = DType::Bool(NonNullable); - let i32_nn = DType::Primitive(PType::I32, NonNullable); - assert_eq!(bool_nn.least_supertype(&i32_nn).unwrap(), i32_nn); - assert_eq!(i32_nn.least_supertype(&bool_nn).unwrap(), i32_nn); - } - - #[test] - fn least_supertype_decimal_widening() { - let d1 = DType::Decimal(DecimalDType::new(10, 2), NonNullable); - let d2 = DType::Decimal(DecimalDType::new(15, 5), NonNullable); - let result = d1.least_supertype(&d2).unwrap(); - // integral digits: max(8, 10) = 10, max scale = 5, precision = 15 - assert_eq!( - result, - DType::Decimal(DecimalDType::new(15, 5), NonNullable) - ); - } - - #[test] - fn least_supertype_incompatible_none() { - let utf8 = DType::Utf8(NonNullable); - let i32_nn = DType::Primitive(PType::I32, NonNullable); - assert!(utf8.least_supertype(&i32_nn).is_none()); - } - - #[test] - fn can_coerce_from_widening() { - let i32_nn = DType::Primitive(PType::I32, NonNullable); - let i64_nn = DType::Primitive(PType::I64, NonNullable); - assert!(i64_nn.can_coerce_from(&i32_nn)); - } - - #[test] - fn can_coerce_from_narrowing_rejected() { - let i32_nn = DType::Primitive(PType::I32, NonNullable); - let i64_nn = DType::Primitive(PType::I64, NonNullable); - assert!(!i32_nn.can_coerce_from(&i64_nn)); - } - - #[test] - fn can_coerce_from_nullability_constraints() { - let i32_nn = DType::Primitive(PType::I32, NonNullable); - let i32_n = DType::Primitive(PType::I32, Nullable); - assert!(i32_n.can_coerce_from(&i32_nn)); - assert!(!i32_nn.can_coerce_from(&i32_n)); - } - - #[test] - fn can_coerce_from_null() { - let i32_n = DType::Primitive(PType::I32, Nullable); - let i32_nn = DType::Primitive(PType::I32, NonNullable); - assert!(i32_n.can_coerce_from(&DType::Null)); - assert!(!i32_nn.can_coerce_from(&DType::Null)); - } - - #[test] - fn are_coercible_mixed() { - let types = [ - DType::Primitive(PType::I32, NonNullable), - DType::Primitive(PType::I64, NonNullable), - ]; - assert!(DType::are_coercible(&types)); - } - - #[test] - fn all_coercible_to_target() { - let types = [ - DType::Primitive(PType::I32, NonNullable), - DType::Primitive(PType::I16, NonNullable), - ]; - let target = DType::Primitive(PType::I64, NonNullable); - assert!(DType::all_coercible_to(&types, &target)); - } - - #[test] - fn coerce_to_supertype_works() { - let types = [ - DType::Primitive(PType::U8, NonNullable), - DType::Primitive(PType::I16, NonNullable), - ]; - let result = DType::coerce_to_supertype(&types).unwrap(); - // U8 + I16: unsigned_signed_supertype max_width=max(1,2)=2 => I32 - assert_eq!(result, vec![DType::Primitive(PType::I32, NonNullable); 2]); - } - - #[test] - fn fsl_widens_element_dtype_when_size_matches() { - let lhs = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::F32, NonNullable)), - 768, - NonNullable, - ); - let rhs = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::F64, NonNullable)), - 768, - NonNullable, - ); - let expected = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::F64, NonNullable)), - 768, - NonNullable, - ); - assert_eq!(lhs.least_supertype(&rhs), Some(expected)); - } - - #[test] - fn fsl_size_mismatch_returns_none() { - let lhs = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::F32, NonNullable)), - 768, - NonNullable, - ); - let rhs = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::F32, NonNullable)), - 1024, - NonNullable, - ); - assert_eq!(lhs.least_supertype(&rhs), None); - } - - #[test] - fn fsl_unions_outer_nullability() { - let lhs = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::F32, NonNullable)), - 4, - NonNullable, - ); - let rhs = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::F64, NonNullable)), - 4, - Nullable, - ); - let expected = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::F64, NonNullable)), - 4, - Nullable, - ); - assert_eq!(lhs.least_supertype(&rhs), Some(expected)); - } - - #[test] - fn fsl_widening_unions_element_nullability() { - let lhs = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::F32, NonNullable)), - 4, - NonNullable, - ); - let rhs = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::F64, Nullable)), - 4, - NonNullable, - ); - let expected = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::F64, Nullable)), - 4, - NonNullable, - ); - assert_eq!(lhs.least_supertype(&rhs), Some(expected)); - } - - #[test] - fn fsl_incompatible_elements_returns_none() { - let lhs = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::F32, NonNullable)), - 4, - NonNullable, - ); - let rhs = DType::FixedSizeList(Arc::new(DType::Utf8(NonNullable)), 4, NonNullable); - assert_eq!(lhs.least_supertype(&rhs), None); - } - - #[test] - fn fsl_can_coerce_from_widening() { - let target = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::F64, NonNullable)), - 4, - NonNullable, - ); - let source = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::F32, NonNullable)), - 4, - NonNullable, - ); - assert!(target.can_coerce_from(&source)); - assert!(!source.can_coerce_from(&target)); - } - - #[test] - fn fsl_same_element_widening_unions_inner_nullability() { - let lhs = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::I32, NonNullable)), - 4, - NonNullable, - ); - let rhs = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::I32, Nullable)), - 4, - NonNullable, - ); - let expected = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::I32, Nullable)), - 4, - NonNullable, - ); - assert_eq!(lhs.least_supertype(&rhs), Some(expected)); - } - - #[test] - fn list_widens_element_dtype() { - let lhs = DType::List( - Arc::new(DType::Primitive(PType::F32, NonNullable)), - NonNullable, - ); - let rhs = DType::List( - Arc::new(DType::Primitive(PType::F64, NonNullable)), - NonNullable, - ); - let expected = DType::List( - Arc::new(DType::Primitive(PType::F64, NonNullable)), - NonNullable, - ); - assert_eq!(lhs.least_supertype(&rhs), Some(expected)); - } - - #[test] - fn list_unions_outer_nullability() { - let lhs = DType::List( - Arc::new(DType::Primitive(PType::F32, NonNullable)), - NonNullable, - ); - let rhs = DType::List( - Arc::new(DType::Primitive(PType::F64, NonNullable)), - Nullable, - ); - let expected = DType::List( - Arc::new(DType::Primitive(PType::F64, NonNullable)), - Nullable, - ); - assert_eq!(lhs.least_supertype(&rhs), Some(expected)); - } - - #[test] - fn list_unions_inner_nullability() { - let lhs = DType::List( - Arc::new(DType::Primitive(PType::I32, NonNullable)), - NonNullable, - ); - let rhs = DType::List( - Arc::new(DType::Primitive(PType::I32, Nullable)), - NonNullable, - ); - let expected = DType::List( - Arc::new(DType::Primitive(PType::I32, Nullable)), - NonNullable, - ); - assert_eq!(lhs.least_supertype(&rhs), Some(expected)); - } - - #[test] - fn list_incompatible_elements_returns_none() { - let lhs = DType::List( - Arc::new(DType::Primitive(PType::F32, NonNullable)), - NonNullable, - ); - let rhs = DType::List(Arc::new(DType::Utf8(NonNullable)), NonNullable); - assert_eq!(lhs.least_supertype(&rhs), None); - } - - #[test] - fn list_can_coerce_from_widening() { - let target = DType::List( - Arc::new(DType::Primitive(PType::F64, NonNullable)), - NonNullable, - ); - let source = DType::List( - Arc::new(DType::Primitive(PType::F32, NonNullable)), - NonNullable, - ); - assert!(target.can_coerce_from(&source)); - assert!(!source.can_coerce_from(&target)); - } - - #[test] - fn fsl_can_coerce_from_rejects_nullable_source_elements() { - let target = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::F64, NonNullable)), - 4, - NonNullable, - ); - let source = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::F32, Nullable)), - 4, - NonNullable, - ); - assert!(!target.can_coerce_from(&source)); - } - - #[test] - fn list_can_coerce_from_rejects_nullable_source_elements() { - let target = DType::List( - Arc::new(DType::Primitive(PType::F64, NonNullable)), - NonNullable, - ); - let source = DType::List( - Arc::new(DType::Primitive(PType::F32, Nullable)), - NonNullable, - ); - assert!(!target.can_coerce_from(&source)); - } - - #[test] - fn fsl_can_coerce_from_allows_widening_nullable_target() { - let target = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::I32, Nullable)), - 4, - NonNullable, - ); - let source = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::I32, NonNullable)), - 4, - NonNullable, - ); - assert!(target.can_coerce_from(&source)); - } - - #[test] - fn fsl_can_coerce_from_size_mismatch_rejected() { - let a = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::F32, NonNullable)), - 4, - NonNullable, - ); - let b = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::F32, NonNullable)), - 8, - NonNullable, - ); - assert!(!a.can_coerce_from(&b)); - assert!(!b.can_coerce_from(&a)); - } - - #[test] - fn least_supertype_integer_decimal() { - let i32_nn = DType::Primitive(PType::I32, NonNullable); - let dec = DType::Decimal(DecimalDType::new(15, 5), NonNullable); - let result = i32_nn.least_supertype(&dec).unwrap(); - // int_dec for I32 = Decimal(10, 0). integral digits = 10. - // dec integral = 15 - 5 = 10. - // max_integral = 10, max_scale = 5, precision = 15 - assert_eq!( - result, - DType::Decimal(DecimalDType::new(15, 5), NonNullable) - ); - } -} diff --git a/vortex-array/src/dtype/decimal/precision.rs b/vortex-array/src/dtype/decimal/precision.rs index 79d47c65462..62c00b80e0d 100644 --- a/vortex-array/src/dtype/decimal/precision.rs +++ b/vortex-array/src/dtype/decimal/precision.rs @@ -93,12 +93,14 @@ impl PrecisionScale { } /// The precision is the number of significant figures that the decimal tracks. + #[allow(clippy::inline_always)] #[inline(always)] pub fn precision(&self) -> u8 { self.precision.get() } /// The scale is the maximum number of digits relative to the decimal point. + #[allow(clippy::inline_always)] #[inline(always)] pub fn scale(&self) -> i8 { self.scale diff --git a/vortex-array/src/dtype/dtype_impl.rs b/vortex-array/src/dtype/dtype_impl.rs index 20ff3d6a182..2cf50600b9a 100644 --- a/vortex-array/src/dtype/dtype_impl.rs +++ b/vortex-array/src/dtype/dtype_impl.rs @@ -3,16 +3,20 @@ use std::fmt::Display; use std::fmt::Formatter; +use std::hash::Hash; +use std::hash::Hasher; use std::sync::Arc; use DType::*; use itertools::Itertools; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_error::vortex_panic; use super::DType; use crate::dtype::FieldDType; use crate::dtype::FieldName; +use crate::dtype::MapDType; use crate::dtype::PType; use crate::dtype::StructFields; use crate::dtype::UnionVariants; @@ -33,7 +37,7 @@ pub trait NativeDType { fn dtype() -> DType; } -/// Assert that the size of DType is 16 bytes. +/// Assert that the size of DType is 24 bytes. #[cfg(not(target_arch = "wasm32"))] const _: [(); size_of::()] = [(); 24]; // FIXME(ngates): should we keep this at 16? @@ -63,6 +67,7 @@ impl DType { | Binary(null) | List(_, null) | FixedSizeList(_, _, null) + | Map(_, null) | Struct(_, null) | Union(_, null) | Variant(null) => matches!(null, Nullability::Nullable), @@ -82,9 +87,12 @@ impl DType { /// Get a new DType with the given nullability (but otherwise the same as `self`). /// - /// [`DType::Null`] has intrinsic nullability and is returned unchanged. + /// Invoking .with_nullability(NonNullable) on DType::Null panics. pub fn with_nullability(&self, nullability: Nullability) -> Self { match self { + Null if nullability == Nullability::NonNullable => { + vortex_panic!("DType::Null cannot be made non-nullable") + } Null => Null, Bool(_) => Bool(nullability), Primitive(pdt, _) => Primitive(*pdt, nullability), @@ -93,6 +101,7 @@ impl DType { Binary(_) => Binary(nullability), List(edt, _) => List(Arc::clone(edt), nullability), FixedSizeList(edt, size, _) => FixedSizeList(Arc::clone(edt), *size, nullability), + Map(map, _) => Map(map.clone(), nullability), Struct(sf, _) => Struct(sf.clone(), nullability), Union(vs, _) => Union(vs.clone(), nullability), Variant(_) => Variant(nullability), @@ -119,22 +128,11 @@ impl DType { (FixedSizeList(lhs_dtype, lhs_size, _), FixedSizeList(rhs_dtype, rhs_size, _)) => { lhs_size == rhs_size && lhs_dtype.eq_ignore_nullability(rhs_dtype) } + (Map(lhs, _), Map(rhs, _)) => lhs.eq_ignore_nullability(rhs), (Struct(lhs_dtype, _), Struct(rhs_dtype, _)) => { - (lhs_dtype.names() == rhs_dtype.names()) - && (lhs_dtype - .fields() - .zip_eq(rhs_dtype.fields()) - .all(|(l, r)| l.eq_ignore_nullability(&r))) - } - (Union(lhs, _), Union(rhs, _)) => { - // Equal `names` implies equal length by FieldNames equality. - lhs.names() == rhs.names() - && lhs.type_ids() == rhs.type_ids() - && lhs - .variants() - .zip_eq(rhs.variants()) - .all(|(l, r)| l.eq_ignore_nullability(&r)) + lhs_dtype.eq_ignore_nullability(rhs_dtype) } + (Union(lhs, _), Union(rhs, _)) => lhs.eq_ignore_nullability(rhs), (Variant(_), Variant(_)) => true, (Extension(lhs_extdtype), Extension(rhs_extdtype)) => { lhs_extdtype.eq_ignore_nullability(rhs_extdtype) @@ -143,6 +141,26 @@ impl DType { } } + /// Hash this dtype using the same equivalence relation as [`Self::eq_ignore_nullability`]. + pub(crate) fn hash_ignore_nullability(&self, state: &mut H) { + std::mem::discriminant(self).hash(state); + + match self { + Null | Bool(_) | Utf8(_) | Binary(_) | Variant(_) => {} + Primitive(ptype, _) => ptype.hash(state), + Decimal(decimal, _) => decimal.hash(state), + List(element, _) => element.hash_ignore_nullability(state), + FixedSizeList(element, size, _) => { + element.hash_ignore_nullability(state); + size.hash(state); + } + Map(map, _) => map.hash_ignore_nullability(state), + Struct(fields, _) => fields.hash_ignore_nullability(state), + Union(variants, _) => variants.hash_ignore_nullability(state), + Extension(ext) => ext.hash_ignore_nullability(state), + } + } + /// Returns `true` if `self` is a subset type of `other, otherwise `false`. /// /// If `self` is nullable, this means that the other `DType` must also be nullable (since a @@ -253,6 +271,11 @@ impl DType { matches!(self, FixedSizeList(..)) } + /// Check if `self` is a [`DType::Map`]. + pub fn is_map(&self) -> bool { + matches!(self, Map(..)) + } + /// Check if `self` is a [`DType::Struct`] pub fn is_struct(&self) -> bool { matches!(self, Struct(_, _)) @@ -277,7 +300,7 @@ impl DType { /// recursive type. pub fn is_nested(&self) -> bool { match self { - List(..) | FixedSizeList(..) | Struct(..) | Union(..) | Variant(..) => true, + List(..) | FixedSizeList(..) | Map(..) | Struct(..) | Union(..) | Variant(..) => true, Extension(ext) => ext.storage_dtype().is_nested(), _ => false, } @@ -296,7 +319,7 @@ impl DType { Decimal(decimal, _) => { Some(DecimalType::smallest_decimal_value_type(decimal).byte_width()) } - Utf8(_) | Binary(_) | List(..) => None, + Utf8(_) | Binary(_) | List(..) | Map(..) => None, FixedSizeList(elem_dtype, list_size, _) => { elem_dtype.element_size().map(|s| s * *list_size as usize) } @@ -310,7 +333,7 @@ impl DType { } Some(sum) } - Union(..) => todo!("TODO(connor)[Union]: unimplemented"), + Union(..) => None, Variant(_) => None, Extension(ext) => ext.storage_dtype().element_size(), } @@ -375,6 +398,24 @@ impl DType { } } + /// Get the [`MapDType`] if `self` is a [`DType::Map`], otherwise `None`. + pub fn as_map_opt(&self) -> Option<&MapDType> { + if let Map(map, _) = self { + Some(map) + } else { + None + } + } + + /// Owned version of [Self::as_map_opt]. + pub fn into_map_opt(self) -> Option { + if let Map(map, _) = self { + Some(map) + } else { + None + } + } + /// Get the inner element dtype if `self` is **either** a [`DType::List`] or a /// [`DType::FixedSizeList`], otherwise returns `None` pub fn as_any_size_list_element_opt(&self) -> Option<&Arc> { @@ -436,7 +477,10 @@ impl DType { } } - /// Get the [`UnionVariants`] if `self` is a [`DType::Union`], otherwise `None`. + /// Get the [`UnionVariants`] if `self` is a [`DType::Union`], otherwise [`None`]. + /// + /// This [`Option`] only represents whether the dtype is a union; Union nullability does not + /// affect the result. pub fn as_union_variants_opt(&self) -> Option<&UnionVariants> { if let Union(uv, _) = self { Some(uv) @@ -445,7 +489,7 @@ impl DType { } } - /// Owned version of [Self::as_union_variants_opt]. + /// Owned version of [`Self::as_union_variants_opt`]. pub fn into_union_variants_opt(self) -> Option { if let Union(uv, _) = self { Some(uv) @@ -476,6 +520,23 @@ impl DType { List(Arc::new(dtype.into()), nullability) } + /// Convenience method for creating a [`DType::Map`]. + /// + /// # Errors + /// + /// Returns an error when the key dtype is nullable. + pub fn map( + key: impl Into, + value: impl Into, + keys_sorted: bool, + nullability: Nullability, + ) -> VortexResult { + Ok(Map( + MapDType::try_new(key.into(), value.into(), keys_sorted)?, + nullability, + )) + } + /// Convenience method for creating a [`DType::Struct`]. pub fn struct_, impl Into)>>( iter: I, @@ -496,6 +557,7 @@ impl Display for DType { Binary(null) => write!(f, "binary{null}"), List(edt, null) => write!(f, "list({edt}){null}"), FixedSizeList(edt, size, null) => write!(f, "fixed_size_list({edt})[{size}]{null}"), + Map(map, null) => write!(f, "{map}{null}"), Struct(sf, null) => write!( f, "{{{}}}{null}", @@ -514,18 +576,29 @@ impl Display for DType { #[cfg(test)] mod tests { + use std::collections::hash_map::DefaultHasher; + use std::hash::Hasher; use std::sync::Arc; + use vortex_error::VortexResult; + use crate::dtype::DType; use crate::dtype::Nullability::NonNullable; use crate::dtype::Nullability::Nullable; use crate::dtype::PType; + use crate::dtype::UnionVariants; use crate::dtype::decimal::DecimalDType; use crate::extension::datetime::Date; use crate::extension::datetime::Time; use crate::extension::datetime::TimeUnit; use crate::extension::datetime::Timestamp; + fn hash_ignore_nullability(dtype: &DType) -> u64 { + let mut hasher = DefaultHasher::new(); + dtype.hash_ignore_nullability(&mut hasher); + hasher.finish() + } + #[test] fn test_ext_dtype_eq_ignore_nullability() { let d1 = DType::Extension(Time::new(TimeUnit::Seconds, Nullable).erased()); @@ -541,6 +614,58 @@ mod tests { assert!(!t1.eq_ignore_nullability(&t2)); } + #[test] + fn test_union_dtype_hash_ignores_variant_nullability() -> VortexResult<()> { + let lhs = DType::Union( + UnionVariants::try_new( + ["int", "string"].into(), + vec![ + DType::Primitive(PType::I32, Nullable), + DType::Utf8(NonNullable), + ], + vec![5, 9], + )?, + NonNullable, + ); + let rhs = DType::Union( + UnionVariants::try_new( + ["int", "string"].into(), + vec![ + DType::Primitive(PType::I32, NonNullable), + DType::Utf8(Nullable), + ], + vec![5, 9], + )?, + NonNullable, + ); + + assert!(lhs.eq_ignore_nullability(&rhs)); + assert_eq!(hash_ignore_nullability(&lhs), hash_ignore_nullability(&rhs)); + + Ok(()) + } + + #[test] + fn test_map_dtype_hash_ignores_entry_nullability() -> VortexResult<()> { + let lhs = DType::map( + DType::Primitive(PType::I32, NonNullable), + DType::Utf8(NonNullable), + true, + NonNullable, + )?; + let rhs = DType::map( + DType::Primitive(PType::I32, NonNullable), + DType::Utf8(Nullable), + true, + Nullable, + )?; + + assert!(lhs.eq_ignore_nullability(&rhs)); + assert_eq!(hash_ignore_nullability(&lhs), hash_ignore_nullability(&rhs)); + + Ok(()) + } + #[test] fn element_size_null() { assert_eq!(DType::Null.element_size(), Some(0)); @@ -627,4 +752,10 @@ mod tests { None ); } + + #[test] + #[should_panic(expected = "DType::Null cannot be made non-nullable")] + fn null_as_nonnullable_panics() { + DType::Null.as_nonnullable(); + } } diff --git a/vortex-array/src/dtype/extension/erased.rs b/vortex-array/src/dtype/extension/erased.rs index c84b198516c..b2f49974247 100644 --- a/vortex-array/src/dtype/extension/erased.rs +++ b/vortex-array/src/dtype/extension/erased.rs @@ -75,6 +75,14 @@ impl ExtDTypeRef { .eq_ignore_nullability(other.storage_dtype()) } + /// Hash this extension dtype using the same equivalence relation as + /// [`Self::eq_ignore_nullability`]. + pub(crate) fn hash_ignore_nullability(&self, state: &mut H) { + self.id().hash(state); + self.0.metadata_hash(state); + self.storage_dtype().hash_ignore_nullability(state); + } + // TODO(connor): We should add a different type that returns something that can be serialized. /// Serialize the metadata into a byte vector. pub fn serialize_metadata(&self) -> VortexResult> { @@ -99,21 +107,6 @@ impl ExtDTypeRef { pub(crate) fn validate_storage_value(&self, storage_value: &ScalarValue) -> VortexResult<()> { self.0.validate_scalar_value(storage_value) } - - /// Can a value of `other` be implicitly coerced into this extension type? - pub fn can_coerce_from(&self, other: &DType) -> bool { - self.0.can_coerce_from(other) - } - - /// Can this extension type be implicitly coerced into `other`? - pub fn can_coerce_to(&self, other: &DType) -> bool { - self.0.can_coerce_to(other) - } - - /// Compute the least supertype of this extension type and another type. - pub fn least_supertype(&self, other: &DType) -> Option { - self.0.least_supertype(other) - } } /// Methods for downcasting type-erased extension dtypes. diff --git a/vortex-array/src/dtype/extension/typed.rs b/vortex-array/src/dtype/extension/typed.rs index 5b176bbb34d..e59d4207860 100644 --- a/vortex-array/src/dtype/extension/typed.rs +++ b/vortex-array/src/dtype/extension/typed.rs @@ -108,21 +108,6 @@ impl ExtDType { V::validate_scalar_value(self, storage_value) } - /// Can a value of `other` be implicitly coerced into this extension type? - pub fn can_coerce_from(&self, other: &DType) -> bool { - V::can_coerce_from(self, other) - } - - /// Can this extension type be implicitly coerced into `other`? - pub fn can_coerce_to(&self, other: &DType) -> bool { - V::can_coerce_to(self, other) - } - - /// Compute the least supertype of this extension type and another type. - pub fn least_supertype(&self, other: &DType) -> Option { - V::least_supertype(self, other) - } - /// Erase the concrete type information, returning a type-erased extension dtype. pub fn erased(self) -> ExtDTypeRef { ExtDTypeRef(Arc::new(self)) @@ -148,9 +133,6 @@ pub(super) trait DynExtDType: 'static + Send + Sync + super::sealed::Sealed { fn validate_scalar_value(&self, storage_value: &ScalarValue) -> VortexResult<()>; fn value_display(&self, f: &mut fmt::Formatter<'_>, storage_value: &ScalarValue) -> fmt::Result; - fn can_coerce_from(&self, other: &DType) -> bool; - fn can_coerce_to(&self, other: &DType) -> bool; - fn least_supertype(&self, other: &DType) -> Option; } /// Blanket impl: thin forwarder to `ExtDType` inherent methods. @@ -220,16 +202,4 @@ impl DynExtDType for ExtDType { ), } } - - fn can_coerce_from(&self, other: &DType) -> bool { - self.can_coerce_from(other) - } - - fn can_coerce_to(&self, other: &DType) -> bool { - self.can_coerce_to(other) - } - - fn least_supertype(&self, other: &DType) -> Option { - self.least_supertype(other) - } } diff --git a/vortex-array/src/dtype/extension/vtable.rs b/vortex-array/src/dtype/extension/vtable.rs index c89f296a22b..ecea9d24a9a 100644 --- a/vortex-array/src/dtype/extension/vtable.rs +++ b/vortex-array/src/dtype/extension/vtable.rs @@ -7,7 +7,6 @@ use std::hash::Hash; use vortex_error::VortexResult; -use crate::dtype::DType; use crate::dtype::extension::ExtDType; use crate::dtype::extension::ExtId; use crate::scalar::ScalarValue; @@ -19,7 +18,8 @@ use crate::scalar::ScalarValue; /// compatibility. /// /// An extension dtype is not a new physical array layout. It is a logical wrapper around a storage -/// [`DType`] plus metadata. Implementations should keep [`validate_dtype`](Self::validate_dtype) +/// [`DType`](crate::dtype::DType) plus metadata. Implementations should keep +/// [`validate_dtype`](Self::validate_dtype) /// strict enough that every valid storage scalar can be interpreted by /// [`unpack_native`](Self::unpack_native). pub trait ExtVTable: 'static + Sized + Send + Sync + Clone + Debug + Eq + Hash { @@ -48,33 +48,6 @@ pub trait ExtVTable: 'static + Sized + Send + Sync + Clone + Debug + Eq + Hash { /// extension metadata. fn validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()>; - /// Can a value of `other` be implicitly widened into this type? (e.g. GeographyType might - /// accept Point, LineString, etc.) - /// - /// Implementors only need to override one of `can_coerce_from` or `can_coerce_to`. We have both - /// so that either side of the coercion can provide the logic. - fn can_coerce_from(ext_dtype: &ExtDType, other: &DType) -> bool { - let _ = (ext_dtype, other); - false - } - - /// Can this type be implicitly widened into `other`? - /// - /// Implementors only need to override one of `can_coerce_from` or `can_coerce_to`. We have both - /// so that either side of the coercion can provide the logic. - fn can_coerce_to(ext_dtype: &ExtDType, other: &DType) -> bool { - let _ = (ext_dtype, other); - false - } - - /// Given two types in a Uniform context, what is their least supertype? - /// - /// Return None if no supertype exists. - fn least_supertype(ext_dtype: &ExtDType, other: &DType) -> Option { - let _ = (ext_dtype, other); - None - } - // Methods related to the extension scalar values. /// Validate the given storage value is compatible with the extension type. diff --git a/vortex-array/src/dtype/map.rs b/vortex-array/src/dtype/map.rs new file mode 100644 index 00000000000..e138cea6545 --- /dev/null +++ b/vortex-array/src/dtype/map.rs @@ -0,0 +1,217 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt; +use std::hash::Hash; +use std::hash::Hasher; +use std::sync::Arc; + +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::dtype::DType; +use crate::dtype::FieldDType; +use crate::dtype::Nullability; + +/// Logical type information for a map's entries. +/// +/// A map has ordered key/value entries. Keys must be non-nullable, while values may be nullable. +/// `keys_sorted` is a producer assertion matching Arrow's map type; it is not validated against +/// data at type construction time. +#[allow( + clippy::derived_hash_with_manual_eq, + reason = "manual PartialEq adds Arc::ptr_eq fast path only" +)] +#[derive(Clone, Eq, Hash)] +pub struct MapDType(Arc); + +impl PartialEq for MapDType { + fn eq(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) || self.0 == other.0 + } +} + +impl fmt::Debug for MapDType { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("MapDType") + .field("key", &self.0.key) + .field("value", &self.0.value) + .field("keys_sorted", &self.0.keys_sorted) + .finish() + } +} + +impl fmt::Display for MapDType { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "map({}, {}, keys_sorted={})", + self.key_dtype(), + self.value_dtype(), + self.keys_sorted() + ) + } +} + +struct MapDTypeInner { + key: FieldDType, + value: FieldDType, + keys_sorted: bool, +} + +impl PartialEq for MapDTypeInner { + fn eq(&self, other: &Self) -> bool { + self.key == other.key && self.value == other.value && self.keys_sorted == other.keys_sorted + } +} + +impl Eq for MapDTypeInner {} + +impl Hash for MapDTypeInner { + fn hash(&self, state: &mut H) { + self.key.hash(state); + self.value.hash(state); + self.keys_sorted.hash(state); + } +} + +impl MapDType { + pub(crate) fn eq_ignore_nullability(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) + || (self.key_dtype().eq_ignore_nullability(&other.key_dtype()) + && self + .value_dtype() + .eq_ignore_nullability(&other.value_dtype()) + && self.keys_sorted() == other.keys_sorted()) + } + + pub(crate) fn hash_ignore_nullability(&self, state: &mut H) { + self.key_dtype().hash_ignore_nullability(state); + self.value_dtype().hash_ignore_nullability(state); + self.keys_sorted().hash(state); + } + + /// Creates a map dtype from its key and value dtypes. + /// + /// # Errors + /// + /// Returns an error when `key` is nullable. Arrow map keys cannot be null. + pub fn try_new(key: DType, value: DType, keys_sorted: bool) -> VortexResult { + Self::try_from_fields(key.into(), value.into(), keys_sorted) + } + + pub(crate) fn try_from_fields( + key: FieldDType, + value: FieldDType, + keys_sorted: bool, + ) -> VortexResult { + vortex_ensure!( + !key.value()?.is_nullable(), + "map key dtype must be non-nullable" + ); + + Ok(Self(Arc::new(MapDTypeInner { + key, + value, + keys_sorted, + }))) + } + + /// Returns the dtype of the map keys. + pub fn key_dtype(&self) -> DType { + self.0 + .key + .value() + .vortex_expect("map key dtype must be valid") + } + + /// Returns the dtype of the map values. + pub fn value_dtype(&self) -> DType { + self.0 + .value + .value() + .vortex_expect("map value dtype must be valid") + } + + /// Returns whether producers assert that keys are sorted within every map value. + pub fn keys_sorted(&self) -> bool { + self.0.keys_sorted + } + + /// Returns the non-nullable `{key, value}` struct dtype used for map entries. + pub fn entries_dtype(&self) -> DType { + DType::struct_( + [("key", self.key_dtype()), ("value", self.value_dtype())], + Nullability::NonNullable, + ) + } +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use crate::dtype::DType; + use crate::dtype::MapDType; + use crate::dtype::Nullability; + use crate::dtype::PType; + + #[test] + fn rejects_nullable_keys() { + let result = MapDType::try_new( + DType::Primitive(PType::I32, Nullability::Nullable), + DType::Utf8(Nullability::Nullable), + false, + ); + + assert!(result.is_err()); + } + + #[test] + fn accepts_nullable_values() -> VortexResult<()> { + let dtype = MapDType::try_new( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + )?; + + assert_eq!( + dtype.entries_dtype(), + DType::struct_( + [ + ( + "key", + DType::Primitive(PType::I32, Nullability::NonNullable) + ), + ("value", DType::Utf8(Nullability::Nullable)), + ], + Nullability::NonNullable, + ) + ); + assert!(dtype.keys_sorted()); + assert_eq!(dtype.to_string(), "map(i32, utf8?, keys_sorted=true)"); + + Ok(()) + } + + #[test] + fn outer_nullability_is_independent_of_map_details() -> VortexResult<()> { + let dtype = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + Nullability::NonNullable, + )?; + let nullable = dtype.as_nullable(); + + assert!(dtype.is_map()); + assert!(!dtype.is_nullable()); + assert!(nullable.is_nullable()); + assert!(dtype.eq_ignore_nullability(&nullable)); + assert_eq!(dtype.as_map_opt(), nullable.as_map_opt()); + assert!(nullable.as_map_opt().unwrap().keys_sorted()); + + Ok(()) + } +} diff --git a/vortex-array/src/dtype/mod.rs b/vortex-array/src/dtype/mod.rs index 8b579c70d39..aaf729ef994 100644 --- a/vortex-array/src/dtype/mod.rs +++ b/vortex-array/src/dtype/mod.rs @@ -16,7 +16,6 @@ #[cfg(feature = "arbitrary")] mod arbitrary; mod bigint; -mod coercion; mod decimal; mod dtype_impl; pub mod extension; @@ -24,6 +23,7 @@ mod f16; mod field; mod field_mask; mod field_names; +mod map; mod native_dtype; mod nullability; mod ptype; @@ -102,6 +102,12 @@ pub enum DType { /// well as a `u32` size that determines the fixed length of each `FixedSizeList` scalar. FixedSizeList(Arc, u32, Nullability), + /// A logical map type. + /// + /// Map keys are non-nullable, values may be nullable, and [`MapDType`] stores Arrow's + /// `keys_sorted` assertion. + Map(MapDType, Nullability), + /// A logical struct type. /// /// A `Struct` type is composed of an ordered list of fields, each with a corresponding name and @@ -147,6 +153,7 @@ impl PartialEq for DType { (Self::FixedSizeList(da, sa, na), Self::FixedSizeList(db, sb, nb)) => { sa == sb && na == nb && (Arc::ptr_eq(da, db) || da == db) } + (Self::Map(ma, na), Self::Map(mb, nb)) => na == nb && ma == mb, // StructFields handles its own Arc::ptr_eq in its PartialEq impl. (Self::Struct(a, na), Self::Struct(b, nb)) => na == nb && a == b, // UnionVariants handles its own Arc::ptr_eq in its PartialEq impl. @@ -163,6 +170,7 @@ impl PartialEq for DType { | (Self::Binary(_), _) | (Self::List(..), _) | (Self::FixedSizeList(..), _) + | (Self::Map(..), _) | (Self::Struct(..), _) | (Self::Union(..), _) | (Self::Variant(_), _) @@ -179,6 +187,7 @@ pub use field::*; pub use field_mask::*; pub use field_names::*; pub use half; +pub use map::*; pub use nullability::*; pub use ptype::*; pub use struct_::*; diff --git a/vortex-array/src/dtype/ptype.rs b/vortex-array/src/dtype/ptype.rs index a038eb328c2..928a068f103 100644 --- a/vortex-array/src/dtype/ptype.rs +++ b/vortex-array/src/dtype/ptype.rs @@ -91,6 +91,32 @@ pub trait UnsignedPType: IntegerPType + Unsigned {} /// Implements [`UnsignedPType`] for all possible `T` that have the correct bounds. impl UnsignedPType for T where T: IntegerPType + Unsigned {} +/// Trait for the integer types that can back a builder's `offsets` and `sizes`: `u32`, `i32`, +/// `u64`, and `i64`. +/// +/// This trait is sealed and cannot be implemented for any other type. Restricting the set of +/// widths keeps the concrete [`ListBuilder`](crate::builders::ListBuilder), +/// [`ListViewBuilder`](crate::builders::ListViewBuilder), and +/// [`VarBinBuilder`](crate::builders::VarBinBuilder) instantiations small enough for +/// [`match_each_list_builder!`](crate::match_each_list_builder) and +/// [`match_each_varbin_builder!`](crate::match_each_varbin_builder) to enumerate them. +pub trait OffsetBuilderPType: IntegerPType + offset_builder_sealed::Sealed {} + +mod offset_builder_sealed { + pub trait Sealed {} +} + +macro_rules! impl_offset_builder_ptype { + ($($T:ty),*) => { + $( + impl offset_builder_sealed::Sealed for $T {} + impl OffsetBuilderPType for $T {} + )* + }; +} + +impl_offset_builder_ptype!(u32, i32, u64, i64); + /// A trait for native Rust types that correspond 1:1 to a PType. /// /// You can use the `match_each_native_ptype` macro to help with writing "generic" code over @@ -691,6 +717,25 @@ macro_rules! match_smallest_offset_type { }}; } +/// Macro to match the smallest [`OffsetBuilderPType`] able to index a given number of elements. +/// +/// Like [`match_smallest_offset_type!`](crate::match_smallest_offset_type), but restricted to the +/// unsigned types valid as builder offsets (`u32` and `u64`). +#[macro_export] +macro_rules! match_smallest_list_offset_type { + ($n_elements:expr, | $offset_type:ident | $body:block) => {{ + let n_elements = $n_elements; + if n_elements <= u32::MAX as usize { + type $offset_type = u32; + $body + } else { + assert!(u64::try_from(n_elements).is_ok()); + type $offset_type = u64; + $body + } + }}; +} + impl PType { /// Returns `true` iff this PType is an unsigned integer type #[inline] diff --git a/vortex-array/src/dtype/serde/flatbuffers.rs b/vortex-array/src/dtype/serde/flatbuffers.rs index a796490f58a..8a1e7e6b611 100644 --- a/vortex-array/src/dtype/serde/flatbuffers.rs +++ b/vortex-array/src/dtype/serde/flatbuffers.rs @@ -21,6 +21,8 @@ use vortex_session::VortexSession; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::FieldDType; +use crate::dtype::FieldNames; +use crate::dtype::MapDType; use crate::dtype::PType; use crate::dtype::StructFields; use crate::dtype::UnionVariants; @@ -68,7 +70,7 @@ impl StructFields { buffer: FlatBuffer, session: VortexSession, ) -> VortexResult { - let names = fb_struct + let names: FieldNames = fb_struct .names() .ok_or_else(|| vortex_err!("failed to parse struct names from flatbuffer"))? .iter() @@ -87,6 +89,14 @@ impl StructFields { }) .collect::>(); + if names.len() != dtypes.len() { + vortex_bail!( + "length mismatch between struct names ({}) and dtypes ({})", + names.len(), + dtypes.len() + ); + } + Ok(StructFields::from_fields(names, dtypes)) } } @@ -128,6 +138,32 @@ impl UnionVariants { } } +impl MapDType { + /// Creates a map dtype from a flatbuffer-defined object and its underlying buffer. + fn from_fb( + fb_map: fbd::Map<'_>, + buffer: FlatBuffer, + session: VortexSession, + ) -> VortexResult { + let key = fb_map + .key_type() + .ok_or_else(|| vortex_err!("failed to parse map key type from flatbuffer"))?; + let value = fb_map + .value_type() + .ok_or_else(|| vortex_err!("failed to parse map value type from flatbuffer"))?; + + MapDType::try_from_fields( + FieldDType::from(ViewedDType::from_fb_loc( + key._tab.loc(), + buffer.clone(), + session.clone(), + )), + FieldDType::from(ViewedDType::from_fb_loc(value._tab.loc(), buffer, session)), + fb_map.keys_sorted(), + ) + } +} + impl DType { /// Create a [`DType`] from a flatbuffer buffer. pub fn from_flatbuffer(buffer: FlatBuffer, session: &VortexSession) -> VortexResult { @@ -220,6 +256,13 @@ impl TryFrom for DType { fb_fixed_size_list.nullable().into(), )) } + fb::Type::Map => { + let fb_map = fb + .type__as_map() + .ok_or_else(|| vortex_err!("failed to parse map from flatbuffer"))?; + let map = MapDType::from_fb(fb_map, vfdt.buffer().clone(), vfdt.session.clone())?; + Ok(Self::Map(map, fb_map.nullable().into())) + } fb::Type::Struct_ => { let fb_struct = fb .type__as_struct_() @@ -270,7 +313,7 @@ impl TryFrom for DType { vortex_err!("failed to parse extension metadata from flatbuffer") })? .bytes(); - let ext_dtype = if let Some(vtable) = vfdt.session.dtypes().registry().find(&id) { + let ext_dtype = if let Some(vtable) = vfdt.session.dtypes().registry().get(&id) { vtable.deserialize(metadata, storage_dtype)? } else if vfdt.session.allows_unknown() { ForeignExtDType::from_parts(id, metadata.to_vec(), storage_dtype)? @@ -357,6 +400,20 @@ impl WriteFlatBuffer for DType { ) .as_union_value() } + Self::Map(map, n) => { + let key_type = Some(map.key_dtype().write_flatbuffer(fbb)?); + let value_type = Some(map.value_dtype().write_flatbuffer(fbb)?); + fb::Map::create( + fbb, + &fb::MapArgs { + key_type, + value_type, + keys_sorted: map.keys_sorted(), + nullable: (*n).into(), + }, + ) + .as_union_value() + } Self::Struct(st, n) => { let names = st .names() @@ -448,6 +505,7 @@ impl WriteFlatBuffer for DType { Self::Binary(_) => fb::Type::Binary, Self::List(..) => fb::Type::List, Self::FixedSizeList(..) => fb::Type::FixedSizeList, + Self::Map(..) => fb::Type::Map, Self::Struct(..) => fb::Type::Struct_, Self::Union(..) => fb::Type::Union, Self::Variant(_) => fb::Type::Variant, @@ -507,8 +565,11 @@ impl TryFrom for PType { mod test { use std::sync::Arc; + use flatbuffers::FlatBufferBuilder; use flatbuffers::root; + use vortex_buffer::ByteBuffer; use vortex_flatbuffers::FlatBuffer; + use vortex_flatbuffers::WriteFlatBuffer; use vortex_flatbuffers::WriteFlatBufferExt; use crate::dtype::DType; @@ -569,6 +630,21 @@ mod test { ), Nullability::NonNullable, )); + let inner_map = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + Nullability::NonNullable, + ) + .unwrap(); + let map = DType::map( + inner_map, + DType::Utf8(Nullability::Nullable), + false, + Nullability::Nullable, + ) + .unwrap(); + roundtrip_dtype(DType::struct_([("map", map)], Nullability::NonNullable)); roundtrip_dtype(DType::Variant(Nullability::Nullable)); } @@ -678,6 +754,48 @@ mod test { assert_eq!(viewed, eager); } + #[test] + fn test_struct_malformed_flatbuffer() { + let mut fbb = FlatBufferBuilder::new(); + let names = fbb.create_vector::>(&[]); + let dtype_offsets = (0..3) + .map(|_| { + DType::Primitive(PType::I32, Nullability::NonNullable) + .write_flatbuffer(&mut fbb) + .unwrap() + }) + .collect::>(); + let dtypes = fbb.create_vector(&dtype_offsets); + + let struct_table = fb::Struct_::create( + &mut fbb, + &fb::Struct_Args { + names: Some(names), + dtypes: Some(dtypes), + nullable: false, + }, + ); + + let dtype = fb::DType::create( + &mut fbb, + &fb::DTypeArgs { + type_type: fb::Type::Struct_, + type_: Some(struct_table.as_union_value()), + }, + ); + fbb.finish_minimal(dtype); + let (vec, start) = fbb.collapse(); + let end = vec.len(); + let buffer = FlatBuffer::align_from(ByteBuffer::from(vec).slice(start..end)); + + let root_fb = root::(&buffer).unwrap(); + let view = ViewedDType::from_fb_loc(root_fb._tab.loc(), buffer, SESSION.clone()); + + let result = DType::try_from(view); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("length mismatch")); + } + /// A malformed flatbuffer (here, `dtypes.len() != type_ids.len()`) must round-trip /// to `Err`, not panic. #[test] diff --git a/vortex-array/src/dtype/serde/mod.rs b/vortex-array/src/dtype/serde/mod.rs index c6837f4ab2d..f26912b5e53 100644 --- a/vortex-array/src/dtype/serde/mod.rs +++ b/vortex-array/src/dtype/serde/mod.rs @@ -203,4 +203,31 @@ mod test { assert_eq!(deserialized, dtype); assert_eq!(deserialized.nullability(), Nullability::Nullable); } + + #[test] + fn test_serde_nested_map_dtype_json_roundtrip() { + let inner_map = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + Nullability::NonNullable, + ) + .unwrap(); + let map = DType::map( + inner_map, + DType::Utf8(Nullability::Nullable), + false, + Nullability::Nullable, + ) + .unwrap(); + let dtype = DType::struct_([("map", map)], Nullability::Nullable); + + let json = serde_json::to_string(&dtype).unwrap(); + let mut deserializer = serde_json::Deserializer::from_str(&json); + let deserialized: DType = DTypeSerde::::new(&SESSION) + .deserialize(&mut deserializer) + .unwrap(); + + assert_eq!(deserialized, dtype); + } } diff --git a/vortex-array/src/dtype/serde/proto.rs b/vortex-array/src/dtype/serde/proto.rs index 8389a90200a..50367a90276 100644 --- a/vortex-array/src/dtype/serde/proto.rs +++ b/vortex-array/src/dtype/serde/proto.rs @@ -10,6 +10,7 @@ use vortex_session::VortexSession; use crate::dtype::DType; use crate::dtype::DecimalDType; +use crate::dtype::MapDType; use crate::dtype::PType; use crate::dtype::StructFields; use crate::dtype::UnionVariants; @@ -77,6 +78,26 @@ impl DType { nullable, )) } + DtypeType::Map(map) => Ok(Self::Map( + MapDType::try_new( + DType::from_proto( + map.key_type + .as_ref() + .ok_or_else(|| vortex_err!(Serde: "Invalid map key type"))? + .as_ref(), + session, + )?, + DType::from_proto( + map.value_type + .as_ref() + .ok_or_else(|| vortex_err!(Serde: "Invalid map value type"))? + .as_ref(), + session, + )?, + map.keys_sorted, + )?, + map.nullable.into(), + )), DtypeType::Struct(s) => Ok(Self::Struct( StructFields::new( s.names.iter().map(|s| s.as_str()).collect(), @@ -116,7 +137,7 @@ impl DType { .ok_or_else(|| vortex_err!("Extension DType missing storage proto"))?, session, )?; - let ext_dtype = if let Some(vtable) = session.dtypes().registry().find(&id) { + let ext_dtype = if let Some(vtable) = session.dtypes().registry().get(&id) { vtable.deserialize(e.metadata(), storage_dtype)? } else if session.allows_unknown() { ForeignExtDType::from_parts(id, e.metadata().to_vec(), storage_dtype)? @@ -165,6 +186,16 @@ impl TryFrom<&DType> for pb::DType { nullable: (*null).into(), })) } + DType::Map(map, null) => { + let key_dtype = map.key_dtype(); + let value_dtype = map.value_dtype(); + DtypeType::Map(Box::new(pb::Map { + key_type: Some(Box::new(Self::try_from(&key_dtype)?)), + value_type: Some(Box::new(Self::try_from(&value_dtype)?)), + keys_sorted: map.keys_sorted(), + nullable: (*null).into(), + })) + } DType::Struct(s, null) => DtypeType::Struct(pb::Struct { names: s.names().iter().map(|s| s.as_ref().to_string()).collect(), dtypes: s @@ -393,6 +424,27 @@ mod tests { } } + #[test] + fn test_nested_map_round_trip() { + let inner_map = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + Nullability::NonNullable, + ) + .unwrap(); + let map = DType::map( + inner_map, + DType::Utf8(Nullability::Nullable), + false, + Nullability::Nullable, + ) + .unwrap(); + let dtype = DType::struct_([("map", map)], Nullability::NonNullable); + + assert_eq!(round_trip_dtype(&dtype), dtype); + } + #[test] fn test_extension_round_trip() { let ext_dtype = diff --git a/vortex-array/src/dtype/serde/serde.rs b/vortex-array/src/dtype/serde/serde.rs index b3870b8c03e..e941fb475bb 100644 --- a/vortex-array/src/dtype/serde/serde.rs +++ b/vortex-array/src/dtype/serde/serde.rs @@ -110,6 +110,14 @@ impl Serialize for DType { state.serialize_field(n)?; state.end() } + DType::Map(map, n) => { + let mut state = serializer.serialize_tuple_variant("DType", 12, "Map", 4)?; + state.serialize_field(&map.key_dtype())?; + state.serialize_field(&map.value_dtype())?; + state.serialize_field(&map.keys_sorted())?; + state.serialize_field(n)?; + state.end() + } DType::Struct(fields, n) => { let mut state = serializer.serialize_tuple_variant("DType", 8, "Struct", 2)?; state.serialize_field(&fields)?; @@ -181,6 +189,7 @@ impl<'de> DeserializeSeed<'de> for DTypeSerde<'_, DType> { "Union", "Variant", "Extension", + "Map", ]; struct DTypeVisitor<'a> { @@ -234,6 +243,9 @@ impl<'de> DeserializeSeed<'de> for DTypeSerde<'_, DType> { "FixedSizeList" => access.newtype_variant_seed(FixedSizeListFieldsSeed { session: self.session, }), + "Map" => access.newtype_variant_seed(MapFieldsSeed { + session: self.session, + }), "Struct" => access.newtype_variant_seed(StructFieldsSeed { session: self.session, }), @@ -264,6 +276,59 @@ impl<'de> DeserializeSeed<'de> for DTypeSerde<'_, DType> { } } +struct MapFieldsSeed<'a> { + session: &'a VortexSession, +} + +impl<'de> DeserializeSeed<'de> for MapFieldsSeed<'_> { + type Value = DType; + + fn deserialize(self, deserializer: D) -> Result + where + D: Deserializer<'de>, + { + struct MapVisitor<'a> { + session: &'a VortexSession, + } + + impl<'de> Visitor<'de> for MapVisitor<'_> { + type Value = DType; + + fn expecting(&self, f: &mut Formatter) -> fmt::Result { + f.write_str("Map tuple (key_dtype, value_dtype, keys_sorted, nullability)") + } + + fn visit_seq(self, mut seq: A) -> Result + where + A: SeqAccess<'de>, + { + let key = seq + .next_element_seed(DTypeSerde::::new(self.session))? + .ok_or_else(|| de::Error::invalid_length(0, &self))?; + let value = seq + .next_element_seed(DTypeSerde::::new(self.session))? + .ok_or_else(|| de::Error::invalid_length(1, &self))?; + let keys_sorted = seq + .next_element()? + .ok_or_else(|| de::Error::invalid_length(2, &self))?; + let nullability = seq + .next_element()? + .ok_or_else(|| de::Error::invalid_length(3, &self))?; + + DType::map(key, value, keys_sorted, nullability) + .map_err(|error| de::Error::custom(error.to_string())) + } + } + + deserializer.deserialize_tuple( + 4, + MapVisitor { + session: self.session, + }, + ) + } +} + // ============================================================================ // Helper seeds for nested DType variants (with session) // ============================================================================ @@ -722,7 +787,7 @@ impl<'de> DeserializeSeed<'de> for DTypeSerde<'_, ExtDTypeRef> { storage_dtype.ok_or_else(|| de::Error::missing_field("storage_dtype"))?; let metadata = metadata.ok_or_else(|| de::Error::missing_field("metadata"))?; - if let Some(vtable) = self.session.dtypes().registry().find(&id) { + if let Some(vtable) = self.session.dtypes().registry().get(&id) { vtable.deserialize(&metadata, storage_dtype).map_err(|e| { de::Error::custom(format!( "failed to deserialize extension dtype {}: {}", diff --git a/vortex-array/src/dtype/session.rs b/vortex-array/src/dtype/session.rs index 2e13b923bc3..bd4a3265d70 100644 --- a/vortex-array/src/dtype/session.rs +++ b/vortex-array/src/dtype/session.rs @@ -6,10 +6,11 @@ use std::any::Any; use std::sync::Arc; +use vortex_session::ArcSwapMap; use vortex_session::SessionExt; use vortex_session::SessionGuard; use vortex_session::SessionVar; -use vortex_session::registry::Registry; +use vortex_session::registry::Id; use crate::dtype::extension::ExtDTypePluginRef; use crate::dtype::extension::ExtVTable; @@ -19,7 +20,7 @@ use crate::extension::datetime::Timestamp; use crate::extension::uuid::Uuid; /// Registry for extension dtypes. -pub type ExtDTypeRegistry = Registry; +pub type ExtDTypeRegistry = ArcSwapMap; /// Session for managing extension dtypes. #[derive(Clone, Debug)] @@ -30,7 +31,7 @@ pub struct DTypeSession { impl Default for DTypeSession { fn default() -> Self { let this = Self { - registry: Registry::default(), + registry: ExtDTypeRegistry::default(), }; // Register built-in temporal extension dtypes @@ -57,7 +58,7 @@ impl DTypeSession { /// Register an extension DType with the Vortex session. pub fn register(&self, vtable: V) { self.registry - .register(vtable.id(), Arc::new(vtable) as ExtDTypePluginRef); + .insert(vtable.id(), Arc::new(vtable) as ExtDTypePluginRef); } /// Return the registry of extension dtypes. diff --git a/vortex-array/src/dtype/struct_.rs b/vortex-array/src/dtype/struct_.rs index 4d276c35205..4b457d84d16 100644 --- a/vortex-array/src/dtype/struct_.rs +++ b/vortex-array/src/dtype/struct_.rs @@ -4,6 +4,7 @@ use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hash; +use std::hash::Hasher; use std::sync::Arc; use std::sync::OnceLock; @@ -29,9 +30,8 @@ pub struct FieldDType { impl From for FieldDType { fn from(value: ViewedDType) -> Self { - Self { - inner: FieldDTypeInner::View(value), - } + let view = FieldDTypeInner::View(value, Arc::new(OnceLock::new())); + Self { inner: view } } } @@ -56,26 +56,38 @@ enum FieldDTypeInner { /// Owned DType instance // TODO(ngates): we should consider making this an Arc. Owned(DType), - /// A view over a flatbuffer, parsed only when accessed. - View(ViewedDType), + /// A view over a flatbuffer, parsed on first access and cached after. + /// We cache it because it's requested for every field for every file. + /// If there are many files with wide schemas (think clickbench), + /// re-parsing each field is costly + /// This form is quite ugly because in some cases we want to panic on + /// parsing but return an error on others + View(ViewedDType, Arc>), } impl PartialEq for FieldDTypeInner { fn eq(&self, other: &Self) -> bool { match (self, other) { (Self::Owned(lhs), Self::Owned(rhs)) => lhs == rhs, - (Self::View(lhs), Self::View(rhs)) => { - let lhs = DType::try_from(lhs.clone()) - .vortex_expect("Failed to parse FieldDType into DType"); - let rhs = DType::try_from(rhs.clone()) - .vortex_expect("Failed to parse FieldDType into DType"); + (Self::View(left_view, left_lock), Self::View(right_view, right_lock)) => { + let lhs = left_lock.get_or_init(|| { + DType::try_from(left_view.clone()) + .vortex_expect("Failed to parse FieldDType into DType") + }); + let rhs = right_lock.get_or_init(|| { + DType::try_from(right_view.clone()) + .vortex_expect("Failed to parse FieldDType into DType") + }); lhs == rhs } - (Self::View(view), Self::Owned(owned)) | (Self::Owned(owned), Self::View(view)) => { - let view = DType::try_from(view.clone()) - .vortex_expect("Failed to parse FieldDType into DType"); - owned == &view + (Self::View(view, lock), Self::Owned(owned)) + | (Self::Owned(owned), Self::View(view, lock)) => { + let view = lock.get_or_init(|| { + DType::try_from(view.clone()) + .vortex_expect("Failed to parse FieldDType into DType") + }); + owned == view } } } @@ -83,14 +95,16 @@ impl PartialEq for FieldDTypeInner { impl Eq for FieldDTypeInner {} impl Hash for FieldDTypeInner { - fn hash(&self, state: &mut H) { + fn hash(&self, state: &mut H) { match self { FieldDTypeInner::Owned(owned) => { owned.hash(state); } - FieldDTypeInner::View(view) => { - let owned = DType::try_from(view.clone()) - .vortex_expect("Failed to parse FieldDType into DType"); + FieldDTypeInner::View(view, lock) => { + let owned = lock.get_or_init(|| { + DType::try_from(view.clone()) + .vortex_expect("Failed to parse FieldDType into DType") + }); owned.hash(state); } } @@ -109,7 +123,13 @@ impl FieldDTypeInner { fn value(&self) -> VortexResult { match &self { FieldDTypeInner::Owned(owned) => Ok(owned.clone()), - FieldDTypeInner::View(view) => DType::try_from(view.clone()), + FieldDTypeInner::View(view, lock) => { + if let Some(dtype) = lock.get() { + return Ok(dtype.clone()); + } + let parsed = DType::try_from(view.clone())?; + Ok(lock.get_or_init(|| parsed).clone()) + } } } } @@ -255,12 +275,33 @@ impl PartialEq for StructFieldsInner { impl Eq for StructFieldsInner {} impl Hash for StructFieldsInner { - fn hash(&self, state: &mut H) { + fn hash(&self, state: &mut H) { self.names.hash(state); self.dtypes.hash(state); } } +impl StructFields { + /// Check if these struct fields are equal, ignoring field dtype nullability recursively. + pub fn eq_ignore_nullability(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) + || (self.0.names == other.0.names + && self + .fields() + .zip_eq(other.fields()) + .all(|(lhs, rhs)| lhs.eq_ignore_nullability(&rhs))) + } + + /// Hash these struct fields using the same equivalence relation as + /// [`Self::eq_ignore_nullability`]. + pub(crate) fn hash_ignore_nullability(&self, state: &mut H) { + self.0.names.hash(state); + for field in self.fields() { + field.hash_ignore_nullability(state); + } + } +} + impl Default for StructFields { fn default() -> Self { Self::empty() @@ -469,12 +510,17 @@ mod test { use insta::assert_snapshot; use itertools::Itertools; + use vortex_error::VortexResult; + use vortex_flatbuffers::FlatBuffer; + use vortex_flatbuffers::WriteFlatBufferExt; + use super::FieldDTypeInner; use crate::dtype::DType; use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; + use crate::dtype::test::SESSION; #[test] fn nullability() { @@ -553,6 +599,36 @@ mod test { assert_eq!(without_a.nfields(), 1); } + #[test] + fn field_dtype_cache() -> VortexResult<()> { + let inner = DType::Struct( + StructFields::from_iter([("x", DType::Bool(Nullability::NonNullable))]), + Nullability::NonNullable, + ); + let dtype = DType::Struct( + StructFields::from_iter([("a", inner)]), + Nullability::NonNullable, + ); + let buffer = FlatBuffer::from(dtype.write_flatbuffer_bytes()?); + let parsed = DType::from_flatbuffer(buffer, &SESSION)?; + let fields = parsed.as_struct_fields_opt().unwrap(); + + let field = &fields.0.dtypes[0]; + let FieldDTypeInner::View(_, cell) = &field.inner else { + panic!("flatbuffer-backed field must be a View"); + }; + assert!(cell.get().is_none()); + + let first = field.value()?; + assert!(cell.get().is_some()); + let second = field.value()?; + let (DType::Struct(first, _), DType::Struct(second, _)) = (&first, &second) else { + panic!("field must parse to a struct"); + }; + assert!(Arc::ptr_eq(&first.0, &second.0)); + Ok(()) + } + #[test] fn test_merge() { let child_a = DType::Primitive(PType::I32, Nullability::NonNullable); diff --git a/vortex-array/src/dtype/union.rs b/vortex-array/src/dtype/union.rs index 7ce53155adc..30df46e2d80 100644 --- a/vortex-array/src/dtype/union.rs +++ b/vortex-array/src/dtype/union.rs @@ -3,6 +3,7 @@ use std::fmt; use std::hash::Hash; +use std::hash::Hasher; use std::sync::Arc; use itertools::Itertools; @@ -18,8 +19,9 @@ use crate::dtype::FieldNames; /// Type information for a union array. /// /// A `UnionVariants` describes the possible alternative types for each row of a union, along with a -/// per-variant `u8` type tag. We use the term **variants** (rather than "fields") because a union -/// is a sum type: each row chooses exactly one alternative. +/// per-variant `u8` type tag. A union must contain at least one variant. We use the term +/// **variants** (rather than "fields") because a union is a sum type: each row chooses exactly one +/// alternative. /// /// By default, tag `i` selects the child at offset `i` (`type_ids = [0, 1, ..., N-1]`). Vortex uses /// unsigned tags and supports up to 256 variants. @@ -133,29 +135,37 @@ impl PartialEq for UnionVariantsInner { impl Eq for UnionVariantsInner {} impl Hash for UnionVariantsInner { - fn hash(&self, state: &mut H) { + fn hash(&self, state: &mut H) { self.names.hash(state); self.dtypes.hash(state); self.type_ids.hash(state); } } -impl Default for UnionVariants { - fn default() -> Self { - Self::empty() +impl UnionVariants { + /// Check if these union variants are equal, ignoring variant dtype nullability recursively. + pub fn eq_ignore_nullability(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) + || (self.0.names == other.0.names + && self + .variants() + .zip_eq(other.variants()) + .all(|(lhs, rhs)| lhs.eq_ignore_nullability(&rhs)) + && self.0.type_ids == other.0.type_ids) + } + + /// Hash these union variants using the same equivalence relation as + /// [`Self::eq_ignore_nullability`]. + pub(crate) fn hash_ignore_nullability(&self, state: &mut H) { + self.0.names.hash(state); + for variant in self.variants() { + variant.hash_ignore_nullability(state); + } + self.0.type_ids.hash(state); } } impl UnionVariants { - /// The variants of the empty union. - pub fn empty() -> Self { - Self(Arc::new(UnionVariantsInner::from_fields( - FieldNames::default(), - Arc::from([]), - Arc::from([]), - ))) - } - /// Validate that `names`, `dtypes`, and `type_ids` are mutually consistent. fn validate_shape(names: &FieldNames, n_dtypes: usize, type_ids: &[u8]) -> VortexResult<()> { vortex_ensure_eq!( @@ -172,6 +182,10 @@ impl UnionVariants { names.len(), type_ids.len() ); + vortex_ensure!( + !names.is_empty(), + "union must have at least one variant (for now)" + ); vortex_ensure!( type_ids.iter().all_unique(), @@ -197,8 +211,9 @@ impl UnionVariants { /// /// # Errors /// - /// Returns an error if names, dtypes, or type IDs do not all have the same length, or if there - /// are any duplicate names or type IDs, or if there are more than 256 variants. + /// Returns an error if the union has no variants, if names, dtypes, or type IDs do not all have + /// the same length, if there are any duplicate names or type IDs, or if there are more than 256 + /// variants. pub fn try_new(names: FieldNames, dtypes: Vec, type_ids: Vec) -> VortexResult { Self::validate_shape(&names, dtypes.len(), &type_ids)?; @@ -214,8 +229,8 @@ impl UnionVariants { /// /// # Errors /// - /// `names` and `dtypes` must have the same length, and `names.len()` cannot be more than - /// `u8::MAX as usize + 1` (256). + /// The union must have at least one variant, `names` and `dtypes` must have the same length, and + /// `names.len()` cannot be more than `u8::MAX as usize + 1` (256). pub fn new(names: FieldNames, dtypes: Vec) -> VortexResult { const MAX_VARIANTS: usize = u8::MAX as usize + 1; vortex_ensure!( @@ -241,8 +256,9 @@ impl UnionVariants { /// /// # Errors /// - /// Returns an error if names, dtypes, or type IDs do not all have the same length, or if there - /// are any duplicate names or type IDs, or if there are more than 256 variants. + /// Returns an error if the union has no variants, if names, dtypes, or type IDs do not all have + /// the same length, if there are any duplicate names or type IDs, or if there are more than 256 + /// variants. pub(crate) fn try_from_fields( names: FieldNames, dtypes: Vec, @@ -306,7 +322,7 @@ impl UnionVariants { ) } - /// Get the [`DType`] of a variant by offset. + /// Get the [`DType`] of a variant by offset, or [`None`] if `index` is out of bounds. pub fn variant_by_index(&self, index: usize) -> Option { Some( self.0 @@ -576,11 +592,19 @@ mod tests { } #[test] - fn test_empty() { - let v = UnionVariants::empty(); - assert!(v.is_empty()); - assert_eq!(v.len(), 0); - assert_eq!(v.type_ids(), &[] as &[u8]); - assert!(v.is_consecutive()); + fn test_empty_rejected() { + let error = UnionVariants::new(FieldNames::default(), vec![]).unwrap_err(); + assert!( + error + .to_string() + .contains("union must have at least one variant") + ); + + let error = UnionVariants::try_new(FieldNames::default(), vec![], vec![]).unwrap_err(); + assert!( + error + .to_string() + .contains("union must have at least one variant") + ); } } diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index a49f98b4d55..18afd45e16c 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -17,6 +17,7 @@ use std::fmt; use std::fmt::Display; use std::sync::Arc; use std::sync::LazyLock; +use std::sync::OnceLock; #[cfg(debug_assertions)] use std::sync::atomic::AtomicUsize; #[cfg(debug_assertions)] @@ -38,7 +39,7 @@ use crate::builders::ArrayBuilder; use crate::builders::builder_with_capacity_in; use crate::dtype::DType; use crate::matcher::Matcher; -use crate::memory::HostAllocatorRef; +use crate::memory::BufferAllocatorRef; use crate::memory::MemorySessionExt; use crate::optimizer::ArrayOptimizer; use crate::optimizer::kernels::ArrayKernelsExt; @@ -46,6 +47,7 @@ use crate::optimizer::kernels::ParentExecutionKernels; use crate::optimizer::kernels::execute_parent_key; use crate::stats::ArrayStats; use crate::stats::StatsSet; +use crate::trace_op; /// Returns the maximum number of iterations to attempt when executing an array before giving up and returning /// an error, can be by the `VORTEX_MAX_ITERATIONS` env variables, otherwise defaults to 2^22. @@ -163,6 +165,7 @@ impl ArrayRef { /// partially consumes `current_array`: some slots already live in the builder, so a /// parent rewrite would observe inconsistent state and could discard accumulated builder /// data. + #[allow(clippy::cognitive_complexity)] pub fn execute_until(self, ctx: &mut ExecutionCtx) -> VortexResult { let mut current_array = self; let mut current_builder: Option> = None; @@ -171,23 +174,40 @@ impl ArrayRef { let kernels = execute_parent_kernels.as_ref(); let max_iterations = max_iterations(); - for _ in 0..max_iterations { + trace_op!(record_execute_until_start::(¤t_array)); + + for _iteration in 0..max_iterations { + trace_op!(record_execute_until_iteration( + _iteration, + ¤t_array, + stack + .last() + .map(|frame| (&frame.parent_array, frame.slot_idx)), + current_builder.is_some(), + )); + let is_done = stack .last() .map_or(M::matches as DonePredicate, |frame| frame.done); - if is_done(¤t_array) || AnyCanonical::matches(¤t_array) { + let done_target = is_done(¤t_array); + let done_canonical = AnyCanonical::matches(¤t_array); + trace_op!(record_execute_until_done_check(done_target, done_canonical)); + + if done_target || done_canonical { match stack.pop() { None => { debug_assert!( current_builder.is_none(), "root activation should not retain a builder" ); - ctx.log(format_args!("-> {}", current_array)); + trace_op!(record_execute_until_return(¤t_array)); return Ok(current_array); } Some(frame) => { + let _slot_idx = frame.slot_idx; (current_array, current_builder) = pop_frame(frame, current_array)?; + trace_op!(record_execute_until_pop_frame(_slot_idx, ¤t_array)); continue; } } @@ -207,6 +227,7 @@ impl ArrayRef { && let Some(frame) = stack.last() && let Some(result) = { execute_parent_for_child( + "stack_execute_parent", &frame.parent_array, ¤t_array, frame.slot_idx, @@ -215,42 +236,48 @@ impl ArrayRef { )? } { - ctx.log(format_args!( - "execute_parent (stack) rewrote {} -> {}", - current_array, result - )); let frame = stack.pop().vortex_expect("just peeked"); - current_array = result.optimize_ctx(ctx.session())?; + let optimized = result.optimize_ctx(ctx.session())?; + trace_op!(record_execute_optimized(&result, &optimized)); + current_array = optimized; current_builder = frame.parent_builder; continue; } + if current_builder.is_none() && stack.last().is_some() { + trace_op!(record_execute_parent_none( + "stack_execute_parent", + ¤t_array, + )); + } // Step 2b: execute_parent against current_array's own children. if current_builder.is_none() && let Some(rewritten) = try_execute_parent(¤t_array, kernels, ctx)? { - ctx.log(format_args!( - "execute_parent rewrote {} -> {}", - current_array, rewritten - )); - current_array = rewritten.optimize_ctx(ctx.session())?; + let optimized = rewritten.optimize_ctx(ctx.session())?; + trace_op!(record_execute_optimized(&rewritten, &optimized)); + current_array = optimized; continue; } + if current_builder.is_none() { + trace_op!(record_execute_parent_none( + "child_execute_parent", + ¤t_array, + )); + } - // execute step let expected_len = current_array.len(); let expected_dtype = current_array.dtype().clone(); let stats = current_array.statistics().to_array_stats(); let encoding_id = current_array.encoding_id(); + trace_op!(record_execute_encoding(¤t_array)); let result = current_array.execute_encoding_unchecked(ctx)?; let (array, step) = result.into_parts(); match step { ExecutionStep::ExecuteSlot(i, done) => { let (parent, child) = unsafe { array.take_slot_unchecked(i) }?; - ctx.log(format_args!( - "ExecuteSlot({i}): pushing {}, focusing on {}", - parent, child - )); + + trace_op!(record_execute_slot(i, &parent, &child)); stack.push(StackFrame { parent_array: parent, parent_builder: current_builder.take(), @@ -264,17 +291,18 @@ impl ArrayRef { } ExecutionStep::AppendChild(i) => { if current_builder.is_none() { + trace_op!(record_builder_start(&array)); current_builder = Some(builder_with_capacity_in( - ctx.allocator(), array.dtype(), array.len(), + ctx.allocator(), )); } let (parent, child) = unsafe { array.take_slot_unchecked(i) }?; - ctx.log(format_args!( - "AppendChild({i}): appending {} into builder", - child - )); + + trace_op!(record_append_child(i, &parent, &child)); + trace_op!(record_builder_append(&child)); + // TODO(joe)[7674]: replace with a builder kernel registry so we don't // need to go through the VTable append_to_builder indirection. child.append_to_builder( @@ -286,7 +314,8 @@ impl ArrayRef { current_array = parent; } ExecutionStep::Done => { - ctx.log(format_args!("Done: {}", array)); + let had_builder = current_builder.is_some(); + trace_op!(record_execute_done(&array)); (current_array, current_builder) = finalize_done( array, current_builder, @@ -295,6 +324,9 @@ impl ArrayRef { stats, encoding_id, )?; + if had_builder { + trace_op!(record_builder_finish(¤t_array)); + } } } } @@ -319,6 +351,8 @@ struct StackFrame { #[derive(Debug, Clone)] pub struct ExecutionCtx { session: VortexSession, + // OnceLock avoids cloning the session allocator when a context does not allocate. + allocator: OnceLock, execute_parent_kernels: Arc, #[cfg(debug_assertions)] id: usize, @@ -336,6 +370,7 @@ impl ExecutionCtx { let execute_parent_kernels = session.kernels().execute_parent_snapshot(); Self { session, + allocator: OnceLock::new(), execute_parent_kernels, #[cfg(debug_assertions)] id: { @@ -352,9 +387,15 @@ impl ExecutionCtx { &self.session } - /// Get the session-scoped host allocator for this execution context. - pub fn allocator(&self) -> HostAllocatorRef { - self.session.allocator() + /// Get the allocator for this execution context. + pub fn allocator(&self) -> &BufferAllocatorRef { + self.allocator.get_or_init(|| self.session.allocator()) + } + + /// Set the allocator for this execution context. + pub fn with_allocator(mut self, allocator: BufferAllocatorRef) -> Self { + self.allocator = OnceLock::from(allocator); + self } /// Log an execution step at the current depth. @@ -419,40 +460,49 @@ impl Drop for ExecutionCtx { /// `AppendChild` is returned. impl Executable for ArrayRef { fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + trace_op!(record_single_step_start(&array)); + if let Some(canonical) = array.as_opt::() { - ctx.log(format_args!("-> canonical {}", array)); - return Ok(Canonical::from(canonical).into_array()); + let output = Canonical::from(canonical).into_array(); + trace_op!(record_single_step_applied("canonical", &array, &output)); + return Ok(output); } + trace_op!(record_single_step_phase_none("canonical", &array)); if let Some(reduced) = array.reduce()? { - ctx.log(format_args!("reduce: rewrote {} -> {}", array, reduced)); reduced.statistics().inherit_from(array.statistics()); + trace_op!(record_single_step_applied("reduce", &array, &reduced)); return Ok(reduced); } + trace_op!(record_single_step_phase_none("reduce", &array)); for (slot_idx, slot) in array.slots().iter().enumerate() { let Some(child) = slot else { continue }; if let Some(reduced_parent) = child.reduce_parent(&array, slot_idx)? { - ctx.log(format_args!( - "reduce_parent: slot[{}]({}) rewrote {} -> {}", - slot_idx, - child.encoding_id(), - array, - reduced_parent - )); reduced_parent.statistics().inherit_from(array.statistics()); + trace_op!(record_single_step_applied( + "reduce_parent", + &array, + &reduced_parent, + )); return Ok(reduced_parent); } } + trace_op!(record_single_step_phase_none("reduce_parent", &array)); let execute_parent_kernels = Arc::clone(&ctx.execute_parent_kernels); let kernels = execute_parent_kernels.as_ref(); for (slot_idx, slot) in array.slots().iter().enumerate() { let Some(child) = slot else { continue }; - if let Some(executed_parent) = - execute_parent_for_child(&array, child, slot_idx, kernels, ctx)? - { + if let Some(executed_parent) = execute_parent_for_child( + "single_step_execute_parent", + &array, + child, + slot_idx, + kernels, + ctx, + )? { ctx.log(format_args!( "execute_parent: slot[{}]({}) rewrote {} -> {}", slot_idx, @@ -463,16 +513,22 @@ impl Executable for ArrayRef { executed_parent .statistics() .inherit_from(array.statistics()); + trace_op!(record_single_step_applied( + "execute_parent", + &array, + &executed_parent, + )); return Ok(executed_parent); } } + trace_op!(record_single_step_phase_none("execute_parent", &array)); + trace_op!(record_execute_encoding(&array)); - ctx.log(format_args!("executing {}", array)); let result = array.execute_encoding(ctx)?; let (array, step) = result.into_parts(); match step { ExecutionStep::Done => { - ctx.log(format_args!("-> {}", array)); + trace_op!(record_execute_done(&array)); Ok(array) } ExecutionStep::ExecuteSlot(i, _) => { @@ -484,9 +540,12 @@ impl Executable for ArrayRef { } ExecutionStep::AppendChild(_) => { // Single-step: build the entire parent via the builder path. - let builder = builder_with_capacity_in(ctx.allocator(), array.dtype(), array.len()); + trace_op!(record_builder_start(&array)); + let builder = builder_with_capacity_in(array.dtype(), array.len(), ctx.allocator()); let mut builder = execute_into_builder(array, builder, ctx)?; - Ok(builder.finish()) + let output = builder.finish(); + trace_op!(record_builder_finish(&output)); + Ok(output) } } } @@ -562,6 +621,7 @@ fn finalize_done( } fn execute_parent_for_child( + _phase: &'static str, parent: &ArrayRef, child: &ArrayRef, slot_idx: usize, @@ -570,7 +630,8 @@ fn execute_parent_for_child( ) -> VortexResult> { let key = execute_parent_key(parent.encoding_id(), child.encoding_id()); if let Some(plugins) = kernels.get(&key) { - for plugin in plugins.as_ref() { + #[allow(clippy::unused_enumerate_index)] + for (_plugin_idx, plugin) in plugins.as_ref().iter().enumerate() { if let Some(result) = plugin.execute_parent(child, parent, slot_idx, ctx)? { if cfg!(debug_assertions) { vortex_ensure!( @@ -582,8 +643,23 @@ fn execute_parent_for_child( "Executed parent canonical dtype mismatch" ); } + trace_op!(record_session_execute_parent_applied( + _phase, + parent, + child, + slot_idx, + _plugin_idx, + &result, + )); return Ok(Some(result)); } + trace_op!(record_session_execute_parent_declined( + _phase, + parent, + child, + slot_idx, + _plugin_idx, + )); } } @@ -599,7 +675,7 @@ fn try_execute_parent( for (slot_idx, slot) in array.slots().iter().enumerate() { let Some(child) = slot else { continue }; if let Some(executed_parent) = - execute_parent_for_child(array, child, slot_idx, kernels, ctx)? + execute_parent_for_child("child_execute_parent", array, child, slot_idx, kernels, ctx)? { ctx.log(format_args!( "execute_parent: slot[{}]({}) rewrote {} -> {}", @@ -709,8 +785,9 @@ impl ExecutionResult { /// /// The provided array is the (possibly modified) parent that still needs its slot executed. pub fn execute_slot(array: impl IntoArray, slot_idx: usize) -> Self { + let array = array.into_array(); Self { - array: array.into_array(), + array, step: ExecutionStep::ExecuteSlot(slot_idx, M::matches), } } @@ -719,8 +796,9 @@ impl ExecutionResult { /// activation's canonical builder, and leave the returned parent as the next /// `current_array`. pub fn append_child(array: impl IntoArray, slot_idx: usize) -> Self { + let array = array.into_array(); Self { - array: array.into_array(), + array, step: ExecutionStep::AppendChild(slot_idx), } } @@ -796,7 +874,12 @@ macro_rules! require_opt_child { /// Like [`require_opt_child!`], `$parent` is moved (not cloned) into the early-return path. /// /// ```ignore -/// require_patches!(array, PATCH_INDICES_SLOT, PATCH_VALUES_SLOT, PATCH_CHUNK_OFFSETS_SLOT); +/// require_patches!( +/// array, +/// MySlots::PATCH_INDICES, +/// MySlots::PATCH_VALUES, +/// MySlots::PATCH_CHUNK_OFFSETS +/// ); /// ``` #[macro_export] macro_rules! require_patches { @@ -826,7 +909,7 @@ macro_rules! require_patches { /// Like [`require_opt_child!`], `$parent` is moved (not cloned) into the early-return path. /// /// ```ignore -/// require_validity!(array, VALIDITY_SLOT); +/// require_validity!(array, MySlots::VALIDITY); /// ``` #[macro_export] macro_rules! require_validity { @@ -853,6 +936,8 @@ impl VortexSessionExecute for VortexSession { #[cfg(test)] mod tests { + use static_assertions::assert_impl_all; + use vortex_session::SessionExt; use vortex_session::VortexSession; use super::*; @@ -860,10 +945,15 @@ mod tests { use crate::VortexSessionExecute; use crate::arrays::Bool; use crate::arrays::Primitive; + use crate::memory::BufferAllocatorRef; + use crate::memory::MemorySession; + use crate::memory::MemorySessionExt; use crate::optimizer::kernels::ExecuteParentFn; use crate::optimizer::kernels::KernelSession; use crate::optimizer::kernels::execute_parent_key; + assert_impl_all!(ExecutionCtx: Send, Sync); + fn noop_execute_parent( _child: &ArrayRef, _parent: &ArrayRef, @@ -901,4 +991,27 @@ mod tests { let after_registration = session.create_execution_ctx(); assert!(after_registration.execute_parent_kernels.contains_key(&key)); } + + #[test] + fn execution_ctx_allocator_override() { + let first = BufferAllocatorRef::new(vortex_buffer::StaticBufferAllocator); + let second = BufferAllocatorRef::new(vortex_buffer::StaticBufferAllocator); + let third = BufferAllocatorRef::new(vortex_buffer::StaticBufferAllocator); + let session = VortexSession::empty() + .with::() + .with_allocator(first.clone()); + let ctx = session.create_execution_ctx(); + + session + .get_mut::() + .set_allocator(third.clone()); + + assert!(session.allocator().ptr_eq(&third)); + assert!(ctx.allocator().ptr_eq(&third)); + + let ctx = ctx.with_allocator(second.clone()); + session.get_mut::().set_allocator(first); + + assert!(ctx.allocator().ptr_eq(&second)); + } } diff --git a/vortex-array/src/expr/analysis/annotation.rs b/vortex-array/src/expr/analysis/annotation.rs index 32608f38d1e..f2c63d7c4d7 100644 --- a/vortex-array/src/expr/analysis/annotation.rs +++ b/vortex-array/src/expr/analysis/annotation.rs @@ -8,7 +8,10 @@ use vortex_error::VortexResult; use vortex_utils::aliases::hash_map::HashMap; use vortex_utils::aliases::hash_set::HashSet; +use crate::expr::BoundExpression; +use crate::expr::ExactBoundExpr; use crate::expr::Expression; +use crate::expr::traversal::Node; use crate::expr::traversal::NodeExt; use crate::expr::traversal::NodeVisitor; use crate::expr::traversal::TraversalOrder; @@ -17,28 +20,38 @@ pub trait Annotation: Clone + Hash + Eq {} impl Annotation for A where A: Clone + Hash + Eq {} -pub trait AnnotationFn: Fn(&Expression) -> Vec { +pub trait AnnotationFn: Fn(&N) -> Vec { type Annotation: Annotation; } -impl AnnotationFn for F +impl AnnotationFn for F where A: Annotation, - F: Fn(&Expression) -> Vec, + F: Fn(&N) -> Vec, { type Annotation = A; } -pub type Annotations<'a, A> = HashMap<&'a Expression, HashSet>; +pub type Annotations<'a, A, N = Expression> = HashMap<&'a N, HashSet>; + +/// Annotations keyed by bound-tree identity. +/// +/// Identity keys avoid structurally hashing every node's dtype. That matters when a bound root +/// carries a lazy schema whose structural hash would deserialize every field. +pub type BoundAnnotations = HashMap>; /// Walk the expression tree and annotate each expression with zero or more annotations. /// /// Returns a map of each expression to all annotations that any of its descendent (child) /// expressions are annotated with. -pub fn descendent_annotations( - expr: &Expression, +pub fn descendent_annotations<'a, N, A>( + expr: &'a N, annotate: A, -) -> Annotations<'_, A::Annotation> { +) -> Annotations<'a, A::Annotation, N> +where + N: Node + Eq + Hash, + A: AnnotationFn, +{ let mut visitor = AnnotationVisitor { annotations: Default::default(), annotate, @@ -53,10 +66,11 @@ pub fn descendent_annotations( /// /// Returns a map of each expression to all annotations. Annotations of /// children are not propagated to parents. -pub fn direct_annotations( - expr: &Expression, - annotate: A, -) -> Annotations<'_, A::Annotation> { +pub fn direct_annotations<'a, N, A>(expr: &'a N, annotate: A) -> Annotations<'a, A::Annotation, N> +where + N: Node + Eq + Hash, + A: AnnotationFn, +{ let mut visitor = AnnotationVisitor { annotations: Default::default(), annotate, @@ -66,14 +80,66 @@ pub fn direct_annotations( visitor.annotations } -struct AnnotationVisitor<'a, A: AnnotationFn> { - annotations: Annotations<'a, A::Annotation>, +/// Annotate a bound expression and propagate each annotation to its ancestors. +/// +/// Unlike [`descendent_annotations`], this uses [`ExactBoundExpr`] keys to preserve the cheap +/// identity semantics of an already-bound tree. +pub fn descendent_bound_annotations( + expr: &BoundExpression, + annotate: A, +) -> BoundAnnotations +where + A: AnnotationFn, +{ + bound_annotations(expr, annotate, true) +} + +/// Annotate each bound-expression node without propagating annotations to its ancestors. +/// +/// The returned map uses [`ExactBoundExpr`] keys so lookups do not structurally hash node dtypes. +pub fn direct_bound_annotations( + expr: &BoundExpression, + annotate: A, +) -> BoundAnnotations +where + A: AnnotationFn, +{ + bound_annotations(expr, annotate, false) +} + +fn bound_annotations( + expr: &BoundExpression, + annotate: A, + propagate_up: bool, +) -> BoundAnnotations +where + A: AnnotationFn, +{ + let mut visitor = BoundAnnotationVisitor { + annotations: Default::default(), + annotate, + propagate_up, + }; + expr.accept(&mut visitor).vortex_expect("Infallible"); + visitor.annotations +} + +struct AnnotationVisitor<'a, N, A> +where + N: Node + Eq + Hash, + A: AnnotationFn, +{ + annotations: Annotations<'a, A::Annotation, N>, annotate: A, propagate_up: bool, } -impl<'a, A: AnnotationFn> NodeVisitor<'a> for AnnotationVisitor<'a, A> { - type NodeTy = Expression; +impl<'a, N, A> NodeVisitor<'a> for AnnotationVisitor<'a, N, A> +where + N: Node + Eq + Hash, + A: AnnotationFn, +{ + type NodeTy = N; fn visit_down(&mut self, node: &'a Self::NodeTy) -> VortexResult { let annotations = (self.annotate)(node); @@ -89,20 +155,74 @@ impl<'a, A: AnnotationFn> NodeVisitor<'a> for AnnotationVisitor<'a, A> { } } - fn visit_up(&mut self, node: &'a Expression) -> VortexResult { + fn visit_up(&mut self, node: &'a N) -> VortexResult { if !self.propagate_up { return Ok(TraversalOrder::Continue); } + let child_annotations = node.iter_children(|children| { + children + .filter_map(|child| self.annotations.get(child).cloned()) + .collect::>() + }); + + let annotations = self.annotations.entry(node).or_default(); + child_annotations + .into_iter() + .for_each(|ps| annotations.extend(ps.iter().cloned())); + + Ok(TraversalOrder::Continue) + } +} + +struct BoundAnnotationVisitor +where + A: AnnotationFn, +{ + annotations: BoundAnnotations, + annotate: A, + propagate_up: bool, +} + +impl<'a, A> NodeVisitor<'a> for BoundAnnotationVisitor +where + A: AnnotationFn, +{ + type NodeTy = BoundExpression; + + fn visit_down(&mut self, node: &'a Self::NodeTy) -> VortexResult { + let annotations = (self.annotate)(node); + if annotations.is_empty() { + return Ok(TraversalOrder::Continue); + } + + self.annotations + .entry(ExactBoundExpr(node.clone())) + .or_default() + .extend(annotations); + Ok(TraversalOrder::Skip) + } + + fn visit_up(&mut self, node: &'a Self::NodeTy) -> VortexResult { + if !self.propagate_up { + return Ok(TraversalOrder::Continue); + } + let child_annotations = node .children() .iter() - .filter_map(|c| self.annotations.get(c).cloned()) + .filter_map(|child| { + self.annotations + .get(&ExactBoundExpr(child.clone())) + .cloned() + }) .collect::>(); - - let annotations = self.annotations.entry(node).or_default(); + let annotations = self + .annotations + .entry(ExactBoundExpr(node.clone())) + .or_default(); child_annotations .into_iter() - .for_each(|ps| annotations.extend(ps.iter().cloned())); + .for_each(|child| annotations.extend(child)); Ok(TraversalOrder::Continue) } diff --git a/vortex-array/src/expr/analysis/immediate_access.rs b/vortex-array/src/expr/analysis/immediate_access.rs index 805400be793..6c2e4975a92 100644 --- a/vortex-array/src/expr/analysis/immediate_access.rs +++ b/vortex-array/src/expr/analysis/immediate_access.rs @@ -2,20 +2,15 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexExpect; -use vortex_utils::aliases::hash_set::HashSet; use crate::dtype::FieldName; use crate::dtype::StructFields; +use crate::expr::BoundExpression; use crate::expr::Expression; use crate::expr::analysis::AnnotationFn; -use crate::expr::analysis::Annotations; -use crate::expr::descendent_annotations; use crate::scalar_fn::fns::get_item::GetItem; -use crate::scalar_fn::fns::root::Root; use crate::scalar_fn::fns::select::Select; -pub type FieldAccesses<'a> = Annotations<'a, FieldName>; - /// Returns the "free fields" for this expression node. /// /// A "free field" is a top-level field from the root scope that this expression references—not @@ -30,9 +25,10 @@ pub type FieldAccesses<'a> = Annotations<'a, FieldName>; /// /// # Annotation Rules /// -/// - **[`Select`]**: Returns the included field names if the child is [`Root`]. -/// - **[`GetItem`] on [`Root`]**: Returns `[field_name]` if the child is [`Root`]. -/// - **[`Root`]**: Returns all field names from `scope` (conservative over-approximation). +/// - **[`Select`]**: Returns the included field names if the child is [`Expression::Root`]. +/// - **[`GetItem`] on the root**: Returns `[field_name]` if the child is [`Expression::Root`]. +/// - **[`Expression::Root`]**: Returns all field names from `scope` (conservative +/// over-approximation). /// - **Everything else**: Returns empty (annotations aggregate from children automatically). /// /// # Example @@ -43,10 +39,10 @@ pub type FieldAccesses<'a> = Annotations<'a, FieldName>; /// - The full expression has free fields `{a, d}` (not `b`, only top-level fields are tracked). pub fn make_free_field_annotator( scope: &StructFields, -) -> impl AnnotationFn { +) -> impl AnnotationFn { move |expr: &Expression| { if let Some(selection) = expr.as_opt::() { + if expr.children()[0].is_root() { + return selection + .normalize_to_included_fields(scope.names()) + .vortex_expect("Select fields must be valid for scope") + .into_iter() + .collect(); + } + } else if let Some(field_name) = scalar_fn.as_opt::() + && expr.children()[0].is_root() + { + return vec![field_name.clone()]; + } + + vec![] + } } diff --git a/vortex-array/src/expr/analysis/fallible.rs b/vortex-array/src/expr/analysis/infallible.rs similarity index 61% rename from vortex-array/src/expr/analysis/fallible.rs rename to vortex-array/src/expr/analysis/infallible.rs index b6d45bd5483..fbf84e124f0 100644 --- a/vortex-array/src/expr/analysis/fallible.rs +++ b/vortex-array/src/expr/analysis/infallible.rs @@ -5,11 +5,20 @@ use crate::expr::Expression; use crate::expr::analysis::BooleanLabels; use crate::expr::label_tree; -pub fn label_is_fallible(expr: &Expression) -> BooleanLabels<'_> { +/// Label each expression with whether its entire subtree is infallible. +/// +/// A subtree is infallible only when the node's scalar function and every child subtree are +/// infallible. See [`crate::scalar_fn::ScalarFnVTable::is_infallible`] for the scalar-function +/// contract. +pub fn label_infallible(expr: &Expression) -> BooleanLabels<'_> { label_tree( expr, - |expr| expr.signature().is_fallible(), - |acc, &child| acc | child, + |expr| match expr { + Expression::Scalar { scalar_fn, .. } => scalar_fn.signature().is_infallible(), + // The scope itself cannot fail. + Expression::Root => true, + }, + |acc, &child| acc & child, ) } @@ -26,55 +35,55 @@ mod tests { use crate::scalar_fn::fns::merge::DuplicateHandling; #[test] - fn not_is_not_fallible() { + fn not_is_infallible() { let expr = not(col("x")); - let labels = label_is_fallible(&expr); - assert_eq!(labels.get(&expr), Some(&false)); + let labels = label_infallible(&expr); + assert_eq!(labels.get(&expr), Some(&true)); } #[test] fn checked_add_defaults_to_fallible() { let expr = checked_add(col("a"), col("b")); - let labels = label_is_fallible(&expr); - assert_eq!(labels.get(&expr), Some(&true)); + let labels = label_infallible(&expr); + assert_eq!(labels.get(&expr), Some(&false)); } #[test] - fn eq_not_fallible() { + fn eq_is_infallible() { let expr = eq(col("a"), lit(5)); - let labels = label_is_fallible(&expr); - assert_eq!(labels.get(&expr), Some(&false)); + let labels = label_infallible(&expr); + assert_eq!(labels.get(&expr), Some(&true)); } #[test] fn merge_with_error_handling_is_fallible() { let expr = merge_opts([col("a"), col("b")], DuplicateHandling::Error); - let labels = label_is_fallible(&expr); - assert_eq!(labels.get(&expr), Some(&true)); + let labels = label_infallible(&expr); + assert_eq!(labels.get(&expr), Some(&false)); } #[test] - fn merge_with_rightmost_handling_is_not_fallible() { + fn merge_with_rightmost_handling_is_infallible() { let expr = merge_opts([col("a"), col("b")], DuplicateHandling::RightMost); - let labels = label_is_fallible(&expr); - assert_eq!(labels.get(&expr), Some(&false)); + let labels = label_infallible(&expr); + assert_eq!(labels.get(&expr), Some(&true)); } #[test] fn nested_with_fallible_child() { let child = checked_add(col("a"), col("b")); let expr = not(child.clone()); - let labels = label_is_fallible(&expr); - assert_eq!(labels.get(&child), Some(&true)); - assert_eq!(labels.get(&expr), Some(&true)); + let labels = label_infallible(&expr); + assert_eq!(labels.get(&child), Some(&false)); + assert_eq!(labels.get(&expr), Some(&false)); } #[test] fn nested_without_fallible_child() { let child = is_null(col("x")); let expr = not(child.clone()); - let labels = label_is_fallible(&expr); - assert_eq!(labels.get(&child), Some(&false)); - assert_eq!(labels.get(&expr), Some(&false)); + let labels = label_infallible(&expr); + assert_eq!(labels.get(&child), Some(&true)); + assert_eq!(labels.get(&expr), Some(&true)); } } diff --git a/vortex-array/src/expr/analysis/labeling.rs b/vortex-array/src/expr/analysis/labeling.rs index 721241a14c9..51957f26190 100644 --- a/vortex-array/src/expr/analysis/labeling.rs +++ b/vortex-array/src/expr/analysis/labeling.rs @@ -1,15 +1,26 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::hash::Hash; + use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_utils::aliases::hash_map::HashMap; +use crate::expr::BoundExpression; +use crate::expr::ExactBoundExpr; use crate::expr::Expression; +use crate::expr::traversal::Node; use crate::expr::traversal::NodeExt; use crate::expr::traversal::NodeVisitor; use crate::expr::traversal::TraversalOrder; +/// Boolean labels keyed by each expression node in a tree. +pub type BooleanLabels<'a, N = Expression> = HashMap<&'a N, bool>; + +/// Labels keyed by bound-tree identity. +pub type BoundLabels = HashMap; + /// Label each node in an expression tree using a bottom-up traversal. /// /// This function separates tree labeling into two distinct steps: @@ -29,11 +40,14 @@ use crate::expr::traversal::TraversalOrder; /// - `merge_child`: Mutable function that folds child labels into an accumulator. /// Takes `(self_label, child_label)` and returns the updated accumulator. /// Called once per child, with the initial accumulator being the node's self-label. -pub fn label_tree( - expr: &Expression, - self_label: impl Fn(&Expression) -> L, +pub fn label_tree( + expr: &N, + self_label: impl Fn(&N) -> L, mut merge_child: impl FnMut(L, &L) -> L, -) -> HashMap<&Expression, L> { +) -> HashMap<&N, L> +where + N: Node + Eq + Hash, +{ let mut visitor = LabelingVisitor { labels: Default::default(), self_label, @@ -44,40 +58,98 @@ pub fn label_tree( visitor.labels } -struct LabelingVisitor<'a, 'b, L, F, G> +/// Label each node in a bound expression using identity-keyed lookups. +/// +/// This avoids structurally hashing bound dtypes, which may deserialize a lazy schema. +pub fn label_bound_tree( + expr: &BoundExpression, + self_label: impl Fn(&BoundExpression) -> L, + mut merge_child: impl FnMut(L, &L) -> L, +) -> BoundLabels { + let mut visitor = BoundLabelingVisitor { + labels: Default::default(), + self_label, + merge_child: &mut merge_child, + }; + expr.accept(&mut visitor) + .vortex_expect("BoundLabelingVisitor is infallible"); + visitor.labels +} + +struct LabelingVisitor<'a, 'b, N, L, F, G> where - F: Fn(&Expression) -> L, + N: Node + Eq + Hash, + F: Fn(&N) -> L, G: FnMut(L, &L) -> L, { - labels: HashMap<&'a Expression, L>, + labels: HashMap<&'a N, L>, self_label: F, merge_child: &'b mut G, } -impl<'a, 'b, L: Clone, F, G> NodeVisitor<'a> for LabelingVisitor<'a, 'b, L, F, G> +impl<'a, 'b, N, L: Clone, F, G> NodeVisitor<'a> for LabelingVisitor<'a, 'b, N, L, F, G> where - F: Fn(&Expression) -> L, + N: Node + Eq + Hash, + F: Fn(&N) -> L, G: FnMut(L, &L) -> L, { - type NodeTy = Expression; + type NodeTy = N; fn visit_down(&mut self, _node: &'a Self::NodeTy) -> VortexResult { Ok(TraversalOrder::Continue) } - fn visit_up(&mut self, node: &'a Expression) -> VortexResult { + fn visit_up(&mut self, node: &'a N) -> VortexResult { let self_label = (self.self_label)(node); + let final_label = node.iter_children(|children| { + children.fold(self_label, |acc, child| { + let child_label = self + .labels + .get(child) + .vortex_expect("child must have label"); + (self.merge_child)(acc, child_label) + }) + }); + + self.labels.insert(node, final_label); + + Ok(TraversalOrder::Continue) + } +} + +struct BoundLabelingVisitor<'a, L, F, G> +where + F: Fn(&BoundExpression) -> L, + G: FnMut(L, &L) -> L, +{ + labels: BoundLabels, + self_label: F, + merge_child: &'a mut G, +} + +impl<'node, 'visitor, L: Clone, F, G> NodeVisitor<'node> for BoundLabelingVisitor<'visitor, L, F, G> +where + F: Fn(&BoundExpression) -> L, + G: FnMut(L, &L) -> L, +{ + type NodeTy = BoundExpression; + + fn visit_down(&mut self, _node: &'node Self::NodeTy) -> VortexResult { + Ok(TraversalOrder::Continue) + } + + fn visit_up(&mut self, node: &'node Self::NodeTy) -> VortexResult { + let self_label = (self.self_label)(node); let final_label = node.children().iter().fold(self_label, |acc, child| { let child_label = self .labels - .get(child) + .get(&ExactBoundExpr(child.clone())) .vortex_expect("child must have label"); (self.merge_child)(acc, child_label) }); - - self.labels.insert(node, final_label); - + self.labels + .insert(ExactBoundExpr(node.clone()), final_label); Ok(TraversalOrder::Continue) } } diff --git a/vortex-array/src/expr/analysis/mod.rs b/vortex-array/src/expr/analysis/mod.rs index f5208a31be8..d47c909761d 100644 --- a/vortex-array/src/expr/analysis/mod.rs +++ b/vortex-array/src/expr/analysis/mod.rs @@ -2,16 +2,15 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors pub mod annotation; -mod fallible; pub mod immediate_access; +mod infallible; mod labeling; -mod null_sensitive; mod referenced_field_paths; +mod strict; pub use annotation::*; -pub use fallible::label_is_fallible; pub use immediate_access::*; +pub use infallible::label_infallible; pub use labeling::*; -pub use null_sensitive::BooleanLabels; -pub use null_sensitive::label_null_sensitive; pub use referenced_field_paths::referenced_field_paths; +pub use strict::label_strict; diff --git a/vortex-array/src/expr/analysis/null_sensitive.rs b/vortex-array/src/expr/analysis/null_sensitive.rs deleted file mode 100644 index 45944dc1bcb..00000000000 --- a/vortex-array/src/expr/analysis/null_sensitive.rs +++ /dev/null @@ -1,62 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use vortex_utils::aliases::hash_map::HashMap; - -use super::labeling::label_tree; -use crate::expr::Expression; - -pub type BooleanLabels<'a> = HashMap<&'a Expression, bool>; - -/// Label each expression in the tree with whether it is null-sensitive. -/// -/// See [`crate::scalar_fn::ScalarFnVTable::is_null_sensitive`] for a definition of null sensitivity. -/// This function operates on a tree of expressions, not just a single expression. -pub fn label_null_sensitive(expr: &Expression) -> BooleanLabels<'_> { - label_tree( - expr, - |expr| expr.signature().is_null_sensitive(), - |acc, &child| acc | child, - ) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::expr::col; - use crate::expr::eq; - use crate::expr::is_null; - use crate::expr::lit; - - #[test] - fn test_null_sensitive_with_is_null() { - let expr = is_null(col("col1")); - let labels = label_null_sensitive(&expr); - - // The root expression should be null-sensitive - assert_eq!(labels.get(&expr), Some(&true)); - } - - #[test] - fn test_null_sensitive_without_is_null() { - let expr = eq(col("col1"), lit(5)); - let labels = label_null_sensitive(&expr); - - // Since the default is conservative (true), all expressions are sensitive - assert_eq!(labels.get(&expr), Some(&true)); - } - - #[test] - fn test_null_sensitive_nested() { - let left = eq(col("col1"), lit(5)); - let right = is_null(col("col2")); - let expr = eq(left.clone(), right.clone()); - - let labels = label_null_sensitive(&expr); - - // With conservative defaults, all are sensitive - assert_eq!(labels.get(&left), Some(&true)); - assert_eq!(labels.get(&right), Some(&true)); - assert_eq!(labels.get(&expr), Some(&true)); - } -} diff --git a/vortex-array/src/expr/analysis/referenced_field_paths.rs b/vortex-array/src/expr/analysis/referenced_field_paths.rs index 5239e4a691c..48ef285c20a 100644 --- a/vortex-array/src/expr/analysis/referenced_field_paths.rs +++ b/vortex-array/src/expr/analysis/referenced_field_paths.rs @@ -4,17 +4,15 @@ use vortex_error::VortexResult; use vortex_error::vortex_err; -use crate::dtype::DType; use crate::dtype::Field; use crate::dtype::FieldPath; use crate::dtype::FieldPathSet; -use crate::expr::Expression; +use crate::expr::BoundExpression; use crate::expr::traversal::FoldDownContext; use crate::expr::traversal::FoldUp; use crate::expr::traversal::NodeExt; use crate::expr::traversal::NodeFolderContext; use crate::scalar_fn::fns::get_item::GetItem; -use crate::scalar_fn::fns::root::Root; use crate::scalar_fn::fns::select::Select; /// Returns the rooted field paths referenced by an expression. @@ -24,47 +22,13 @@ use crate::scalar_fn::fns::select::Select; /// expression is represented by [`FieldPath::root`], which conservatively selects all fields. /// Scalar functions other than `GetItem` and `Select` conservatively reference each complete child /// output. -pub fn referenced_field_paths(expr: &Expression, scope: &DType) -> VortexResult { - // Validate the whole expression so plain GetItem paths and Select paths behave consistently. - expr.return_dtype(scope)?; - +pub fn referenced_field_paths(expr: &BoundExpression) -> VortexResult { let mut collector = ReferencedFieldPaths { - scope, field_paths: FieldPathSet::default(), }; expr.clone() .fold_context(&vec![FieldPath::root()], &mut collector)?; - let field_paths = collector.field_paths; - - // The top-level field of every referenced path must be one of the immediately accessed scope - // fields: this analysis only refines *which nested fields* are read, never which top-level - // fields. `FieldPath::root()` stands in for "all fields", so it expands to the whole scope. - #[cfg(debug_assertions)] - if let Some(scope_fields) = scope.as_struct_fields_opt() { - use vortex_utils::aliases::hash_set::HashSet; - - use crate::dtype::FieldName; - use crate::expr::analysis::immediate_access::immediate_scope_access; - - let referenced_heads: HashSet = if field_paths.iter().any(FieldPath::is_root) { - scope_fields.names().iter().cloned().collect() - } else { - field_paths - .iter() - .filter_map(|path| match path.parts().first() { - Some(Field::Name(name)) => Some(name.clone()), - _ => None, - }) - .collect() - }; - debug_assert_eq!( - referenced_heads, - immediate_scope_access(expr, scope_fields), - "referenced field path heads must match the immediately accessed scope fields" - ); - } - - Ok(field_paths) + Ok(collector.field_paths) } /// Threads the set of currently-requested field paths down the expression tree, narrowing it at @@ -78,22 +42,21 @@ pub fn referenced_field_paths(expr: &Expression, scope: &DType) -> VortexResult< /// column projection). Any other function is opaque—we cannot assume it preserves a field's /// provenance—so its children conservatively re-request the whole scope, which is what keeps an /// expression like `f($).x` reading every field of `$` rather than just `x`. -struct ReferencedFieldPaths<'a> { - scope: &'a DType, +struct ReferencedFieldPaths { field_paths: FieldPathSet, } -impl NodeFolderContext for ReferencedFieldPaths<'_> { - type NodeTy = Expression; +impl NodeFolderContext for ReferencedFieldPaths { + type NodeTy = BoundExpression; type Result = (); type Context = Vec; fn visit_down( &mut self, requested: &Self::Context, - node: &Expression, + node: &BoundExpression, ) -> VortexResult> { - if node.is::() { + if node.is_root() { self.field_paths.extend( requested .iter() @@ -102,7 +65,10 @@ impl NodeFolderContext for ReferencedFieldPaths<'_> { return Ok(FoldDownContext::Skip(())); } - if let Some(field_name) = node.as_opt::() { + if let Some(field_name) = node + .as_scalar() + .and_then(|scalar_fn| scalar_fn.as_opt::()) + { let appended = requested .iter() .map(|path| path.clone().push(Field::Name(field_name.clone()))) @@ -112,9 +78,12 @@ impl NodeFolderContext for ReferencedFieldPaths<'_> { // Keep requested paths whose head is included, expanding a whole-scope request into one // path per included field. - if let Some(selection) = node.as_opt::()) + { + let child_fields = node.children()[0] + .dtype() .as_struct_fields_opt() .ok_or_else(|| vortex_err!("Select child is not a struct"))?; let included_fields = selection.normalize_to_included_fields(child_fields.names())?; @@ -146,7 +115,7 @@ impl NodeFolderContext for ReferencedFieldPaths<'_> { fn visit_up( &mut self, - _node: Expression, + _node: BoundExpression, _requested: &Self::Context, _children: Vec<()>, ) -> VortexResult> { @@ -159,9 +128,11 @@ mod tests { use vortex_utils::aliases::hash_set::HashSet; use super::*; + use crate::dtype::DType; use crate::dtype::Nullability::NonNullable; use crate::dtype::PType::I32; use crate::dtype::StructFields; + use crate::expr::Expression; use crate::expr::get_item; use crate::expr::pack; use crate::expr::root; @@ -183,7 +154,7 @@ mod tests { /// Collects the prefix-minimal field paths referenced by `expr` against [`scope`]. fn referenced(expr: &Expression) -> VortexResult> { - Ok(referenced_field_paths(expr, &scope())? + Ok(referenced_field_paths(&expr.bind(&scope())?)? .into_iter() .collect()) } @@ -259,6 +230,9 @@ mod tests { #[test] fn invalid_get_item_path_returns_error() { - assert!(referenced_field_paths(&get_item("missing", root()), &scope()).is_err()); + let result = get_item("missing", root()) + .bind(&scope()) + .and_then(|expr| referenced_field_paths(&expr)); + assert!(result.is_err()); } } diff --git a/vortex-array/src/expr/analysis/strict.rs b/vortex-array/src/expr/analysis/strict.rs new file mode 100644 index 00000000000..01cc71d91c3 --- /dev/null +++ b/vortex-array/src/expr/analysis/strict.rs @@ -0,0 +1,60 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use super::BooleanLabels; +use super::labeling::label_tree; +use crate::expr::Expression; + +/// Label each expression with whether its entire subtree is strict. +/// +/// A subtree is strict only when the node's scalar function and every child subtree are strict. +/// See [`crate::scalar_fn::ScalarFnVTable::is_strict`] for the scalar-function contract. +pub fn label_strict(expr: &Expression) -> BooleanLabels<'_> { + label_tree( + expr, + |expr| match expr { + Expression::Scalar { scalar_fn, .. } => scalar_fn.signature().is_strict(), + // Vacuously strict. + Expression::Root => true, + }, + |acc, &child| acc & child, + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::expr::col; + use crate::expr::eq; + use crate::expr::is_null; + use crate::expr::lit; + + #[test] + fn test_non_strict_with_is_null() { + let expr = is_null(col("col1")); + let labels = label_strict(&expr); + + assert_eq!(labels.get(&expr), Some(&false)); + } + + #[test] + fn test_strict_expression() { + let expr = eq(lit(4), lit(5)); + let labels = label_strict(&expr); + + assert_eq!(labels.get(&expr), Some(&true)); + } + + #[test] + fn test_non_strict_child_makes_parent_subtree_non_strict() { + let left = eq(lit(4), lit(5)); + let right = is_null(col("col2")); + let expr = eq(left.clone(), right.clone()); + + let labels = label_strict(&expr); + + assert_eq!(labels.get(&left), Some(&true)); + assert_eq!(labels.get(&right), Some(&false)); + assert_eq!(labels.get(&expr), Some(&false)); + } +} diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs new file mode 100644 index 00000000000..4bd276e191d --- /dev/null +++ b/vortex-array/src/expr/bound_expression.rs @@ -0,0 +1,449 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt; +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hash; +use std::hash::Hasher; +use std::sync::Arc; + +use itertools::Itertools; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_session::VortexSession; + +use crate::dtype::DType; +use crate::expr::Expression; +use crate::expr::display::DisplayTreeExpr; +use crate::expr::scope::Scope; +use crate::expr::traversal::TraversalOrder; +use crate::expr::traversal::pre_order_visit_down; +use crate::scalar_fn::ScalarFnRef; +use crate::scalar_fn::ScalarFnVTable; +use crate::stats::rewrite::StatsRewriteCtx; + +/// An [`Expression`] that has been type-checked against a [`Scope`]. +/// +/// Every node carries its own dtype, so reading one is a field access rather than a walk of the +/// subtree. Holding a `BoundExpression` is proof that the whole tree type-checked. +/// +/// Binding is purely logical: it deals only in [`DType`]s and never sees an array, a length, or an +/// encoding. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub enum BoundExpression { + /// A scalar function applied to bound children. + Scalar { + /// The dtype this node evaluates to. + dtype: DType, + /// The scalar function for this node. + scalar_fn: ScalarFnRef, + /// The bound children, in argument order. + /// + /// Sharing keeps clones cheap even though the iterative [`Drop`] implementation prevents + /// consumers from destructuring a `BoundExpression` by value. + children: Arc>, + }, + /// The scope itself. Its dtype is the scope's root dtype. + Root { + /// The dtype this node evaluates to. + dtype: DType, + }, +} + +/// A bound-expression wrapper that compares shared tree identity instead of structure. +#[derive(Clone, Debug)] +pub struct ExactBoundExpr(pub BoundExpression); + +impl PartialEq for ExactBoundExpr { + fn eq(&self, other: &Self) -> bool { + match (&self.0, &other.0) { + ( + BoundExpression::Root { dtype: lhs_dtype }, + BoundExpression::Root { dtype: rhs_dtype }, + ) => lhs_dtype == rhs_dtype, + ( + BoundExpression::Scalar { + dtype: lhs_dtype, + scalar_fn: lhs_fn, + children: lhs_children, + }, + BoundExpression::Scalar { + dtype: rhs_dtype, + scalar_fn: rhs_fn, + children: rhs_children, + }, + ) => { + lhs_fn == rhs_fn + && Arc::ptr_eq(lhs_children, rhs_children) + && lhs_dtype == rhs_dtype + } + _ => false, + } + } +} + +impl Eq for ExactBoundExpr {} + +impl Hash for ExactBoundExpr { + fn hash(&self, state: &mut H) { + // DType differences are resolved by equality. Omitting the potentially lazy dtype keeps + // identity-keyed cache lookups from deserializing an entire schema just to compute a hash. + match &self.0 { + BoundExpression::Root { .. } => state.write_u8(0), + BoundExpression::Scalar { + scalar_fn, + children, + .. + } => { + state.write_u8(1); + scalar_fn.hash(state); + Arc::as_ptr(children).hash(state); + } + } + } +} + +impl BoundExpression { + /// Create a bound root expression with the given dtype. + pub fn new_root(dtype: DType) -> Self { + Self::Root { dtype } + } + + /// Create a bound scalar node from a scalar function and already-bound children. + pub fn try_new( + scalar_fn: ScalarFnRef, + children: impl IntoIterator, + ) -> VortexResult { + Self::try_new_vec(scalar_fn, children.into_iter().collect()) + } + + fn try_new_vec(scalar_fn: ScalarFnRef, children: Vec) -> VortexResult { + vortex_ensure!( + scalar_fn.signature().arity().matches(children.len()), + "Expression arity mismatch: expected {} children but got {}", + scalar_fn.signature().arity(), + children.len() + ); + + let arg_dtypes = children + .iter() + .map(|child| child.dtype().clone()) + .collect_vec(); + let dtype = scalar_fn.return_dtype(&arg_dtypes)?; + + Ok(Self::Scalar { + dtype, + scalar_fn, + children: children.into(), + }) + } + + /// Rebuild this node with new bound children, recomputing its dtype. + pub fn with_children( + self, + children: impl IntoIterator, + ) -> VortexResult { + let children = Vec::from_iter(children); + let BoundExpression::Scalar { scalar_fn, .. } = &self else { + vortex_ensure!( + children.is_empty(), + "Root expression cannot have {} children", + children.len() + ); + return Ok(self); + }; + + Self::try_new_vec(scalar_fn.clone(), children) + } + + /// The dtype this expression evaluates to. + pub fn dtype(&self) -> &DType { + match self { + Self::Scalar { dtype, .. } | Self::Root { dtype } => dtype, + } + } + + /// The bound children of this node, in argument order. Empty for [`BoundExpression::Root`]. + pub fn children(&self) -> &[BoundExpression] { + match self { + Self::Scalar { children, .. } => children.as_slice(), + Self::Root { .. } => &[], + } + } + + /// Return the child at `index`. + pub fn child(&self, index: usize) -> &BoundExpression { + &self.children()[index] + } + + /// The scalar function for this node, or `None` if it is the scope root. + pub fn as_scalar(&self) -> Option<&ScalarFnRef> { + match self { + Self::Scalar { scalar_fn, .. } => Some(scalar_fn), + Self::Root { .. } => None, + } + } + + /// Return whether this node uses the given scalar-function vtable. + pub fn is(&self) -> bool { + self.as_scalar().is_some_and(ScalarFnRef::is::) + } + + /// Return whether this expression tree contains a node using the given scalar-function vtable. + pub fn contains(&self) -> VortexResult { + let mut contains = false; + pre_order_visit_down(self, |node| { + if node.is::() { + contains = true; + return Ok(TraversalOrder::Stop); + } + Ok(TraversalOrder::Continue) + })?; + Ok(contains) + } + + /// Return the typed scalar-function options when this node uses the given vtable. + pub fn as_opt(&self) -> Option<&V::Options> { + self.as_scalar().and_then(ScalarFnRef::as_opt::) + } + + /// Return the typed scalar-function options for this node. + /// + /// # Panics + /// + /// Panics when this node is the scope root or uses a different scalar-function vtable. + pub fn as_(&self) -> &V::Options { + self.as_opt::() + .vortex_expect("Bound expression options type mismatch") + } + + /// Whether this node is the scope root. + pub fn is_root(&self) -> bool { + matches!(self, Self::Root { .. }) + } + + /// Return whether every scope root in this expression has `dtype`. + /// + /// Expressions without a scope root, such as literals, match every dtype. + pub fn is_root_bound_to(&self, dtype: &DType) -> bool { + let mut is_bound_to = true; + pre_order_visit_down(self, |node| { + if node.is_root() && node.dtype() != dtype { + is_bound_to = false; + return Ok(TraversalOrder::Stop); + } + Ok(TraversalOrder::Continue) + }) + .vortex_expect("bound expression traversal cannot not fail"); + is_bound_to + } + + /// Return an expression that proves this predicate is definitely false from statistics. + pub fn falsify(&self, session: &VortexSession) -> VortexResult> { + StatsRewriteCtx::new(session).falsify(self) + } + + /// Return an expression that proves this predicate is definitely true from statistics. + pub fn satisfy(&self, session: &VortexSession) -> VortexResult> { + StatsRewriteCtx::new(session).satisfy(self) + } + + /// Display the bound expression as a formatted tree structure. + pub fn display_tree(&self) -> impl Display { + DisplayTreeExpr(self) + } +} + +impl Display for BoundExpression { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self { + Self::Scalar { scalar_fn, .. } => scalar_fn.fmt_sql(self, f), + Self::Root { .. } => f.write_str("$"), + } + } +} + +impl Expression { + /// Bind this expression against a root dtype, type-checking every node in a single walk. + /// + /// The returned tree carries a dtype on each node, so callers needing types at more than one + /// node should bind once and read fields rather than calling + /// [`return_dtype`](Expression::return_dtype) repeatedly. + pub fn bind(&self, dtype: &DType) -> VortexResult { + self.bind_scope(&Scope::new(dtype.clone())) + } + + /// Bind this expression against an explicit [`Scope`]. + pub fn bind_scope(&self, scope: &Scope) -> VortexResult { + if self.is_root() { + return Ok(BoundExpression::new_root(scope.root().clone())); + } + + let children: Vec<_> = self + .children() + .iter() + .map(|child| child.bind_scope(scope)) + .try_collect()?; + let scalar_fn = self + .as_scalar() + .vortex_expect("root was handled above, so this is a scalar node"); + BoundExpression::try_new(scalar_fn.clone(), children) + } +} + +/// Iterative drop to avoid stack overflows on deep trees. +impl Drop for BoundExpression { + fn drop(&mut self) { + let Self::Scalar { children, .. } = self else { + return; + }; + let Some(children) = Arc::get_mut(children) else { + return; + }; + + let mut to_drop = std::mem::take(children); + while let Some(mut child) = to_drop.pop() { + if let BoundExpression::Scalar { children, .. } = &mut child + && let Some(grandchildren) = Arc::get_mut(children) + { + to_drop.append(grandchildren); + } + } + } +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use super::*; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::expr::col; + use crate::expr::eq; + use crate::expr::lit; + use crate::expr::root; + use crate::expr::test_harness::struct_dtype; + use crate::scalar_fn::fns::literal::Literal; + + fn scope() -> Scope { + Scope::new(struct_dtype()) + } + + #[test] + fn root_binds_to_the_scope() -> VortexResult<()> { + let bound = root().bind_scope(&scope())?; + assert!(bound.is_root()); + assert_eq!(bound.dtype(), &struct_dtype()); + assert_eq!(bound, BoundExpression::new_root(struct_dtype())); + Ok(()) + } + + #[test] + fn every_node_carries_its_dtype() -> VortexResult<()> { + let expr = eq(col("a"), lit(1_i32)); + let bound = expr.bind_scope(&scope())?; + + assert_eq!(bound.dtype(), &DType::Bool(Nullability::NonNullable)); + + let lhs = &bound.children()[0]; + assert_eq!( + lhs.dtype(), + &DType::Primitive(PType::I32, Nullability::NonNullable) + ); + assert_eq!(lhs.children()[0].dtype(), &struct_dtype()); + Ok(()) + } + + #[test] + fn bind_agrees_with_return_dtype() -> VortexResult<()> { + for expr in [root(), col("a"), eq(col("a"), lit(1_i32)), lit(true)] { + assert_eq!( + expr.bind(&struct_dtype())?.dtype(), + &expr.return_dtype(&struct_dtype())?, + "disagreement for {expr}" + ); + } + Ok(()) + } + + #[test] + fn contains_scalar_function() -> VortexResult<()> { + let bound = eq(col("a"), lit(1_i32)).bind_scope(&scope())?; + assert!(bound.contains::()?); + assert!(!root().bind_scope(&scope())?.contains::()?); + Ok(()) + } + + #[test] + fn bound_to_checks_every_root() -> VortexResult<()> { + let dtype = struct_dtype(); + let bound = eq(col("a"), col("a")).bind(&dtype)?; + assert!(bound.is_root_bound_to(&dtype)); + assert!(!bound.is_root_bound_to(&DType::Bool(Nullability::NonNullable))); + assert!( + lit(true) + .bind(&dtype)? + .is_root_bound_to(&DType::Bool(Nullability::NonNullable)) + ); + Ok(()) + } + + #[test] + fn bound_display_matches_unbound() -> VortexResult<()> { + for expr in [root(), col("a"), eq(col("a"), lit(1_i32)), lit(true)] { + let bound = expr.bind_scope(&scope())?; + assert_eq!(bound.to_string(), expr.to_string()); + assert_eq!( + bound.display_tree().to_string(), + expr.display_tree().to_string() + ); + } + Ok(()) + } + + #[test] + fn clone_shares_children() -> VortexResult<()> { + let bound = eq(col("a"), lit(1_i32)).bind_scope(&scope())?; + let cloned = bound.clone(); + + let ( + BoundExpression::Scalar { children: a, .. }, + BoundExpression::Scalar { children: b, .. }, + ) = (&bound, &cloned) + else { + unreachable!("eq is a scalar node") + }; + assert!(Arc::ptr_eq(a, b)); + Ok(()) + } + + #[test] + fn repeated_subtree_is_bound_per_occurrence() -> VortexResult<()> { + let shared = col("a"); + let bound = eq(shared.clone(), shared).bind_scope(&scope())?; + let children = bound.children(); + assert_eq!(children[0].dtype(), children[1].dtype()); + Ok(()) + } + + #[test] + fn structural_and_exact_equality_are_distinct() -> VortexResult<()> { + let expr = eq(col("a"), lit(1_i32)); + let bound = expr.bind_scope(&scope())?; + let independently_bound = expr.bind_scope(&scope())?; + + assert_eq!(bound, independently_bound); + assert_eq!(ExactBoundExpr(bound.clone()), ExactBoundExpr(bound.clone())); + assert_ne!(ExactBoundExpr(bound), ExactBoundExpr(independently_bound)); + Ok(()) + } + + #[test] + fn binding_reports_a_type_error() { + let expr = eq(col("a"), lit("nope")); + assert!(expr.bind_scope(&scope()).is_err()); + } +} diff --git a/vortex-array/src/expr/display.rs b/vortex-array/src/expr/display.rs index 492df401b9f..250d7063834 100644 --- a/vortex-array/src/expr/display.rs +++ b/vortex-array/src/expr/display.rs @@ -1,47 +1,137 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::fmt; use std::fmt::Display; use std::fmt::Formatter; -use std::ops::Deref; +use vortex_utils::tree::TreeDisplayAdapter; +use vortex_utils::tree::write_branch_tree; + +use crate::expr::BoundExpression; use crate::expr::Expression; -use crate::scalar_fn::ScalarFnRef; +use crate::scalar_fn::ChildName; pub enum DisplayFormat { Compact, Tree, } -pub struct DisplayTreeExpr<'a>(pub &'a Expression); - -impl Display for DisplayTreeExpr<'_> { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - pub use termtree::Tree; - fn make_tree(expr: &Expression) -> Result, std::fmt::Error> { - let scalar_fn: &ScalarFnRef = expr.deref(); - let node_name = format!("{}", scalar_fn); - - // Get child names for display purposes - let child_names = (0..expr.children().len()).map(|i| expr.signature().child_name(i)); - let children = expr.children(); - - let child_trees: Result>, std::fmt::Error> = children - .iter() - .zip(child_names) - .map(|(child, name)| { - let child_tree = make_tree(child)?; - Ok::, std::fmt::Error>( - Tree::new(format!("{}: {}", name, child_tree.root)) - .with_leaves(child_tree.leaves), - ) - }) - .collect(); - - Ok(Tree::new(node_name).with_leaves(child_trees?)) +/// Read-only expression-tree interface used by scalar functions for SQL-style formatting. +/// +/// Both [`Expression`] and [`BoundExpression`] implement this interface, allowing scalar +/// functions to format either representation without converting between them. +pub trait ExprDisplay: Display { + /// Return the child at `index`. + fn display_child(&self, index: usize) -> &dyn ExprDisplay; + + /// Return the number of children in this node. + fn display_children_count(&self) -> usize; +} + +impl ExprDisplay for Expression { + fn display_child(&self, index: usize) -> &dyn ExprDisplay { + Expression::child(self, index) + } + + fn display_children_count(&self) -> usize { + self.children().len() + } +} + +impl ExprDisplay for BoundExpression { + fn display_child(&self, index: usize) -> &dyn ExprDisplay { + &self.children()[index] + } + + fn display_children_count(&self) -> usize { + self.children().len() + } +} + +trait DisplayTreeNode: Sized { + fn tree_children(&self) -> &[Self]; + + fn tree_child_name(&self, index: usize) -> ChildName; + + fn fmt_tree_node(&self, f: &mut Formatter<'_>) -> fmt::Result; +} + +/// Tree-display label for the scope root. +const ROOT_DISPLAY: &str = "vortex.root()"; + +impl DisplayTreeNode for Expression { + fn tree_children(&self) -> &[Self] { + Expression::children(self) + } + + fn tree_child_name(&self, index: usize) -> ChildName { + match self { + Expression::Scalar { scalar_fn, .. } => scalar_fn.signature().child_name(index), + Expression::Root => unreachable!("the scope root has no children"), } + } + + fn fmt_tree_node(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self { + Expression::Scalar { scalar_fn, .. } => Display::fmt(scalar_fn, f), + Expression::Root => write!(f, "{ROOT_DISPLAY}"), + } + } +} + +impl DisplayTreeNode for BoundExpression { + fn tree_children(&self) -> &[Self] { + BoundExpression::children(self) + } + + fn tree_child_name(&self, index: usize) -> ChildName { + match self { + BoundExpression::Scalar { scalar_fn, .. } => scalar_fn.signature().child_name(index), + BoundExpression::Root { .. } => unreachable!("the scope root has no children"), + } + } + + fn fmt_tree_node(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self { + BoundExpression::Scalar { scalar_fn, .. } => Display::fmt(scalar_fn, f), + BoundExpression::Root { .. } => write!(f, "{ROOT_DISPLAY}"), + } + } +} + +pub struct DisplayTreeExpr<'a, T: ?Sized = Expression>(pub &'a T); + +impl TreeDisplayAdapter for DisplayTreeExpr<'_, T> { + type Context = (); + type Node = T; + + fn write_node( + &self, + node: &Self::Node, + _context: &Self::Context, + formatter: &mut Formatter<'_>, + ) -> fmt::Result { + node.fmt_tree_node(formatter) + } + + fn visit_children( + &self, + node: &Self::Node, + visit: &mut dyn FnMut(&str, &Self::Node, bool) -> fmt::Result, + ) -> fmt::Result { + let children = node.tree_children(); + for (index, child) in children.iter().enumerate() { + let child_name = node.tree_child_name(index); + visit(child_name.as_ref(), child, index + 1 == children.len())?; + } + Ok(()) + } +} - write!(f, "{}", make_tree(self.0)?) +impl Display for DisplayTreeExpr<'_, T> { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write_branch_tree(self, self.0, &mut (), f) } } diff --git a/vortex-array/src/expr/expression.rs b/vortex-array/src/expr/expression.rs index 41cd61a369d..3ae9adb72d5 100644 --- a/vortex-array/src/expr/expression.rs +++ b/vortex-array/src/expr/expression.rs @@ -1,42 +1,46 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::cell::Cell; use std::fmt; use std::fmt::Debug; use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hash; -use std::ops::Deref; use std::sync::Arc; use itertools::Itertools; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use vortex_session::VortexSession; use crate::dtype::DType; use crate::expr::display::DisplayTreeExpr; +use crate::expr::traversal::TraversalOrder; +use crate::expr::traversal::pre_order_visit_down; use crate::scalar_fn::ScalarFnRef; -use crate::scalar_fn::fns::root::Root; +use crate::scalar_fn::ScalarFnVTable; + +/// An empty child slice, returned by [`Expression::children`] for childless variants. +const NO_CHILDREN: &[Expression] = &[]; /// A node in a Vortex expression tree. /// -/// Expressions represent scalar computations that can be performed on data. Each -/// expression consists of an encoding (vtable), heap-allocated metadata, and child expressions. +/// Most nodes are a scalar function applied to child expressions. [`Expression::Root`] is the scope +/// itself: a language primitive rather than a registered function, because its dtype comes from the +/// scope rather than from children and it is not executable. A [`ScalarFnVTable`] can answer neither +/// of those, so `Root` is a variant instead. #[derive(Clone, Debug, PartialEq, Eq, Hash)] -pub struct Expression { - /// The scalar fn for this node. - scalar_fn: ScalarFnRef, - /// Any children of this expression. - children: Arc>, -} - -impl Deref for Expression { - type Target = ScalarFnRef; - - fn deref(&self) -> &Self::Target { - &self.scalar_fn - } +pub enum Expression { + /// A scalar function applied to child expressions. + Scalar { + /// The scalar fn for this node. + scalar_fn: ScalarFnRef, + /// Any children of this expression. + children: Arc>, + }, + /// The full scope of the expression evaluation. + Root, } impl Expression { @@ -54,90 +58,114 @@ impl Expression { children.len() ); - Ok(Self { + Ok(Self::Scalar { scalar_fn, children: children.into(), }) } - /// Returns the scalar fn vtable for this expression. - pub fn scalar_fn(&self) -> &ScalarFnRef { - &self.scalar_fn + /// Whether this expression is the scope root. + pub fn is_root(&self) -> bool { + matches!(self, Self::Root) + } + + /// Returns the scalar fn for this expression, or `None` if it is not a scalar node. + pub fn as_scalar(&self) -> Option<&ScalarFnRef> { + match self { + Self::Scalar { scalar_fn, .. } => Some(scalar_fn), + Self::Root => None, + } + } + + /// Whether this expression's scalar fn is of the given vtable type. + pub fn is(&self) -> bool { + self.as_scalar().is_some_and(|sf| sf.is::()) + } + + /// The typed options for this expression if its scalar fn matches the given vtable type. + pub fn as_opt(&self) -> Option<&V::Options> { + self.as_scalar().and_then(|sf| sf.as_opt::()) + } + + /// The typed options for this expression. + /// + /// # Panics + /// + /// Panics if the vtable type does not match. + pub fn as_(&self) -> &V::Options { + self.as_opt::() + .vortex_expect("Expression options type mismatch") } /// Returns the children of this expression. - pub fn children(&self) -> &Arc> { - &self.children + pub fn children(&self) -> &[Expression] { + match self { + Self::Scalar { children, .. } => children.as_slice(), + Self::Root => NO_CHILDREN, + } } /// Returns the n'th child of this expression. pub fn child(&self, n: usize) -> &Expression { - &self.children[n] + &self.children()[n] } /// Replace the children of this expression with the provided new children. pub fn with_children( - mut self, + self, children: impl IntoIterator, ) -> VortexResult { let children = Vec::from_iter(children); - vortex_ensure!( - self.signature().arity().matches(children.len()), - "Expression arity mismatch: expected {} children but got {}", - self.signature().arity(), - children.len() - ); - self.children = Arc::new(children); - Ok(self) + match &self { + Self::Root => { + vortex_ensure!( + children.is_empty(), + "Expression arity mismatch: root expects 0 children but got {}", + children.len() + ); + Ok(Self::Root) + } + Self::Scalar { scalar_fn, .. } => { + vortex_ensure!( + scalar_fn.signature().arity().matches(children.len()), + "Expression arity mismatch: expected {} children but got {}", + scalar_fn.signature().arity(), + children.len() + ); + Ok(Self::Scalar { + scalar_fn: scalar_fn.clone(), + children: children.into(), + }) + } + } } /// Computes the return dtype of this expression given the input dtype. pub fn return_dtype(&self, scope: &DType) -> VortexResult { - if self.is::() { - return Ok(scope.clone()); + match self { + Self::Root => Ok(scope.clone()), + Self::Scalar { + scalar_fn, + children, + } => { + let dtypes: Vec<_> = children + .iter() + .map(|c| c.return_dtype(scope)) + .try_collect()?; + scalar_fn.return_dtype(&dtypes) + } } - - let dtypes: Vec<_> = self - .children - .iter() - .map(|c| c.return_dtype(scope)) - .try_collect()?; - self.scalar_fn.return_dtype(&dtypes) } /// Returns a new expression representing the validity mask output of this expression. /// /// The returned expression evaluates to a non-nullable boolean array. pub fn validity(&self) -> VortexResult { - self.scalar_fn.validity(self) - } - - /// Returns an expression that proves this predicate is definitely false from stats. - /// - /// `scope` is the dtype of the row this expression evaluates over. - /// - /// If the returned expression evaluates to `true` for a stats scope, this expression is - /// guaranteed to be false for every row in that scope. `false` and `null` are unknown. - pub fn falsify( - &self, - scope: &DType, - session: &VortexSession, - ) -> VortexResult> { - crate::stats::rewrite::StatsRewriteCtx::new(session, scope).falsify(self) - } - - /// Returns an expression that proves this predicate is definitely true from stats. - /// - /// `scope` is the dtype of the row this expression evaluates over. - /// - /// If the returned expression evaluates to `true` for a stats scope, this expression is - /// guaranteed to be true for every row in that scope. `false` and `null` are unknown. - pub fn satisfy( - &self, - scope: &DType, - session: &VortexSession, - ) -> VortexResult> { - crate::stats::rewrite::StatsRewriteCtx::new(session, scope).satisfy(self) + match self { + // The scope is exactly as valid as itself. + Self::Root => Ok(Self::Root), + Self::Scalar { scalar_fn, .. } => scalar_fn.validity(self), + } } /// Format the expression as a compact string. @@ -145,7 +173,10 @@ impl Expression { /// Since this is a recursive formatter, it is exposed on the public Expression type. /// See fmt_data that is only implemented on the vtable trait. pub fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result { - self.scalar_fn().fmt_sql(self, f) + match self { + Self::Root => write!(f, "$"), + Self::Scalar { scalar_fn, .. } => scalar_fn.fmt_sql(self, f), + } } /// Display the expression as a formatted tree structure. @@ -203,6 +234,30 @@ impl Expression { pub fn display_tree(&self) -> impl Display { DisplayTreeExpr(self) } + + /// Returns true if this expression contains expression E inside. + /// + /// # Example + /// + /// ```rust + /// # use vortex_array::scalar_fn::fns::literal::Literal; + /// # use vortex_array::expr::{eq, lit, root}; + /// let expression = &eq(root(), lit(3u64)); + /// assert!(expression.contains::().unwrap()); + /// let expression = root(); + /// assert!(!expression.contains::().unwrap()); + /// ``` + pub fn contains(&self) -> VortexResult { + let mut contains = false; + pre_order_visit_down(self, |node| { + if node.is::() { + contains = true; + return Ok(TraversalOrder::Stop); + } + Ok(TraversalOrder::Continue) + })?; + Ok(contains) + } } /// The default display implementation for expressions uses the 'SQL'-style format. @@ -212,17 +267,103 @@ impl Display for Expression { } } -/// Iterative drop for expression to avoid stack overflows. +// Switch to iterative cleanup after this many recursive drops. +const MAX_DROP_DEPTH: u32 = 32; + +thread_local! { + static DROP_DEPTH: Cell = const { Cell::new(0) }; +} + +// Increments the thread-local drop depth until the guard is dropped. +struct DropDepthGuard; + +impl DropDepthGuard { + fn enter() -> Option { + DROP_DEPTH.with(|depth| { + let current = depth.get(); + (current < MAX_DROP_DEPTH).then(|| { + depth.set(current + 1); + Self + }) + }) + } +} + +impl Drop for DropDepthGuard { + fn drop(&mut self) { + DROP_DEPTH.with(|depth| depth.set(depth.get() - 1)); + } +} + impl Drop for Expression { fn drop(&mut self) { - if let Some(children) = Arc::get_mut(&mut self.children) { - let mut children_to_drop = std::mem::take(children); + let Self::Scalar { children, .. } = self else { + return; + }; + let Some(children) = Arc::get_mut(children) else { + return; + }; + if children.is_empty() { + return; + } + + let mut children_to_drop = std::mem::take(children); - while let Some(mut child) = children_to_drop.pop() { - if let Some(expr_children) = Arc::get_mut(&mut child.children) { - children_to_drop.append(expr_children); + match DropDepthGuard::enter() { + Some(_guard) => drop(children_to_drop), + None => { + while let Some(mut child) = children_to_drop.pop() { + if let Self::Scalar { children, .. } = &mut child + && let Some(expr_children) = Arc::get_mut(children) + { + children_to_drop.append(expr_children); + } } } } } } + +#[cfg(test)] +mod tests { + use std::thread; + + use super::*; + use crate::expr::lit; + use crate::expr::not; + + fn deep_expression(depth: usize) -> Expression { + let mut expr = lit(true); + for _ in 0..depth { + expr = not(expr); + } + expr + } + + #[test] + fn deep_expression_drops_within_a_small_stack() -> VortexResult<()> { + const DEPTH: usize = 100_000; + const STACK_SIZE: usize = 256 * 1024; + + let dropper = thread::Builder::new() + .stack_size(STACK_SIZE) + .spawn(|| drop(deep_expression(DEPTH)))?; + + assert!( + dropper.join().is_ok(), + "dropping a tree of depth {DEPTH} exhausted a {STACK_SIZE} byte stack" + ); + + Ok(()) + } + + #[test] + fn shallow_expression_keeps_shared_children() { + let expr = not(lit(true)); + let shared = expr.clone(); + + drop(expr); + + assert_eq!(shared.children().len(), 1); + } +} diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index da70427a0f7..fb8bfe227aa 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -4,9 +4,9 @@ //! Factory functions for creating [`Expression`]s from scalar function vtables. use std::sync::Arc; -use std::sync::LazyLock; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_error::vortex_panic; use vortex_utils::iter::ReduceBalancedIterExt; @@ -15,6 +15,7 @@ use crate::dtype::DType; use crate::dtype::FieldName; use crate::dtype::FieldNames; use crate::dtype::Nullability; +use crate::expr::BoundExpression; use crate::expr::Expression; use crate::scalar::Scalar; use crate::scalar::ScalarValue; @@ -49,7 +50,6 @@ use crate::scalar_fn::fns::operators::CompareOperator; use crate::scalar_fn::fns::operators::Operator; use crate::scalar_fn::fns::pack::Pack; use crate::scalar_fn::fns::pack::PackOptions; -use crate::scalar_fn::fns::root::Root; use crate::scalar_fn::fns::select::FieldSelection; use crate::scalar_fn::fns::select::Select; use crate::scalar_fn::fns::variant_get::VariantGet; @@ -57,24 +57,22 @@ use crate::scalar_fn::fns::variant_get::VariantGetOptions; use crate::scalar_fn::fns::variant_get::VariantPath; use crate::scalar_fn::fns::zip::Zip; -static ROOT: LazyLock = LazyLock::new(|| { - Root.try_new_expr(EmptyOptions, vec![]) - .vortex_expect("Creating root() shouldn't fail") -}); - /// Creates an expression that references the root scope. /// /// Returns the entire input array as passed to the expression evaluator. /// This is commonly used as the starting point for field access and other operations. pub fn root() -> Expression { - ROOT.clone() + Expression::Root +} + +/// Creates a bound expression that references a root scope with the given dtype. +pub fn bound_root(dtype: DType) -> BoundExpression { + BoundExpression::new_root(dtype) } /// Return whether the expression is a root expression. pub fn is_root(expr: &Expression) -> bool { - // root doesn't have any children, and scalar_fns have distinct ids - // so we should almost always hit this eq check - (expr.scalar_fn().id() == ROOT.scalar_fn().id()) || expr.is::() + expr.is_root() } // ---- Literal ---- @@ -100,6 +98,13 @@ pub fn lit(value: impl Into) -> Expression { Literal.new_expr(value.into(), []) } +/// Creates a bound literal expression. +pub fn bound_lit(value: impl Into) -> BoundExpression { + Literal + .try_new_bound_expr(value.into(), []) + .vortex_expect("literal expressions are always well-typed") +} + // ---- GetItem / Col ---- /// Creates an expression that accesses a field from the root array. @@ -114,6 +119,11 @@ pub fn col(field: impl Into) -> Expression { GetItem.new_expr(field.into(), vec![root()]) } +/// Creates a bound expression that accesses a field from a root scope with the given dtype. +pub fn bound_col(field: impl Into, scope: DType) -> BoundExpression { + bound_get_item(field, bound_root(scope)) +} + /// Creates an expression that extracts a named field from a struct expression. /// /// Accesses the specified field from the result of the child expression. @@ -126,6 +136,13 @@ pub fn get_item(field: impl Into, child: Expression) -> Expression { GetItem.new_expr(field.into(), vec![child]) } +/// Creates a bound expression that extracts a named field from a struct expression. +pub fn bound_get_item(field: impl Into, child: BoundExpression) -> BoundExpression { + GetItem + .try_new_bound_expr(field.into(), [child]) + .vortex_expect("get-item expressions must reference a field in the child dtype") +} + // ---- VariantGet ---- /// Creates an expression that extracts a path from a Variant expression. @@ -140,6 +157,17 @@ pub fn variant_get( VariantGet.new_expr(VariantGetOptions::new(path.into(), dtype), vec![child]) } +/// Creates a bound expression that extracts a path from a Variant expression. +pub fn bound_variant_get( + child: BoundExpression, + path: impl Into, + dtype: Option, +) -> BoundExpression { + VariantGet + .try_new_bound_expr(VariantGetOptions::new(path.into(), dtype), [child]) + .vortex_expect("variant-get expressions require a Variant child") +} + // ---- CaseWhen ---- /// Creates a CASE WHEN expression with one WHEN/THEN pair and an ELSE value. @@ -155,6 +183,21 @@ pub fn case_when( CaseWhen.new_expr(options, [condition, then_value, else_value]) } +/// Creates a bound CASE WHEN expression with one WHEN/THEN pair and an ELSE value. +pub fn bound_case_when( + condition: BoundExpression, + then_value: BoundExpression, + else_value: BoundExpression, +) -> BoundExpression { + let options = CaseWhenOptions { + num_when_then_pairs: 1, + has_else: true, + }; + CaseWhen + .try_new_bound_expr(options, [condition, then_value, else_value]) + .vortex_expect("case expressions must have boolean conditions and matching branch dtypes") +} + /// Creates a CASE WHEN expression with one WHEN/THEN pair and no ELSE value. pub fn case_when_no_else(condition: Expression, then_value: Expression) -> Expression { let options = CaseWhenOptions { @@ -164,6 +207,20 @@ pub fn case_when_no_else(condition: Expression, then_value: Expression) -> Expre CaseWhen.new_expr(options, [condition, then_value]) } +/// Creates a bound CASE WHEN expression with one WHEN/THEN pair and no ELSE value. +pub fn bound_case_when_no_else( + condition: BoundExpression, + then_value: BoundExpression, +) -> BoundExpression { + let options = CaseWhenOptions { + num_when_then_pairs: 1, + has_else: false, + }; + CaseWhen + .try_new_bound_expr(options, [condition, then_value]) + .vortex_expect("case expressions must have boolean conditions") +} + /// Creates an n-ary CASE WHEN expression from WHEN/THEN pairs and an optional ELSE value. pub fn nested_case_when( when_then_pairs: Vec<(Expression, Expression)>, @@ -194,8 +251,58 @@ pub fn nested_case_when( CaseWhen.new_expr(options, children) } +/// Creates a bound n-ary CASE WHEN expression from WHEN/THEN pairs and an optional ELSE value. +pub fn bound_nested_case_when( + when_then_pairs: Vec<(BoundExpression, BoundExpression)>, + else_value: Option, +) -> BoundExpression { + assert!( + !when_then_pairs.is_empty(), + "nested_case_when requires at least one when/then pair" + ); + + let Ok(num_when_then_pairs) = u32::try_from(when_then_pairs.len()) else { + vortex_panic!("nested_case_when has too many when/then pairs"); + }; + let has_else = else_value.is_some(); + let mut children = Vec::with_capacity(when_then_pairs.len() * 2 + usize::from(has_else)); + for (condition, then_value) in when_then_pairs { + children.push(condition); + children.push(then_value); + } + if let Some(else_expr) = else_value { + children.push(else_expr); + } + + let options = CaseWhenOptions { + num_when_then_pairs, + has_else, + }; + CaseWhen + .try_new_bound_expr(options, children) + .vortex_expect("case expressions must have boolean conditions and matching branch dtypes") +} + // ---- Binary operators ---- +/// Creates a binary expression with the given operator. +pub fn binary(operator: Operator, lhs: Expression, rhs: Expression) -> Expression { + Binary + .try_new_expr(operator, [lhs, rhs]) + .vortex_expect("Failed to create binary expression") +} + +/// Creates a bound binary expression with the given operator. +pub fn bound_binary( + operator: Operator, + lhs: BoundExpression, + rhs: BoundExpression, +) -> BoundExpression { + Binary + .try_new_bound_expr(operator, [lhs, rhs]) + .vortex_expect("binary expressions must have compatible operand dtypes") +} + /// Create a new [`Binary`] using the [`Eq`](Operator::Eq) operator. /// /// ## Example usage @@ -223,6 +330,11 @@ pub fn eq(lhs: Expression, rhs: Expression) -> Expression { .vortex_expect("Failed to create Eq binary expression") } +/// Creates a bound equality expression. +pub fn bound_eq(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { + bound_binary(Operator::Eq, lhs, rhs) +} + /// Create a new [`Binary`] using the [`NotEq`](Operator::NotEq) operator. /// /// ## Example usage @@ -250,6 +362,11 @@ pub fn not_eq(lhs: Expression, rhs: Expression) -> Expression { .vortex_expect("Failed to create NotEq binary expression") } +/// Creates a bound inequality expression. +pub fn bound_not_eq(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { + bound_binary(Operator::NotEq, lhs, rhs) +} + /// Create a new [`Binary`] using the [`Gte`](Operator::Gte) operator. /// /// ## Example usage @@ -277,6 +394,11 @@ pub fn gt_eq(lhs: Expression, rhs: Expression) -> Expression { .vortex_expect("Failed to create Gte binary expression") } +/// Creates a bound greater-than-or-equal expression. +pub fn bound_gt_eq(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { + bound_binary(Operator::Gte, lhs, rhs) +} + /// Create a new [`Binary`] using the [`Gt`](Operator::Gt) operator. /// /// ## Example usage @@ -304,6 +426,11 @@ pub fn gt(lhs: Expression, rhs: Expression) -> Expression { .vortex_expect("Failed to create Gt binary expression") } +/// Creates a bound greater-than expression. +pub fn bound_gt(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { + bound_binary(Operator::Gt, lhs, rhs) +} + /// Create a new [`Binary`] using the [`Lte`](Operator::Lte) operator. /// /// ## Example usage @@ -331,6 +458,11 @@ pub fn lt_eq(lhs: Expression, rhs: Expression) -> Expression { .vortex_expect("Failed to create Lte binary expression") } +/// Creates a bound less-than-or-equal expression. +pub fn bound_lt_eq(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { + bound_binary(Operator::Lte, lhs, rhs) +} + /// Create a new [`Binary`] using the [`Lt`](Operator::Lt) operator. /// /// ## Example usage @@ -358,6 +490,11 @@ pub fn lt(lhs: Expression, rhs: Expression) -> Expression { .vortex_expect("Failed to create Lt binary expression") } +/// Creates a bound less-than expression. +pub fn bound_lt(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { + bound_binary(Operator::Lt, lhs, rhs) +} + /// Create a new [`Binary`] using the [`Or`](Operator::Or) operator. /// /// ## Example usage @@ -383,6 +520,11 @@ pub fn or(lhs: Expression, rhs: Expression) -> Expression { .vortex_expect("Failed to create Or binary expression") } +/// Creates a bound boolean OR expression. +pub fn bound_or(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { + bound_binary(Operator::Or, lhs, rhs) +} + /// Collects a list of `or`ed values into a single expression using a balanced tree. /// /// This creates a balanced binary tree to avoid deep nesting that could cause @@ -396,6 +538,14 @@ where iter.into_iter().reduce_balanced(or) } +/// Collects bound expressions into a balanced tree of boolean OR expressions. +pub fn bound_or_collect(iter: I) -> Option +where + I: IntoIterator, +{ + iter.into_iter().reduce_balanced(bound_or) +} + /// Create a new [`Binary`] using the [`And`](Operator::And) operator. /// /// ## Example usage @@ -421,6 +571,11 @@ pub fn and(lhs: Expression, rhs: Expression) -> Expression { .vortex_expect("Failed to create And binary expression") } +/// Creates a bound boolean AND expression. +pub fn bound_and(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { + bound_binary(Operator::And, lhs, rhs) +} + /// Collects a list of `and`ed values into a single expression using a balanced tree. /// /// This creates a balanced binary tree to avoid deep nesting that could cause @@ -434,6 +589,30 @@ where iter.into_iter().reduce_balanced(and) } +/// Collects bound expressions into a balanced tree of boolean AND expressions. +pub fn bound_and_collect(iter: I) -> Option +where + I: IntoIterator, +{ + iter.into_iter().reduce_balanced(bound_and) +} + +/// The conjunction of an expression's child validities — i.e. the validity of a scalar function +/// whose result is null exactly when any operand is null. +/// +/// This is the `ScalarFnVTable::validity` for kernels that propagate nulls and never produce a +/// null from non-null inputs (comparisons, arithmetic, most spatial and tensor operations). Returning it lets +/// the planner derive the output's null mask without executing the kernel. Yields `None` when the +/// expression has no children. +pub fn union_child_validities(expression: &Expression) -> VortexResult> { + let child_validities = expression + .children() + .iter() + .map(Expression::validity) + .collect::>>()?; + Ok(and_collect(child_validities)) +} + /// Create a new [`Binary`] using the [`Add`](Operator::Add) operator. /// /// ## Example usage @@ -458,6 +637,11 @@ pub fn checked_add(lhs: Expression, rhs: Expression) -> Expression { .vortex_expect("Failed to create Add binary expression") } +/// Creates a bound checked-add expression. +pub fn bound_checked_add(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { + bound_binary(Operator::Add, lhs, rhs) +} + // ---- Not ---- /// Creates an expression that logically inverts boolean values. @@ -472,6 +656,12 @@ pub fn not(operand: Expression) -> Expression { Not.new_expr(EmptyOptions, vec![operand]) } +/// Creates a bound expression that logically inverts boolean values. +pub fn bound_not(operand: BoundExpression) -> BoundExpression { + Not.try_new_bound_expr(EmptyOptions, [operand]) + .vortex_expect("not expressions require a boolean operand") +} + // ---- Between ---- /// Creates an expression that checks if values are between two bounds. @@ -500,6 +690,18 @@ pub fn between( .vortex_expect("Failed to create Between expression") } +/// Creates a bound expression that checks if values are between two bounds. +pub fn bound_between( + arr: BoundExpression, + lower: BoundExpression, + upper: BoundExpression, + options: BetweenOptions, +) -> BoundExpression { + Between + .try_new_bound_expr(options, [arr, lower, upper]) + .vortex_expect("between expressions require compatible operand dtypes") +} + // ---- Select ---- /// Creates an expression that selects (includes) specific fields from an array. @@ -515,6 +717,13 @@ pub fn select(field_names: impl Into, child: Expression) -> Expressi .vortex_expect("Failed to create Select expression") } +/// Creates a bound expression that selects specific fields from a struct expression. +pub fn bound_select(field_names: impl Into, child: BoundExpression) -> BoundExpression { + Select + .try_new_bound_expr(FieldSelection::Include(field_names.into()), [child]) + .vortex_expect("select expressions require fields from a struct child") +} + /// Creates an expression that excludes specific fields from an array. /// /// Projects all fields except the specified ones from the input struct expression. @@ -529,6 +738,16 @@ pub fn select_exclude(fields: impl Into, child: Expression) -> Expre .vortex_expect("Failed to create Select expression") } +/// Creates a bound expression that excludes specific fields from a struct expression. +pub fn bound_select_exclude( + fields: impl Into, + child: BoundExpression, +) -> BoundExpression { + Select + .try_new_bound_expr(FieldSelection::Exclude(fields.into()), [child]) + .vortex_expect("select expressions require fields from a struct child") +} + // ---- Pack ---- /// Creates an expression that packs values into a struct with named fields. @@ -555,6 +774,25 @@ pub fn pack( ) } +/// Creates a bound expression that packs values into a struct with named fields. +pub fn bound_pack( + elements: impl IntoIterator, BoundExpression)>, + nullability: Nullability, +) -> BoundExpression { + let (names, values): (Vec<_>, Vec<_>) = elements + .into_iter() + .map(|(name, value)| (name.into(), value)) + .unzip(); + Pack.try_new_bound_expr( + PackOptions { + names: names.into(), + nullability, + }, + values, + ) + .vortex_expect("pack expressions must have one name per child") +} + // ---- Cast ---- /// Creates an expression that casts values to a target data type. @@ -571,6 +809,12 @@ pub fn cast(child: Expression, target: DType) -> Expression { .vortex_expect("Failed to create Cast expression") } +/// Creates a bound expression that casts values to a target dtype. +pub fn bound_cast(child: BoundExpression, target: DType) -> BoundExpression { + Cast.try_new_bound_expr(target, [child]) + .vortex_expect("cast expressions require a supported source and target dtype") +} + // ---- FillNull ---- /// Creates an expression that replaces null values with a fill value. @@ -583,6 +827,13 @@ pub fn fill_null(child: Expression, fill_value: Expression) -> Expression { FillNull.new_expr(EmptyOptions, [child, fill_value]) } +/// Creates a bound expression that replaces null values with a fill value. +pub fn bound_fill_null(child: BoundExpression, fill_value: BoundExpression) -> BoundExpression { + FillNull + .try_new_bound_expr(EmptyOptions, [child, fill_value]) + .vortex_expect("fill-null expressions require compatible child and fill dtypes") +} + // ---- IsNull ---- /// Creates an expression that checks for null values. @@ -597,6 +848,13 @@ pub fn is_null(child: Expression) -> Expression { IsNull.new_expr(EmptyOptions, vec![child]) } +/// Creates a bound expression that checks for null values. +pub fn bound_is_null(child: BoundExpression) -> BoundExpression { + IsNull + .try_new_bound_expr(EmptyOptions, [child]) + .vortex_expect("is-null expressions are always well-typed") +} + // ---- IsNotNull ---- /// Creates an expression that checks for non-null values. @@ -611,6 +869,13 @@ pub fn is_not_null(child: Expression) -> Expression { IsNotNull.new_expr(EmptyOptions, vec![child]) } +/// Creates a bound expression that checks for non-null values. +pub fn bound_is_not_null(child: BoundExpression) -> BoundExpression { + IsNotNull + .try_new_bound_expr(EmptyOptions, [child]) + .vortex_expect("is-not-null expressions are always well-typed") +} + // ---- Like ---- /// Creates a SQL LIKE expression. @@ -624,6 +889,11 @@ pub fn like(child: Expression, pattern: Expression) -> Expression { ) } +/// Creates a bound SQL LIKE expression. +pub fn bound_like(child: BoundExpression, pattern: BoundExpression) -> BoundExpression { + bound_like_with_options(child, pattern, false, false) +} + /// Creates a case-insensitive SQL ILIKE expression. pub fn ilike(child: Expression, pattern: Expression) -> Expression { Like.new_expr( @@ -635,6 +905,11 @@ pub fn ilike(child: Expression, pattern: Expression) -> Expression { ) } +/// Creates a bound case-insensitive SQL ILIKE expression. +pub fn bound_ilike(child: BoundExpression, pattern: BoundExpression) -> BoundExpression { + bound_like_with_options(child, pattern, false, true) +} + /// Creates a negated SQL NOT LIKE expression. pub fn not_like(child: Expression, pattern: Expression) -> Expression { Like.new_expr( @@ -646,6 +921,11 @@ pub fn not_like(child: Expression, pattern: Expression) -> Expression { ) } +/// Creates a bound negated SQL NOT LIKE expression. +pub fn bound_not_like(child: BoundExpression, pattern: BoundExpression) -> BoundExpression { + bound_like_with_options(child, pattern, true, false) +} + /// Creates a negated case-insensitive SQL NOT ILIKE expression. pub fn not_ilike(child: Expression, pattern: Expression) -> Expression { Like.new_expr( @@ -657,6 +937,27 @@ pub fn not_ilike(child: Expression, pattern: Expression) -> Expression { ) } +/// Creates a bound negated case-insensitive SQL NOT ILIKE expression. +pub fn bound_not_ilike(child: BoundExpression, pattern: BoundExpression) -> BoundExpression { + bound_like_with_options(child, pattern, true, true) +} + +fn bound_like_with_options( + child: BoundExpression, + pattern: BoundExpression, + negated: bool, + case_insensitive: bool, +) -> BoundExpression { + Like.try_new_bound_expr( + LikeOptions { + negated, + case_insensitive, + }, + [child, pattern], + ) + .vortex_expect("like expressions require UTF-8 or binary operands") +} + // ---- Mask ---- /// Creates a mask expression that applies the given boolean mask to the input array. @@ -664,6 +965,12 @@ pub fn mask(array: Expression, mask: Expression) -> Expression { Mask.new_expr(EmptyOptions, [array, mask]) } +/// Creates a bound mask expression. +pub fn bound_mask(array: BoundExpression, mask: BoundExpression) -> BoundExpression { + Mask.try_new_bound_expr(EmptyOptions, [array, mask]) + .vortex_expect("mask expressions require a boolean mask") +} + // ---- Merge ---- /// Creates an expression that merges struct expressions into a single struct. @@ -682,6 +989,11 @@ pub fn merge(elements: impl IntoIterator>) -> Expre Merge.new_expr(DuplicateHandling::default(), values) } +/// Creates a bound expression that merges struct expressions. +pub fn bound_merge(elements: impl IntoIterator) -> BoundExpression { + bound_merge_opts(elements, DuplicateHandling::default()) +} + /// Creates a merge expression with explicit duplicate handling. pub fn merge_opts( elements: impl IntoIterator>, @@ -692,6 +1004,16 @@ pub fn merge_opts( Merge.new_expr(duplicate_handling, values) } +/// Creates a bound merge expression with explicit duplicate handling. +pub fn bound_merge_opts( + elements: impl IntoIterator, + duplicate_handling: DuplicateHandling, +) -> BoundExpression { + Merge + .try_new_bound_expr(duplicate_handling, elements) + .vortex_expect("merge expressions require non-nullable struct children") +} + // ---- Zip ---- /// Creates a zip expression that conditionally selects between two arrays. @@ -704,8 +1026,33 @@ pub fn zip_expr(mask: Expression, if_true: Expression, if_false: Expression) -> Zip.new_expr(EmptyOptions, [if_true, if_false, mask]) } +/// Creates a bound zip expression that conditionally selects between two arrays. +pub fn bound_zip_expr( + mask: BoundExpression, + if_true: BoundExpression, + if_false: BoundExpression, +) -> BoundExpression { + Zip.try_new_bound_expr(EmptyOptions, [if_true, if_false, mask]) + .vortex_expect("zip expressions require a boolean mask and compatible value dtypes") +} + // ---- Dynamic ---- +/// Creates a dynamic comparison expression from its complete options. +pub fn dynamic_with_options(options: DynamicComparisonExpr, lhs: Expression) -> Expression { + DynamicComparison.new_expr(options, [lhs]) +} + +/// Creates a bound dynamic comparison expression from its complete options. +pub fn bound_dynamic_with_options( + options: DynamicComparisonExpr, + lhs: BoundExpression, +) -> BoundExpression { + DynamicComparison + .try_new_bound_expr(options, [lhs]) + .vortex_expect("dynamic comparisons require a compatible left-hand dtype") +} + /// Creates a dynamic comparison expression. pub fn dynamic( operator: CompareOperator, @@ -714,7 +1061,28 @@ pub fn dynamic( default: bool, lhs: Expression, ) -> Expression { - DynamicComparison.new_expr( + dynamic_with_options( + DynamicComparisonExpr { + operator, + rhs: Arc::new(Rhs { + value: Arc::new(rhs_value), + dtype: rhs_dtype, + }), + default, + }, + lhs, + ) +} + +/// Creates a bound dynamic comparison expression. +pub fn bound_dynamic( + operator: CompareOperator, + rhs_value: impl Fn() -> Option + Send + Sync + 'static, + rhs_dtype: DType, + default: bool, + lhs: BoundExpression, +) -> BoundExpression { + bound_dynamic_with_options( DynamicComparisonExpr { operator, rhs: Arc::new(Rhs { @@ -723,7 +1091,7 @@ pub fn dynamic( }), default, }, - [lhs], + lhs, ) } @@ -741,6 +1109,13 @@ pub fn list_contains(list: Expression, value: Expression) -> Expression { ListContains.new_expr(EmptyOptions, [list, value]) } +/// Creates a bound expression that checks if a value is contained in a list. +pub fn bound_list_contains(list: BoundExpression, value: BoundExpression) -> BoundExpression { + ListContains + .try_new_bound_expr(EmptyOptions, [list, value]) + .vortex_expect("list-contains expressions require a compatible list and value dtype") +} + // ---- ByteLength ---- /// Creates an expression that computes the byte length of each element. @@ -754,6 +1129,13 @@ pub fn byte_length(input: Expression) -> Expression { ByteLength.new_expr(EmptyOptions, [input]) } +/// Creates a bound expression that computes each element's byte length. +pub fn bound_byte_length(input: BoundExpression) -> BoundExpression { + ByteLength + .try_new_bound_expr(EmptyOptions, [input]) + .vortex_expect("byte-length expressions require a variable-length binary child") +} + // ---- ExtStorage ---- /// Creates an expression that extracts the storage values from an extension array. @@ -766,6 +1148,13 @@ pub fn ext_storage(input: Expression) -> Expression { ExtStorage.new_expr(EmptyOptions, [input]) } +/// Creates a bound expression that extracts an extension array's storage values. +pub fn bound_ext_storage(input: BoundExpression) -> BoundExpression { + ExtStorage + .try_new_bound_expr(EmptyOptions, [input]) + .vortex_expect("extension-storage expressions require an extension child") +} + // ---- ListLength ---- /// Creates an expression that computes the number of elements in each list @@ -780,6 +1169,13 @@ pub fn list_length(input: Expression) -> Expression { ListLength.new_expr(EmptyOptions, [input]) } +/// Creates a bound expression that computes the number of elements in each list. +pub fn bound_list_length(input: BoundExpression) -> BoundExpression { + ListLength + .try_new_bound_expr(EmptyOptions, [input]) + .vortex_expect("list-length expressions require a list child") +} + // ---- ListSum ---- /// Creates an expression that sums the elements of each list for `List` and @@ -798,8 +1194,78 @@ pub fn list_sum(input: Expression) -> Expression { ListSum.new_expr(NumericalAggregateOpts::default(), [input]) } +/// Creates a bound expression that sums the elements of each list. +pub fn bound_list_sum(input: BoundExpression) -> BoundExpression { + ListSum + .try_new_bound_expr(NumericalAggregateOpts::default(), [input]) + .vortex_expect("list-sum expressions require a numeric list child") +} + /// Creates a [`list_sum`] expression with explicit [`NumericalAggregateOpts`], controlling /// whether NaN float elements are skipped (the default) or poison the list's sum to NaN. pub fn list_sum_opts(input: Expression, options: NumericalAggregateOpts) -> Expression { ListSum.new_expr(options, [input]) } + +/// Creates a bound list-sum expression with explicit aggregate options. +pub fn bound_list_sum_opts( + input: BoundExpression, + options: NumericalAggregateOpts, +) -> BoundExpression { + ListSum + .try_new_bound_expr(options, [input]) + .vortex_expect("list-sum expressions require a numeric list child") +} + +/// Constructors for expressions whose children have already been bound and type-checked. +/// +/// These mirror the constructors in [`crate::expr`] and panic when the supplied children do not +/// form a well-typed expression. Use [`BoundExpression::try_new`] when construction must be +/// fallible. +pub mod bound { + pub use super::bound_and as and; + pub use super::bound_and_collect as and_collect; + pub use super::bound_between as between; + pub use super::bound_binary as binary; + pub use super::bound_byte_length as byte_length; + pub use super::bound_case_when as case_when; + pub use super::bound_case_when_no_else as case_when_no_else; + pub use super::bound_cast as cast; + pub use super::bound_checked_add as checked_add; + pub use super::bound_col as col; + pub use super::bound_dynamic as dynamic; + pub use super::bound_dynamic_with_options as dynamic_with_options; + pub use super::bound_eq as eq; + pub use super::bound_ext_storage as ext_storage; + pub use super::bound_fill_null as fill_null; + pub use super::bound_get_item as get_item; + pub use super::bound_gt as gt; + pub use super::bound_gt_eq as gt_eq; + pub use super::bound_ilike as ilike; + pub use super::bound_is_not_null as is_not_null; + pub use super::bound_is_null as is_null; + pub use super::bound_like as like; + pub use super::bound_list_contains as list_contains; + pub use super::bound_list_length as list_length; + pub use super::bound_list_sum as list_sum; + pub use super::bound_list_sum_opts as list_sum_opts; + pub use super::bound_lit as lit; + pub use super::bound_lt as lt; + pub use super::bound_lt_eq as lt_eq; + pub use super::bound_mask as mask; + pub use super::bound_merge as merge; + pub use super::bound_merge_opts as merge_opts; + pub use super::bound_nested_case_when as nested_case_when; + pub use super::bound_not as not; + pub use super::bound_not_eq as not_eq; + pub use super::bound_not_ilike as not_ilike; + pub use super::bound_not_like as not_like; + pub use super::bound_or as or; + pub use super::bound_or_collect as or_collect; + pub use super::bound_pack as pack; + pub use super::bound_root as root; + pub use super::bound_select as select; + pub use super::bound_select_exclude as select_exclude; + pub use super::bound_variant_get as variant_get; + pub use super::bound_zip_expr as zip_expr; +} diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index 4331b2817f0..59fd21c46ee 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -12,7 +12,7 @@ //! //! Each node references a scalar function defined by a //! [`ScalarFnVTable`](crate::scalar_fn::ScalarFnVTable). The vtable declares the function signature, -//! properties such as null-sensitivity, and the logic that executes it over input arrays. Built-in +//! properties such as strictness, and the logic that executes it over input arrays. Built-in //! functions live in [`crate::scalar_fn`]; integration and plugin crates supply additional, //! use-case-specific functions. //! @@ -24,12 +24,12 @@ //! together, before any data is materialized. The deferred tree is executed toward canonical form //! only when a result is actually required. //! -//! # Typing and coercion +//! # Type checking //! //! Expressions are strictly typed: an input array's dtype must match the function signature exactly, -//! so callers perform any required type coercion themselves before building the expression (see the -//! [`transform`] passes). The one relaxation is null-coercion — for example, equality may compare a -//! `u32` against a `u32?`, but never a `u32` against an `i32`. +//! so callers perform any required casts themselves before building the expression. The one +//! relaxation is nullability—for example, equality may compare a `u32` against a `u32?`, but never +//! a `u32` against an `i32`. //! //! Filter expressions are decomposed into independent conjuncts with [`split_conjunction`] so that //! scans can evaluate and reorder the most selective predicates first. @@ -56,6 +56,7 @@ pub mod aliases; pub mod analysis; #[cfg(feature = "arbitrary")] pub mod arbitrary; +pub mod bound_expression; pub mod display; pub(crate) mod expression; mod exprs; @@ -63,13 +64,63 @@ pub(crate) mod field; pub mod forms; mod optimize; pub mod proto; +pub mod scope; pub mod stats; pub mod transform; pub mod traversal; pub use analysis::*; +pub use bound_expression::*; pub use expression::*; -pub use exprs::*; +pub use exprs::and; +pub use exprs::and_collect; +pub use exprs::between; +pub use exprs::binary; +pub use exprs::bound; +pub use exprs::byte_length; +pub use exprs::case_when; +pub use exprs::case_when_no_else; +pub use exprs::cast; +pub use exprs::checked_add; +pub use exprs::col; +pub use exprs::dynamic; +pub use exprs::dynamic_with_options; +pub use exprs::eq; +pub use exprs::ext_storage; +pub use exprs::fill_null; +pub use exprs::get_item; +pub use exprs::gt; +pub use exprs::gt_eq; +pub use exprs::ilike; +pub use exprs::is_not_null; +pub use exprs::is_null; +pub use exprs::is_root; +pub use exprs::like; +pub use exprs::list_contains; +pub use exprs::list_length; +pub use exprs::list_sum; +pub use exprs::list_sum_opts; +pub use exprs::lit; +pub use exprs::lt; +pub use exprs::lt_eq; +pub use exprs::mask; +pub use exprs::merge; +pub use exprs::merge_opts; +pub use exprs::nested_case_when; +pub use exprs::not; +pub use exprs::not_eq; +pub use exprs::not_ilike; +pub use exprs::not_like; +pub use exprs::or; +pub use exprs::or_collect; +pub use exprs::pack; +pub use exprs::root; +pub use exprs::select; +pub use exprs::select_exclude; +pub use exprs::union_child_validities; +pub use exprs::variant_get; +pub use exprs::zip_expr; +pub use scope::*; pub trait VortexExprExt { /// Accumulate all field references from this expression and its children in a set @@ -106,19 +157,41 @@ fn split_inner(expr: &Expression, exprs: &mut Vec) { } /// An expression wrapper that performs pointer equality on child expressions. -#[derive(Clone)] +#[derive(Clone, Debug)] pub struct ExactExpr(pub Expression); impl PartialEq for ExactExpr { fn eq(&self, other: &Self) -> bool { - self.0.scalar_fn() == other.0.scalar_fn() - && Arc::ptr_eq(self.0.children(), other.0.children()) + match (&self.0, &other.0) { + (Expression::Root, Expression::Root) => true, + ( + Expression::Scalar { + scalar_fn: lhs_fn, + children: lhs_children, + }, + Expression::Scalar { + scalar_fn: rhs_fn, + children: rhs_children, + }, + ) => lhs_fn == rhs_fn && Arc::ptr_eq(lhs_children, rhs_children), + _ => false, + } } } impl Eq for ExactExpr {} impl Hash for ExactExpr { fn hash(&self, state: &mut H) { - self.0.hash(state); + match &self.0 { + Expression::Root => state.write_u8(0), + Expression::Scalar { + scalar_fn, + children, + } => { + state.write_u8(1); + scalar_fn.hash(state); + Arc::as_ptr(children).hash(state); + } + } } } @@ -148,6 +221,13 @@ pub mod test_harness { #[cfg(test)] mod tests { + use std::collections::hash_map::RandomState; + use std::hash::BuildHasher; + + use vortex_array::expr::eq; + use vortex_array::expr::lit; + use vortex_array::expr::root; + use super::*; use crate::dtype::DType; use crate::dtype::FieldNames; @@ -155,21 +235,21 @@ mod tests { use crate::dtype::PType; use crate::dtype::StructFields; use crate::expr::and; + use crate::expr::bound; + use crate::expr::case_when; use crate::expr::col; - use crate::expr::eq; use crate::expr::get_item; use crate::expr::gt; use crate::expr::gt_eq; - use crate::expr::lit; use crate::expr::lt; use crate::expr::lt_eq; use crate::expr::not; use crate::expr::not_eq; use crate::expr::or; - use crate::expr::root; use crate::expr::select; use crate::expr::select_exclude; use crate::scalar::Scalar; + use crate::scalar_fn::fns::literal::Literal; #[test] fn basic_expr_split_test() { @@ -189,6 +269,62 @@ mod tests { assert_eq!(conjunction.len(), 2, "Conjunction is {conjunction:?}"); } + #[test] + fn exact_expr_hash_consistent_with_eq() { + let state = RandomState::new(); + let expr = eq(get_item("col1", root()), lit(1)); + + // Clones share the children Arc, so they are equal and must hash equally. + let a = ExactExpr(expr.clone()); + let b = ExactExpr(expr); + assert_eq!(a, b); + assert_eq!(state.hash_one(&a), state.hash_one(&b)); + + // Structurally identical expressions built separately are distinct keys. + let rebuilt = ExactExpr(eq(get_item("col1", root()), lit(1))); + assert_ne!(a, rebuilt); + } + + #[test] + fn bound_constructors_preserve_order_and_types() -> vortex_error::VortexResult<()> { + let value_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let scope = DType::Struct( + StructFields::from_iter([("value", value_dtype.clone())]), + Nullability::NonNullable, + ); + + let root = bound::root(scope.clone()); + let value = bound::get_item("value", root); + let literal = bound::lit(5i32); + let condition = bound::gt(value.clone(), literal.clone()); + assert_eq!(condition.dtype(), &DType::Bool(Nullability::NonNullable)); + assert_eq!(condition.children(), &[value.clone(), literal.clone()]); + + let case = bound::case_when(condition.clone(), value.clone(), literal.clone()); + assert_eq!(case.dtype(), &value_dtype); + assert_eq!(case.children(), &[condition.clone(), value, literal]); + + let packed = bound::pack( + [("condition", condition.clone()), ("value", case.clone())], + Nullability::NonNullable, + ); + assert_eq!(packed.children(), &[condition, case.clone()]); + assert_eq!( + packed.dtype(), + &DType::Struct( + StructFields::from_iter([ + ("condition", DType::Bool(Nullability::NonNullable)), + ("value", value_dtype), + ]), + Nullability::NonNullable, + ) + ); + + let unbound = case_when(gt(col("value"), lit(5i32)), col("value"), lit(5i32)); + assert_eq!(unbound.bind(&scope)?, case); + Ok(()) + } + #[test] fn expr_display() { assert_eq!(col("a").to_string(), "$.a"); @@ -287,4 +423,12 @@ mod tests { "{dog: 32u32, cat: \"rufus\"}" ); } + + #[test] + fn expr_contains() { + let expression = &eq(root(), lit(3u64)); + assert!(expression.contains::().unwrap()); + let expression = root(); + assert!(!expression.contains::().unwrap()); + } } diff --git a/vortex-array/src/expr/optimize.rs b/vortex-array/src/expr/optimize.rs index 27959a96070..3e625324a18 100644 --- a/vortex-array/src/expr/optimize.rs +++ b/vortex-array/src/expr/optimize.rs @@ -1,25 +1,18 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::any::Any; use std::cell::RefCell; -use std::ops::Deref; -use std::sync::Arc; use itertools::Itertools; -use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_err; use vortex_utils::aliases::hash_map::HashMap; use crate::dtype::DType; use crate::expr::Expression; use crate::expr::transform::match_between::find_between; -use crate::scalar_fn::ReduceCtx; -use crate::scalar_fn::ReduceNode; -use crate::scalar_fn::ReduceNodeRef; -use crate::scalar_fn::ScalarFnRef; +use crate::scalar_fn::ExpressionReduceNode; use crate::scalar_fn::SimplifyCtx; -use crate::scalar_fn::fns::root::Root; impl Expression { /// Optimize the root expression node only, iterating to convergence. @@ -27,30 +20,46 @@ impl Expression { /// This applies optimization rules repeatedly until no more changes occur: /// 1. `simplify_untyped` - type-independent simplifications /// 2. `simplify` - type-aware simplifications - /// 3. `reduce` - abstract reduction rules via `ReduceNode`/`ReduceCtx` + /// 3. `reduce` - abstract reduction rules via `ReduceNode` pub fn optimize(&self, scope: &DType) -> VortexResult { - let cache = SimplifyCache { - scope, - dtype_cache: RefCell::new(HashMap::new()), - }; - Ok(self - .clone() - .try_optimize(scope, &cache)? - .unwrap_or_else(|| self.clone())) + let cache = SimplifyCache::new(scope); + Ok(self.try_optimize(&cache)?.unwrap_or_else(|| self.clone())) } - /// Try to optimize the root expression node only, returning None if no optimizations applied. - fn try_optimize( + /// Apply this node's own untyped simplification rule, if it has one. + /// + /// Non-scalar nodes carry no rules, so they never simplify. + fn simplify_untyped_node(&self) -> VortexResult> { + match self { + Expression::Scalar { scalar_fn, .. } => scalar_fn.simplify_untyped(self), + Expression::Root => Ok(None), + } + } + + /// Apply this node's own type-aware simplification rule, if it has one. + fn simplify_node(&self, ctx: &dyn SimplifyCtx) -> VortexResult> { + match self { + Expression::Scalar { scalar_fn, .. } => scalar_fn.simplify(self, ctx), + Expression::Root => Ok(None), + } + } + + /// Apply this node's own abstract reduction rule, if it has one. + fn reduce_node<'a>( &self, - scope: &DType, - cache: &SimplifyCache<'_>, - ) -> VortexResult> { - let reduce_ctx = ExpressionReduceCtx { - scope: scope.clone(), - }; + node: &ExpressionReduceNode<'a>, + ) -> VortexResult>> { + match self { + Expression::Scalar { scalar_fn, .. } => scalar_fn.reduce_expression(node), + Expression::Root => Ok(None), + } + } - let mut current = self.clone(); - let mut any_optimizations = false; + /// Try to optimize the root expression node only, returning None if no optimizations applied. + fn try_optimize(&self, cache: &SimplifyCache<'_>) -> VortexResult> { + // Copy-on-write: `current` stays None until a rule fires, so unchanged nodes (the common + // case) are never cloned. + let mut current: Option = None; let mut loop_counter = 0; loop { @@ -61,37 +70,33 @@ impl Expression { } loop_counter += 1; + let expr = current.as_ref().unwrap_or(self); let mut changed = false; // Try simplify_untyped - if let Some(simplified) = current.scalar_fn().simplify_untyped(¤t)? { - current = simplified; + if let Some(simplified) = expr.simplify_untyped_node()? { + current = Some(simplified); changed = true; - any_optimizations = true; } // Try simplify (typed) - if let Some(simplified) = current.scalar_fn().simplify(¤t, cache)? { - current = simplified; + let expr = current.as_ref().unwrap_or(self); + if let Some(simplified) = expr.simplify_node(cache)? { + current = Some(simplified); changed = true; - any_optimizations = true; } - // Try reduce via ReduceNode/ReduceCtx - let reduce_node = ExpressionReduceNode { - expression: current.clone(), - scope: scope.clone(), + // Try reduce via ReduceNode. The node borrows the expression and scope, so + // constructing it is free; the block scopes the borrows so `current` can be updated. + let reduced = { + let expr = current.as_ref().unwrap_or(self); + let reduce_node = ExpressionReduceNode::new(expr, cache.scope); + expr.reduce_node(&reduce_node)? + .map(ExpressionReduceNode::into_expression) }; - if let Some(reduced) = current.scalar_fn().reduce(&reduce_node, &reduce_ctx)? { - let reduced_expr = reduced - .as_any() - .downcast_ref::() - .vortex_expect("ReduceNode not an ExpressionReduceNode") - .expression - .clone(); - current = reduced_expr; + if let Some(reduced_expr) = reduced { + current = Some(reduced_expr); changed = true; - any_optimizations = true; } if !changed { @@ -99,11 +104,7 @@ impl Expression { } } - if any_optimizations { - Ok(Some(current)) - } else { - Ok(None) - } + Ok(current) } /// Optimize the entire expression tree recursively. @@ -118,11 +119,8 @@ impl Expression { /// Try to optimize the entire expression tree recursively. pub fn try_optimize_recursive(&self, scope: &DType) -> VortexResult> { - let cache = SimplifyCache { - scope, - dtype_cache: RefCell::new(HashMap::new()), - }; - let result = self.try_optimize_recursive_inner(scope, &cache)?; + let cache = SimplifyCache::new(scope); + let result = self.try_optimize_recursive_inner(&cache)?; // Apply the between optimization once at the top level only. // TODO(ngates): remove the "between" optimization, or rewrite it to not always convert @@ -132,92 +130,34 @@ impl Expression { fn try_optimize_recursive_inner( &self, - scope: &DType, cache: &SimplifyCache<'_>, ) -> VortexResult> { - let mut current = self.clone(); - let mut any_optimizations = false; - // First optimize the root - if let Some(optimized) = current.clone().try_optimize(scope, cache)? { - current = optimized; - any_optimizations = true; - } - - // Then recursively optimize children - let mut new_children = Vec::with_capacity(current.children().len()); - let mut any_child_optimized = false; - for child in current.children().iter() { - if let Some(optimized) = child.try_optimize_recursive_inner(scope, cache)? { - new_children.push(optimized); - any_child_optimized = true; - } else { + let mut current = self.try_optimize(cache)?; + + // Then recursively optimize children. The new children vector is only allocated once a + // child actually changes, so fully-optimized subtrees cost no allocations. + let expr = current.as_ref().unwrap_or(self); + let children = expr.children(); + let mut new_children: Option> = None; + for (idx, child) in children.iter().enumerate() { + if let Some(optimized) = child.try_optimize_recursive_inner(cache)? { + new_children + .get_or_insert_with(|| children[..idx].to_vec()) + .push(optimized); + } else if let Some(new_children) = new_children.as_mut() { new_children.push(child.clone()); } } - if any_child_optimized { - current = current.with_children(new_children)?; - any_optimizations = true; + if let Some(new_children) = new_children { + let updated = expr.clone().with_children(new_children)?; // After updating children, try to optimize root again - if let Some(optimized) = current.clone().try_optimize(scope, cache)? { - current = optimized; - } + current = Some(updated.try_optimize(cache)?.unwrap_or(updated)); } - if any_optimizations { - Ok(Some(current)) - } else { - Ok(None) - } - } - - /// Simplify the expression, returning a potentially new expression. - /// - /// Deprecated: Use [`Expression::optimize_recursive`] instead, which iterates to convergence. - #[deprecated(note = "Use Expression::optimize_recursive instead")] - pub fn simplify(&self, scope: &DType) -> VortexResult { - self.optimize_recursive(scope) - } - - /// Simplify the expression without type information. - /// - /// Deprecated: Use [`Expression::optimize_recursive`] instead. - #[deprecated(note = "Use Expression::optimize_recursive instead")] - pub fn simplify_untyped(&self) -> VortexResult { - // For backwards compat, do a single bottom-up pass of untyped simplification - fn inner(expr: &Expression) -> VortexResult> { - let children: Vec<_> = expr.children().iter().map(inner).try_collect()?; - - if children.iter().any(|c| c.is_some()) { - let new_children: Vec<_> = children - .into_iter() - .zip(expr.children().iter()) - .map(|(new_c, old_c)| new_c.unwrap_or_else(|| old_c.clone())) - .collect(); - - let new_expr = expr.clone().with_children(new_children)?; - Ok(Some( - new_expr - .scalar_fn() - .simplify_untyped(&new_expr)? - .unwrap_or(new_expr), - )) - } else { - expr.scalar_fn().simplify_untyped(expr) - } - } - - let simplified = self - .scalar_fn() - .simplify_untyped(self)? - .unwrap_or_else(|| self.clone()); - - let simplified = inner(&simplified)?.unwrap_or(simplified); - let simplified = find_between(simplified); - - Ok(simplified) + Ok(current) } } @@ -226,10 +166,19 @@ struct SimplifyCache<'a> { dtype_cache: RefCell>, } +impl<'a> SimplifyCache<'a> { + fn new(scope: &'a DType) -> Self { + Self { + scope, + dtype_cache: RefCell::new(HashMap::new()), + } + } +} + impl SimplifyCtx for SimplifyCache<'_> { fn return_dtype(&self, expr: &Expression) -> VortexResult { // If the expression is "root", return the scope dtype - if expr.is::() { + if expr.is_root() { return Ok(self.scope.clone()); } @@ -243,7 +192,10 @@ impl SimplifyCtx for SimplifyCache<'_> { .iter() .map(|c| self.return_dtype(c)) .try_collect()?; - let dtype = expr.deref().return_dtype(&input_dtypes)?; + let dtype = expr + .as_scalar() + .ok_or_else(|| vortex_err!("cannot type a non-scalar expression: {expr}"))? + .return_dtype(&input_dtypes)?; self.dtype_cache .borrow_mut() .insert(expr.clone(), dtype.clone()); @@ -252,79 +204,24 @@ impl SimplifyCtx for SimplifyCache<'_> { } } -struct ExpressionReduceNode { - expression: Expression, - scope: DType, -} - -impl ReduceNode for ExpressionReduceNode { - fn as_any(&self) -> &dyn Any { - self - } - - fn node_dtype(&self) -> VortexResult { - self.expression.return_dtype(&self.scope) - } - - fn scalar_fn(&self) -> Option<&ScalarFnRef> { - Some(self.expression.scalar_fn()) - } - - fn child(&self, idx: usize) -> ReduceNodeRef { - Arc::new(ExpressionReduceNode { - expression: self.expression.child(idx).clone(), - scope: self.scope.clone(), - }) - } - - fn child_count(&self) -> usize { - self.expression.children().len() - } -} - -struct ExpressionReduceCtx { - scope: DType, -} -impl ReduceCtx for ExpressionReduceCtx { - fn new_node( - &self, - scalar_fn: ScalarFnRef, - children: &[ReduceNodeRef], - ) -> VortexResult { - let expression = Expression::try_new( - scalar_fn, - children - .iter() - .map(|c| { - c.as_any() - .downcast_ref::() - .vortex_expect("ReduceNode not an ExpressionReduceNode") - .expression - .clone() - }) - .collect::>(), - )?; - - Ok(Arc::new(ExpressionReduceNode { - expression, - scope: self.scope.clone(), - })) - } -} - #[cfg(test)] mod tests { use vortex_error::VortexResult; + use vortex_error::vortex_err; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; + use crate::expr::cast; use crate::expr::eq; use crate::expr::get_item; use crate::expr::lit; + use crate::expr::lt_eq; use crate::expr::or; use crate::expr::root; + use crate::scalar::Scalar; + use crate::scalar_fn::fns::literal::Literal; #[test] fn optimize_or_chain_correctness() -> VortexResult<()> { @@ -347,4 +244,32 @@ mod tests { assert!(s.contains("2i32") || s.contains('2'), "expected 2 in {s}"); Ok(()) } + + #[test] + fn optimize_folds_cast_of_literal_in_comparison() -> VortexResult<()> { + let expr = lt_eq( + get_item("x", root()), + cast( + lit(3i32), + DType::Primitive(PType::F64, Nullability::NonNullable), + ), + ); + let scope = DType::Struct( + StructFields::new( + ["x"].into(), + vec![DType::Primitive(PType::F64, Nullability::NonNullable)], + ), + Nullability::NonNullable, + ); + let optimized = expr.optimize_recursive(&scope)?; + + // Prune rules pattern-match a bare Literal on the comparison RHS; a cast wrapper + // silently disables pruning. + let rhs = optimized + .child(1) + .as_opt::() + .ok_or_else(|| vortex_err!("expected a bare literal RHS, got {optimized}"))?; + assert_eq!(rhs, &Scalar::primitive(3.0f64, Nullability::NonNullable)); + Ok(()) + } } diff --git a/vortex-array/src/expr/proto.rs b/vortex-array/src/expr/proto.rs index 0a8d71f2954..4f544fecafb 100644 --- a/vortex-array/src/expr/proto.rs +++ b/vortex-array/src/expr/proto.rs @@ -3,6 +3,7 @@ use itertools::Itertools; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_proto::expr as pb; use vortex_session::VortexSession; @@ -17,20 +18,36 @@ pub trait ExprSerializeProtoExt { fn serialize_proto(&self) -> VortexResult; } +/// The wire id for [`Expression::Root`], retained from when `Root` was a scalar function so that +/// already-serialized expressions keep round-tripping. +pub(crate) const ROOT_ID: &str = "vortex.root"; + impl ExprSerializeProtoExt for Expression { fn serialize_proto(&self) -> VortexResult { + let Some(scalar_fn) = self.as_scalar() else { + return Ok(pb::Expr { + id: ROOT_ID.to_string(), + children: vec![], + metadata: Some(vec![]), + }); + }; + let children = self .children() .iter() .map(|child| child.serialize_proto()) .try_collect()?; - let metadata = self.options().serialize()?.ok_or_else(|| { - vortex_err!("Expression '{}' is not serializable: {}", self.id(), self) + let metadata = scalar_fn.options().serialize()?.ok_or_else(|| { + vortex_err!( + "Expression '{}' is not serializable: {}", + scalar_fn.id(), + self + ) })?; Ok(pb::Expr { - id: self.id().to_string(), + id: scalar_fn.id().to_string(), children, metadata: Some(metadata), }) @@ -39,6 +56,16 @@ impl ExprSerializeProtoExt for Expression { impl Expression { pub fn from_proto(expr: &pb::Expr, session: &VortexSession) -> VortexResult { + // Root is not a registered scalar fn, so it must be resolved before the registry lookup. + if expr.id == ROOT_ID { + vortex_ensure!( + expr.children.is_empty(), + "root expression must have no children, got {}", + expr.children.len() + ); + return Ok(Expression::Root); + } + #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] let expr_id = ScalarFnId::new(expr.id.as_str()); let children = expr @@ -47,7 +74,7 @@ impl Expression { .map(|e| Expression::from_proto(e, session)) .collect::>>()?; - let scalar_fn = if let Some(vtable) = session.scalar_fns().registry().find(&expr_id) { + let scalar_fn = if let Some(vtable) = session.scalar_fns().registry().get(&expr_id) { vtable.deserialize(expr.metadata(), session)? } else if session.allows_unknown() { ForeignScalarFnVTable::make_scalar_fn(expr_id, expr.metadata().to_vec(), children.len()) @@ -126,7 +153,10 @@ mod tests { }; let expr = Expression::from_proto(&expr_proto, &session).unwrap(); - assert_eq!(expr.id().as_ref(), "vortex.test.foreign_scalar_fn"); + assert_eq!( + expr.as_scalar().map(|f| f.id().as_ref().to_string()), + Some("vortex.test.foreign_scalar_fn".to_string()) + ); let roundtrip = expr.serialize_proto().unwrap(); assert_eq!(roundtrip.id, expr_proto.id); diff --git a/vortex-array/src/expr/scope.rs b/vortex-array/src/expr/scope.rs new file mode 100644 index 00000000000..04de3202cf5 --- /dev/null +++ b/vortex-array/src/expr/scope.rs @@ -0,0 +1,45 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use crate::dtype::DType; + +/// The context an [`Expression`](crate::expr::Expression) is bound against. +/// +/// Today a scope is just the dtype that [`root`](crate::expr::root) resolves to. It is an opaque +/// struct rather than a bare [`DType`] so that lexical bindings can be added later without changing +/// [`Expression::bind_scope`](crate::expr::Expression::bind_scope)'s signature. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Scope { + root: DType, +} + +impl Scope { + /// Create a scope in which `root` resolves to the given dtype. + pub fn new(root: DType) -> Self { + Self { root } + } + + /// The dtype that `root` resolves to. + pub fn root(&self) -> &DType { + &self.root + } +} + +impl From for Scope { + fn from(root: DType) -> Self { + Self::new(root) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::dtype::Nullability; + + #[test] + fn root_round_trips() { + let dtype = DType::Bool(Nullability::Nullable); + assert_eq!(Scope::new(dtype.clone()).root(), &dtype); + assert_eq!(Scope::from(dtype.clone()).root(), &dtype); + } +} diff --git a/vortex-array/src/expr/stats/bound.rs b/vortex-array/src/expr/stats/bound.rs index c3b62033f4b..bf7852e829c 100644 --- a/vortex-array/src/expr/stats/bound.rs +++ b/vortex-array/src/expr/stats/bound.rs @@ -139,9 +139,7 @@ impl PartialOrd for LowerBound { fn partial_cmp(&self, other: &T) -> Option { match &self.0 { Exact(lhs) => lhs.partial_cmp(other), - Inexact(lhs) => lhs - .partial_cmp(other) - .and_then(|o| if o == Ordering::Less { None } else { Some(o) }), + Inexact(lhs) => lhs.partial_cmp(other).filter(|&o| o != Ordering::Less), Absent => None, } } @@ -241,13 +239,7 @@ impl PartialOrd for UpperBound { fn partial_cmp(&self, other: &T) -> Option { match &self.0 { Exact(lhs) => lhs.partial_cmp(other), - Inexact(lhs) => lhs.partial_cmp(other).and_then(|o| { - if o == Ordering::Greater { - None - } else { - Some(o) - } - }), + Inexact(lhs) => lhs.partial_cmp(other).filter(|&o| o != Ordering::Greater), Absent => None, } } diff --git a/vortex-array/src/expr/transform/bound_partition.rs b/vortex-array/src/expr/transform/bound_partition.rs new file mode 100644 index 00000000000..3272be862bc --- /dev/null +++ b/vortex-array/src/expr/transform/bound_partition.rs @@ -0,0 +1,474 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hash; + +use itertools::Itertools; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::dtype::DType; +use crate::dtype::FieldName; +use crate::dtype::FieldNames; +use crate::dtype::Nullability; +use crate::dtype::StructFields; +use crate::expr::BoundExpression; +use crate::expr::ExactBoundExpr; +use crate::expr::analysis::Annotation; +use crate::expr::analysis::AnnotationFn; +use crate::expr::analysis::BoundAnnotations; +use crate::expr::analysis::descendent_bound_annotations; +use crate::expr::bound::get_item; +use crate::expr::bound::pack; +use crate::expr::traversal::NodeExt; +use crate::expr::traversal::NodeRewriter; +use crate::expr::traversal::Transformed; +use crate::expr::traversal::TraversalOrder; + +/// Partition an expression into sub-expressions that are uniquely associated with an annotation. +/// A root expression is also returned that can be used to recombine the results of the partitions +/// into the result of the original expression. +/// +/// ## Note +/// +/// This function currently respects the validity of each field in the scope, but the not validity +/// of the scope itself. The fix would be for the returned `BoundPartitionedExpr` to include a +/// partition expression for computing the validity, or to include that expression as part of the +/// root. +/// +/// See . +pub fn partition_bound>( + expr: BoundExpression, + annotate_fn: A, +) -> VortexResult> +where + A::Annotation: Display, + FieldName: From, +{ + // Annotate each expression with the annotations that any of its descendent expressions have. + let annotations = descendent_bound_annotations(&expr, annotate_fn); + partition_bound_annotations(expr, annotations) +} + +/// Partition an already-annotated bound expression tree. +/// +/// Prefer [`partition_bound`] when annotations can be derived by an [`AnnotationFn`]. +pub fn partition_bound_annotations( + expr: BoundExpression, + annotations: BoundAnnotations, +) -> VortexResult> +where + A: Display + Clone + Eq + Hash, + FieldName: From, +{ + let mut collector = PartitionCollector::::new(&annotations); + expr.clone().rewrite(&mut collector)?; + + let mut partitions = Vec::with_capacity(collector.sub_expressions.len()); + let mut partition_annotations = Vec::with_capacity(collector.sub_expressions.len()); + + for (annotation, exprs) in collector.sub_expressions { + // We pack all sub-expressions for the same annotation into a single expression. + let names: FieldNames = exprs + .iter() + .enumerate() + .map(|(idx, _)| PartitionCollector::field_name(&annotation, idx)) + .collect(); + let expr = pack(names.into_iter().zip(exprs), Nullability::NonNullable); + + partitions.push(expr); + partition_annotations.push(annotation); + } + + let partition_names = partition_annotations + .iter() + .map(|id| FieldName::from(id.clone())) + .collect::(); + let root_scope = partition_root_dtype(&partition_names, &partitions); + let mut rewriter = PartitionRootRewriter::new(&annotations, root_scope); + let root = expr.rewrite(&mut rewriter)?.value; + + Ok(BoundPartitionedExpr { + root, + partitions: partitions.into_boxed_slice(), + partition_names, + partition_annotations: partition_annotations.into_boxed_slice(), + }) +} + +/// The result of partitioning a bound expression. +/// +/// The root and partitions remain bound so callers can cache and reuse their shared tree identity +/// without an unbind/rebind round trip. +#[derive(Debug)] +pub struct BoundPartitionedExpr { + /// The root expression used to re-assemble the results. + pub root: BoundExpression, + /// The partition expressions themselves. + pub partitions: Box<[BoundExpression]>, + /// The field name of each partition as referenced in the root expression. + pub partition_names: FieldNames, + /// The annotation associated with each partition. + pub partition_annotations: Box<[A]>, +} + +impl Display for BoundPartitionedExpr { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "root: {} {{{}}}", + self.root, + self.partition_names + .iter() + .zip(self.partitions.iter()) + .map(|(name, partition)| format!("{name}: {partition}")) + .join(", ") + ) + } +} + +impl BoundPartitionedExpr +where + FieldName: From, +{ + /// Return the partition for a given field, if it exists. + // FIXME(ngates): this should return an iterator since an annotation may have multiple partitions. + pub fn find_partition(&self, id: &A) -> Option<&BoundExpression> { + let id = FieldName::from(id.clone()); + self.partition_names + .iter() + .position(|field| field == id) + .map(|idx| &self.partitions[idx]) + } + + /// Replace the partition expressions and update every root dtype in the recombination tree. + pub fn replace_partitions(&mut self, partitions: Box<[BoundExpression]>) -> VortexResult<()> { + vortex_ensure!( + partitions.len() == self.partition_names.len(), + "Expected {} partitions, got {}", + self.partition_names.len(), + partitions.len() + ); + + let root_dtype = partition_root_dtype(&self.partition_names, &partitions); + let root = replace_root_dtype(self.root.clone(), root_dtype)?; + self.partitions = partitions; + self.root = root; + Ok(()) + } +} + +#[derive(Debug)] +struct PartitionCollector<'a, A: Annotation> { + annotations: &'a BoundAnnotations, + sub_expressions: HashMap>, +} + +impl<'a, A: Annotation + Display> PartitionCollector<'a, A> { + fn new(annotations: &'a BoundAnnotations) -> Self { + Self { + sub_expressions: HashMap::new(), + annotations, + } + } + + /// Each annotation may be associated with multiple sub-expressions, so we need to + /// a unique name for each sub-expression. + fn field_name(annotation: &A, idx: usize) -> FieldName { + format!("{annotation}_{idx}").into() + } +} + +impl NodeRewriter for PartitionCollector<'_, A> +where + FieldName: From, +{ + type NodeTy = BoundExpression; + + fn visit_down(&mut self, node: Self::NodeTy) -> VortexResult> { + match self.annotations.get(&ExactBoundExpr(node.clone())) { + // If this expression only accesses a single field, then we can skip the children + Some(annotations) if annotations.len() == 1 => { + let annotation = annotations + .iter() + .next() + .vortex_expect("expected one field"); + let sub_exprs = self.sub_expressions.entry(annotation.clone()).or_default(); + sub_exprs.push(node.clone()); + Ok(Transformed { + value: node, + changed: false, + order: TraversalOrder::Skip, + }) + } + + // Otherwise, continue traversing. + _ => Ok(Transformed::no(node)), + } + } + + fn visit_up(&mut self, node: Self::NodeTy) -> VortexResult> { + Ok(Transformed::no(node)) + } +} + +struct PartitionRootRewriter<'a, A: Annotation> { + annotations: &'a BoundAnnotations, + partition_offsets: HashMap, + root_dtype: DType, +} + +impl<'a, A: Annotation> PartitionRootRewriter<'a, A> { + fn new(annotations: &'a BoundAnnotations, root_dtype: DType) -> Self { + Self { + annotations, + partition_offsets: HashMap::new(), + root_dtype, + } + } +} + +impl NodeRewriter for PartitionRootRewriter<'_, A> +where + FieldName: From, +{ + type NodeTy = BoundExpression; + + fn visit_down(&mut self, node: Self::NodeTy) -> VortexResult> { + let Some(annotations) = self.annotations.get(&ExactBoundExpr(node.clone())) else { + return Ok(Transformed::no(node)); + }; + if annotations.len() != 1 { + return Ok(Transformed::no(node)); + } + + let annotation = annotations + .iter() + .next() + .vortex_expect("expected one annotation"); + let offset = self + .partition_offsets + .entry(annotation.clone()) + .or_default(); + let field_name = PartitionCollector::field_name(annotation, *offset); + *offset += 1; + + let partition = get_item( + FieldName::from(annotation.clone()), + BoundExpression::new_root(self.root_dtype.clone()), + ); + let value = get_item(field_name, partition); + + Ok(Transformed { + value, + changed: true, + order: TraversalOrder::Skip, + }) + } +} + +fn partition_root_dtype(names: &FieldNames, partitions: &[BoundExpression]) -> DType { + DType::Struct( + StructFields::new( + names.clone(), + partitions + .iter() + .map(|partition| partition.dtype().clone()) + .collect(), + ), + Nullability::NonNullable, + ) +} + +fn replace_root_dtype(expr: BoundExpression, root_dtype: DType) -> VortexResult { + Ok(expr + .transform_down(|node| { + if node.is_root() { + Ok(Transformed { + value: BoundExpression::new_root(root_dtype.clone()), + changed: true, + order: TraversalOrder::Skip, + }) + } else { + Ok(Transformed::no(node)) + } + })? + .into_inner()) +} + +#[cfg(test)] +mod tests { + use rstest::fixture; + use rstest::rstest; + + use super::*; + use crate::dtype::DType; + use crate::dtype::Nullability::NonNullable; + use crate::dtype::Nullability::Nullable; + use crate::dtype::PType::I32; + use crate::dtype::StructFields; + use crate::expr::analysis::make_bound_free_field_annotator; + use crate::expr::and; + use crate::expr::col; + use crate::expr::get_item; + use crate::expr::lit; + use crate::expr::merge; + use crate::expr::pack; + use crate::expr::root; + use crate::expr::transform::replace::replace_root_fields; + + #[fixture] + fn dtype() -> DType { + DType::Struct( + StructFields::from_iter([ + ( + "a", + DType::Struct( + StructFields::from_iter([("x", I32.into()), ("y", DType::from(I32))]), + NonNullable, + ), + ), + ("b", I32.into()), + ("c", I32.into()), + ]), + NonNullable, + ) + } + + fn partition_by_field( + expr: BoundExpression, + dtype: &DType, + ) -> VortexResult> { + let fields = dtype.as_struct_fields_opt().unwrap(); + partition_bound(expr, make_bound_free_field_annotator(fields)) + } + + #[rstest] + fn test_expr_top_level_ref(dtype: DType) { + let fields = dtype.as_struct_fields_opt().unwrap(); + + let expr = root(); + let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap(); + + // An un-expanded root expression is annotated by all fields, but since it is a single node + assert_eq!(partitioned.partitions.len(), 0); + assert_eq!(partitioned.root, root().bind(&dtype).unwrap()); + + // Instead, callers must expand the root expression themselves. + let expr = replace_root_fields(expr, fields); + let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap(); + + assert_eq!(partitioned.partitions.len(), fields.names().len()); + } + + #[rstest] + fn test_expr_top_level_ref_get_item_and_split(dtype: DType) { + let expr = get_item("y", get_item("a", root())); + + let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap(); + let root_dtype = + partition_root_dtype(&partitioned.partition_names, &partitioned.partitions); + assert_eq!( + partitioned.root, + get_item("a_0", get_item("a", root())) + .bind(&root_dtype) + .unwrap() + ); + } + + #[rstest] + fn test_expr_top_level_ref_get_item_and_split_pack(dtype: DType) { + let expr = pack( + [ + ("x", get_item("x", get_item("a", root()))), + ("y", get_item("y", get_item("a", root()))), + ("c", get_item("c", root())), + ], + NonNullable, + ); + let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap(); + + let split_a = partitioned.find_partition(&"a".into()).unwrap(); + assert_eq!( + split_a, + &pack( + [ + ("a_0", get_item("x", get_item("a", root()))), + ("a_1", get_item("y", get_item("a", root()))) + ], + NonNullable + ) + .bind(&dtype) + .unwrap() + ); + } + + #[rstest] + fn test_expr_top_level_ref_get_item_add(dtype: DType) { + let expr = and(get_item("y", get_item("a", root())), lit(1)); + let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap(); + + // Whole expr is a single split + assert_eq!(partitioned.partitions.len(), 1); + } + + #[rstest] + fn test_expr_top_level_ref_get_item_add_cannot_split(dtype: DType) { + let expr = and(get_item("y", get_item("a", root())), get_item("b", root())); + let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap(); + + // One for id.a and id.b + assert_eq!(partitioned.partitions.len(), 2); + } + + #[rstest] + fn test_expr_merge(dtype: DType) { + let expr = merge([col("a"), pack([("b", col("b"))], NonNullable)]); + + let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap(); + let expected = merge([get_item("a_0", col("a")), get_item("b_0", col("b"))]); + let root_dtype = + partition_root_dtype(&partitioned.partition_names, &partitioned.partitions); + assert_eq!( + partitioned.root, + expected.bind(&root_dtype).unwrap(), + "{} {}", + partitioned.root, + expected + ); + + assert_eq!(partitioned.partitions.len(), 2); + + let part_a = partitioned.find_partition(&"a".into()).unwrap(); + let expected_a = pack([("a_0", col("a"))], NonNullable); + assert_eq!( + part_a, + &expected_a.bind(&dtype).unwrap(), + "{part_a} {expected_a}" + ); + + let part_b = partitioned.find_partition(&"b".into()).unwrap(); + let expected_b = pack([("b_0", pack([("b", col("b"))], NonNullable))], NonNullable); + assert_eq!( + part_b, + &expected_b.bind(&dtype).unwrap(), + "{part_b} {expected_b}" + ); + } + + #[rstest] + fn replacing_partitions_refreshes_root_dtype(dtype: DType) -> VortexResult<()> { + let mut partitioned = partition_by_field(col("b").bind(&dtype)?, &dtype)?; + let field_dtype = DType::Primitive(I32, Nullable); + let replacement = pack([("b_0", root())], NonNullable).bind(&field_dtype)?; + + partitioned.replace_partitions(vec![replacement].into_boxed_slice())?; + + assert_eq!(partitioned.root.dtype(), &field_dtype); + Ok(()) + } +} diff --git a/vortex-array/src/expr/transform/coerce.rs b/vortex-array/src/expr/transform/coerce.rs deleted file mode 100644 index 1b6e9acd661..00000000000 --- a/vortex-array/src/expr/transform/coerce.rs +++ /dev/null @@ -1,194 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Expression-level type coercion pass. - -use vortex_error::VortexResult; - -use crate::dtype::DType; -use crate::expr::Expression; -use crate::expr::cast; -use crate::expr::traversal::NodeExt; -use crate::expr::traversal::Transformed; -use crate::scalar_fn::fns::literal::Literal; -use crate::scalar_fn::fns::root::Root; - -/// Rewrite an expression tree to insert casts where a scalar function's `coerce_args` demands -/// a different type than what the child currently produces. -/// -/// The rewrite is bottom-up: children are coerced first, then each parent node checks whether -/// its children match the coerced argument types. -pub fn coerce_expression(expr: Expression, scope: &DType) -> VortexResult { - // We capture scope by reference for the closure. - let scope = scope.clone(); - expr.transform_up(|node| { - // Leaf nodes (Root, Literal) have no children to coerce. - if node.is::() || node.is::() || node.children().is_empty() { - return Ok(Transformed::no(node)); - } - - // Compute the current child return types. - let child_dtypes: Vec = node - .children() - .iter() - .map(|c| c.return_dtype(&scope)) - .collect::>()?; - - // Ask the scalar function what types it wants. - let coerced_dtypes = node.scalar_fn().coerce_args(&child_dtypes)?; - - // If nothing changed, skip. - if child_dtypes == coerced_dtypes { - return Ok(Transformed::no(node)); - } - - // Build new children, inserting casts where needed. - let new_children: Vec = node - .children() - .iter() - .zip(coerced_dtypes.iter()) - .map(|(child, target)| { - let child_dtype = child.return_dtype(&scope)?; - if child_dtype.eq_ignore_nullability(target) - && child_dtype.nullability() == target.nullability() - { - Ok(child.clone()) - } else { - Ok(cast(child.clone(), target.clone())) - } - }) - .collect::>()?; - - let new_expr = node.with_children(new_children)?; - Ok(Transformed::yes(new_expr)) - }) - .map(|t| t.into_inner()) -} - -#[cfg(test)] -mod tests { - use vortex_error::VortexResult; - - use crate::dtype::DType; - use crate::dtype::Nullability::NonNullable; - use crate::dtype::PType; - use crate::dtype::StructFields; - use crate::expr::col; - use crate::expr::lit; - use crate::expr::transform::coerce::coerce_expression; - use crate::scalar::Scalar; - use crate::scalar_fn::ScalarFnVTableExt; - use crate::scalar_fn::fns::binary::Binary; - use crate::scalar_fn::fns::cast::Cast; - use crate::scalar_fn::fns::operators::Operator; - - fn test_scope() -> DType { - DType::Struct( - StructFields::new( - ["x", "y"].into(), - vec![ - DType::Primitive(PType::I32, NonNullable), - DType::Primitive(PType::I64, NonNullable), - ], - ), - NonNullable, - ) - } - - #[test] - fn mixed_type_comparison_inserts_cast() -> VortexResult<()> { - let scope = test_scope(); - // x (I32) < y (I64) => should cast x to I64 - let expr = Binary.new_expr(Operator::Lt, [col("x"), col("y")]); - let coerced = coerce_expression(expr, &scope)?; - - // The LHS child should now be a cast expression - assert!(coerced.child(0).is::()); - // The coerced LHS should return I64 - assert_eq!( - coerced.child(0).return_dtype(&scope)?, - DType::Primitive(PType::I64, NonNullable) - ); - // The RHS should be unchanged - assert!(!coerced.child(1).is::()); - Ok(()) - } - - #[test] - fn same_type_comparison_no_cast() -> VortexResult<()> { - let scope = test_scope(); - // x (I32) < x (I32) => no cast needed - let expr = Binary.new_expr(Operator::Lt, [col("x"), col("x")]); - let coerced = coerce_expression(expr, &scope)?; - - // Neither child should be a cast - assert!(!coerced.child(0).is::()); - assert!(!coerced.child(1).is::()); - Ok(()) - } - - #[test] - fn mixed_type_arithmetic_coerces_both() -> VortexResult<()> { - let scope = DType::Struct( - StructFields::new( - ["a", "b"].into(), - vec![ - DType::Primitive(PType::U8, NonNullable), - DType::Primitive(PType::I32, NonNullable), - ], - ), - NonNullable, - ); - // a (U8) + b (I32) => both should be coerced to I32 - // U8 + I32: unsigned_signed_supertype(U8, I32) => max(1,4)=4 => I64 - let expr = Binary.new_expr(Operator::Add, [col("a"), col("b")]); - let coerced = coerce_expression(expr, &scope)?; - - // LHS (U8) should be cast - assert!(coerced.child(0).is::()); - // Both should return the same supertype - let lhs_dt = coerced.child(0).return_dtype(&scope)?; - let rhs_dt = coerced.child(1).return_dtype(&scope)?; - assert_eq!(lhs_dt, rhs_dt); - Ok(()) - } - - #[test] - fn boolean_operators_no_coercion() -> VortexResult<()> { - let scope = DType::Struct( - StructFields::new( - ["p", "q"].into(), - vec![DType::Bool(NonNullable), DType::Bool(NonNullable)], - ), - NonNullable, - ); - let expr = Binary.new_expr(Operator::And, [col("p"), col("q")]); - let coerced = coerce_expression(expr, &scope)?; - - assert!(!coerced.child(0).is::()); - assert!(!coerced.child(1).is::()); - Ok(()) - } - - #[test] - fn literal_coercion() -> VortexResult<()> { - let scope = DType::Struct( - StructFields::new( - ["x"].into(), - vec![DType::Primitive(PType::I64, NonNullable)], - ), - NonNullable, - ); - // x (I64) + 1i32 => literal should be cast to I64 - let expr = Binary.new_expr(Operator::Add, [col("x"), lit(Scalar::from(1i32))]); - let coerced = coerce_expression(expr, &scope)?; - - // The RHS (literal) should be cast to I64 - assert!(coerced.child(1).is::()); - assert_eq!( - coerced.child(1).return_dtype(&scope)?, - DType::Primitive(PType::I64, NonNullable) - ); - Ok(()) - } -} diff --git a/vortex-array/src/expr/transform/match_between.rs b/vortex-array/src/expr/transform/match_between.rs index 56f03dbac4d..7cb427596f4 100644 --- a/vortex-array/src/expr/transform/match_between.rs +++ b/vortex-array/src/expr/transform/match_between.rs @@ -125,7 +125,18 @@ fn is_strict_comparison(op: Operator) -> Option { #[cfg(test)] mod tests { + use vortex_buffer::buffer; + use vortex_error::VortexResult; + use super::find_between; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::BoolArray; + use crate::arrays::StructArray; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; use crate::expr::and; use crate::expr::between; use crate::expr::col; @@ -134,9 +145,42 @@ mod tests { use crate::expr::lit; use crate::expr::lt; use crate::expr::lt_eq; + use crate::scalar::Scalar; use crate::scalar_fn::fns::between::BetweenOptions; use crate::scalar_fn::fns::between::StrictComparison; + /// A null literal bound must not change the values of the rewritten expression. Kleene `AND` + /// keeps a row false when the surviving comparison is false, so the rewrite cannot null it. + #[test] + fn test_null_literal_bound_is_value_preserving() -> VortexResult<()> { + let session = array_session(); + let ctx = &mut session.create_execution_ctx(); + let data = StructArray::from_fields(&[("x", buffer![10, 1].into_array())])?.into_array(); + + let null_lit = lit(Scalar::null(DType::Primitive( + PType::I32, + Nullability::Nullable, + ))); + let expr = and(gt_eq(col("x"), null_lit), lt_eq(col("x"), lit(5i32))); + + let before = data + .clone() + .apply(&expr)? + .execute::(ctx)? + .opt_bool_vec(ctx); + + let after = data + .apply(&find_between(expr))? + .execute::(ctx)? + .opt_bool_vec(ctx); + + // Row 0 is false rather than null because `$.x <= 5` falsifies it on its own. + assert_eq!(before, [Some(false), None]); + assert_eq!(before, after); + + Ok(()) + } + #[test] fn test_bad_match() { // An impossible expression diff --git a/vortex-array/src/expr/transform/mod.rs b/vortex-array/src/expr/transform/mod.rs index 6013241d971..713029497ba 100644 --- a/vortex-array/src/expr/transform/mod.rs +++ b/vortex-array/src/expr/transform/mod.rs @@ -1,12 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! A collection of transformations that can be applied to a [`crate::expr::Expression`]. -mod coerce; +//! Transformations for [`crate::expr::Expression`] and [`crate::expr::BoundExpression`] trees. +mod bound_partition; pub(crate) mod match_between; mod partition; mod replace; -pub use coerce::*; +pub use bound_partition::*; pub use partition::*; pub use replace::*; diff --git a/vortex-array/src/expr/traversal/mod.rs b/vortex-array/src/expr/traversal/mod.rs index b000439082c..0f55600afe5 100644 --- a/vortex-array/src/expr/traversal/mod.rs +++ b/vortex-array/src/expr/traversal/mod.rs @@ -17,12 +17,12 @@ pub use fold::FoldDownContext; pub use fold::FoldUp; pub use fold::NodeFolder; pub use fold::NodeFolderContext; -use itertools::Itertools; pub use references::ReferenceCollector; pub use visitor::pre_order_visit_down; pub use visitor::pre_order_visit_up; use vortex_error::VortexResult; +use crate::expr::BoundExpression; use crate::expr::Expression; use crate::expr::traversal::fold::NodeFolderContextWrapper; @@ -499,19 +499,14 @@ impl Node for Expression { &'a self, mut f: F, ) -> VortexResult { - self.children().as_ref().apply_elements(&mut f) + self.children().apply_ref_elements(&mut f) } fn map_children VortexResult>>( self, f: F, ) -> VortexResult> { - let transformed = self - .children() - .iter() - .cloned() - .collect_vec() - .map_elements(f)?; + let transformed = self.children().to_vec().map_elements(f)?; if transformed.changed { Ok(Transformed { @@ -533,6 +528,83 @@ impl Node for Expression { } } +impl Node for BoundExpression { + fn apply_children<'a, F: FnMut(&'a Self) -> VortexResult>( + &'a self, + mut f: F, + ) -> VortexResult { + let BoundExpression::Scalar { children, .. } = self else { + return Ok(TraversalOrder::Continue); + }; + + for child in children.iter() { + match f(child)? { + TraversalOrder::Continue | TraversalOrder::Skip => {} + TraversalOrder::Stop => return Ok(TraversalOrder::Stop), + } + } + + Ok(TraversalOrder::Continue) + } + + fn map_children VortexResult>>( + self, + mut f: F, + ) -> VortexResult> { + let BoundExpression::Scalar { children, .. } = &self else { + return Ok(Transformed::no(self)); + }; + + let mut order = TraversalOrder::Continue; + // Stays `None` until a child actually changes. Most nodes of a rewritten tree are + // untouched. + let mut rewritten: Option> = None; + + for (index, child) in children.iter().enumerate() { + let value = match order { + TraversalOrder::Continue | TraversalOrder::Skip => { + let result = f(child.clone())?; + order = result.order; + if result.changed && rewritten.is_none() { + let mut prefix = Vec::with_capacity(children.len()); + prefix.extend_from_slice(&children[..index]); + rewritten = Some(prefix); + } + result.value + } + TraversalOrder::Stop => child.clone(), + }; + + if let Some(rewritten) = &mut rewritten { + rewritten.push(value); + } + } + + match rewritten { + Some(rewritten) => Ok(Transformed { + value: self.with_children(rewritten)?, + order, + changed: true, + }), + None => Ok(Transformed::no(self)), + } + } + + fn iter_children(&self, f: impl FnOnce(&mut dyn Iterator) -> T) -> T { + match self { + BoundExpression::Scalar { children, .. } => f(&mut children.iter()), + BoundExpression::Root { .. } => f(&mut std::iter::empty()), + } + } + + fn children_count(&self) -> usize { + match self { + BoundExpression::Scalar { children, .. } => children.len(), + BoundExpression::Root { .. } => 0, + } + } +} + #[cfg(test)] mod tests { use vortex_error::VortexResult; diff --git a/vortex-array/src/expression.rs b/vortex-array/src/expression.rs index d6b8739abf6..d0590f2bf58 100644 --- a/vortex-array/src/expression.rs +++ b/vortex-array/src/expression.rs @@ -2,22 +2,49 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::ArrayRef; use crate::IntoArray; use crate::arrays::ConstantArray; use crate::arrays::ScalarFnArray; +use crate::expr::BoundExpression; use crate::expr::Expression; use crate::optimizer::ArrayOptimizer; use crate::scalar_fn::fns::literal::Literal; -use crate::scalar_fn::fns::root::Root; impl ArrayRef { + /// Apply a bound expression to this array, producing a new array in constant time. + pub fn apply_bound(self, expr: &BoundExpression) -> VortexResult { + let BoundExpression::Scalar { + scalar_fn, + children, + .. + } = expr + else { + return Ok(self); + }; + + if let Some(scalar) = scalar_fn.as_opt::() { + return Ok(ConstantArray::new(scalar.clone(), self.len()).into_array()); + } + + let children: Vec<_> = children + .iter() + .map(|child| self.clone().apply_bound(child)) + .try_collect()?; + + let array = + ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, self.len())?.into_array(); + + array.optimize() + } + /// Apply the expression to this array, producing a new array in constant time. pub fn apply(self, expr: &Expression) -> VortexResult { // If the expression is a root, return self. - if expr.is::() { + if expr.is_root() { return Ok(self); } @@ -34,9 +61,11 @@ impl ArrayRef { .try_collect()?; // And wrap the scalar function up in an array. + let scalar_fn = expr + .as_scalar() + .vortex_expect("root and literal were handled above, so this is a scalar node"); let array = - ScalarFnArray::try_new_with_len(expr.scalar_fn().clone(), children, self.len())? - .into_array(); + ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, self.len())?.into_array(); // Optimize the resulting array's root. array.optimize() diff --git a/vortex-array/src/extension/datetime/date.rs b/vortex-array/src/extension/datetime/date.rs index 9eda23dd2d1..8cbe5b019c2 100644 --- a/vortex-array/src/extension/datetime/date.rs +++ b/vortex-array/src/extension/datetime/date.rs @@ -109,29 +109,6 @@ impl ExtVTable for Date { Ok(()) } - fn can_coerce_from(ext_dtype: &ExtDType, other: &DType) -> bool { - let DType::Extension(other_ext) = other else { - return false; - }; - let Some(other_unit) = other_ext.metadata_opt::() else { - return false; - }; - let our_unit = ext_dtype.metadata(); - // We can coerce from other if our unit is finer (<=) and nullability is compatible. - our_unit <= other_unit && (ext_dtype.storage_dtype().is_nullable() || !other.is_nullable()) - } - - fn least_supertype(ext_dtype: &ExtDType, other: &DType) -> Option { - let DType::Extension(other_ext) = other else { - return None; - }; - let other_unit = other_ext.metadata_opt::()?; - let our_unit = ext_dtype.metadata(); - let finest = (*our_unit).min(*other_unit); - let union_null = ext_dtype.storage_dtype().nullability() | other.nullability(); - Some(DType::Extension(Date::new(finest, union_null).erased())) - } - fn unpack_native<'a>( ext_dtype: &'a ExtDType, storage_value: &'a ScalarValue, @@ -192,27 +169,6 @@ mod tests { assert_eq!(format!("{}", scalar.as_extension()), "1971-01-01"); } - #[test] - fn least_supertype_date_units() { - use crate::dtype::Nullability::NonNullable; - - let days = DType::Extension(Date::new(TimeUnit::Days, NonNullable).erased()); - let ms = DType::Extension(Date::new(TimeUnit::Milliseconds, NonNullable).erased()); - let expected = DType::Extension(Date::new(TimeUnit::Milliseconds, NonNullable).erased()); - assert_eq!(days.least_supertype(&ms).unwrap(), expected); - assert_eq!(ms.least_supertype(&days).unwrap(), expected); - } - - #[test] - fn can_coerce_from_date() { - use crate::dtype::Nullability::NonNullable; - - let days = DType::Extension(Date::new(TimeUnit::Days, NonNullable).erased()); - let ms = DType::Extension(Date::new(TimeUnit::Milliseconds, NonNullable).erased()); - assert!(ms.can_coerce_from(&days)); - assert!(!days.can_coerce_from(&ms)); - } - #[test] fn deserialize_empty_metadata_returns_error() { use crate::dtype::extension::ExtVTable; diff --git a/vortex-array/src/extension/datetime/time.rs b/vortex-array/src/extension/datetime/time.rs index 88455a121b5..125845e758c 100644 --- a/vortex-array/src/extension/datetime/time.rs +++ b/vortex-array/src/extension/datetime/time.rs @@ -110,28 +110,6 @@ impl ExtVTable for Time { Ok(()) } - fn can_coerce_from(ext_dtype: &ExtDType, other: &DType) -> bool { - let DType::Extension(other_ext) = other else { - return false; - }; - let Some(other_unit) = other_ext.metadata_opt::

(ctx, [needle])? else { - return Ok(None); - }; - let Some(value_max) = max(needle, ctx) else { return Ok(None); }; @@ -463,19 +446,8 @@ fn list_contains_falsify( return Ok(None); }; - // A null element cannot make `list_contains` true — the equality form maps - // it to `false` through its own null fill — so it is not a value a zone has - // to be checked against. Dropping nulls rather than abandoning the whole - // list is what lets `x IN (1, 2, NULL, 3)` prune at all. - let mut sorted: Vec<&Scalar> = elements.iter().filter(|e| !e.is_null()).collect(); - - // Sorting is what lets the whole list be described by a handful of - // intervals instead of one term per element, and the intervals are only - // sound if that order is the one the statistics comparison uses. Whether - // values order at all is a property of their dtype, which every element of - // a list shares, so one incomparable pair means none of them order and the - // list proves nothing. Floats do order here: they compare totally, so `NaN` - // and the two signed zeros have positions rather than being incomparable. + let mut sorted: Vec<&crate::scalar::Scalar> = + elements.iter().filter(|e| !e.is_null()).collect(); sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(Ordering::Equal)); if sorted .windows(2) @@ -483,26 +455,12 @@ fn list_contains_falsify( { return Ok(None); } - - // Equal values have nothing between them, so they cannot bound a gap. - // Collapsing them before the boundaries are chosen is what stops a list of - // mostly duplicates from spending every boundary inside one repeated value: - // `[0; 32] + [1000]` has exactly one gap, and sampling the sorted list - // directly would place all 31 boundaries inside the run of zeroes and miss - // it. sorted.dedup_by(|a, b| (**a).partial_cmp(&**b) == Some(Ordering::Equal)); - // Every element was null, so nothing orders and nothing is provable. let (Some(lowest), Some(highest)) = (sorted.first(), sorted.last()) else { return Ok(None); }; - // A scope is proven free of the list when it lies wholly outside the list's - // range, or wholly inside one of the gaps between two adjacent list values — - // nothing orders between an adjacent pair, so a row strictly between them - // cannot be in the list. The gaps are what make a clustered list prunable: - // a list split into a low group and a high group leaves most of the column - // *inside* its range, where the outer bound alone proves nothing. let mut terms = Vec::with_capacity((sorted.len() + 1).min(FALSIFY_MAX_GAP_TERMS + 2)); terms.push(lt(value_max.clone(), lit((*lowest).clone()))); terms.push(gt(value_min.clone(), lit((*highest).clone()))); @@ -517,10 +475,8 @@ fn list_contains_falsify( )); } - // The two outer-range terms are pushed unconditionally, so there is always - // something to reduce. let value_predicate = or_collect(terms).vortex_expect("terms is never empty"); - Ok(Some(guard.apply(value_predicate))) + with_non_nan_guards::

(ctx, [needle], value_predicate) } #[derive(Debug)] @@ -533,9 +489,9 @@ impl StatsRewriteRule for DynamicComparisonNanCountStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { dynamic_comparison_falsify::(expr, ctx) } } @@ -550,17 +506,17 @@ impl StatsRewriteRule for DynamicComparisonAllNonNanStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { dynamic_comparison_falsify::(expr, ctx) } } fn dynamic_comparison_falsify( - expr: &Expression, + expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>, -) -> VortexResult> { +) -> VortexResult> { let dynamic = expr.as_::(); let lhs = expr.child(0); @@ -574,47 +530,47 @@ fn dynamic_comparison_falsify( return Ok(None); }; - let value_predicate = DynamicComparison.new_expr( + let value_predicate = dynamic_with_options( DynamicComparisonExpr { operator, rhs: Arc::clone(&dynamic.rhs), default: !dynamic.default, }, - [lhs_stat], + lhs_stat, ); with_non_nan_guards::

(ctx, [lhs], value_predicate) } -fn min(expr: &Expression, ctx: &StatsRewriteCtx<'_>) -> Option { +fn min(expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>) -> Option { stat_expr(expr, Stat::Min, ctx) } -fn max(expr: &Expression, ctx: &StatsRewriteCtx<'_>) -> Option { +fn max(expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>) -> Option { stat_expr(expr, Stat::Max, ctx) } -fn null_count(expr: &Expression, ctx: &StatsRewriteCtx<'_>) -> Option { +fn null_count(expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>) -> Option { stat_expr(expr, Stat::NullCount, ctx) } -fn all_null(expr: &Expression) -> Expression { +fn all_null(expr: &BoundExpression) -> BoundExpression { stat_fn(expr.clone(), AllNull.bind(AggregateEmptyOptions)) } -fn all_non_null(expr: &Expression) -> Expression { +fn all_non_null(expr: &BoundExpression) -> BoundExpression { stat_fn(expr.clone(), AllNonNull.bind(AggregateEmptyOptions)) } enum NanCheck { NotNeeded, - Check(Expression), + Check(BoundExpression), Unavailable, } trait NonNanProof { const EMIT_UNGUARDED_REWRITES: bool; - fn check(ctx: &StatsRewriteCtx<'_>, expr: &Expression) -> VortexResult; + fn check(ctx: &StatsRewriteCtx<'_>, expr: &BoundExpression) -> VortexResult; } struct NanCountProof; @@ -622,7 +578,7 @@ struct NanCountProof; impl NonNanProof for NanCountProof { const EMIT_UNGUARDED_REWRITES: bool = true; - fn check(ctx: &StatsRewriteCtx<'_>, expr: &Expression) -> VortexResult { + fn check(ctx: &StatsRewriteCtx<'_>, expr: &BoundExpression) -> VortexResult { non_nan_check(ctx, expr, |expr| { match stat_expr(expr, Stat::NaNCount, ctx) { Some(nan_count) => NanCheck::Check(eq(nan_count, lit(0u64))), @@ -637,7 +593,7 @@ struct AllNonNanProof; impl NonNanProof for AllNonNanProof { const EMIT_UNGUARDED_REWRITES: bool = false; - fn check(ctx: &StatsRewriteCtx<'_>, expr: &Expression) -> VortexResult { + fn check(ctx: &StatsRewriteCtx<'_>, expr: &BoundExpression) -> VortexResult { non_nan_check(ctx, expr, |expr| { NanCheck::Check(stat_fn(expr.clone(), AllNonNan.bind(AggregateEmptyOptions))) }) @@ -649,14 +605,14 @@ impl NonNanProof for AllNonNanProof { // from float to non-float still needs a proof about the float source values. fn non_nan_check( ctx: &StatsRewriteCtx<'_>, - expr: &Expression, - proof: impl FnOnce(&Expression) -> NanCheck, + expr: &BoundExpression, + proof: impl FnOnce(&BoundExpression) -> NanCheck, ) -> VortexResult { if let Some(scalar) = expr.as_opt::() { - let Some(value) = scalar.as_primitive_opt() else { + if !scalar.dtype().is_float() { return Ok(NanCheck::NotNeeded); - }; - return Ok(if value.is_nan() { + } + return Ok(if scalar.as_primitive().is_nan() { NanCheck::Check(lit(false)) } else { NanCheck::NotNeeded @@ -679,10 +635,14 @@ fn non_nan_check( } fn has_nans(dtype: &DType) -> bool { - matches!(dtype, DType::Primitive(ptype, _) if ptype.is_float()) + dtype.is_float() } -fn stat_expr(expr: &Expression, stat: Stat, ctx: &StatsRewriteCtx<'_>) -> Option { +fn stat_expr( + expr: &BoundExpression, + stat: Stat, + ctx: &StatsRewriteCtx<'_>, +) -> Option { if let Some(literal) = literal_stat(expr, stat) { return Some(literal); } @@ -709,36 +669,11 @@ fn stat_expr(expr: &Expression, stat: Stat, ctx: &StatsRewriteCtx<'_>) -> Option .then(|| stat_fn(expr.clone(), aggregate_fn)) } -/// The NaN guard a rewrite carries, resolved independently of the value -/// predicate it will guard. -enum NanGuard { - /// The value predicate has to be conjoined with this check. - Check(Expression), - /// No possible NaN-bearing expression remains, so the value predicate is - /// already guarded and stands alone. - Unguarded, -} - -impl NanGuard { - fn apply(self, value_predicate: Expression) -> Expression { - match self { - Self::Check(check) => and(check, value_predicate), - Self::Unguarded => value_predicate, - } - } -} - -/// Resolves the guard for `exprs`, or `None` when this proof cannot guard them -/// and the rule must therefore emit nothing. -/// -/// Worth resolving before the value predicate is built rather than after: two -/// rules are registered per function and differ only in this proof, so for any -/// given needle one of them commonly emits nothing. Asking first is what stops -/// it from building a predicate it will discard. -fn non_nan_guard<'a, P: NonNanProof>( +fn with_non_nan_guards<'a, P: NonNanProof>( ctx: &StatsRewriteCtx<'_>, - exprs: impl IntoIterator, -) -> VortexResult> { + exprs: impl IntoIterator, + value_predicate: BoundExpression, +) -> VortexResult> { let mut nan_checks = Vec::new(); for expr in exprs { match P::check(ctx, expr)? { @@ -747,25 +682,19 @@ fn non_nan_guard<'a, P: NonNanProof>( NanCheck::Unavailable => return Ok(None), } } - - Ok(match and_collect(nan_checks) { - Some(check) => Some(NanGuard::Check(check)), - // Only one registered rule emits an unguarded rewrite, so a non-float - // comparison is not rewritten twice. - None if P::EMIT_UNGUARDED_REWRITES => Some(NanGuard::Unguarded), + let nan_predicate = and_collect(nan_checks); + + Ok(match nan_predicate { + Some(nan_check) => Some(and(nan_check, value_predicate)), + // No possible NaN-bearing expression remains, so the value predicate is + // already guarded. Only one registered rule emits this unguarded + // rewrite so non-float comparisons are not duplicated. + None if P::EMIT_UNGUARDED_REWRITES => Some(value_predicate), None => None, }) } -fn with_non_nan_guards<'a, P: NonNanProof>( - ctx: &StatsRewriteCtx<'_>, - exprs: impl IntoIterator, - value_predicate: Expression, -) -> VortexResult> { - Ok(non_nan_guard::

(ctx, exprs)?.map(|guard| guard.apply(value_predicate))) -} - -fn literal_stat(expr: &Expression, stat: Stat) -> Option { +fn literal_stat(expr: &BoundExpression, stat: Stat) -> Option { let scalar = expr.as_opt::()?; match stat { Stat::Min | Stat::Max => Some(lit(scalar.clone())), @@ -787,11 +716,11 @@ fn literal_stat(expr: &Expression, stat: Stat) -> Option { } fn cast_stat( - expr: &Expression, + expr: &BoundExpression, dtype: &DType, stat: Stat, ctx: &StatsRewriteCtx<'_>, -) -> Option { +) -> Option { match stat { Stat::Min | Stat::Max => stat_expr(expr, stat, ctx).map(|stat| cast(stat, dtype.clone())), Stat::NaNCount | Stat::Sum | Stat::UncompressedSizeInBytes => stat_expr(expr, stat, ctx), @@ -799,30 +728,30 @@ fn cast_stat( } } -fn stat_fn(expr: Expression, aggregate_fn: AggregateFnRef) -> Expression { - StatFn.new_expr(StatOptions::new(aggregate_fn), [expr]) +fn stat_fn(expr: BoundExpression, aggregate_fn: AggregateFnRef) -> BoundExpression { + stat(expr, aggregate_fn) } #[cfg(test)] mod tests { use std::sync::Arc; use std::sync::LazyLock; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; use vortex_error::VortexResult; use vortex_session::VortexSession; - use super::StatFn; - use super::StatOptions; - use super::all_non_null; - use super::all_null; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTableExt; use crate::aggregate_fn::EmptyOptions as AggregateEmptyOptions; use crate::aggregate_fn::fns::all_non_nan::AllNonNan; + use crate::array_session; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; + use crate::expr::BoundExpression; use crate::expr::Expression; use crate::expr::and; use crate::expr::between; @@ -843,15 +772,23 @@ mod tests { use crate::expr::stats::Stat; use crate::scalar::Scalar; use crate::scalar_fn::EmptyOptions; + use crate::scalar_fn::ScalarFnId; + use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; use crate::scalar_fn::fns::between::BetweenOptions; use crate::scalar_fn::fns::between::StrictComparison; + use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::dynamic::DynamicComparison; use crate::scalar_fn::fns::dynamic::DynamicComparisonExpr; use crate::scalar_fn::fns::operators::CompareOperator; use crate::scalar_fn::internal::row_count::RowCount; + use crate::stats::expr::StatFn; + use crate::stats::expr::StatOptions; + use crate::stats::rewrite::StatsRewriteCtx; + use crate::stats::rewrite::StatsRewriteRule; + use crate::stats::session::StatsSessionExt; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = LazyLock::new(array_session); fn stat(expr: Expression, stat: Stat) -> Expression { let aggregate_fn = stat.aggregate_fn().expect("stat should have aggregate fn"); @@ -883,12 +820,30 @@ mod tests { ) } - fn falsify(expr: &Expression) -> VortexResult> { - expr.falsify(&test_scope(), &SESSION) + fn falsify(expr: &Expression) -> VortexResult> { + expr.bind(&test_scope())?.falsify(&SESSION) + } + + fn satisfy(expr: &Expression) -> VortexResult> { + expr.bind(&test_scope())?.satisfy(&SESSION) + } + + fn bind_expected(expr: Option) -> VortexResult> { + expr.map(|expr| expr.bind(&test_scope())).transpose() + } + + fn all_null(expr: &Expression) -> Expression { + crate::stats::all_null(expr.clone()) } - fn satisfy(expr: &Expression) -> VortexResult> { - expr.satisfy(&test_scope(), &SESSION) + fn all_non_null(expr: &Expression) -> Expression { + crate::stats::all_non_null(expr.clone()) + } + + macro_rules! assert_rewrite_eq { + ($actual:expr, $expected:expr) => { + assert_eq!($actual, bind_expected($expected)?) + }; } fn nan_guarded(expr: Expression, value_predicate: Expression) -> Expression { @@ -907,13 +862,13 @@ mod tests { #[test] fn rewrites_comparison_falsifier() -> VortexResult<()> { let expr = gt(col("a"), lit(10)); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(lt_eq(stat(col("a"), Stat::Max), lit(10))) ); let expr = eq(col("a"), col("b")); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(or( gt(stat(col("a"), Stat::Min), stat(col("b"), Stat::Max)), @@ -922,7 +877,7 @@ mod tests { ); let expr = eq(col("s"), col("t")); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(or( gt(stat(col("s"), Stat::Min), stat(col("t"), Stat::Max)), @@ -935,13 +890,60 @@ mod tests { #[test] fn rewrites_boolean_falsifiers() -> VortexResult<()> { let expr = and(gt(col("a"), lit(10)), lt(col("a"), lit(50))); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(or( lt_eq(stat(col("a"), Stat::Max), lit(10)), gt_eq(stat(col("a"), Stat::Min), lit(50)), )) ); + + let expr = or(gt(col("a"), lit(10)), lt(col("a"), lit(5))); + assert_rewrite_eq!( + falsify(&expr)?, + Some(and( + lt_eq(stat(col("a"), Stat::Max), lit(10)), + gt_eq(stat(col("a"), Stat::Min), lit(5)), + )) + ); + Ok(()) + } + + /// Counts how many times the stats rewrite visits a `Binary` node. + #[derive(Debug)] + struct BinaryVisitCounter(Arc); + + impl StatsRewriteRule for BinaryVisitCounter { + fn scalar_fn_id(&self) -> ScalarFnId { + Binary.id() + } + + fn falsify( + &self, + _expr: &BoundExpression, + _ctx: &StatsRewriteCtx<'_>, + ) -> VortexResult> { + self.0.fetch_add(1, Ordering::Relaxed); + Ok(None) + } + } + + #[test] + fn or_chain_falsify_visits_each_node_once() -> VortexResult<()> { + let session = array_session(); + let visits = Arc::new(AtomicUsize::new(0)); + session + .stats() + .register_rewrite(BinaryVisitCounter(Arc::clone(&visits))); + + let expr = (1..16).fold(eq(col("a"), lit(0)), |chain, i| { + or(chain, eq(col("a"), lit(i))) + }); + let falsifier = expr.bind(&test_scope())?.falsify(&session)?; + assert!(falsifier.is_some()); + + // One visit per `Binary` node: 16 comparisons plus 15 `or`s. + assert_eq!(visits.load(Ordering::Relaxed), 31); Ok(()) } @@ -957,7 +959,7 @@ mod tests { }, ); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(or( gt(lit(10), stat(col("a"), Stat::Max)), @@ -969,7 +971,7 @@ mod tests { #[test] fn rewrites_null_falsifiers() -> VortexResult<()> { - assert_eq!( + assert_rewrite_eq!( falsify(&is_null(col("a")))?, Some(or( eq(stat(col("a"), Stat::NullCount), lit(0u64)), @@ -977,7 +979,7 @@ mod tests { )) ); - assert_eq!( + assert_rewrite_eq!( falsify(&is_not_null(col("a")))?, Some(or( eq( @@ -992,7 +994,7 @@ mod tests { #[test] fn rewrites_null_satisfiers() -> VortexResult<()> { - assert_eq!( + assert_rewrite_eq!( satisfy(&is_null(col("a")))?, Some(or( eq( @@ -1003,7 +1005,7 @@ mod tests { )) ); - assert_eq!( + assert_rewrite_eq!( satisfy(&is_not_null(col("a")))?, Some(or( eq(stat(col("a"), Stat::NullCount), lit(0u64)), @@ -1026,7 +1028,7 @@ mod tests { // outside the list's range, or wholly inside a gap between two adjacent // list values. For a list this dense the gaps cover every element, so // it proves exactly what a term per element would. - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(or( or( @@ -1051,7 +1053,7 @@ mod tests { #[test] fn rewrites_like_falsifier() -> VortexResult<()> { let expr = like(col("s"), lit("prefix%")); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(or( gt_eq(stat(col("s"), Stat::Min), lit("prefiy")), @@ -1060,7 +1062,7 @@ mod tests { ); let expr = like(col("s"), lit(r"\%%")); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(or( gt_eq(stat(col("s"), Stat::Min), lit("&")), @@ -1069,7 +1071,7 @@ mod tests { ); let expr = like(col("s"), lit("pref%ix%")); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(or( gt_eq(stat(col("s"), Stat::Min), lit("preg")), @@ -1078,7 +1080,7 @@ mod tests { ); let expr = like(col("s"), lit("pref_ix_")); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(or( gt_eq(stat(col("s"), Stat::Min), lit("preg")), @@ -1087,7 +1089,7 @@ mod tests { ); let expr = like(col("s"), lit("exact")); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(or( gt(stat(col("s"), Stat::Min), lit("exact")), @@ -1096,7 +1098,7 @@ mod tests { ); let expr = like(col("s"), lit("%suffix")); - assert_eq!(falsify(&expr)?, None); + assert_rewrite_eq!(falsify(&expr)?, None); Ok(()) } @@ -1111,7 +1113,7 @@ mod tests { ); let dynamic = expr.as_::(); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(DynamicComparison.new_expr( DynamicComparisonExpr { @@ -1130,7 +1132,7 @@ mod tests { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let expr = gt(cast(col("f"), dtype.clone()), lit(5i32)); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(nan_guarded( col("f"), @@ -1148,8 +1150,8 @@ mod tests { nested_struct_dtype(), vec![Scalar::primitive(1.0f32, Nullability::Nullable)], ); - assert_eq!(falsify(<_eq(col("n"), lit(struct_scalar.clone())))?, None); - assert_eq!(falsify(&eq(col("n"), lit(struct_scalar)))?, None); + assert_rewrite_eq!(falsify(<_eq(col("n"), lit(struct_scalar.clone())))?, None); + assert_rewrite_eq!(falsify(&eq(col("n"), lit(struct_scalar)))?, None); Ok(()) } @@ -1158,7 +1160,7 @@ mod tests { let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); let expr = eq(cast(col("a"), dtype.clone()), lit(42i64)); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(or( gt(cast(stat(col("a"), Stat::Min), dtype.clone()), lit(42i64)), diff --git a/vortex-array/src/stats/stats_set.rs b/vortex-array/src/stats/stats_set.rs index cb4af0b2c7f..d8a27788f85 100644 --- a/vortex-array/src/stats/stats_set.rs +++ b/vortex-array/src/stats/stats_set.rs @@ -443,30 +443,33 @@ impl MutTypedStatsSetRef<'_, '_> { other.get_scalar_bound::(), ) { (Some(m1), Some(m2)) => { - // If the combine sum is exact, then we can sum them. - if let Some(scalar_value) = - m1.zip(m2).as_exact().and_then(|(s1, s2)| match s1.dtype() { - DType::Primitive(..) => s1 - .as_primitive() - .checked_add(&s2.as_primitive()) - .and_then(|pscalar| pscalar.pvalue().map(ScalarValue::Primitive)), - DType::Decimal(..) => s1 - .as_decimal() - .checked_binary_numeric( - &s2.as_decimal(), - crate::scalar::NumericOperator::Add, - ) - .map(|scalar| { - ScalarValue::Decimal( - scalar - .decimal_value() - .vortex_expect("no decimal value in scalar"), - ) - }), - _ => None, - }) - { - self.set(Stat::Sum, Precision::Exact(scalar_value)); + // If the combined sum is exact, then we can sum them. + let merged = m1.zip(m2).as_exact().and_then(|(s1, s2)| match s1.dtype() { + DType::Primitive(..) => s1 + .as_primitive() + .checked_add(&s2.as_primitive()) + .and_then(|pscalar| pscalar.pvalue().map(ScalarValue::Primitive)), + // Add widens the result precision, so the merged sum is only exact as a + // stat if it still fits the summands' own decimal type. + DType::Decimal(decimal_dtype, _) => s1 + .as_decimal() + .checked_binary_numeric( + &s2.as_decimal(), + crate::scalar::NumericOperator::Add, + ) + .ok() + .flatten() + .and_then(|scalar| scalar.as_decimal().decimal_value()) + .filter(|value| value.fits_in_precision(*decimal_dtype)) + .map(ScalarValue::Decimal), + _ => None, + }); + + match merged { + Some(scalar_value) => self.set(Stat::Sum, Precision::Exact(scalar_value)), + // An overflow, an inexact bound or a non-numeric dtype leaves the combined + // sum unknown. Keeping this set's own sum would report a partial total. + None => self.clear(Stat::Sum), } } _ => self.clear(Stat::Sum), @@ -565,13 +568,19 @@ mod test { use crate::array_session; use crate::arrays::PrimitiveArray; use crate::dtype::DType; + use crate::dtype::DecimalDType; + use crate::dtype::MAX_PRECISION; + use crate::dtype::NativeDecimalType; use crate::dtype::Nullability; use crate::dtype::PType; + use crate::dtype::i256; use crate::expr::stats::IsConstant; use crate::expr::stats::Precision; use crate::expr::stats::Stat; use crate::expr::stats::StatsProvider; use crate::expr::stats::StatsProviderExt; + use crate::scalar::DecimalValue; + use crate::scalar::ScalarValue; use crate::stats::StatsSet; use crate::stats::stats_set::Scalar; @@ -631,6 +640,32 @@ mod test { ); } + #[test] + fn merge_sums_overflow_clears() { + // A sum that cannot be combined exactly leaves the merged set with no Sum at all; + // retaining this set's own sum would report a partial total as the combined one. + let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); + let merged = StatsSet::of(Stat::Sum, Precision::exact(i64::MAX)) + .merge_ordered(&StatsSet::of(Stat::Sum, Precision::exact(i64::MAX)), &dtype); + + assert!(merged.get(Stat::Sum).is_absent()); + } + + #[test] + fn merge_decimal_sums_out_of_precision_clears() { + // A decimal(70, 0) array's Sum stat is itself decimal(76, 0): `sum_decimal_dtype` adds + // ten digits of headroom but saturates at MAX_PRECISION, so at this input precision there + // is none left and two maximal sums cannot be combined exactly. + let dtype = DType::Decimal(DecimalDType::new(70, 0), Nullability::NonNullable); + let max = ScalarValue::Decimal(DecimalValue::I256( + ::MAX_BY_PRECISION[usize::from(MAX_PRECISION)], + )); + let merged = StatsSet::of(Stat::Sum, Precision::exact(max.clone())) + .merge_ordered(&StatsSet::of(Stat::Sum, Precision::exact(max)), &dtype); + + assert!(merged.get(Stat::Sum).is_absent()); + } + #[test] fn merge_constant() { let first = StatsSet::from_iter([ diff --git a/vortex-array/src/test_harness.rs b/vortex-array/src/test_harness/mod.rs similarity index 97% rename from vortex-array/src/test_harness.rs rename to vortex-array/src/test_harness/mod.rs index 98be3a85ce4..d3a6b829e80 100644 --- a/vortex-array/src/test_harness.rs +++ b/vortex-array/src/test_harness/mod.rs @@ -12,6 +12,9 @@ use crate::ExecutionCtx; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; +#[cfg(not(codspeed))] +pub mod trace; + /// Check that a named metadata matches its previous versioning. /// /// Goldenfile takes care of checking for equality against a checked-in file. diff --git a/vortex-array/src/test_harness/trace/mod.rs b/vortex-array/src/test_harness/trace/mod.rs new file mode 100644 index 00000000000..1e053cf5a7a --- /dev/null +++ b/vortex-array/src/test_harness/trace/mod.rs @@ -0,0 +1,1259 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Snapshot-friendly tracing harness for the optimizer and executor. +//! +//! # What this records +//! +//! [`trace_op`] runs a closure with a thread-local recorder installed. While the recorder is +//! active, calls to the `trace_op!` macro inside the optimizer and executor push +//! structured events into the recorder. The recorder produces a [`TraceDisplay`] that renders +//! as a deterministic, hierarchical text trace suitable for `insta` snapshot assertions. +//! +//! Events cover: +//! +//! - **Optimization**: optimize/recursive-optimize entry, fixpoint loop iterations, applied +//! reduce rules, applied parent-reduce rules. +//! - **Execution**: `execute_until` iterations, single-step entries, parent kernel attempts and +//! matches, slot transitions, builder start/append/finish, and the eventual canonical output. +//! +//! Despite the name `trace_op`, the harness is *not* a generic logging facility: it is closely +//! coupled to the optimizer/executor state machines so that the resulting trace is stable enough +//! to commit as a snapshot. +//! +//! # When to use it +//! +//! Use [`trace_op`] to write a regression test that asserts on the sequence of optimizer +//! rewrites or executor steps an array goes through. Typical scenarios: +//! +//! - A reduce rule should fire exactly once on a specific input shape. +//! - A parent kernel should be tried in a specific order and the first match should win. +//! - The executor should walk into a slot, finish it, and pop back to the parent without +//! building a canonical intermediate. +//! - A chunked array should drive the builder path rather than the stack path. +//! +//! Two resolutions are available: +//! +//! - [`TraceResolution::ExecutedOnly`] (default) — only events that actually fired (rule +//! rewrites that matched, kernels that succeeded, execution steps that ran). Optimizer +//! passes that produced no change are elided. +//! - [`TraceResolution::Attempts`] — also records declined rule attempts, kernels that did +//! not match, and per-loop bookkeeping. Use this when ordering or fall-through matters. +//! +//! # Cost and scope +//! +//! - Capture is thread-local. Worker threads spawned inside `f` do not inherit the recorder. +//! - Nested captures return an error so that unrelated traces never merge. +//! - In release builds and CodSpeed benchmark builds, every `trace_op!` invocation is compiled +//! away by the macro's `cfg` gating; this module is then unused. See +//! `trace_op!` for the gating rules. +//! +//! # Example +//! +//! ```ignore +//! use vortex_array::test_harness::trace::trace_op; +//! +//! let traced = trace_op(|| filter_array.optimize())?; +//! assert!(traced.output.is::()); +//! insta::assert_snapshot!(traced.trace.to_string(), @r" +//! optimize root=vortex.filter(i32, len=4) session=false +//! reduce TrivialFilterRule: vortex.filter(i32, len=4) -> vortex.primitive(i32, len=4) +//! done output=vortex.primitive(i32, len=4) +//! "); +//! ``` + +use std::cell::Cell; +use std::cell::RefCell; +use std::fmt; +use std::fmt::Debug; +use std::fmt::Display; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::ArrayRef; + +/// Controls how much rule and kernel resolution detail is captured. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum TraceResolution { + /// Record only the operations that actually executed. + #[default] + ExecutedOnly, + /// Also record rule and kernel attempts that matched but declined, or did not match. + Attempts, +} + +/// Options for [`trace_op_with`]. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct TraceOptions { + /// The amount of rule and kernel resolution detail to include. + pub resolution: TraceResolution, +} + +/// The result of a traced operation. +#[derive(Clone, Debug)] +pub struct Traced { + /// The value returned by the traced closure. + pub output: T, + /// A stable, snapshot-friendly rendering of optimizer and execution activity. + pub trace: TraceDisplay, +} + +/// A stable, snapshot-friendly trace. +#[derive(Clone, Debug, Default)] +pub struct TraceDisplay { + options: TraceOptions, + events: Vec, +} + +impl Display for TraceDisplay { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let hidden_events = self.hidden_events(); + let mut optimize_depth = 0usize; + let mut wrote_event = false; + + for (idx, event) in self.events.iter().enumerate() { + if hidden_events[idx] { + continue; + } + + if event.closes_before(self.options.resolution) { + optimize_depth = optimize_depth.saturating_sub(1); + } + + if event.is_hidden(self.options.resolution) { + continue; + } + + if wrote_event { + writeln!(f)?; + } else { + wrote_event = true; + } + + write_indent( + f, + optimize_depth + event.relative_indent(self.options.resolution, optimize_depth > 0), + )?; + event.fmt_line(f, self.options.resolution)?; + + if event.opens_after(self.options.resolution) { + optimize_depth += 1; + } + if event.closes_after(self.options.resolution) { + optimize_depth = optimize_depth.saturating_sub(1); + } + } + Ok(()) + } +} + +impl TraceDisplay { + fn hidden_events(&self) -> Vec { + let mut hidden = vec![false; self.events.len()]; + if self.options.resolution != TraceResolution::ExecutedOnly { + return hidden; + } + + let mut optimize_stack = Vec::new(); + for (idx, event) in self.events.iter().enumerate() { + match event { + TraceEvent::OptimizeStart { .. } => optimize_stack.push(idx), + TraceEvent::OptimizeDone { changed, .. } => { + let Some(start) = optimize_stack.pop() else { + continue; + }; + if !changed { + hidden[start..=idx].fill(true); + } + } + _ => {} + } + } + hidden + } +} + +fn write_indent(f: &mut fmt::Formatter<'_>, depth: usize) -> fmt::Result { + for _ in 0..depth { + f.write_str(" ")?; + } + Ok(()) +} + +/// Run `f` while capturing a trace of the optimizer and executor work it performs. +/// +/// `f` typically invokes an operation that drives the executor or optimizer, such as +/// [`ArrayOptimizer::optimize`][crate::optimizer::ArrayOptimizer::optimize] or +/// [`ArrayRef::execute`][crate::ArrayRef::execute]. While `f` runs, +/// the optimizer and executor emit structured events via the `trace_op!` +/// macro into a thread-local recorder. When `f` returns, the recorder is finalized and +/// returned alongside the closure's output as a [`Traced`]. +/// +/// The default resolution ([`TraceResolution::ExecutedOnly`]) records the rule rewrites, +/// parent kernels, execution steps, and builder activity that actually executed. Optimizer +/// passes that produced no change are hidden from the rendered trace. Use [`trace_op_with`] +/// with [`TraceResolution::Attempts`] when a test needs to assert on declined rule attempts, +/// kernels that did not match, or other fall-through detail. +/// +/// # Examples +/// +/// ```ignore +/// let traced = trace_op(|| filter_array.optimize())?; +/// assert!(traced.output.is::()); +/// insta::assert_snapshot!(traced.trace.to_string(), @r" +/// optimize root=vortex.filter(i32, len=4) session=false +/// reduce TrivialFilterRule: vortex.filter(i32, len=4) -> vortex.primitive(i32, len=4) +/// done output=vortex.primitive(i32, len=4) +/// "); +/// ``` +/// +/// # Errors +/// +/// Returns whatever error `f` produces. Returns an error if a recorder is already active on +/// the current thread — nested traces are not supported. +pub fn trace_op(f: impl FnOnce() -> VortexResult) -> VortexResult> { + trace_op_with(TraceOptions::default(), f) +} + +/// Run `f` while capturing a trace using `options`. +/// +/// See [`trace_op`] for the common case. Use this when you need to override the default +/// [`TraceResolution`] to capture declined rules and unmatched kernels. +/// +/// Trace capture is thread-local and intentionally does not propagate to worker threads. Nested +/// trace captures return an error so tests do not accidentally merge unrelated traces. +pub fn trace_op_with( + options: TraceOptions, + f: impl FnOnce() -> VortexResult, +) -> VortexResult> { + let interest = TraceInterest::from(options.resolution); + ACTIVE_TRACE.with(|active| { + let mut active = active.borrow_mut(); + if active.is_some() { + return Err(vortex_err!("trace_op captures cannot be nested")); + } + *active = Some(TraceRecorder::new(options)); + Ok(()) + })?; + TRACE_INTEREST.with(|trace_interest| trace_interest.set(interest)); + ACTIVE_TRACE_COUNT.fetch_add(1, Ordering::Relaxed); + if interest == TraceInterest::Attempts { + ATTEMPTS_TRACE_COUNT.fetch_add(1, Ordering::Relaxed); + } + + let guard = ActiveTraceGuard { interest }; + let output = f(); + let recorder = ACTIVE_TRACE.with(|active| { + active + .borrow_mut() + .take() + .vortex_expect("trace recorder must be installed") + }); + drop(guard); + + output.map(|output| Traced { + output, + trace: recorder.finish(), + }) +} + +/// Returns true when the current thread has an active trace recorder. +#[allow(clippy::inline_always)] +#[inline(always)] +pub(crate) fn is_active() -> bool { + if ACTIVE_TRACE_COUNT.load(Ordering::Relaxed) == 0 { + return false; + } + TRACE_INTEREST.with(|interest| interest.get().is_active()) +} + +#[allow(clippy::inline_always)] +#[inline(always)] +fn attempts_enabled() -> bool { + if ATTEMPTS_TRACE_COUNT.load(Ordering::Relaxed) == 0 { + return false; + } + TRACE_INTEREST.with(|interest| interest.get() == TraceInterest::Attempts) +} + +#[derive(Clone, Copy, Debug)] +enum TraceSource { + Static, + Session(usize), +} + +impl Display for TraceSource { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + TraceSource::Static => f.write_str("static"), + TraceSource::Session(idx) => write!(f, "session[{idx}]"), + } + } +} + +#[derive(Clone, Copy, Debug)] +enum AttemptOutcome { + Declined, + NoMatch, +} + +impl Display for AttemptOutcome { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + AttemptOutcome::Declined => f.write_str("declined"), + AttemptOutcome::NoMatch => f.write_str("no-match"), + } + } +} + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +enum TraceInterest { + #[default] + Off, + ExecutedOnly, + Attempts, +} + +impl TraceInterest { + #[inline] + fn is_active(self) -> bool { + self != Self::Off + } +} + +impl From for TraceInterest { + fn from(resolution: TraceResolution) -> Self { + match resolution { + TraceResolution::ExecutedOnly => Self::ExecutedOnly, + TraceResolution::Attempts => Self::Attempts, + } + } +} + +/// Snapshot-friendly wrapper around [`ArrayRef`] that renders the encoding, dtype, and length +/// using the canonical [`Display`] format (`vortex.primitive(i32, len=4)`). +/// +/// Carries a clone of the [`ArrayRef`] instead of duplicating fields, +/// so trace events stay small and share the same rendering as every other `{array}` print in +/// the codebase. +#[derive(Clone, Debug)] +pub(crate) struct ArraySummary(ArrayRef); + +impl ArraySummary { + pub(crate) fn new(array: &ArrayRef) -> Self { + Self(array.clone()) + } +} + +impl Display for ArraySummary { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + Display::fmt(&self.0, f) + } +} + +pub(crate) fn record_optimize_start(root: &ArrayRef, session: bool) { + record(TraceEvent::OptimizeStart { + root: ArraySummary::new(root), + session, + }); +} + +pub(crate) fn record_optimize_loop_start(array: &ArrayRef) { + if !attempts_enabled() { + return; + } + record(TraceEvent::OptimizeLoopStart { + array: ArraySummary::new(array), + }); +} + +pub(crate) fn record_optimize_loop_end() { + if !attempts_enabled() { + return; + } + record(TraceEvent::OptimizeLoopEnd); +} + +pub(crate) fn record_optimize_reduce_none(array: &ArrayRef) { + if !attempts_enabled() { + return; + } + record(TraceEvent::PhaseNone { + indent: 0, + phase: "reduce", + subject: "array", + array: ArraySummary::new(array), + }); +} + +pub(crate) fn record_optimize_parent_reduce_none(array: &ArrayRef) { + if !attempts_enabled() { + return; + } + record(TraceEvent::PhaseNone { + indent: 0, + phase: "reduce_parent", + subject: "array", + array: ArraySummary::new(array), + }); +} + +pub(crate) fn record_optimize_done(output: &ArrayRef, changed: bool) { + record(TraceEvent::OptimizeDone { + output: ArraySummary::new(output), + changed, + }); +} + +pub(crate) fn record_optimize_recursive_start(root: &ArrayRef) { + record(TraceEvent::OptimizeRecursiveStart { + root: ArraySummary::new(root), + }); +} + +pub(crate) fn record_optimize_recursive_slot(slot_idx: usize, input: &ArrayRef, output: &ArrayRef) { + record(TraceEvent::OptimizeRecursiveSlot { + slot_idx, + input: ArraySummary::new(input), + output: ArraySummary::new(output), + }); +} + +pub(crate) fn record_reduce_applied(array: &ArrayRef, rule: &dyn Debug, output: &ArrayRef) { + record(TraceEvent::ReduceApplied { + array: ArraySummary::new(array), + rule: compact_label(rule), + output: ArraySummary::new(output), + }); +} + +pub(crate) fn record_reduce_declined(array: &ArrayRef, rule: &dyn Debug) { + if !attempts_enabled() { + return; + } + record(TraceEvent::ReduceAttempt { + array: ArraySummary::new(array), + rule: compact_label(rule), + outcome: AttemptOutcome::Declined, + }); +} + +pub(crate) fn record_session_parent_reduce_applied( + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, + plugin_idx: usize, + output: &ArrayRef, +) { + record_parent_reduce_applied( + parent, + child, + slot_idx, + TraceSource::Session(plugin_idx), + "reduce_parent_fn", + output, + ); +} + +pub(crate) fn record_session_parent_reduce_declined( + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, + plugin_idx: usize, +) { + record_parent_reduce_attempt( + parent, + child, + slot_idx, + TraceSource::Session(plugin_idx), + "reduce_parent_fn", + AttemptOutcome::Declined, + ); +} + +pub(crate) fn record_static_parent_reduce_no_match( + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, + rule: &dyn Debug, +) { + record_parent_reduce_attempt( + parent, + child, + slot_idx, + TraceSource::Static, + compact_label(rule), + AttemptOutcome::NoMatch, + ); +} + +pub(crate) fn record_static_parent_reduce_applied( + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, + rule: &dyn Debug, + output: &ArrayRef, +) { + record_parent_reduce_applied( + parent, + child, + slot_idx, + TraceSource::Static, + compact_label(rule), + output, + ); +} + +pub(crate) fn record_static_parent_reduce_declined( + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, + rule: &dyn Debug, +) { + record_parent_reduce_attempt( + parent, + child, + slot_idx, + TraceSource::Static, + compact_label(rule), + AttemptOutcome::Declined, + ); +} + +fn record_parent_reduce_attempt( + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, + source: TraceSource, + rule: impl Into, + outcome: AttemptOutcome, +) { + if !attempts_enabled() { + return; + } + record(TraceEvent::ParentReduceAttempt { + parent: ArraySummary::new(parent), + child: ArraySummary::new(child), + slot_idx, + source, + rule: rule.into(), + outcome, + }); +} + +fn record_parent_reduce_applied( + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, + source: TraceSource, + rule: impl Into, + output: &ArrayRef, +) { + record(TraceEvent::ParentReduceApplied { + parent: ArraySummary::new(parent), + child: ArraySummary::new(child), + slot_idx, + source, + rule: rule.into(), + output: ArraySummary::new(output), + }); +} + +pub(crate) fn record_execute_until_start(root: &ArrayRef) { + record(TraceEvent::ExecuteUntilStart { + target: short_type_name::(), + root: ArraySummary::new(root), + }); +} + +pub(crate) fn record_execute_until_iteration( + iteration: usize, + current: &ArrayRef, + stack_parent: Option<(&ArrayRef, usize)>, + builder_active: bool, +) { + record(TraceEvent::ExecuteUntilIteration { + iteration, + current: ArraySummary::new(current), + stack_parent: stack_parent.map(|(array, slot_idx)| (ArraySummary::new(array), slot_idx)), + builder_active, + }); +} + +pub(crate) fn record_execute_until_done_check(target: bool, canonical: bool) { + if !attempts_enabled() { + return; + } + record(TraceEvent::ExecuteUntilDoneCheck { target, canonical }); +} + +pub(crate) fn record_execute_until_return(output: &ArrayRef) { + record(TraceEvent::ExecuteUntilReturn { + output: ArraySummary::new(output), + }); +} + +pub(crate) fn record_execute_until_pop_frame(slot_idx: usize, output: &ArrayRef) { + record(TraceEvent::ExecuteUntilPopFrame { + slot_idx, + output: ArraySummary::new(output), + }); +} + +pub(crate) fn record_session_execute_parent_applied( + phase: &'static str, + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, + plugin_idx: usize, + output: &ArrayRef, +) { + record_execute_parent_applied( + phase, + parent, + child, + slot_idx, + TraceSource::Session(plugin_idx), + "execute_parent_fn", + output, + ); +} + +pub(crate) fn record_session_execute_parent_declined( + phase: &'static str, + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, + plugin_idx: usize, +) { + record_execute_parent_attempt( + phase, + parent, + child, + slot_idx, + TraceSource::Session(plugin_idx), + "execute_parent_fn", + AttemptOutcome::Declined, + ); +} + +fn record_execute_parent_attempt( + phase: &'static str, + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, + source: TraceSource, + kernel: impl Into, + outcome: AttemptOutcome, +) { + if !attempts_enabled() { + return; + } + record(TraceEvent::ExecuteParentAttempt { + phase, + parent: ArraySummary::new(parent), + child: ArraySummary::new(child), + slot_idx, + source, + kernel: kernel.into(), + outcome, + }); +} + +fn record_execute_parent_applied( + phase: &'static str, + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, + source: TraceSource, + kernel: impl Into, + output: &ArrayRef, +) { + record(TraceEvent::ExecuteParentApplied { + phase, + parent: ArraySummary::new(parent), + child: ArraySummary::new(child), + slot_idx, + source, + kernel: kernel.into(), + output: ArraySummary::new(output), + }); +} + +pub(crate) fn record_execute_parent_none(phase: &'static str, current: &ArrayRef) { + if !attempts_enabled() { + return; + } + record(TraceEvent::PhaseNone { + indent: 2, + phase, + subject: "current", + array: ArraySummary::new(current), + }); +} + +pub(crate) fn record_execute_optimized(input: &ArrayRef, output: &ArrayRef) { + let changed = !ArrayRef::ptr_eq(input, output); + if !changed && !attempts_enabled() { + return; + } + record(TraceEvent::ExecuteOptimized { + input: ArraySummary::new(input), + output: ArraySummary::new(output), + changed, + }); +} + +pub(crate) fn record_execute_encoding(array: &ArrayRef) { + if !attempts_enabled() { + return; + } + record(TraceEvent::ExecuteEncoding { + array: ArraySummary::new(array), + }); +} + +pub(crate) fn record_execute_slot(slot_idx: usize, parent: &ArrayRef, child: &ArrayRef) { + record(TraceEvent::SlotTransition { + step: "ExecuteSlot", + slot_idx, + parent: ArraySummary::new(parent), + child: ArraySummary::new(child), + }); +} + +pub(crate) fn record_builder_start(array: &ArrayRef) { + record(TraceEvent::BuilderEvent { + action: "start", + subject: "array", + array: ArraySummary::new(array), + }); +} + +pub(crate) fn record_append_child(slot_idx: usize, parent: &ArrayRef, child: &ArrayRef) { + record(TraceEvent::SlotTransition { + step: "AppendChild", + slot_idx, + parent: ArraySummary::new(parent), + child: ArraySummary::new(child), + }); +} + +pub(crate) fn record_builder_append(child: &ArrayRef) { + record(TraceEvent::BuilderEvent { + action: "append", + subject: "child", + array: ArraySummary::new(child), + }); +} + +pub(crate) fn record_execute_done(array: &ArrayRef) { + record(TraceEvent::ExecuteDone { + array: ArraySummary::new(array), + }); +} + +pub(crate) fn record_builder_finish(output: &ArrayRef) { + record(TraceEvent::BuilderEvent { + action: "finish", + subject: "output", + array: ArraySummary::new(output), + }); +} + +pub(crate) fn record_single_step_start(array: &ArrayRef) { + record(TraceEvent::SingleStepStart { + array: ArraySummary::new(array), + }); +} + +pub(crate) fn record_single_step_phase_none(phase: &'static str, array: &ArrayRef) { + if !attempts_enabled() { + return; + } + record(TraceEvent::PhaseNone { + indent: 1, + phase, + subject: "array", + array: ArraySummary::new(array), + }); +} + +pub(crate) fn record_single_step_applied(phase: &'static str, input: &ArrayRef, output: &ArrayRef) { + record(TraceEvent::SingleStepApplied { + phase, + input: ArraySummary::new(input), + output: ArraySummary::new(output), + }); +} + +fn record(event: TraceEvent) { + ACTIVE_TRACE.with(|active| { + if let Some(recorder) = active.borrow_mut().as_mut() { + recorder.events.push(event); + } + }); +} + +fn compact_label(value: &dyn Debug) -> String { + let label = format!("{value:?}"); + if let Some(label) = adapter_field(&label, "rule") { + return label.to_string(); + } + if let Some(label) = adapter_field(&label, "kernel") { + return label.to_string(); + } + label +} + +fn adapter_field<'a>(label: &'a str, field: &str) -> Option<&'a str> { + let marker = format!("{field}: "); + let start = label.find(&marker)? + marker.len(); + let rest = &label[start..]; + let end = rest.rfind(" }")?; + Some(&rest[..end]) +} + +fn short_type_name() -> String { + std::any::type_name::() + .rsplit("::") + .next() + .vortex_expect("type names are never empty") + .to_string() +} + +thread_local! { + static TRACE_INTEREST: Cell = const { Cell::new(TraceInterest::Off) }; + static ACTIVE_TRACE: RefCell> = const { RefCell::new(None) }; +} + +static ACTIVE_TRACE_COUNT: AtomicUsize = AtomicUsize::new(0); +static ATTEMPTS_TRACE_COUNT: AtomicUsize = AtomicUsize::new(0); + +struct ActiveTraceGuard { + interest: TraceInterest, +} + +impl Drop for ActiveTraceGuard { + fn drop(&mut self) { + if self.interest == TraceInterest::Attempts { + ATTEMPTS_TRACE_COUNT.fetch_sub(1, Ordering::Relaxed); + } + ACTIVE_TRACE_COUNT.fetch_sub(1, Ordering::Relaxed); + TRACE_INTEREST.with(|interest| interest.set(TraceInterest::Off)); + ACTIVE_TRACE.with(|active| { + active.borrow_mut().take(); + }); + } +} + +#[derive(Debug)] +struct TraceRecorder { + options: TraceOptions, + events: Vec, +} + +impl TraceRecorder { + fn new(options: TraceOptions) -> Self { + Self { + options, + events: Vec::new(), + } + } + + fn finish(self) -> TraceDisplay { + TraceDisplay { + options: self.options, + events: self.events, + } + } +} + +#[derive(Clone, Debug)] +enum TraceEvent { + OptimizeStart { + root: ArraySummary, + session: bool, + }, + OptimizeLoopStart { + array: ArraySummary, + }, + OptimizeLoopEnd, + OptimizeDone { + output: ArraySummary, + changed: bool, + }, + OptimizeRecursiveStart { + root: ArraySummary, + }, + OptimizeRecursiveSlot { + slot_idx: usize, + input: ArraySummary, + output: ArraySummary, + }, + ReduceAttempt { + array: ArraySummary, + rule: String, + outcome: AttemptOutcome, + }, + ReduceApplied { + array: ArraySummary, + rule: String, + output: ArraySummary, + }, + ParentReduceAttempt { + parent: ArraySummary, + child: ArraySummary, + slot_idx: usize, + source: TraceSource, + rule: String, + outcome: AttemptOutcome, + }, + ParentReduceApplied { + parent: ArraySummary, + child: ArraySummary, + slot_idx: usize, + source: TraceSource, + rule: String, + output: ArraySummary, + }, + ExecuteUntilStart { + target: String, + root: ArraySummary, + }, + ExecuteUntilIteration { + iteration: usize, + current: ArraySummary, + stack_parent: Option<(ArraySummary, usize)>, + builder_active: bool, + }, + ExecuteUntilDoneCheck { + target: bool, + canonical: bool, + }, + ExecuteUntilReturn { + output: ArraySummary, + }, + ExecuteUntilPopFrame { + slot_idx: usize, + output: ArraySummary, + }, + ExecuteParentAttempt { + phase: &'static str, + parent: ArraySummary, + child: ArraySummary, + slot_idx: usize, + source: TraceSource, + kernel: String, + outcome: AttemptOutcome, + }, + ExecuteParentApplied { + phase: &'static str, + parent: ArraySummary, + child: ArraySummary, + slot_idx: usize, + source: TraceSource, + kernel: String, + output: ArraySummary, + }, + PhaseNone { + indent: usize, + phase: &'static str, + subject: &'static str, + array: ArraySummary, + }, + ExecuteOptimized { + input: ArraySummary, + output: ArraySummary, + changed: bool, + }, + ExecuteEncoding { + array: ArraySummary, + }, + SlotTransition { + step: &'static str, + slot_idx: usize, + parent: ArraySummary, + child: ArraySummary, + }, + BuilderEvent { + action: &'static str, + subject: &'static str, + array: ArraySummary, + }, + ExecuteDone { + array: ArraySummary, + }, + SingleStepStart { + array: ArraySummary, + }, + SingleStepApplied { + phase: &'static str, + input: ArraySummary, + output: ArraySummary, + }, +} + +impl TraceEvent { + fn is_hidden(&self, resolution: TraceResolution) -> bool { + match resolution { + TraceResolution::Attempts => matches!(self, TraceEvent::OptimizeLoopEnd), + TraceResolution::ExecutedOnly => matches!( + self, + TraceEvent::OptimizeLoopStart { .. } + | TraceEvent::OptimizeLoopEnd + | TraceEvent::PhaseNone { .. } + | TraceEvent::ExecuteUntilDoneCheck { .. } + | TraceEvent::ExecuteEncoding { .. } + | TraceEvent::ExecuteOptimized { changed: false, .. } + | TraceEvent::ExecuteParentAttempt { .. } + | TraceEvent::ReduceAttempt { .. } + | TraceEvent::ParentReduceAttempt { .. } + ), + } + } + + fn opens_after(&self, resolution: TraceResolution) -> bool { + match resolution { + TraceResolution::Attempts => matches!( + self, + TraceEvent::OptimizeStart { .. } | TraceEvent::OptimizeLoopStart { .. } + ), + TraceResolution::ExecutedOnly => matches!(self, TraceEvent::OptimizeStart { .. }), + } + } + + fn closes_before(&self, resolution: TraceResolution) -> bool { + match resolution { + TraceResolution::Attempts => matches!(self, TraceEvent::OptimizeLoopEnd), + TraceResolution::ExecutedOnly => false, + } + } + + fn closes_after(&self, _resolution: TraceResolution) -> bool { + matches!(self, TraceEvent::OptimizeDone { .. }) + } + + fn relative_indent(&self, _resolution: TraceResolution, in_optimize_scope: bool) -> usize { + match self { + TraceEvent::OptimizeStart { .. } + | TraceEvent::OptimizeLoopStart { .. } + | TraceEvent::OptimizeDone { .. } => 0, + TraceEvent::ReduceAttempt { .. } + | TraceEvent::ReduceApplied { .. } + | TraceEvent::ParentReduceAttempt { .. } + | TraceEvent::ParentReduceApplied { .. } + if in_optimize_scope => + { + 0 + } + TraceEvent::PhaseNone { indent, .. } => *indent, + TraceEvent::ReduceAttempt { .. } + | TraceEvent::ReduceApplied { .. } + | TraceEvent::ParentReduceAttempt { .. } + | TraceEvent::ParentReduceApplied { .. } + | TraceEvent::ExecuteUntilDoneCheck { .. } + | TraceEvent::ExecuteUntilPopFrame { .. } + | TraceEvent::ExecuteParentAttempt { .. } + | TraceEvent::ExecuteParentApplied { .. } + | TraceEvent::ExecuteOptimized { .. } + | TraceEvent::ExecuteEncoding { .. } + | TraceEvent::SlotTransition { .. } + | TraceEvent::BuilderEvent { .. } + | TraceEvent::ExecuteDone { .. } => 2, + TraceEvent::OptimizeRecursiveSlot { .. } + | TraceEvent::ExecuteUntilIteration { .. } + | TraceEvent::ExecuteUntilReturn { .. } + | TraceEvent::SingleStepApplied { .. } => 1, + TraceEvent::OptimizeLoopEnd + | TraceEvent::OptimizeRecursiveStart { .. } + | TraceEvent::ExecuteUntilStart { .. } + | TraceEvent::SingleStepStart { .. } => 0, + } + } + + fn fmt_line(&self, f: &mut fmt::Formatter<'_>, resolution: TraceResolution) -> fmt::Result { + match self { + TraceEvent::OptimizeStart { root, session } => { + write!(f, "optimize root={root} session={session}") + } + TraceEvent::OptimizeLoopStart { array } => { + write!(f, "loop input={array}") + } + TraceEvent::OptimizeLoopEnd => Ok(()), + TraceEvent::OptimizeDone { output, changed } => match resolution { + TraceResolution::Attempts => write!(f, "done output={output} changed={changed}"), + TraceResolution::ExecutedOnly => write!(f, "done output={output}"), + }, + TraceEvent::OptimizeRecursiveStart { root } => { + write!(f, "optimize_recursive root={root}") + } + TraceEvent::OptimizeRecursiveSlot { + slot_idx, + input, + output, + } => write!(f, "recursive slot={slot_idx} input={input} output={output}"), + TraceEvent::ReduceAttempt { + array, + rule, + outcome, + } => write!( + f, + "reduce attempt array={array} source=static rule={rule} outcome={outcome}" + ), + TraceEvent::ReduceApplied { + array, + rule, + output, + } => match resolution { + TraceResolution::Attempts => write!( + f, + "reduce applied array={array} source=static rule={rule} output={output}" + ), + TraceResolution::ExecutedOnly => { + write!(f, "reduce {rule}: {array} -> {output}") + } + }, + TraceEvent::ParentReduceAttempt { + parent, + child, + slot_idx, + source, + rule, + outcome, + } => write!( + f, + "reduce_parent attempt slot={slot_idx} parent={parent} child={child} source={source} rule={rule} outcome={outcome}" + ), + TraceEvent::ParentReduceApplied { + parent, + child, + slot_idx, + source, + rule, + output, + } => match resolution { + TraceResolution::Attempts => write!( + f, + "reduce_parent applied slot={slot_idx} parent={parent} child={child} source={source} rule={rule} output={output}" + ), + TraceResolution::ExecutedOnly => write!( + f, + "reduce_parent {source}:{rule} slot={slot_idx} parent={parent} child={child} -> {output}" + ), + }, + TraceEvent::ExecuteUntilStart { target, root } => { + write!(f, "execute_until target={target} root={root}") + } + TraceEvent::ExecuteUntilIteration { + iteration, + current, + stack_parent, + builder_active, + } => { + write!(f, "iter {iteration} current={current}")?; + if let Some((parent, slot_idx)) = stack_parent { + write!(f, " stack_parent={parent} slot={slot_idx}")?; + } + write!(f, " builder_active={builder_active}") + } + TraceEvent::ExecuteUntilDoneCheck { target, canonical } => { + write!(f, "done_check target={target} canonical={canonical}") + } + TraceEvent::ExecuteUntilReturn { output } => { + write!(f, "return output={output}") + } + TraceEvent::ExecuteUntilPopFrame { slot_idx, output } => { + write!(f, "pop_frame slot={slot_idx} output={output}") + } + TraceEvent::ExecuteParentAttempt { + phase, + parent, + child, + slot_idx, + source, + kernel, + outcome, + } => write!( + f, + "{phase} attempt slot={slot_idx} parent={parent} child={child} source={source} kernel={kernel} outcome={outcome}" + ), + TraceEvent::ExecuteParentApplied { + phase, + parent, + child, + slot_idx, + source, + kernel, + output, + } => match resolution { + TraceResolution::Attempts => write!( + f, + "{phase} applied slot={slot_idx} parent={parent} child={child} source={source} kernel={kernel} output={output}" + ), + TraceResolution::ExecutedOnly => write!( + f, + "{phase} {source}:{kernel} slot={slot_idx} parent={parent} child={child} -> {output}" + ), + }, + TraceEvent::PhaseNone { + phase, + subject, + array, + .. + } => { + write!(f, "{phase} none {subject}={array}") + } + TraceEvent::ExecuteOptimized { + input, + output, + changed, + } => match resolution { + TraceResolution::Attempts => write!( + f, + "optimize_ctx input={input} output={output} changed={changed}" + ), + TraceResolution::ExecutedOnly => write!(f, "optimize_ctx {input} -> {output}"), + }, + TraceEvent::ExecuteEncoding { array } => { + write!(f, "execute encoding={array}") + } + TraceEvent::SlotTransition { + step, + slot_idx, + parent, + child, + } => write!(f, "{step} slot={slot_idx} parent={parent} child={child}"), + TraceEvent::ExecuteDone { array } => { + write!(f, "Done array={array}") + } + TraceEvent::BuilderEvent { + action, + subject, + array, + } => { + write!(f, "builder {action} {subject}={array}") + } + TraceEvent::SingleStepStart { array } => { + write!(f, "execute_step input={array}") + } + TraceEvent::SingleStepApplied { + phase, + input, + output, + } => match resolution { + TraceResolution::Attempts => { + write!(f, "{phase} applied input={input} output={output}") + } + TraceResolution::ExecutedOnly => write!(f, "{phase} {input} -> {output}"), + }, + } + } +} + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/test_harness/trace/tests.rs b/vortex-array/src/test_harness/trace/tests.rs new file mode 100644 index 00000000000..bfe99c3857b --- /dev/null +++ b/vortex-array/src/test_harness/trace/tests.rs @@ -0,0 +1,814 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hasher; + +use rstest::fixture; +use rstest::rstest; +use smallvec::smallvec; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_panic; +use vortex_mask::Mask; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ArrayEq; +use crate::ArrayHash; +use crate::ArrayRef; +use crate::Canonical; +use crate::EqMode; +use crate::ExecutionCtx; +use crate::ExecutionResult; +use crate::IntoArray; +use crate::VTable; +use crate::array::Array; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayView; +use crate::array::vtable::NotSupported; +use crate::array::vtable::ValidityVTable; +use crate::array::vtable::with_empty_buffers; +use crate::arrays::BoolArray; +use crate::arrays::ChunkedArray; +use crate::arrays::ConstantArray; +use crate::arrays::DictArray; +use crate::arrays::Filter; +use crate::arrays::FilterArray; +use crate::arrays::Primitive; +use crate::arrays::PrimitiveArray; +use crate::arrays::StructArray; +use crate::arrays::VarBinViewArray; +use crate::arrays::filter::FilterArraySlotsExt; +use crate::assert_arrays_eq; +use crate::buffer::BufferHandle; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::kernel::ExecuteParentKernel; +use crate::matcher::Matcher; +use crate::optimizer::ArrayOptimizer; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::Binary; +use crate::scalar_fn::fns::like::Like; +use crate::scalar_fn::fns::like::LikeOptions; +use crate::scalar_fn::fns::operators::Operator; +use crate::serde::ArrayChildren; +use crate::session::ArraySession; +use crate::test_harness::trace::TraceOptions; +use crate::test_harness::trace::TraceResolution; +use crate::test_harness::trace::trace_op; +use crate::test_harness::trace::trace_op_with; +use crate::validity::Validity; + +#[fixture] +fn stack_parent_fixture() -> VortexResult { + stack_parent(stack_child()?) +} + +/// Build a session with the `StackChild` parent kernels registered. +/// +/// The declining kernel is registered first so strict trace snapshots can assert that both +/// session kernels are attempted in registration order. +#[fixture] +fn stack_parent_session() -> VortexSession { + let session = VortexSession::empty().with::(); + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(StackParent.id(), StackChild, StackDeclineKernel); + kernels.register_execute_parent_kernel(StackParent.id(), StackChild, StackParentKernel); + drop(kernels); + session +} + +fn stack_child() -> VortexResult { + Ok( + Array::try_from_parts(ArrayParts::new(StackChild, test_dtype(), 3, StackChildData))? + .into_array(), + ) +} + +fn stack_parent(child: ArrayRef) -> VortexResult { + Ok(Array::try_from_parts( + ArrayParts::new( + StackParent, + child.dtype().clone(), + child.len(), + StackParentData, + ) + .with_slots(smallvec![Some(child)]), + )? + .into_array()) +} + +fn test_dtype() -> DType { + DType::Primitive(PType::I32, Nullability::NonNullable) +} + +#[derive(Clone, Debug)] +struct StackParent; + +#[derive(Clone, Debug)] +struct StackParentData; + +impl Display for StackParentData { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str("stack-parent") + } +} + +impl ArrayHash for StackParentData { + fn array_hash(&self, _state: &mut H, _eq_mode: EqMode) {} +} + +impl ArrayEq for StackParentData { + fn array_eq(&self, _other: &Self, _eq_mode: EqMode) -> bool { + true + } +} + +impl ValidityVTable for StackParent { + fn validity(_array: ArrayView<'_, StackParent>) -> VortexResult { + Ok(Validity::NonNullable) + } +} + +impl VTable for StackParent { + type TypedArrayData = StackParentData; + type OperationsVTable = NotSupported; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.test.stack-parent"); + *ID + } + + fn validate( + &self, + _data: &Self::TypedArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure!(dtype == &test_dtype(), "unexpected stack parent dtype"); + vortex_ensure!(len == 3, "unexpected stack parent length"); + vortex_ensure!(slots.len() == 1, "stack parent must have one child slot"); + let Some(child) = &slots[0] else { + vortex_bail!("stack parent child slot is missing"); + }; + vortex_ensure!(child.dtype() == dtype, "stack parent child dtype mismatch"); + vortex_ensure!(child.len() == len, "stack parent child length mismatch"); + Ok(()) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + vortex_panic!("StackParent buffer index {idx} out of bounds") + } + + fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option { + None + } + + fn serialize( + _array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + Ok(None) + } + + fn deserialize( + &self, + _dtype: &DType, + _len: usize, + _metadata: &[u8], + _buffers: &[BufferHandle], + _children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_bail!("StackParent cannot be deserialized") + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + match idx { + 0 => "child".to_string(), + _ => vortex_panic!("StackParent slot index {idx} out of bounds"), + } + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + let Some(child) = array.slots()[0].as_ref() else { + vortex_bail!("stack parent child slot is missing"); + }; + if !child.is::() { + return Ok(ExecutionResult::execute_slot::(array, 0)); + } + + Ok(ExecutionResult::done(child.clone())) + } +} + +#[derive(Clone, Debug)] +struct StackChild; + +#[derive(Clone, Debug)] +struct StackChildData; + +impl Display for StackChildData { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str("stack-child") + } +} + +impl ArrayHash for StackChildData { + fn array_hash(&self, _state: &mut H, _eq_mode: EqMode) {} +} + +impl ArrayEq for StackChildData { + fn array_eq(&self, _other: &Self, _eq_mode: EqMode) -> bool { + true + } +} + +impl ValidityVTable for StackChild { + fn validity(_array: ArrayView<'_, StackChild>) -> VortexResult { + Ok(Validity::NonNullable) + } +} + +impl VTable for StackChild { + type TypedArrayData = StackChildData; + type OperationsVTable = NotSupported; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.test.stack-child"); + *ID + } + + fn validate( + &self, + _data: &Self::TypedArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure!(dtype == &test_dtype(), "unexpected stack child dtype"); + vortex_ensure!(len == 3, "unexpected stack child length"); + vortex_ensure!(slots.is_empty(), "stack child must not have slots"); + Ok(()) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + vortex_panic!("StackChild buffer index {idx} out of bounds") + } + + fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option { + None + } + + fn serialize( + _array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + Ok(None) + } + + fn deserialize( + &self, + _dtype: &DType, + _len: usize, + _metadata: &[u8], + _buffers: &[BufferHandle], + _children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_bail!("StackChild cannot be deserialized") + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + vortex_panic!("StackChild slot index {idx} out of bounds") + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + debug_assert!(array.slots().is_empty()); + Ok(ExecutionResult::done(PrimitiveArray::from_iter([ + 99i32, 99, 99, + ]))) + } +} + +#[derive(Debug)] +struct StackDeclineKernel; + +impl ExecuteParentKernel for StackDeclineKernel { + type Parent = StackParent; + + fn execute_parent( + &self, + _array: ArrayView<'_, StackChild>, + _parent: ::Match<'_>, + _child_idx: usize, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(None) + } +} + +#[derive(Debug)] +struct StackParentKernel; + +impl ExecuteParentKernel for StackParentKernel { + type Parent = StackParent; + + fn execute_parent( + &self, + _array: ArrayView<'_, StackChild>, + parent: ::Match<'_>, + child_idx: usize, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if parent + .slots() + .get(child_idx) + .is_some_and(|slot| slot.is_none()) + { + return Ok(Some(PrimitiveArray::from_iter([1i32, 2, 3]).into_array())); + } + + Ok(None) + } +} + +#[test] +fn trace_optimize_reduce_fixpoint() -> VortexResult<()> { + let values = PrimitiveArray::from_iter([0i32, 1, 2, 3]).into_array(); + let filter = FilterArray::try_new(values.clone(), Mask::new_true(values.len()))?.into_array(); + + let traced = trace_op(|| filter.optimize())?; + + assert!(traced.output.is::()); + assert_arrays_eq!(traced.output, values, &mut execution_ctx()); + insta::assert_snapshot!(traced.trace.to_string(), @r" +optimize root=vortex.filter(i32, len=4) session=false + reduce TrivialFilterRule: vortex.filter(i32, len=4) -> vortex.primitive(i32, len=4) + done output=vortex.primitive(i32, len=4) +"); + + Ok(()) +} + +#[test] +fn trace_optimize_parent_reduce_fixpoint_attempts() -> VortexResult<()> { + let values = PrimitiveArray::from_iter([0i32, 1, 2, 3, 4, 5]).into_array(); + let inner = FilterArray::try_new( + values, + Mask::from_iter([true, false, true, true, false, true]), + )? + .into_array(); + let outer = + FilterArray::try_new(inner, Mask::from_iter([false, true, true, false]))?.into_array(); + + let traced = trace_op_with( + TraceOptions { + resolution: TraceResolution::ExecutedOnly, + }, + || outer.optimize(), + )?; + + let optimized_filter = traced.output.as_::(); + assert!(optimized_filter.child().is::()); + assert_arrays_eq!( + traced.output, + PrimitiveArray::from_iter([2i32, 3]), + &mut execution_ctx() + ); + insta::assert_snapshot!(traced.trace.to_string(), @r" + optimize root=vortex.filter(i32, len=2) session=false + reduce_parent static:FilterReduceAdaptor(Filter) slot=0 parent=vortex.filter(i32, len=2) child=vortex.filter(i32, len=4) -> vortex.filter(i32, len=2) + done output=vortex.filter(i32, len=2) + "); + + let mut ctx = ExecutionCtx::new(VortexSession::empty().with::()); + let traced = trace_op_with( + TraceOptions { + resolution: TraceResolution::ExecutedOnly, + }, + || { + outer + .execute::(&mut ctx) + .map(IntoArray::into_array) + }, + )?; + + // The trace below only records the shape of the execution; pin the values it produced too. + assert_arrays_eq!( + traced.output, + PrimitiveArray::from_iter([2i32, 3]), + &mut execution_ctx() + ); + insta::assert_snapshot!(traced.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.filter(i32, len=2) + iter 0 current=vortex.filter(i32, len=2) builder_active=false + Done array=vortex.slice(i32, len=2) + iter 1 current=vortex.slice(i32, len=2) builder_active=false + ExecuteSlot slot=0 parent=vortex.slice(i32, len=2) child=vortex.filter(i32, len=4) + iter 2 current=vortex.filter(i32, len=4) stack_parent=vortex.slice(i32, len=2) slot=0 builder_active=false + Done array=vortex.primitive(i32, len=4) + iter 3 current=vortex.primitive(i32, len=4) stack_parent=vortex.slice(i32, len=2) slot=0 builder_active=false + pop_frame slot=0 output=vortex.slice(i32, len=2) + iter 4 current=vortex.slice(i32, len=2) builder_active=false + optimize root=vortex.slice(i32, len=2) session=false + reduce_parent static:SliceReduceAdaptor(Primitive) slot=0 parent=vortex.slice(i32, len=2) child=vortex.primitive(i32, len=4) -> vortex.primitive(i32, len=2) + done output=vortex.primitive(i32, len=2) + Done array=vortex.primitive(i32, len=2) + iter 5 current=vortex.primitive(i32, len=2) builder_active=false + return output=vortex.primitive(i32, len=2) + "); + + Ok(()) +} + +/// A filter whose mask is not one contiguous run cannot be answered as a slice, so it has to +/// execute through its child. +/// +/// The test above happens to build a contiguous combined mask, which short-circuits before the +/// child is reached; without this case no trace would cover an executed filter at all. +#[test] +fn trace_execute_filter_with_scattered_mask() -> VortexResult<()> { + let values = PrimitiveArray::from_iter([0i32, 1, 2, 3, 4, 5]).into_array(); + let inner = FilterArray::try_new( + values, + Mask::from_iter([true, false, true, true, false, true]), + )? + .into_array(); + // Keeps ranks 0 and 2 of the inner selection — values 0 and 3, which are two runs, not one. + let outer = + FilterArray::try_new(inner, Mask::from_iter([true, false, true, false]))?.into_array(); + + let mut ctx = ExecutionCtx::new(VortexSession::empty().with::()); + let traced = trace_op_with( + TraceOptions { + resolution: TraceResolution::ExecutedOnly, + }, + || { + outer + .execute::(&mut ctx) + .map(IntoArray::into_array) + }, + )?; + + assert_arrays_eq!( + traced.output, + PrimitiveArray::from_iter([0i32, 3]), + &mut execution_ctx() + ); + insta::assert_snapshot!(traced.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.filter(i32, len=2) + iter 0 current=vortex.filter(i32, len=2) builder_active=false + ExecuteSlot slot=0 parent=vortex.filter(i32, len=2) child=vortex.filter(i32, len=4) + iter 1 current=vortex.filter(i32, len=4) stack_parent=vortex.filter(i32, len=2) slot=0 builder_active=false + Done array=vortex.primitive(i32, len=4) + iter 2 current=vortex.primitive(i32, len=4) stack_parent=vortex.filter(i32, len=2) slot=0 builder_active=false + pop_frame slot=0 output=vortex.filter(i32, len=2) + iter 3 current=vortex.filter(i32, len=2) builder_active=false + Done array=vortex.primitive(i32, len=2) + iter 4 current=vortex.primitive(i32, len=2) builder_active=false + return output=vortex.primitive(i32, len=2) + "); + + Ok(()) +} + +#[rstest] +fn trace_execution_stack_parent_kernel_attempts( + stack_parent_fixture: VortexResult, + stack_parent_session: VortexSession, +) -> VortexResult<()> { + let mut ctx = ExecutionCtx::new(stack_parent_session); + let parent = stack_parent_fixture?; + + let traced = trace_op_with( + TraceOptions { + resolution: TraceResolution::Attempts, + }, + || parent.execute::(&mut ctx), + )?; + + assert_arrays_eq!( + traced.output, + PrimitiveArray::from_iter([1i32, 2, 3]), + &mut ctx + ); + insta::assert_snapshot!(traced.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.test.stack-parent(i32, len=3) + iter 0 current=vortex.test.stack-parent(i32, len=3) builder_active=false + done_check target=false canonical=false + child_execute_parent attempt slot=0 parent=vortex.test.stack-parent(i32, len=3) child=vortex.test.stack-child(i32, len=3) source=session[0] kernel=execute_parent_fn outcome=declined + child_execute_parent attempt slot=0 parent=vortex.test.stack-parent(i32, len=3) child=vortex.test.stack-child(i32, len=3) source=session[1] kernel=execute_parent_fn outcome=declined + child_execute_parent none current=vortex.test.stack-parent(i32, len=3) + execute encoding=vortex.test.stack-parent(i32, len=3) + ExecuteSlot slot=0 parent=vortex.test.stack-parent(i32, len=3) child=vortex.test.stack-child(i32, len=3) + iter 1 current=vortex.test.stack-child(i32, len=3) stack_parent=vortex.test.stack-parent(i32, len=3) slot=0 builder_active=false + done_check target=false canonical=false + stack_execute_parent attempt slot=0 parent=vortex.test.stack-parent(i32, len=3) child=vortex.test.stack-child(i32, len=3) source=session[0] kernel=execute_parent_fn outcome=declined + stack_execute_parent applied slot=0 parent=vortex.test.stack-parent(i32, len=3) child=vortex.test.stack-child(i32, len=3) source=session[1] kernel=execute_parent_fn output=vortex.primitive(i32, len=3) + optimize root=vortex.primitive(i32, len=3) session=true + loop input=vortex.primitive(i32, len=3) + reduce none array=vortex.primitive(i32, len=3) + reduce_parent none array=vortex.primitive(i32, len=3) + done output=vortex.primitive(i32, len=3) changed=false + optimize_ctx input=vortex.primitive(i32, len=3) output=vortex.primitive(i32, len=3) changed=false + iter 2 current=vortex.primitive(i32, len=3) builder_active=false + done_check target=true canonical=true + return output=vortex.primitive(i32, len=3) + "); + + Ok(()) +} + +#[test] +fn trace_execution_chunked_append_child_flow() -> VortexResult<()> { + let chunks = vec![ + PrimitiveArray::from_iter([1i32, 2]).into_array(), + PrimitiveArray::from_iter([3i32]).into_array(), + PrimitiveArray::from_iter([4i32, 5]).into_array(), + ]; + let dtype = chunks[0].dtype().clone(); + let chunked = ChunkedArray::try_new(chunks, dtype)?.into_array(); + let mut ctx = ExecutionCtx::new(VortexSession::empty().with::()); + + let traced = trace_op(|| { + chunked + .execute::(&mut ctx) + .map(IntoArray::into_array) + })?; + + assert_arrays_eq!( + traced.output, + PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]), + &mut ctx + ); + insta::assert_snapshot!(traced.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.chunked(i32, len=5) + iter 0 current=vortex.chunked(i32, len=5) builder_active=false + builder start array=vortex.chunked(i32, len=5) + AppendChild slot=1 parent=vortex.chunked(i32, len=5) child=vortex.primitive(i32, len=2) + builder append child=vortex.primitive(i32, len=2) + iter 1 current=vortex.chunked(i32, len=5) builder_active=true + AppendChild slot=2 parent=vortex.chunked(i32, len=5) child=vortex.primitive(i32, len=1) + builder append child=vortex.primitive(i32, len=1) + iter 2 current=vortex.chunked(i32, len=5) builder_active=true + AppendChild slot=3 parent=vortex.chunked(i32, len=5) child=vortex.primitive(i32, len=2) + builder append child=vortex.primitive(i32, len=2) + iter 3 current=vortex.chunked(i32, len=5) builder_active=true + Done array=vortex.primitive(i32, len=0) + builder finish output=vortex.primitive(i32, len=5) + iter 4 current=vortex.primitive(i32, len=5) builder_active=false + return output=vortex.primitive(i32, len=5) + "); + + Ok(()) +} + +/// A dictionary of strings: codes `[0, 1, 2, 1, 0, 2]` over values +/// `["alpha", "beta", "gamma"]`. +fn dict_of_strings() -> VortexResult { + Ok(DictArray::try_new( + PrimitiveArray::from_iter([0u32, 1, 2, 1, 0, 2]).into_array(), + VarBinViewArray::from_iter_str(["alpha", "beta", "gamma"]).into_array(), + )? + .into_array()) +} + +fn execution_ctx() -> ExecutionCtx { + ExecutionCtx::new(VortexSession::empty().with::()) +} + +#[test] +fn trace_take_on_chunked() -> VortexResult<()> { + let chunked = ChunkedArray::try_new( + vec![ + PrimitiveArray::from_iter([10i32, 11]).into_array(), + PrimitiveArray::from_iter([12i32, 13, 14]).into_array(), + ], + DType::Primitive(PType::I32, Nullability::NonNullable), + )? + .into_array(); + + // A take is expressed as a `DictArray` whose codes are the take indices. + let indices = PrimitiveArray::from_iter([4u64, 0, 2, 2]).into_array(); + let take = DictArray::try_new(indices, chunked)?.into_array(); + + // No reduce rule rewrites a take over a chunked array: the work all happens at execution + // time, so the optimizer trace is empty. + let traced = trace_op(|| take.optimize())?; + insta::assert_snapshot!(traced.trace.to_string(), @""); + + let optimized = traced.output; + let traced = trace_op(|| { + optimized + .execute::(&mut execution_ctx()) + .map(IntoArray::into_array) + })?; + + assert_arrays_eq!( + traced.output, + PrimitiveArray::from_iter([14i32, 10, 12, 12]), + &mut execution_ctx() + ); + insta::assert_snapshot!(traced.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.dict(i32, len=4) + iter 0 current=vortex.dict(i32, len=4) builder_active=false + child_execute_parent session[0]:execute_parent_fn slot=1 parent=vortex.dict(i32, len=4) child=vortex.chunked(i32, len=5) -> vortex.dict(i32, len=4) + iter 1 current=vortex.dict(i32, len=4) builder_active=false + child_execute_parent session[0]:execute_parent_fn slot=1 parent=vortex.dict(i32, len=4) child=vortex.primitive(i32, len=4) -> vortex.primitive(i32, len=4) + iter 2 current=vortex.primitive(i32, len=4) builder_active=false + return output=vortex.primitive(i32, len=4) + "); + + Ok(()) +} + +#[test] +fn trace_filter_on_struct_with_complex_children() -> VortexResult<()> { + let names = DictArray::try_new( + PrimitiveArray::from_iter([0u32, 1, 2, 1, 0]).into_array(), + VarBinViewArray::from_iter_str(["alpha", "beta", "gamma"]).into_array(), + )? + .into_array(); + let scores = ChunkedArray::try_new( + vec![ + PrimitiveArray::from_iter([1i64, 2]).into_array(), + PrimitiveArray::from_iter([3i64, 4, 5]).into_array(), + ], + DType::Primitive(PType::I64, Nullability::NonNullable), + )? + .into_array(); + let struct_ = StructArray::from_fields(&[("name", names), ("score", scores)])?.into_array(); + + let filtered = + FilterArray::try_new(struct_, Mask::from_iter([true, false, true, true, false]))? + .into_array(); + + let traced = trace_op(|| filtered.optimize())?; + insta::assert_snapshot!(traced.trace.to_string(), @" + optimize root=vortex.filter({name=utf8, score=i64}, len=3) session=false + optimize root=vortex.filter(utf8, len=3) session=false + reduce_parent static:FilterReduceAdaptor(Dict) slot=0 parent=vortex.filter(utf8, len=3) child=vortex.dict(utf8, len=5) -> vortex.dict(utf8, len=3) + done output=vortex.dict(utf8, len=3) + reduce FilterStructRule: vortex.filter({name=utf8, score=i64}, len=3) -> vortex.struct({name=utf8, score=i64}, len=3) + done output=vortex.struct({name=utf8, score=i64}, len=3) + "); + + let optimized = traced.output; + let traced = trace_op(|| { + optimized + .execute::(&mut execution_ctx()) + .map(IntoArray::into_array) + })?; + + let expected = StructArray::from_fields(&[ + ( + "name", + VarBinViewArray::from_iter_str(["alpha", "gamma", "beta"]).into_array(), + ), + ( + "score", + PrimitiveArray::from_iter([1i64, 3, 4]).into_array(), + ), + ])? + .into_array(); + assert_arrays_eq!(traced.output, expected, &mut execution_ctx()); + insta::assert_snapshot!(traced.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.struct({name=utf8, score=i64}, len=3) + iter 0 current=vortex.struct({name=utf8, score=i64}, len=3) builder_active=false + return output=vortex.struct({name=utf8, score=i64}, len=3) + "); + + Ok(()) +} + +#[test] +fn trace_compare_on_dict() -> VortexResult<()> { + let dict = DictArray::try_new( + PrimitiveArray::from_iter([0u32, 1, 2, 1, 0]).into_array(), + PrimitiveArray::from_iter([10i32, 20, 30]).into_array(), + )? + .into_array(); + let rhs = ConstantArray::new(Scalar::from(20i32), dict.len()).into_array(); + + let compared = Binary::try_new(dict, rhs, Operator::Eq)?.into_array(); + + let traced = trace_op(|| compared.optimize())?; + insta::assert_snapshot!(traced.trace.to_string(), @" + optimize root=vortex.binary(bool, len=5) session=false + reduce_parent static:DictionaryScalarFnValuesPushDownRule slot=0 parent=vortex.binary(bool, len=5) child=vortex.dict(i32, len=5) -> vortex.dict(bool, len=5) + done output=vortex.dict(bool, len=5) + "); + + let optimized = traced.output; + let traced = trace_op(|| { + optimized + .execute::(&mut execution_ctx()) + .map(IntoArray::into_array) + })?; + + assert_arrays_eq!( + traced.output, + BoolArray::from_iter([false, true, false, true, false]), + &mut execution_ctx() + ); + insta::assert_snapshot!(traced.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.dict(bool, len=5) + iter 0 current=vortex.dict(bool, len=5) builder_active=false + ExecuteSlot slot=1 parent=vortex.dict(bool, len=5) child=vortex.binary(bool, len=3) + iter 1 current=vortex.binary(bool, len=3) stack_parent=vortex.dict(bool, len=5) slot=1 builder_active=false + Done array=vortex.bool(bool, len=3) + iter 2 current=vortex.bool(bool, len=3) stack_parent=vortex.dict(bool, len=5) slot=1 builder_active=false + pop_frame slot=1 output=vortex.dict(bool, len=5) + iter 3 current=vortex.dict(bool, len=5) builder_active=false + child_execute_parent session[0]:execute_parent_fn slot=1 parent=vortex.dict(bool, len=5) child=vortex.bool(bool, len=3) -> vortex.bool(bool, len=5) + iter 4 current=vortex.bool(bool, len=5) builder_active=false + return output=vortex.bool(bool, len=5) + "); + + Ok(()) +} + +#[test] +fn trace_like_on_dict() -> VortexResult<()> { + let strings = dict_of_strings()?; + let pattern = ConstantArray::new(Scalar::from("b%"), strings.len()).into_array(); + + let like = Like::try_new( + strings, + pattern, + LikeOptions { + negated: false, + case_insensitive: false, + }, + )? + .into_array(); + + let traced = trace_op(|| like.optimize())?; + insta::assert_snapshot!(traced.trace.to_string(), @" + optimize root=vortex.like(bool, len=6) session=false + reduce_parent static:LikeReduceAdaptor(Dict) slot=0 parent=vortex.like(bool, len=6) child=vortex.dict(utf8, len=6) -> vortex.dict(bool, len=6) + done output=vortex.dict(bool, len=6) + "); + + let optimized = traced.output; + let traced = trace_op(|| { + optimized + .execute::(&mut execution_ctx()) + .map(IntoArray::into_array) + })?; + + assert_arrays_eq!( + traced.output, + BoolArray::from_iter([false, true, false, true, false, false]), + &mut execution_ctx() + ); + insta::assert_snapshot!(traced.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.dict(bool, len=6) + iter 0 current=vortex.dict(bool, len=6) builder_active=false + ExecuteSlot slot=1 parent=vortex.dict(bool, len=6) child=vortex.like(bool, len=3) + iter 1 current=vortex.like(bool, len=3) stack_parent=vortex.dict(bool, len=6) slot=1 builder_active=false + Done array=vortex.bool(bool, len=3) + iter 2 current=vortex.bool(bool, len=3) stack_parent=vortex.dict(bool, len=6) slot=1 builder_active=false + pop_frame slot=1 output=vortex.dict(bool, len=6) + iter 3 current=vortex.dict(bool, len=6) builder_active=false + child_execute_parent session[0]:execute_parent_fn slot=1 parent=vortex.dict(bool, len=6) child=vortex.bool(bool, len=3) -> vortex.bool(bool, len=6) + iter 4 current=vortex.bool(bool, len=6) builder_active=false + return output=vortex.bool(bool, len=6) + "); + + Ok(()) +} diff --git a/vortex-array/src/trace_macros.rs b/vortex-array/src/trace_macros.rs new file mode 100644 index 00000000000..fd66f38b201 --- /dev/null +++ b/vortex-array/src/trace_macros.rs @@ -0,0 +1,38 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Dispatch macro used by the optimizer and executor to emit trace events. +//! +//! See [`test_harness::trace`][crate::test_harness::trace] for the harness that consumes the +//! events recorded here. + +/// Dispatch a trace event from inside the optimizer or executor. +/// +/// This macro is the only call-site interface for the trace harness. It is designed so that +/// release builds (and CodSpeed benchmark builds) compile every invocation away to nothing, +/// while debug-test builds with an active recorder route the call into +/// [`test_harness::trace`][crate::test_harness::trace]. +/// +/// # Gating +/// +/// Each invocation expands to a `cfg`-gated branch: +/// +/// - When `test` OR `_test-harness` is on, AND `codspeed` is off, the body is compiled in and +/// guarded by a runtime check for an active recorder. +/// - In all other configurations, the body is replaced with `()`. +/// +/// The macro is named `trace_op` because it records *operations* (optimizer rewrites, parent +/// kernels, execution steps, builder activity) rather than array values. The events it emits +/// describe what work the optimizer/executor did, not the contents of any array. +macro_rules! trace_op { + ($event:ident $(::<$($generic:ty),*>)?($($arg:expr),* $(,)?)) => {{ + #[cfg(all(any(test, feature = "_test-harness"), not(codspeed)))] + { + if $crate::test_harness::trace::is_active() { + $crate::test_harness::trace::$event $(::<$($generic),*>)?($($arg),*); + } + } + }}; +} + +pub(crate) use trace_op; diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs index 3bb47eda0ae..a2c10e30ca7 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -28,7 +28,6 @@ use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; -use crate::arrays::scalar_fn::ScalarFnFactoryExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; @@ -264,6 +263,11 @@ impl Validity { } } + /// Execute this validity into a [`Mask`] of the given length. + /// + /// Resolving [`Validity::Array`] can execute and scan up to `length` values. Optimistic fast + /// paths that do not need the exact mask should use [`Self::definitely_no_nulls`] or + /// [`Self::definitely_all_null`] and delay this work until necessary. #[inline] pub fn execute_mask(&self, length: usize, ctx: &mut ExecutionCtx) -> VortexResult { match self { @@ -322,8 +326,8 @@ impl Validity { | (Validity::AllValid, Validity::AllValid) => Validity::AllValid, // Here we actually have to do some work (Validity::Array(lhs), Validity::Array(rhs)) => Validity::Array( - Binary - .try_new_array(lhs.len(), Operator::And, [lhs, rhs])? + Binary::try_new(lhs, rhs, Operator::And)? + .into_array() .optimize()?, ), }) @@ -586,10 +590,7 @@ impl Validity { Some(Validity::Array( unsafe { ChunkedArray::new_unchecked( - validities - .into_iter() - .map(|(v, len)| v.to_array(len)) - .collect(), + validities.into_iter().map(|(v, len)| v.to_array(len)), DType::Bool(Nullability::NonNullable), ) } diff --git a/vortex-arrow/Cargo.toml b/vortex-arrow/Cargo.toml index 9d69a06fd7d..11f2308451c 100644 --- a/vortex-arrow/Cargo.toml +++ b/vortex-arrow/Cargo.toml @@ -28,6 +28,7 @@ arrow-schema = { workspace = true, features = ["canonical_extension_types"] } arrow-select = { workspace = true } itertools = { workspace = true } num-traits = { workspace = true } +simdutf8 = { workspace = true } tracing = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true, features = ["arrow"] } @@ -40,6 +41,9 @@ vortex-session = { workspace = true } divan = { workspace = true } rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-fsst = { workspace = true } +vortex-onpair = { workspace = true } +vortex-zstd = { workspace = true } [[bench]] name = "to_arrow" diff --git a/vortex-arrow/benches/to_arrow.rs b/vortex-arrow/benches/to_arrow.rs index 1cec274c28e..598dbcf8adf 100644 --- a/vortex-arrow/benches/to_arrow.rs +++ b/vortex-arrow/benches/to_arrow.rs @@ -6,20 +6,31 @@ use std::sync::Arc; use std::sync::LazyLock; +use arrow_schema::DataType; +use arrow_schema::Field; use divan::Bencher; +use divan::counter::ItemsCount; use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; +use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::FilterArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::builders::VarBinBuilder; +use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; +use vortex_array::session::ArraySessionExt; #[expect( deprecated, reason = "benchmark comparing deprecated method with new one" @@ -28,14 +39,26 @@ use vortex_arrow::ArrowArrayExecutor; use vortex_arrow::ArrowSessionExt; #[allow(deprecated)] use vortex_arrow::dtype::ToArrowType as _; +use vortex_fsst::fsst_compress; +use vortex_fsst::fsst_train_compressor; +use vortex_mask::Mask; +use vortex_onpair::DEFAULT_CONFIG; +use vortex_onpair::onpair_compress; use vortex_session::VortexSession; +use vortex_zstd::Zstd; fn main() { LazyLock::force(&SESSION); divan::main(); } -static SESSION: LazyLock = LazyLock::new(array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = array_session(); + vortex_fsst::initialize(&session); + vortex_onpair::initialize(&session); + session.arrays().register(Zstd); + session +}); fn schema() -> DType { let fields = StructFields::from_iter([ @@ -98,6 +121,254 @@ fn ArrowExportVTable_to_arrow_field(bencher: Bencher) { .bench_values(|dtype| SESSION.arrow().to_arrow_field("", &dtype).unwrap()) } +#[derive(Clone, Copy, Debug)] +enum StringEncoding { + View, + Fsst, + OnPair, + Zstd, + Dict, + DictFsst, + DictZstd, + FilterFsst, + FilterZstd, + FilterDictFsst, + ChunkedFsst, + /// Every third row null, so the export walks a partial validity mask rather than an all-valid + /// one and has to interleave nulls with the decoded values. + NullableFsst, + NullableZstd, + NullableDict, +} + +const STRING_ENCODINGS: &[StringEncoding] = &[ + StringEncoding::View, + StringEncoding::Fsst, + StringEncoding::OnPair, + StringEncoding::Zstd, + StringEncoding::Dict, + StringEncoding::DictFsst, + StringEncoding::DictZstd, + StringEncoding::FilterFsst, + StringEncoding::FilterZstd, + StringEncoding::FilterDictFsst, + StringEncoding::ChunkedFsst, + StringEncoding::NullableFsst, + StringEncoding::NullableZstd, + StringEncoding::NullableDict, +]; + +/// Encodings whose `append_to_builder` the builder benchmarks reach directly. +/// +/// The Arrow export cannot stand in for these: `execute_until` stops at the first canonical array, +/// so a bare FSST/OnPair/Zstd root is canonicalized to `VarBinView` before any builder sees it. +/// Only `Chunked`, `Constant` and `VarBin` roots reach an encoding's own `append_to_builder` that +/// way, whereas the scan machinery appends encoded arrays into a builder directly. +const BUILDER_STRING_ENCODINGS: &[StringEncoding] = &[ + StringEncoding::View, + StringEncoding::Fsst, + StringEncoding::OnPair, + StringEncoding::Zstd, + StringEncoding::Dict, + StringEncoding::ChunkedFsst, + StringEncoding::NullableFsst, + StringEncoding::NullableZstd, + StringEncoding::NullableDict, +]; + +const OFFSET_STRING_ROWS: usize = 100_000; +const OFFSET_STRING_CHUNKS: usize = 4; +const DICTIONARY_SIZE: usize = 2_048; + +fn structured_strings(len: usize) -> VarBinViewArray { + let values = (0..len) + .map(|index| format!("https://example.com/common/path/{index:06}/shared-suffix")) + .collect::>(); + VarBinViewArray::from_iter_str(values.iter().map(String::as_str)) +} + +fn nullable_structured_strings(len: usize) -> VarBinViewArray { + let values = (0..len) + .map(|index| { + (!index.is_multiple_of(3)) + .then(|| format!("https://example.com/common/path/{index:06}/shared-suffix")) + }) + .collect::>(); + VarBinViewArray::from_iter( + values.iter().map(|value| value.as_deref()), + DType::Utf8(Nullability::Nullable), + ) +} + +fn dictionary_values() -> VarBinViewArray { + structured_strings(DICTIONARY_SIZE) +} + +fn dictionary_codes() -> ArrayRef { + PrimitiveArray::from_iter( + (0..OFFSET_STRING_ROWS).map(|index| u16::try_from(index % DICTIONARY_SIZE).unwrap()), + ) + .into_array() +} + +fn half_rows_mask() -> Mask { + Mask::from_iter((0..OFFSET_STRING_ROWS).map(|index| index.is_multiple_of(2))) +} + +fn filtered(array: ArrayRef) -> ArrayRef { + // Keep Filter as a lazy intermediate so benchmark setup cannot optimize it away. + FilterArray::new(array, half_rows_mask()).into_array() +} + +fn chunked_fsst(ctx: &mut ExecutionCtx) -> ArrayRef { + let source = structured_strings(OFFSET_STRING_ROWS).into_array(); + let compressor = fsst_train_compressor(&source, ctx).unwrap(); + let chunk_size = OFFSET_STRING_ROWS / OFFSET_STRING_CHUNKS; + let chunks = (0..OFFSET_STRING_CHUNKS).map(|chunk_index| { + let start = chunk_index * chunk_size; + let end = if chunk_index + 1 == OFFSET_STRING_CHUNKS { + OFFSET_STRING_ROWS + } else { + start + chunk_size + }; + let chunk = source.slice(start..end).unwrap(); + fsst_compress(&chunk, &compressor, ctx) + .unwrap() + .into_array() + }); + ChunkedArray::try_new(chunks, source.dtype().clone()) + .unwrap() + .into_array() +} + +fn fsst(source: ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { + let compressor = fsst_train_compressor(&source, ctx).unwrap(); + fsst_compress(&source, &compressor, ctx) + .unwrap() + .into_array() +} + +fn string_array(encoding: StringEncoding) -> ArrayRef { + let mut ctx = SESSION.create_execution_ctx(); + match encoding { + StringEncoding::View => structured_strings(OFFSET_STRING_ROWS).into_array(), + StringEncoding::Fsst => fsst( + structured_strings(OFFSET_STRING_ROWS).into_array(), + &mut ctx, + ), + StringEncoding::OnPair => onpair_compress( + &structured_strings(OFFSET_STRING_ROWS).into_array(), + DEFAULT_CONFIG, + &mut ctx, + ) + .unwrap(), + StringEncoding::Zstd => { + let source = structured_strings(OFFSET_STRING_ROWS); + Zstd::from_var_bin_view_without_dict(&source, 3, 8_192, &mut ctx) + .unwrap() + .into_array() + } + StringEncoding::Dict => { + DictArray::try_new(dictionary_codes(), dictionary_values().into_array()) + .unwrap() + .into_array() + } + StringEncoding::DictFsst => { + let values = fsst(dictionary_values().into_array(), &mut ctx); + DictArray::try_new(dictionary_codes(), values) + .unwrap() + .into_array() + } + StringEncoding::DictZstd => { + let values = dictionary_values(); + let compressed_values = + Zstd::from_var_bin_view_without_dict(&values, 3, 8_192, &mut ctx) + .unwrap() + .into_array(); + DictArray::try_new(dictionary_codes(), compressed_values) + .unwrap() + .into_array() + } + StringEncoding::FilterFsst => filtered(string_array(StringEncoding::Fsst)), + StringEncoding::FilterZstd => filtered(string_array(StringEncoding::Zstd)), + StringEncoding::FilterDictFsst => filtered(string_array(StringEncoding::DictFsst)), + StringEncoding::ChunkedFsst => chunked_fsst(&mut ctx), + StringEncoding::NullableFsst => fsst( + nullable_structured_strings(OFFSET_STRING_ROWS).into_array(), + &mut ctx, + ), + StringEncoding::NullableZstd => { + let source = nullable_structured_strings(OFFSET_STRING_ROWS); + Zstd::from_var_bin_view_without_dict(&source, 3, 8_192, &mut ctx) + .unwrap() + .into_array() + } + StringEncoding::NullableDict => { + // Nulls live in the dictionary rather than the codes, so the export has to combine + // the two validities. + let values = nullable_structured_strings(DICTIONARY_SIZE); + DictArray::try_new(dictionary_codes(), values.into_array()) + .unwrap() + .into_array() + } + } +} + +/// End-to-end export to Arrow `Utf8`, which is served through a `VarBinBuilder`. +#[divan::bench(args = STRING_ENCODINGS)] +fn offset_string_export(bencher: Bencher, encoding: StringEncoding) { + let array = string_array(encoding); + let field = Field::new("value", DataType::Utf8, array.dtype().is_nullable()); + + bencher + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) + .input_counter(|(array, _)| ItemsCount::new(array.len())) + .bench_values(|(array, mut ctx)| { + SESSION + .arrow() + .execute_arrow(array, Some(&field), &mut ctx) + .unwrap() + }); +} + +/// Appends an encoded array straight into an offset builder. +#[divan::bench(args = BUILDER_STRING_ENCODINGS)] +fn append_to_varbin_builder(bencher: Bencher, encoding: StringEncoding) { + let array = string_array(encoding); + + bencher + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) + .input_counter(|(array, _)| ItemsCount::new(array.len())) + .bench_values(|(array, mut ctx)| { + let mut builder = VarBinBuilder::::with_capacity_in( + array.dtype().clone(), + array.len(), + ctx.allocator(), + ); + array.append_to_builder(&mut builder, &mut ctx).unwrap(); + builder.finish_into_varbin() + }); +} + +/// Appends an encoded array straight into a view builder. +#[divan::bench(args = BUILDER_STRING_ENCODINGS)] +fn append_to_view_builder(bencher: Bencher, encoding: StringEncoding) { + let array = string_array(encoding); + + bencher + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) + .input_counter(|(array, _)| ItemsCount::new(array.len())) + .bench_values(|(array, mut ctx)| { + let mut builder = VarBinViewBuilder::with_capacity_in( + array.dtype().clone(), + array.len(), + ctx.allocator().clone(), + ); + array.append_to_builder(&mut builder, &mut ctx).unwrap(); + builder.finish_into_varbinview() + }); +} + #[divan::bench] fn to_arrow_array(bencher: Bencher) { bencher diff --git a/vortex-arrow/src/convert.rs b/vortex-arrow/src/convert.rs index 8193ff9ac6f..90fcdbbc3f1 100644 --- a/vortex-arrow/src/convert.rs +++ b/vortex-arrow/src/convert.rs @@ -1,6 +1,15 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Canonical (non-plugin) conversions of Arrow arrays into Vortex arrays. +//! +//! Each `from_arrow_*` function converts one Arrow array shape; [`from_arrow_dyn`] dispatches on +//! the Arrow [`DataType`]. These functions have no awareness of Arrow extension types — the +//! plugin-aware entry points are the [`ArrowSession`](crate::ArrowSession) methods, which fall +//! back to these conversions for non-extension data. The deprecated [`FromArrowArray`] impls are +//! thin shims over these functions and will eventually be removed. +#![allow(deprecated)] + use std::sync::Arc; use arrow_array::AnyDictionaryArray; @@ -66,16 +75,22 @@ use vortex_array::arrays::DictArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::MapArray; use vortex_array::arrays::NullArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::Struct; use vortex_array::arrays::StructArray; use vortex_array::arrays::TemporalArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::FieldNames; use vortex_array::dtype::IntegerPType; +use vortex_array::dtype::MapDType; use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::i256; use vortex_array::extension::datetime::TimeUnit; @@ -92,7 +107,7 @@ use vortex_error::vortex_err; use vortex_error::vortex_panic; use crate::FromArrowArray; -use crate::dtype::FromArrowType; +use crate::dtype::from_arrow_time_unit; /// Zero-copy conversion of Arrow buffers into non-nullable Vortex arrays. /// @@ -147,13 +162,27 @@ where } } +/// Zero-copy conversion of an Arrow numeric primitive array into a Vortex array. +/// +/// Use [`from_arrow_temporal`] for Arrow temporal arrays, which carry a logical dtype beyond +/// their physical primitive storage. +pub fn from_arrow_primitive( + value: &ArrowPrimitiveArray, + nullable: bool, +) -> VortexResult +where + T::Native: NativePType, +{ + let buffer = Buffer::from_arrow_scalar_buffer(value.values().clone()); + let validity = nulls(value.nulls(), nullable)?; + Ok(PrimitiveArray::new(buffer, validity).into_array()) +} + macro_rules! impl_from_arrow_primitive { ($T:path) => { impl FromArrowArray<&ArrowPrimitiveArray<$T>> for ArrayRef { fn from_arrow(value: &ArrowPrimitiveArray<$T>, nullable: bool) -> VortexResult { - let buffer = Buffer::from_arrow_scalar_buffer(value.values().clone()); - let validity = nulls(value.nulls(), nullable)?; - Ok(PrimitiveArray::new(buffer, validity).into_array()) + from_arrow_primitive(value, nullable) } } }; @@ -171,56 +200,87 @@ impl_from_arrow_primitive!(Float16Type); impl_from_arrow_primitive!(Float32Type); impl_from_arrow_primitive!(Float64Type); +/// Zero-copy conversion of an Arrow `Decimal32` array into a Vortex decimal array. +pub fn from_arrow_decimal32( + array: &ArrowPrimitiveArray, + nullable: bool, +) -> VortexResult { + let decimal_type = DecimalDType::new(array.precision(), array.scale()); + let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone()); + let validity = nulls(array.nulls(), nullable)?; + Ok(DecimalArray::new(buffer, decimal_type, validity).into_array()) +} + impl FromArrowArray<&ArrowPrimitiveArray> for ArrayRef { fn from_arrow( array: &ArrowPrimitiveArray, nullable: bool, ) -> VortexResult { - let decimal_type = DecimalDType::new(array.precision(), array.scale()); - let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone()); - let validity = nulls(array.nulls(), nullable)?; - Ok(DecimalArray::new(buffer, decimal_type, validity).into_array()) + from_arrow_decimal32(array, nullable) } } +/// Zero-copy conversion of an Arrow `Decimal64` array into a Vortex decimal array. +pub fn from_arrow_decimal64( + array: &ArrowPrimitiveArray, + nullable: bool, +) -> VortexResult { + let decimal_type = DecimalDType::new(array.precision(), array.scale()); + let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone()); + let validity = nulls(array.nulls(), nullable)?; + Ok(DecimalArray::new(buffer, decimal_type, validity).into_array()) +} + impl FromArrowArray<&ArrowPrimitiveArray> for ArrayRef { fn from_arrow( array: &ArrowPrimitiveArray, nullable: bool, ) -> VortexResult { - let decimal_type = DecimalDType::new(array.precision(), array.scale()); - let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone()); - let validity = nulls(array.nulls(), nullable)?; - Ok(DecimalArray::new(buffer, decimal_type, validity).into_array()) + from_arrow_decimal64(array, nullable) } } +/// Zero-copy conversion of an Arrow `Decimal128` array into a Vortex decimal array. +pub fn from_arrow_decimal128( + array: &ArrowPrimitiveArray, + nullable: bool, +) -> VortexResult { + let decimal_type = DecimalDType::new(array.precision(), array.scale()); + let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone()); + let validity = nulls(array.nulls(), nullable)?; + Ok(DecimalArray::new(buffer, decimal_type, validity).into_array()) +} + impl FromArrowArray<&ArrowPrimitiveArray> for ArrayRef { fn from_arrow( array: &ArrowPrimitiveArray, nullable: bool, ) -> VortexResult { - let decimal_type = DecimalDType::new(array.precision(), array.scale()); - let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone()); - let validity = nulls(array.nulls(), nullable)?; - Ok(DecimalArray::new(buffer, decimal_type, validity).into_array()) + from_arrow_decimal128(array, nullable) } } +/// Zero-copy conversion of an Arrow `Decimal256` array into a Vortex decimal array. +pub fn from_arrow_decimal256( + array: &ArrowPrimitiveArray, + nullable: bool, +) -> VortexResult { + let decimal_type = DecimalDType::new(array.precision(), array.scale()); + let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone()); + // SAFETY: Our i256 implementation has the same bit-pattern representation of the + // arrow_buffer::i256 type. It is safe to treat values held inside the buffer as values + // of either type. + let buffer = unsafe { std::mem::transmute::, Buffer>(buffer) }; + let validity = nulls(array.nulls(), nullable)?; + Ok(DecimalArray::new(buffer, decimal_type, validity).into_array()) +} + impl FromArrowArray<&ArrowPrimitiveArray> for ArrayRef { fn from_arrow( array: &ArrowPrimitiveArray, nullable: bool, ) -> VortexResult { - let decimal_type = DecimalDType::new(array.precision(), array.scale()); - let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone()); - // SAFETY: Our i256 implementation has the same bit-pattern representation of the - // arrow_buffer::i256 type. It is safe to treat values held inside the buffer as values - // of either type. - let buffer = - unsafe { std::mem::transmute::, Buffer>(buffer) }; - let validity = nulls(array.nulls(), nullable)?; - Ok(DecimalArray::new(buffer, decimal_type, validity).into_array()) + from_arrow_decimal256(array, nullable) } } @@ -231,7 +291,7 @@ macro_rules! impl_from_arrow_temporal { value: &ArrowPrimitiveArray<$T>, nullable: bool, ) -> vortex_error::VortexResult { - temporal_array(value, nullable) + from_arrow_temporal(value, nullable) } } }; @@ -253,7 +313,9 @@ impl_from_arrow_temporal!(Time64NanosecondType); impl_from_arrow_temporal!(Date32Type); impl_from_arrow_temporal!(Date64Type); -fn temporal_array( +/// Conversion of an Arrow temporal array (timestamp/date/time) into a Vortex temporal +/// extension array. +pub fn from_arrow_temporal( value: &ArrowPrimitiveArray, nullable: bool, ) -> VortexResult @@ -268,13 +330,13 @@ where Ok(match value.data_type() { DataType::Timestamp(time_unit, tz) => { - TemporalArray::new_timestamp(arr, TimeUnit::from_arrow(time_unit), tz.clone()).into() + TemporalArray::new_timestamp(arr, from_arrow_time_unit(*time_unit), tz.clone()).into() } DataType::Time32(time_unit) => { - TemporalArray::new_time(arr, TimeUnit::from_arrow(time_unit)).into() + TemporalArray::new_time(arr, from_arrow_time_unit(*time_unit)).into() } DataType::Time64(time_unit) => { - TemporalArray::new_time(arr, TimeUnit::from_arrow(time_unit)).into() + TemporalArray::new_time(arr, from_arrow_time_unit(*time_unit)).into() } DataType::Date32 => TemporalArray::new_date(arr, TimeUnit::Days).into(), DataType::Date64 => TemporalArray::new_date(arr, TimeUnit::Milliseconds).into(), @@ -284,68 +346,92 @@ where }) } +/// Zero-copy conversion of an Arrow (large) string/binary array into a Vortex `VarBin` array. +pub fn from_arrow_bytes( + value: &GenericByteArray, + nullable: bool, +) -> VortexResult +where + ::Offset: IntegerPType, +{ + let dtype = match T::DATA_TYPE { + DataType::Binary | DataType::LargeBinary => DType::Binary(nullable.into()), + DataType::Utf8 | DataType::LargeUtf8 => DType::Utf8(nullable.into()), + dt => vortex_panic!("Invalid data type for ByteArray: {dt}"), + }; + // SAFETY: Arrow arrays are already validated (valid UTF-8, valid offsets, correct validity). + Ok(unsafe { + VarBinArray::new_unchecked( + value.offsets().clone().into_array(), + ByteBuffer::from_arrow_buffer(value.values().clone(), Alignment::of::()), + dtype, + nulls(value.nulls(), nullable)?, + ) + } + .into_array()) +} + impl FromArrowArray<&GenericByteArray> for ArrayRef where ::Offset: IntegerPType, { fn from_arrow(value: &GenericByteArray, nullable: bool) -> VortexResult { - let dtype = match T::DATA_TYPE { - DataType::Binary | DataType::LargeBinary => DType::Binary(nullable.into()), - DataType::Utf8 | DataType::LargeUtf8 => DType::Utf8(nullable.into()), - dt => vortex_panic!("Invalid data type for ByteArray: {dt}"), - }; - // SAFETY: Arrow arrays are already validated (valid UTF-8, valid offsets, correct validity). - Ok(unsafe { - VarBinArray::new_unchecked( - value.offsets().clone().into_array(), - ByteBuffer::from_arrow_buffer(value.values().clone(), Alignment::of::()), - dtype, - nulls(value.nulls(), nullable)?, - ) - } - .into_array()) + from_arrow_bytes(value, nullable) } } +/// Zero-copy conversion of an Arrow string/binary view array into a Vortex `VarBinView` array. +pub fn from_arrow_byte_view( + value: &GenericByteViewArray, + nullable: bool, +) -> VortexResult { + let dtype = match T::DATA_TYPE { + DataType::BinaryView => DType::Binary(nullable.into()), + DataType::Utf8View => DType::Utf8(nullable.into()), + dt => vortex_panic!("Invalid data type for ByteViewArray: {dt}"), + }; + + let views_buffer = Buffer::from_byte_buffer( + Buffer::from_arrow_scalar_buffer(value.views().clone()).into_byte_buffer(), + ); + + // SAFETY: arrow-rs ByteViewArray already checks the same invariants, we inherit those + // guarantees by zero-copy constructing from one. + Ok(unsafe { + VarBinViewArray::new_unchecked( + views_buffer, + Arc::from( + value + .data_buffers() + .iter() + .map(|b| ByteBuffer::from_arrow_buffer(b.clone(), Alignment::of::())) + .collect::>(), + ), + dtype, + nulls(value.nulls(), nullable)?, + ) + .into_array() + }) +} + impl FromArrowArray<&GenericByteViewArray> for ArrayRef { fn from_arrow(value: &GenericByteViewArray, nullable: bool) -> VortexResult { - let dtype = match T::DATA_TYPE { - DataType::BinaryView => DType::Binary(nullable.into()), - DataType::Utf8View => DType::Utf8(nullable.into()), - dt => vortex_panic!("Invalid data type for ByteViewArray: {dt}"), - }; - - let views_buffer = Buffer::from_byte_buffer( - Buffer::from_arrow_scalar_buffer(value.views().clone()).into_byte_buffer(), - ); - - // SAFETY: arrow-rs ByteViewArray already checks the same invariants, we inherit those - // guarantees by zero-copy constructing from one. - Ok(unsafe { - VarBinViewArray::new_unchecked( - views_buffer, - Arc::from( - value - .data_buffers() - .iter() - .map(|b| ByteBuffer::from_arrow_buffer(b.clone(), Alignment::of::())) - .collect::>(), - ), - dtype, - nulls(value.nulls(), nullable)?, - ) - .into_array() - }) + from_arrow_byte_view(value, nullable) } } +/// Zero-copy conversion of an Arrow boolean array into a Vortex `Bool` array. +pub fn from_arrow_boolean(value: &ArrowBooleanArray, nullable: bool) -> VortexResult { + Ok(BoolArray::new( + value.values().clone().into(), + nulls(value.nulls(), nullable)?, + ) + .into_array()) +} + impl FromArrowArray<&ArrowBooleanArray> for ArrayRef { fn from_arrow(value: &ArrowBooleanArray, nullable: bool) -> VortexResult { - Ok(BoolArray::new( - value.values().clone().into(), - nulls(value.nulls(), nullable)?, - ) - .into_array()) + from_arrow_boolean(value, nullable) } } @@ -397,118 +483,222 @@ pub(crate) fn remove_nulls(data: arrow_data::ArrayData) -> VortexResult VortexResult { + Ok(StructArray::try_new( + value.column_names().iter().copied().collect(), + value + .columns() + .iter() + .zip(value.fields()) + .map(|(c, field)| { + // Arrow pushes down nulls, even into non-nullable fields. So we strip them + // out here because Vortex is a little more strict. + if c.null_count() > 0 && !field.is_nullable() { + let stripped = make_array(remove_nulls(c.into_data())?); + from_arrow_dyn(stripped.as_ref(), false) + } else { + from_arrow_dyn(c.as_ref(), field.is_nullable()) + } + }) + .collect::>>()?, + value.len(), + nulls(value.nulls(), nullable)?, + )? + .into_array()) +} + impl FromArrowArray<&ArrowStructArray> for ArrayRef { fn from_arrow(value: &ArrowStructArray, nullable: bool) -> VortexResult { - Ok(StructArray::try_new( - value.column_names().iter().copied().collect(), - value - .columns() - .iter() - .zip(value.fields()) - .map(|(c, field)| { - // Arrow pushes down nulls, even into non-nullable fields. So we strip them - // out here because Vortex is a little more strict. - if c.null_count() > 0 && !field.is_nullable() { - let stripped = make_array(remove_nulls(c.into_data())?); - Self::from_arrow(stripped.as_ref(), false) - } else { - Self::from_arrow(c.as_ref(), field.is_nullable()) - } - }) - .collect::>>()?, - value.len(), - nulls(value.nulls(), nullable)?, - )? - .into_array()) + from_arrow_struct(value, nullable) } } +/// Conversion of an Arrow (large) list array into a Vortex `List` array. +pub fn from_arrow_list( + value: &GenericListArray, + nullable: bool, +) -> VortexResult { + // Extract the validity of the underlying element array. + let elements_are_nullable = match value.data_type() { + DataType::List(field) => field.is_nullable(), + DataType::LargeList(field) => field.is_nullable(), + dt => vortex_panic!("Invalid data type for ListArray: {dt}"), + }; + + let elements = from_arrow_dyn(value.values().as_ref(), elements_are_nullable)?; + + // `offsets` are always non-nullable. + let offsets = value.offsets().clone().into_array(); + let nulls = nulls(value.nulls(), nullable)?; + + Ok(ListArray::try_new(elements, offsets, nulls)?.into_array()) +} + impl FromArrowArray<&GenericListArray> for ArrayRef { fn from_arrow(value: &GenericListArray, nullable: bool) -> VortexResult { - // Extract the validity of the underlying element array. - let elements_are_nullable = match value.data_type() { - DataType::List(field) => field.is_nullable(), - DataType::LargeList(field) => field.is_nullable(), - dt => vortex_panic!("Invalid data type for ListArray: {dt}"), - }; + from_arrow_list(value, nullable) + } +} - let elements = Self::from_arrow(value.values().as_ref(), elements_are_nullable)?; +/// Conversion of an Arrow (large) list-view array into a Vortex `ListView` array. +pub fn from_arrow_list_view( + array: &GenericListViewArray, + nullable: bool, +) -> VortexResult { + // Extract the validity of the underlying element array. + let elements_are_nullable = match array.data_type() { + DataType::ListView(field) => field.is_nullable(), + DataType::LargeListView(field) => field.is_nullable(), + dt => vortex_panic!("Invalid data type for ListViewArray: {dt}"), + }; - // `offsets` are always non-nullable. - let offsets = value.offsets().clone().into_array(); - let nulls = nulls(value.nulls(), nullable)?; + let elements = from_arrow_dyn(array.values().as_ref(), elements_are_nullable)?; - Ok(ListArray::try_new(elements, offsets, nulls)?.into_array()) - } + // `offsets` and `sizes` are always non-nullable. + let offsets = array.offsets().clone().into_array(); + let sizes = array.sizes().clone().into_array(); + let nulls = nulls(array.nulls(), nullable)?; + + Ok(ListViewArray::try_new(elements, offsets, sizes, nulls)?.into_array()) } impl FromArrowArray<&GenericListViewArray> for ArrayRef { fn from_arrow(array: &GenericListViewArray, nullable: bool) -> VortexResult { - // Extract the validity of the underlying element array. - let elements_are_nullable = match array.data_type() { - DataType::ListView(field) => field.is_nullable(), - DataType::LargeListView(field) => field.is_nullable(), - dt => vortex_panic!("Invalid data type for ListViewArray: {dt}"), - }; - - let elements = Self::from_arrow(array.values().as_ref(), elements_are_nullable)?; + from_arrow_list_view(array, nullable) + } +} - // `offsets` and `sizes` are always non-nullable. - let offsets = array.offsets().clone().into_array(); - let sizes = array.sizes().clone().into_array(); - let nulls = nulls(array.nulls(), nullable)?; +/// Conversion of an Arrow fixed-size list array into a Vortex `FixedSizeList` array. +pub fn from_arrow_fixed_size_list( + array: &ArrowFixedSizeListArray, + nullable: bool, +) -> VortexResult { + let DataType::FixedSizeList(field, list_size) = array.data_type() else { + vortex_panic!("Invalid data type for ListArray: {}", array.data_type()); + }; - Ok(ListViewArray::try_new(elements, offsets, sizes, nulls)?.into_array()) - } + Ok(FixedSizeListArray::try_new( + from_arrow_dyn(array.values().as_ref(), field.is_nullable())?, + *list_size as u32, + nulls(array.nulls(), nullable)?, + array.len(), + )? + .into_array()) } impl FromArrowArray<&ArrowFixedSizeListArray> for ArrayRef { fn from_arrow(array: &ArrowFixedSizeListArray, nullable: bool) -> VortexResult { - let DataType::FixedSizeList(field, list_size) = array.data_type() else { - vortex_panic!("Invalid data type for ListArray: {}", array.data_type()); - }; - - Ok(FixedSizeListArray::try_new( - Self::from_arrow(array.values().as_ref(), field.is_nullable())?, - *list_size as u32, - nulls(array.nulls(), nullable)?, - array.len(), - )? - .into_array()) + from_arrow_fixed_size_list(array, nullable) } } +pub(crate) fn map_from_arrow_parts( + entries: ArrayRef, + offsets: &OffsetBuffer, + arrow_nulls: Option<&NullBuffer>, + keys_sorted: bool, + nullable: bool, +) -> VortexResult { + let DType::Struct(struct_dtype, Nullability::NonNullable) = entries.dtype() else { + vortex_bail!( + "Arrow map entries must import as non-nullable struct, got {}", + entries.dtype() + ); + }; + vortex_ensure!( + struct_dtype.nfields() == 2, + "Arrow map entries struct must contain exactly two fields" + ); + + let key_dtype = struct_dtype + .field_by_index(0) + .ok_or_else(|| vortex_err!("Arrow map entries struct missing key field"))?; + vortex_ensure!( + !key_dtype.is_nullable(), + "Arrow map key field must be non-nullable" + ); + let value_dtype = struct_dtype + .field_by_index(1) + .ok_or_else(|| vortex_err!("Arrow map entries struct missing value field"))?; + let map_dtype = MapDType::try_new(key_dtype, value_dtype, keys_sorted)?; + let entries_struct = entries.as_::(); + let entries = StructArray::try_new( + FieldNames::from(["key", "value"]), + vec![ + entries_struct.unmasked_field(0).clone(), + entries_struct.unmasked_field(1).clone(), + ], + entries.len(), + Validity::NonNullable, + )? + .into_array(); + + let len = offsets.len() - 1; + let sizes = Buffer::::from_iter(offsets.windows(2).map(|window| window[1] - window[0])) + .into_array(); + let offsets = offsets.inner().slice(0, len).into_array(); + let validity = nulls(arrow_nulls, nullable)?; + let entries = ListViewArray::try_new(entries, offsets, sizes, validity)?; + // SAFETY: Arrow Map offsets are sorted list offsets. Vortex sizes are adjacent offset + // differences, so every view's end equals the next view's start. Sliced Arrow arrays may + // leave leading or trailing entries unreferenced, which zero-copy-to-list permits. + let entries = unsafe { entries.with_zero_copy_to_list(true) }; + + Ok(MapArray::try_new(map_dtype, entries)?.into_array()) +} + +/// Conversion of an Arrow map array into a Vortex `Map` array. +pub fn from_arrow_map(array: &ArrowMapArray, nullable: bool) -> VortexResult { + let DataType::Map(_, keys_sorted) = array.data_type() else { + vortex_panic!("Invalid data type for MapArray: {}", array.data_type()); + }; + let entries = from_arrow_struct(array.entries(), false)?; + map_from_arrow_parts( + entries, + array.offsets(), + array.nulls(), + *keys_sorted, + nullable, + ) +} + impl FromArrowArray<&ArrowMapArray> for ArrayRef { - fn from_arrow(value: &ArrowMapArray, nullable: bool) -> VortexResult { - // Arrow's Map is logically List> with i32 offsets, and Vortex - // aliases it to exactly that (see `DType::try_from_arrow`). Arrow requires both the - // entries struct and the keys within it to be non-nullable, so the entries convert - // as non-nullable and only the map itself carries validity. - let entries = Self::from_arrow(value.entries() as &dyn ArrowArray, false)?; - let offsets = value.offsets().clone().into_array(); - let nulls = nulls(value.nulls(), nullable)?; - - Ok(ListArray::try_new(entries, offsets, nulls)?.into_array()) + fn from_arrow(array: &ArrowMapArray, nullable: bool) -> VortexResult { + from_arrow_map(array, nullable) } } +/// Conversion of an Arrow null array into a Vortex `Null` array. +pub fn from_arrow_null(value: &ArrowNullArray, nullable: bool) -> VortexResult { + vortex_ensure!( + nullable, + "Cannot convert an Arrow NullArray into a non-nullable Vortex array" + ); + Ok(NullArray::new(value.len()).into_array()) +} + impl FromArrowArray<&ArrowNullArray> for ArrayRef { fn from_arrow(value: &ArrowNullArray, nullable: bool) -> VortexResult { - vortex_ensure!( - nullable, - "Cannot convert an Arrow NullArray into a non-nullable Vortex array" - ); - Ok(NullArray::new(value.len()).into_array()) + from_arrow_null(value, nullable) } } +/// Conversion of an Arrow dictionary array into a Vortex `Dict` array. +pub fn from_arrow_dictionary( + array: &DictionaryArray, + nullable: bool, +) -> VortexResult { + let keys = AnyDictionaryArray::keys(array); + let keys = from_arrow_dyn(keys, keys.is_nullable())?; + let values = from_arrow_dyn(array.values().as_ref(), nullable)?; + // SAFETY: we assume that Arrow has checked the invariants on construction. + Ok(unsafe { DictArray::new_unchecked(keys, values) }) +} + impl FromArrowArray<&DictionaryArray> for DictArray { fn from_arrow(array: &DictionaryArray, nullable: bool) -> VortexResult { - let keys = AnyDictionaryArray::keys(array); - let keys = ArrayRef::from_arrow(keys, keys.is_nullable())?; - let values = ArrayRef::from_arrow(array.values().as_ref(), nullable)?; - // SAFETY: we assume that Arrow has checked the invariants on construction. - Ok(unsafe { DictArray::new_unchecked(keys, values) }) + from_arrow_dictionary(array, nullable) } } @@ -534,145 +724,157 @@ pub(crate) fn nulls(nulls: Option<&NullBuffer>, nullable: bool) -> VortexResult< } } -impl FromArrowArray<&dyn ArrowArray> for ArrayRef { - fn from_arrow(array: &dyn ArrowArray, nullable: bool) -> VortexResult { - match array.data_type() { - DataType::Boolean => Self::from_arrow(array.as_boolean(), nullable), - DataType::UInt8 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::UInt16 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::UInt32 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::UInt64 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::Int8 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::Int16 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::Int32 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::Int64 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::Float16 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::Float32 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::Float64 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::Utf8 => Self::from_arrow(array.as_string::(), nullable), - DataType::LargeUtf8 => Self::from_arrow(array.as_string::(), nullable), - DataType::Binary => Self::from_arrow(array.as_binary::(), nullable), - DataType::LargeBinary => Self::from_arrow(array.as_binary::(), nullable), - DataType::BinaryView => Self::from_arrow(array.as_binary_view(), nullable), - DataType::Utf8View => Self::from_arrow(array.as_string_view(), nullable), - DataType::Struct(_) => Self::from_arrow(array.as_struct(), nullable), - DataType::List(_) => Self::from_arrow(array.as_list::(), nullable), - DataType::LargeList(_) => Self::from_arrow(array.as_list::(), nullable), - DataType::ListView(_) => Self::from_arrow(array.as_list_view::(), nullable), - DataType::LargeListView(_) => Self::from_arrow(array.as_list_view::(), nullable), - DataType::FixedSizeList(..) => Self::from_arrow(array.as_fixed_size_list(), nullable), - DataType::Map(..) => Self::from_arrow(array.as_map(), nullable), - DataType::Null => Self::from_arrow(as_null_array(array), nullable), - DataType::Timestamp(u, _) => match u { - ArrowTimeUnit::Second => { - Self::from_arrow(array.as_primitive::(), nullable) - } - ArrowTimeUnit::Millisecond => { - Self::from_arrow(array.as_primitive::(), nullable) - } - ArrowTimeUnit::Microsecond => { - Self::from_arrow(array.as_primitive::(), nullable) - } - ArrowTimeUnit::Nanosecond => { - Self::from_arrow(array.as_primitive::(), nullable) - } - }, - DataType::Date32 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::Date64 => Self::from_arrow(array.as_primitive::(), nullable), - DataType::Time32(u) => match u { - ArrowTimeUnit::Second => { - Self::from_arrow(array.as_primitive::(), nullable) - } - ArrowTimeUnit::Millisecond => { - Self::from_arrow(array.as_primitive::(), nullable) - } - ArrowTimeUnit::Microsecond | ArrowTimeUnit::Nanosecond => unreachable!(), - }, - DataType::Time64(u) => match u { - ArrowTimeUnit::Microsecond => { - Self::from_arrow(array.as_primitive::(), nullable) - } - ArrowTimeUnit::Nanosecond => { - Self::from_arrow(array.as_primitive::(), nullable) - } - ArrowTimeUnit::Second | ArrowTimeUnit::Millisecond => unreachable!(), - }, - DataType::Decimal32(..) => { - Self::from_arrow(array.as_primitive::(), nullable) +/// Canonical conversion of any supported Arrow array into a Vortex array, dispatching on the +/// Arrow [`DataType`]. +pub fn from_arrow_dyn(array: &dyn ArrowArray, nullable: bool) -> VortexResult { + match array.data_type() { + DataType::Boolean => from_arrow_boolean(array.as_boolean(), nullable), + DataType::UInt8 => from_arrow_primitive(array.as_primitive::(), nullable), + DataType::UInt16 => from_arrow_primitive(array.as_primitive::(), nullable), + DataType::UInt32 => from_arrow_primitive(array.as_primitive::(), nullable), + DataType::UInt64 => from_arrow_primitive(array.as_primitive::(), nullable), + DataType::Int8 => from_arrow_primitive(array.as_primitive::(), nullable), + DataType::Int16 => from_arrow_primitive(array.as_primitive::(), nullable), + DataType::Int32 => from_arrow_primitive(array.as_primitive::(), nullable), + DataType::Int64 => from_arrow_primitive(array.as_primitive::(), nullable), + DataType::Float16 => from_arrow_primitive(array.as_primitive::(), nullable), + DataType::Float32 => from_arrow_primitive(array.as_primitive::(), nullable), + DataType::Float64 => from_arrow_primitive(array.as_primitive::(), nullable), + DataType::Utf8 => from_arrow_bytes(array.as_string::(), nullable), + DataType::LargeUtf8 => from_arrow_bytes(array.as_string::(), nullable), + DataType::Binary => from_arrow_bytes(array.as_binary::(), nullable), + DataType::LargeBinary => from_arrow_bytes(array.as_binary::(), nullable), + DataType::BinaryView => from_arrow_byte_view(array.as_binary_view(), nullable), + DataType::Utf8View => from_arrow_byte_view(array.as_string_view(), nullable), + DataType::Struct(_) => from_arrow_struct(array.as_struct(), nullable), + DataType::List(_) => from_arrow_list(array.as_list::(), nullable), + DataType::LargeList(_) => from_arrow_list(array.as_list::(), nullable), + DataType::ListView(_) => from_arrow_list_view(array.as_list_view::(), nullable), + DataType::LargeListView(_) => from_arrow_list_view(array.as_list_view::(), nullable), + DataType::FixedSizeList(..) => { + from_arrow_fixed_size_list(array.as_fixed_size_list(), nullable) + } + DataType::Map(..) => from_arrow_map(array.as_map(), nullable), + DataType::Null => from_arrow_null(as_null_array(array), nullable), + DataType::Timestamp(u, _) => match u { + ArrowTimeUnit::Second => { + from_arrow_temporal(array.as_primitive::(), nullable) + } + ArrowTimeUnit::Millisecond => { + from_arrow_temporal(array.as_primitive::(), nullable) + } + ArrowTimeUnit::Microsecond => { + from_arrow_temporal(array.as_primitive::(), nullable) + } + ArrowTimeUnit::Nanosecond => { + from_arrow_temporal(array.as_primitive::(), nullable) } - DataType::Decimal64(..) => { - Self::from_arrow(array.as_primitive::(), nullable) + }, + DataType::Date32 => from_arrow_temporal(array.as_primitive::(), nullable), + DataType::Date64 => from_arrow_temporal(array.as_primitive::(), nullable), + DataType::Time32(u) => match u { + ArrowTimeUnit::Second => { + from_arrow_temporal(array.as_primitive::(), nullable) } - DataType::Decimal128(..) => { - Self::from_arrow(array.as_primitive::(), nullable) + ArrowTimeUnit::Millisecond => { + from_arrow_temporal(array.as_primitive::(), nullable) } - DataType::Decimal256(..) => { - Self::from_arrow(array.as_primitive::(), nullable) + ArrowTimeUnit::Microsecond | ArrowTimeUnit::Nanosecond => unreachable!(), + }, + DataType::Time64(u) => match u { + ArrowTimeUnit::Microsecond => { + from_arrow_temporal(array.as_primitive::(), nullable) } - DataType::Dictionary(key_type, _) => match key_type.as_ref() { - DataType::Int8 => Ok(DictArray::from_arrow( - array.as_dictionary::(), - nullable, - )? - .into_array()), - DataType::Int16 => Ok(DictArray::from_arrow( - array.as_dictionary::(), - nullable, - )? - .into_array()), - DataType::Int32 => Ok(DictArray::from_arrow( - array.as_dictionary::(), - nullable, - )? - .into_array()), - DataType::Int64 => Ok(DictArray::from_arrow( - array.as_dictionary::(), - nullable, - )? - .into_array()), - DataType::UInt8 => Ok(DictArray::from_arrow( - array.as_dictionary::(), - nullable, - )? - .into_array()), - DataType::UInt16 => Ok(DictArray::from_arrow( - array.as_dictionary::(), - nullable, - )? - .into_array()), - DataType::UInt32 => Ok(DictArray::from_arrow( - array.as_dictionary::(), - nullable, - )? - .into_array()), - DataType::UInt64 => Ok(DictArray::from_arrow( - array.as_dictionary::(), - nullable, - )? - .into_array()), - key_dt => vortex_bail!("Unsupported dictionary key type: {key_dt}"), - }, - dt => vortex_bail!("Array encoding not implemented for Arrow data type {dt}"), + ArrowTimeUnit::Nanosecond => { + from_arrow_temporal(array.as_primitive::(), nullable) + } + ArrowTimeUnit::Second | ArrowTimeUnit::Millisecond => unreachable!(), + }, + DataType::Decimal32(..) => { + from_arrow_decimal32(array.as_primitive::(), nullable) + } + DataType::Decimal64(..) => { + from_arrow_decimal64(array.as_primitive::(), nullable) + } + DataType::Decimal128(..) => { + from_arrow_decimal128(array.as_primitive::(), nullable) } + DataType::Decimal256(..) => { + from_arrow_decimal256(array.as_primitive::(), nullable) + } + DataType::Dictionary(key_type, _) => match key_type.as_ref() { + DataType::Int8 => Ok(from_arrow_dictionary( + array.as_dictionary::(), + nullable, + )? + .into_array()), + DataType::Int16 => Ok(from_arrow_dictionary( + array.as_dictionary::(), + nullable, + )? + .into_array()), + DataType::Int32 => Ok(from_arrow_dictionary( + array.as_dictionary::(), + nullable, + )? + .into_array()), + DataType::Int64 => Ok(from_arrow_dictionary( + array.as_dictionary::(), + nullable, + )? + .into_array()), + DataType::UInt8 => Ok(from_arrow_dictionary( + array.as_dictionary::(), + nullable, + )? + .into_array()), + DataType::UInt16 => Ok(from_arrow_dictionary( + array.as_dictionary::(), + nullable, + )? + .into_array()), + DataType::UInt32 => Ok(from_arrow_dictionary( + array.as_dictionary::(), + nullable, + )? + .into_array()), + DataType::UInt64 => Ok(from_arrow_dictionary( + array.as_dictionary::(), + nullable, + )? + .into_array()), + key_dt => vortex_bail!("Unsupported dictionary key type: {key_dt}"), + }, + dt => vortex_bail!("Array encoding not implemented for Arrow data type {dt}"), + } +} + +impl FromArrowArray<&dyn ArrowArray> for ArrayRef { + fn from_arrow(array: &dyn ArrowArray, nullable: bool) -> VortexResult { + from_arrow_dyn(array, nullable) } } +/// Canonical conversion of an Arrow [`RecordBatch`] into a Vortex struct array. +pub fn from_arrow_batch(batch: &RecordBatch, nullable: bool) -> VortexResult { + from_arrow_struct(&arrow_array::StructArray::from(batch.clone()), nullable) +} + impl FromArrowArray for ArrayRef { fn from_arrow(array: RecordBatch, nullable: bool) -> VortexResult { - ArrayRef::from_arrow(&arrow_array::StructArray::from(array), nullable) + from_arrow_batch(&array, nullable) } } impl FromArrowArray<&RecordBatch> for ArrayRef { fn from_arrow(array: &RecordBatch, nullable: bool) -> VortexResult { - Self::from_arrow(array.clone(), nullable) + from_arrow_batch(array, nullable) } } #[cfg(test)] mod tests { use std::sync::Arc; - use std::sync::LazyLock; use arrow_array::Array as ArrowArray; use arrow_array::BinaryArray; @@ -711,6 +913,8 @@ mod tests { use arrow_array::builder::Int32Builder; use arrow_array::builder::LargeListBuilder; use arrow_array::builder::ListBuilder; + use arrow_array::builder::MapBuilder as ArrowMapBuilder; + use arrow_array::builder::StringBuilder; use arrow_array::builder::StringViewBuilder; use arrow_array::new_null_array; use arrow_array::types::ArrowPrimitiveType; @@ -725,7 +929,6 @@ mod tests { use arrow_schema::Schema; use rstest::rstest; use vortex_array::ArrayRef; - use vortex_array::VortexSessionExecute as _; use vortex_array::arrays::Decimal; use vortex_array::arrays::FixedSizeList; use vortex_array::arrays::List; @@ -734,18 +937,21 @@ mod tests { use vortex_array::arrays::Struct; use vortex_array::arrays::VarBinView; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; - use vortex_array::arrays::list::ListArrayExt; - use vortex_array::arrays::listview::ListViewArrayExt; + use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; + use vortex_array::arrays::list::ListArraySlotsExt; + use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::extension::datetime::TimeUnit; use vortex_array::extension::datetime::Timestamp; + use vortex_error::VortexResult; - use crate::ArrowSessionExt as _; use crate::FromArrowArray as _; use crate::IntoVortexArray as _; + use crate::convert::from_arrow_batch; + use crate::convert::from_arrow_dyn; #[rstest] #[case::i8( @@ -1532,6 +1738,40 @@ mod tests { ArrayRef::from_arrow(null_struct_array_with_non_nullable_field.as_ref(), true).unwrap(); } + #[test] + fn dyn_dispatch_imports_map() -> VortexResult<()> { + let mut builder = ArrowMapBuilder::new(None, StringBuilder::new(), Int32Builder::new()); + builder.keys().append_value("joe"); + builder.values().append_value(1); + builder.append(true)?; + builder.append(false)?; + let arrow = builder.finish(); + + let expected = DType::map( + DType::Utf8(Nullability::NonNullable), + DType::Primitive(PType::I32, Nullability::Nullable), + false, + Nullability::Nullable, + )?; + + // Dispatching on the Arrow data type must cover maps... + assert_eq!(from_arrow_dyn(&arrow, true)?.dtype(), &expected); + + // ...including maps reached by recursing into a container. + let batch = RecordBatch::try_from_iter_with_nullable([( + "maps", + Arc::new(arrow) as Arc, + true, + )])?; + let imported = from_arrow_batch(&batch, false)?; + assert_eq!( + imported.as_::().unmasked_field(0).dtype(), + &expected + ); + + Ok(()) + } + #[test] fn non_nullable_request_rejects_nulls() { // Requesting `nullable = false` on an Arrow array that physically contains nulls is a @@ -1588,53 +1828,4 @@ mod tests { ArrayRef::from_arrow(null_struct_array_with_non_nullable_field.as_ref(), true).is_err() ); } - - /// Arrow `Map` has no Vortex `DType`; it is aliased to `List>` on import - /// and rebuilt on export. Both halves have to be present for a `Map` column to survive a - /// Vortex file, so this asserts the full round trip rather than either direction alone. - #[test] - fn map_array_roundtrips_through_the_list_alias() { - use arrow_array::MapArray; - use arrow_array::builder::MapBuilder; - use arrow_array::builder::StringBuilder; - use vortex_session::VortexSession; - - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); - - let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new()); - builder.keys().append_value("a"); - builder.values().append_value(1); - builder.keys().append_value("b"); - builder.values().append_value(2); - builder.append(true).expect("first map entry should append"); - builder.append(false).expect("null map entry should append"); - builder.keys().append_value("c"); - builder.values().append_value(3); - builder.append(true).expect("third map entry should append"); - let arrow_map = builder.finish(); - - // Import: the map becomes a Vortex list of key/value structs. - let vortex_array = - ArrayRef::from_arrow(&arrow_map, true).expect("map should convert to a vortex array"); - assert_eq!(vortex_array.len(), 3); - let list = vortex_array.as_::(); - assert_eq!(list.elements().len(), 3); - assert_eq!(list.elements().as_::().names().len(), 2); - - // Export: asking for the original Arrow `Map` type rebuilds a `MapArray`, not a list. - let field = Field::new("m", arrow_map.data_type().clone(), true); - let back = SESSION - .arrow() - .execute_arrow( - vortex_array, - Some(&field), - &mut SESSION.create_execution_ctx(), - ) - .expect("vortex list should convert back to an arrow map"); - let back = back - .as_any() - .downcast_ref::() - .expect("export should produce a MapArray"); - assert_eq!(back, &arrow_map); - } } diff --git a/vortex-arrow/src/datum.rs b/vortex-arrow/src/datum.rs index 5019642c452..953a9d8d714 100644 --- a/vortex-arrow/src/datum.rs +++ b/vortex-arrow/src/datum.rs @@ -18,6 +18,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_panic; use crate::ArrowSessionExt; +#[allow(deprecated)] use crate::FromArrowArray; /// A wrapper around a generic Arrow array that can be used as a Datum in Arrow compute. @@ -109,6 +110,7 @@ impl ArrowDatum for Datum { note = "Relies on the hidden global `legacy_session()`; use `from_arrow_columnar` with an explicit `ExecutionCtx` instead" )] #[allow(clippy::disallowed_methods)] +#[allow(deprecated)] pub fn from_arrow_array_with_len(array: A, len: usize, nullable: bool) -> VortexResult where ArrayRef: FromArrowArray, @@ -142,19 +144,20 @@ where /// This is useful for compute functions that delegate to Arrow using [Datum], /// which will return a scalar (length 1 Arrow array) if the input array is constant. /// +/// The array is imported through the [`ArrowSession`](crate::ArrowSession) of `ctx`, so nested +/// extension fields reach their registered import plugins. +/// /// # Error /// /// The provided array must have length `len` or `1`. -pub fn from_arrow_columnar( - array: A, +pub fn from_arrow_columnar( + array: ArrowArrayRef, len: usize, nullable: bool, ctx: &mut ExecutionCtx, -) -> VortexResult -where - ArrayRef: FromArrowArray, -{ - let array = ArrayRef::from_arrow(array, nullable)?; +) -> VortexResult { + let session = ctx.session().clone(); + let array = session.arrow().from_arrow_array(array, nullable)?; if array.len() == len { return Ok(array); } diff --git a/vortex-arrow/src/dtype.rs b/vortex-arrow/src/dtype.rs index 01ddc898f49..f58cdf1cafb 100644 --- a/vortex-arrow/src/dtype.rs +++ b/vortex-arrow/src/dtype.rs @@ -6,9 +6,12 @@ //! Apache Arrow's type system includes physical information, which could lead to ambiguities as //! Vortex treats encodings as separate from logical types. //! -//! [`DType::to_arrow_schema`] and its sibling [`DType::to_arrow_dtype`] use a simple algorithm, -//! where every logical type is encoded in its simplest corresponding Arrow type. This reflects the -//! reality that most compute engines don't make use of the entire type range arrow-rs supports. +//! The conversions in this module are "naive": every logical type is encoded in its simplest +//! corresponding Arrow type, and Arrow extension types (other than the builtin temporal types and +//! Parquet Variant) are not understood. The authoritative, plugin-aware conversion entry point is +//! [`ArrowSession`](crate::ArrowSession) — prefer its `to_arrow_schema` / `to_arrow_datatype` / +//! `from_arrow_schema` / `from_arrow_field` / `from_arrow_datatype` methods over the deprecated +//! traits in this module. //! //! For this reason, it's recommended to do as much computation as possible within Vortex, and then //! materialize an Arrow ArrayRef at the very end of the processing chain. @@ -38,73 +41,66 @@ use vortex_array::extension::datetime::Timestamp; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_error::vortex_panic; /// Trait for converting Arrow types to Vortex types. +#[deprecated( + note = "Use `ArrowSession` (`from_arrow_schema`, `from_arrow_field`, `from_arrow_datatype`) instead" +)] pub trait FromArrowType: Sized { /// Convert the Arrow type to a Vortex type. + #[deprecated( + note = "Use `ArrowSession` (`from_arrow_schema`, `from_arrow_field`, `from_arrow_datatype`) instead" + )] fn from_arrow(value: T) -> Self; } -/// Trait for converting Vortex types to Arrow types. +/// Trait for converting Arrow types to Vortex types. +#[deprecated( + note = "Use `ArrowSession` (`from_arrow_schema`, `from_arrow_field`, `from_arrow_datatype`) instead" +)] pub trait TryFromArrowType: Sized { /// Convert the Arrow type to a Vortex type. + #[deprecated( + note = "Use `ArrowSession` (`from_arrow_schema`, `from_arrow_field`, `from_arrow_datatype`) instead" + )] fn try_from_arrow(value: T) -> VortexResult; } -impl TryFromArrowType<&DataType> for PType { - fn try_from_arrow(value: &DataType) -> VortexResult { - match value { - DataType::Int8 => Ok(Self::I8), - DataType::Int16 => Ok(Self::I16), - DataType::Int32 => Ok(Self::I32), - DataType::Int64 => Ok(Self::I64), - DataType::UInt8 => Ok(Self::U8), - DataType::UInt16 => Ok(Self::U16), - DataType::UInt32 => Ok(Self::U32), - DataType::UInt64 => Ok(Self::U64), - DataType::Float16 => Ok(Self::F16), - DataType::Float32 => Ok(Self::F32), - DataType::Float64 => Ok(Self::F64), - _ => Err(vortex_err!( - "Arrow datatype {:?} cannot be converted to ptype", - value - )), - } - } -} - -impl TryFromArrowType<&DataType> for DecimalDType { - fn try_from_arrow(value: &DataType) -> VortexResult { - match value { - DataType::Decimal32(precision, scale) - | DataType::Decimal64(precision, scale) - | DataType::Decimal128(precision, scale) - | DataType::Decimal256(precision, scale) => Self::try_new(*precision, *scale), - - _ => Err(vortex_err!( - "Arrow datatype {:?} cannot be converted to DecimalDType", - value - )), - } - } -} +/// Extension trait converting Vortex [`DType`]s into Arrow schemas and data types. +/// +/// This mirrors inherent methods that lived on [`DType`] before Arrow interoperability moved +/// into this crate. +#[deprecated(note = "Use `ArrowSession::to_arrow_schema` / `to_arrow_datatype` instead")] +pub trait ToArrowType { + /// Convert a Vortex [`DType`] into an Arrow [`Schema`]. + /// + /// This method is not plugin-aware and strips any `ARROW:extension:name` metadata for + /// non-builtin extensions (only `arrow.parquet.variant` is special-cased here). Use the + /// session method when you need round-trippable extension metadata. + #[deprecated(note = "Use `ArrowSession::to_arrow_schema` instead")] + fn to_arrow_schema(&self) -> VortexResult; -impl FromArrowType<&ArrowTimeUnit> for TimeUnit { - fn from_arrow(value: &ArrowTimeUnit) -> Self { - Self::from_arrow(*value) - } + /// Returns the Arrow [`DataType`] that best corresponds to this Vortex [`DType`]. + /// + /// This method has no awareness of registered Arrow extension plugins, so any + /// [`DType::Extension`] outside the builtin temporal set will fail or silently lose its + /// `ARROW:extension:name` metadata. The session methods recurse through containers + /// and dispatch plugins at every extension node. + #[deprecated(note = "Use `ArrowSession::to_arrow_datatype` instead")] + fn to_arrow_dtype(&self) -> VortexResult; } -impl FromArrowType for TimeUnit { - fn from_arrow(value: ArrowTimeUnit) -> Self { - match value { - ArrowTimeUnit::Second => Self::Seconds, - ArrowTimeUnit::Millisecond => Self::Milliseconds, - ArrowTimeUnit::Microsecond => Self::Microseconds, - ArrowTimeUnit::Nanosecond => Self::Nanoseconds, - } +/// Convert an Arrow [`ArrowTimeUnit`] to a Vortex [`TimeUnit`]. +pub(crate) fn from_arrow_time_unit(value: ArrowTimeUnit) -> TimeUnit { + match value { + ArrowTimeUnit::Second => TimeUnit::Seconds, + ArrowTimeUnit::Millisecond => TimeUnit::Milliseconds, + ArrowTimeUnit::Microsecond => TimeUnit::Microseconds, + ArrowTimeUnit::Nanosecond => TimeUnit::Nanoseconds, } } @@ -113,7 +109,7 @@ impl FromArrowType for TimeUnit { /// # Errors /// /// Returns an error for units with no Arrow equivalent (e.g. [`TimeUnit::Days`]). -pub fn to_arrow_time_unit(value: TimeUnit) -> VortexResult { +pub(crate) fn to_arrow_time_unit(value: TimeUnit) -> VortexResult { Ok(match value { TimeUnit::Seconds => ArrowTimeUnit::Second, TimeUnit::Milliseconds => ArrowTimeUnit::Millisecond, @@ -123,207 +119,197 @@ pub fn to_arrow_time_unit(value: TimeUnit) -> VortexResult { }) } -impl FromArrowType for DType { - fn from_arrow(value: SchemaRef) -> Self { - Self::from_arrow(value.as_ref()) +/// Naive conversion of an Arrow [`DataType`] to the Vortex [`PType`] it stores. +pub(crate) fn ptype_from_arrow(value: &DataType) -> VortexResult { + match value { + DataType::Int8 => Ok(PType::I8), + DataType::Int16 => Ok(PType::I16), + DataType::Int32 => Ok(PType::I32), + DataType::Int64 => Ok(PType::I64), + DataType::UInt8 => Ok(PType::U8), + DataType::UInt16 => Ok(PType::U16), + DataType::UInt32 => Ok(PType::U32), + DataType::UInt64 => Ok(PType::U64), + DataType::Float16 => Ok(PType::F16), + DataType::Float32 => Ok(PType::F32), + DataType::Float64 => Ok(PType::F64), + _ => Err(vortex_err!( + "Arrow datatype {:?} cannot be converted to ptype", + value + )), } } -impl TryFromArrowType for DType { - fn try_from_arrow(value: SchemaRef) -> VortexResult { - Self::try_from_arrow(value.as_ref()) +/// Naive conversion of an Arrow decimal [`DataType`] to a [`DecimalDType`]. +pub(crate) fn decimal_dtype_from_arrow(value: &DataType) -> VortexResult { + match value { + DataType::Decimal32(precision, scale) + | DataType::Decimal64(precision, scale) + | DataType::Decimal128(precision, scale) + | DataType::Decimal256(precision, scale) => DecimalDType::try_new(*precision, *scale), + + _ => Err(vortex_err!( + "Arrow datatype {:?} cannot be converted to DecimalDType", + value + )), } } -impl FromArrowType<&Schema> for DType { - fn from_arrow(value: &Schema) -> Self { - Self::try_from_arrow(value).vortex_expect("arrow schema to dtype") +/// Naive conversion of an Arrow [`DataType`] to a Vortex [`DType`], with no awareness of Arrow +/// extension plugins. +pub(crate) fn from_arrow_data_type( + data_type: &DataType, + nullability: Nullability, +) -> VortexResult { + if data_type.is_integer() || data_type.is_floating() { + return Ok(DType::Primitive(ptype_from_arrow(data_type)?, nullability)); } -} - -impl TryFromArrowType<&Schema> for DType { - fn try_from_arrow(value: &Schema) -> VortexResult { - Ok(Self::Struct( - StructFields::try_from_arrow(value.fields())?, - Nullability::NonNullable, // Must match From for Array - )) - } -} - -impl FromArrowType<&Fields> for StructFields { - fn from_arrow(value: &Fields) -> Self { - Self::try_from_arrow(value).vortex_expect("arrow fields to struct fields") - } -} -impl TryFromArrowType<&Fields> for StructFields { - fn try_from_arrow(value: &Fields) -> VortexResult { - value - .into_iter() - .map(|f| { - Ok(( - FieldName::from(f.name().as_str()), - DType::try_from_arrow(f.as_ref())?, - )) - }) - .collect::>() - } -} - -impl FromArrowType<(&DataType, Nullability)> for DType { - fn from_arrow(value: (&DataType, Nullability)) -> Self { - Self::try_from_arrow(value).vortex_expect("arrow data type to dtype") - } -} - -impl TryFromArrowType<(&DataType, Nullability)> for DType { - fn try_from_arrow((data_type, nullability): (&DataType, Nullability)) -> VortexResult { - if data_type.is_integer() || data_type.is_floating() { - return Ok(DType::Primitive( - PType::try_from_arrow(data_type)?, - nullability, - )); + Ok(match data_type { + DataType::Null => DType::Null, + DataType::Decimal32(precision, scale) + | DataType::Decimal64(precision, scale) + | DataType::Decimal128(precision, scale) + | DataType::Decimal256(precision, scale) => { + DType::Decimal(DecimalDType::new(*precision, *scale), nullability) } - - Ok(match data_type { - DataType::Null => DType::Null, - DataType::Decimal32(precision, scale) - | DataType::Decimal64(precision, scale) - | DataType::Decimal128(precision, scale) - | DataType::Decimal256(precision, scale) => { - DType::Decimal(DecimalDType::new(*precision, *scale), nullability) - } - DataType::Boolean => DType::Bool(nullability), - DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => DType::Utf8(nullability), - DataType::Binary | DataType::LargeBinary | DataType::BinaryView => { - DType::Binary(nullability) - } - DataType::Date32 => DType::Extension(Date::new(TimeUnit::Days, nullability).erased()), - DataType::Date64 => { - DType::Extension(Date::new(TimeUnit::Milliseconds, nullability).erased()) - } - DataType::Time32(unit) => { - DType::Extension(Time::new(TimeUnit::from_arrow(unit), nullability).erased()) - } - DataType::Time64(unit) => { - DType::Extension(Time::new(TimeUnit::from_arrow(unit), nullability).erased()) - } - DataType::Timestamp(unit, tz) => DType::Extension( - Timestamp::new_with_tz(TimeUnit::from_arrow(unit), tz.clone(), nullability) - .erased(), - ), - DataType::List(e) - | DataType::LargeList(e) - | DataType::ListView(e) - | DataType::LargeListView(e) => { - DType::List(Arc::new(Self::try_from_arrow(e.as_ref())?), nullability) - } - DataType::Map(entries_field, _ordered) => { - // Map is logically List>. - // The entries_field contains a Struct type with the key and value fields. - DType::List( - Arc::new(Self::from_arrow(entries_field.as_ref())), - nullability, - ) - } - DataType::FixedSizeList(e, size) => DType::FixedSizeList( - Arc::new(Self::try_from_arrow(e.as_ref())?), - *size as u32, + DataType::Boolean => DType::Bool(nullability), + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => DType::Utf8(nullability), + DataType::Binary | DataType::LargeBinary | DataType::BinaryView => { + DType::Binary(nullability) + } + DataType::Date32 => DType::Extension(Date::new(TimeUnit::Days, nullability).erased()), + DataType::Date64 => { + DType::Extension(Date::new(TimeUnit::Milliseconds, nullability).erased()) + } + DataType::Time32(unit) => { + DType::Extension(Time::new(from_arrow_time_unit(*unit), nullability).erased()) + } + DataType::Time64(unit) => { + DType::Extension(Time::new(from_arrow_time_unit(*unit), nullability).erased()) + } + DataType::Timestamp(unit, tz) => DType::Extension( + Timestamp::new_with_tz(from_arrow_time_unit(*unit), tz.clone(), nullability).erased(), + ), + DataType::List(e) + | DataType::LargeList(e) + | DataType::ListView(e) + | DataType::LargeListView(e) => { + DType::List(Arc::new(from_arrow_field_naive(e.as_ref())?), nullability) + } + DataType::FixedSizeList(e, size) => DType::FixedSizeList( + Arc::new(from_arrow_field_naive(e.as_ref())?), + *size as u32, + nullability, + ), + DataType::Struct(f) => DType::Struct(from_arrow_fields_naive(f)?, nullability), + DataType::Map(entries, keys_sorted) => { + vortex_ensure!( + !entries.is_nullable(), + "Arrow map entries field must be non-nullable" + ); + let DataType::Struct(fields) = entries.data_type() else { + vortex_bail!( + "Arrow map entries field must have Struct type, got {:?}", + entries.data_type() + ); + }; + vortex_ensure_eq!( + fields.len(), + 2, + InvalidArgument: "Arrow map entries struct must contain exactly two fields" + ); + let key = &fields[0]; + let value = &fields[1]; + vortex_ensure!( + !key.is_nullable(), + "Arrow map key field must be non-nullable" + ); + DType::map( + from_arrow_field_naive(key.as_ref())?, + from_arrow_field_naive(value.as_ref())?, + *keys_sorted, nullability, - ), - DataType::Struct(f) => DType::Struct(StructFields::try_from_arrow(f)?, nullability), - DataType::Dictionary(_, value_type) => { - Self::try_from_arrow((value_type.as_ref(), nullability))? - } - DataType::RunEndEncoded(_, value_type) => { - Self::try_from_arrow((value_type.data_type(), nullability))? - } - _ => vortex_bail!("Arrow data type not supported: {data_type:?}"), - }) - } + )? + } + DataType::Dictionary(_, value_type) => { + from_arrow_data_type(value_type.as_ref(), nullability)? + } + DataType::RunEndEncoded(_, value_type) => { + from_arrow_data_type(value_type.data_type(), nullability)? + } + _ => vortex_bail!("Arrow data type not supported: {data_type:?}"), + }) } -impl FromArrowType<&Field> for DType { - fn from_arrow(field: &Field) -> Self { - Self::try_from_arrow(field).vortex_expect("arrow field to dtype") +/// Naive conversion of an Arrow [`Field`] to a Vortex [`DType`]. Only the builtin +/// `arrow.parquet.variant` extension metadata is recognized. +pub(crate) fn from_arrow_field_naive(field: &Field) -> VortexResult { + if field + .metadata() + .get("ARROW:extension:name") + .map(|s| s.as_str()) + == Some("arrow.parquet.variant") + { + return Ok(DType::Variant(field.is_nullable().into())); } + from_arrow_data_type(field.data_type(), field.is_nullable().into()) } -impl TryFromArrowType<&Field> for DType { - fn try_from_arrow(field: &Field) -> VortexResult { - if field - .metadata() - .get("ARROW:extension:name") - .map(|s| s.as_str()) - == Some("arrow.parquet.variant") - { - return Ok(DType::Variant(field.is_nullable().into())); - } - Self::try_from_arrow((field.data_type(), field.is_nullable().into())) - } +/// Naive conversion of Arrow [`Fields`] to Vortex [`StructFields`]. +pub(crate) fn from_arrow_fields_naive(fields: &Fields) -> VortexResult { + fields + .into_iter() + .map(|f| { + Ok(( + FieldName::from(f.name().as_str()), + from_arrow_field_naive(f.as_ref())?, + )) + }) + .collect::>() } -/// Extension trait converting Vortex [`DType`]s into Arrow schemas and data types. -/// -/// This mirrors inherent methods that lived on [`DType`] before Arrow interoperability moved -/// into this crate. -pub trait ToArrowType { - /// Convert a Vortex [`DType`] into an Arrow [`Schema`]. - /// - /// **Prefer `ArrowSession::to_arrow_schema`.** This method is not plugin-aware and - /// strips any `ARROW:extension:name` metadata for non-builtin extensions (only - /// `arrow.parquet.variant` is special-cased here). Use the session method when you - /// need round-trippable extension metadata. - fn to_arrow_schema(&self) -> VortexResult; - - /// Returns the Arrow [`DataType`] that best corresponds to this Vortex [`DType`]. - /// - /// **Prefer `ArrowSession::to_arrow_datatype` (or `to_arrow_field`).** This method - /// has no awareness of registered Arrow extension plugins, so any [`DType::Extension`] - /// outside the builtin temporal set will fail or silently lose its - /// `ARROW:extension:name` metadata. The session methods recurse through containers - /// and dispatch plugins at every extension node. - #[deprecated(note = "Use `ArrowSession::to_arrow_datatype` instead")] - fn to_arrow_dtype(&self) -> VortexResult; +/// Naive conversion of an Arrow [`Schema`] to a Vortex top-level non-nullable struct [`DType`]. +pub(crate) fn from_arrow_schema_naive(schema: &Schema) -> VortexResult { + Ok(DType::Struct( + from_arrow_fields_naive(schema.fields())?, + Nullability::NonNullable, // Must match From for Array + )) } -impl ToArrowType for DType { - fn to_arrow_schema(&self) -> VortexResult { - let DType::Struct(struct_dtype, nullable) = self else { - vortex_bail!("only DType::Struct can be converted to arrow schema"); - }; +/// Naive conversion of a Vortex top-level struct [`DType`] to an Arrow [`Schema`], with the +/// builtin `arrow.parquet.variant` special-case. +fn to_arrow_schema_naive(dtype: &DType) -> VortexResult { + let DType::Struct(struct_dtype, nullable) = dtype else { + vortex_bail!("only DType::Struct can be converted to arrow schema"); + }; - if *nullable != Nullability::NonNullable { - vortex_bail!("top-level struct in Schema must be NonNullable"); - } - - let mut builder = SchemaBuilder::with_capacity(struct_dtype.names().len()); - for (field_name, field_dtype) in struct_dtype.names().iter().zip(struct_dtype.fields()) { - let field = if field_dtype.is_variant() { - let storage = DataType::Struct(variant_storage_fields_minimal()); - Field::new(field_name.as_ref(), storage, field_dtype.is_nullable()).with_metadata( - [( - "ARROW:extension:name".to_owned(), - "arrow.parquet.variant".to_owned(), - )] - .into(), - ) - } else { - Field::new( - field_name.as_ref(), - to_data_type_naive(&field_dtype)?, - field_dtype.is_nullable(), - ) - }; - builder.push(field); - } - - Ok(builder.finish()) + if *nullable != Nullability::NonNullable { + vortex_bail!("top-level struct in Schema must be NonNullable"); } - fn to_arrow_dtype(&self) -> VortexResult { - to_data_type_naive(self) + let mut builder = SchemaBuilder::with_capacity(struct_dtype.names().len()); + for (field_name, field_dtype) in struct_dtype.names().iter().zip(struct_dtype.fields()) { + let field = if field_dtype.is_variant() { + let storage = DataType::Struct(variant_storage_fields_minimal()); + Field::new(field_name.as_ref(), storage, field_dtype.is_nullable()).with_metadata( + [( + "ARROW:extension:name".to_owned(), + "arrow.parquet.variant".to_owned(), + )] + .into(), + ) + } else { + Field::new( + field_name.as_ref(), + to_data_type_naive(&field_dtype)?, + field_dtype.is_nullable(), + ) + }; + builder.push(field); } + + Ok(builder.finish()) } /// Naive conversion from a Vortex `DType` to the nearest Arrow physical data type. @@ -376,6 +362,17 @@ pub(crate) fn to_data_type_naive(dtype: &DType) -> VortexResult { )), *size as i32, ), + DType::Map(map_dtype, _) => { + let key = Field::new("key", to_data_type_naive(&map_dtype.key_dtype())?, false); + let value_dtype = map_dtype.value_dtype(); + let value = Field::new( + "value", + to_data_type_naive(&value_dtype)?, + value_dtype.is_nullable(), + ); + let entries = Field::new_struct("entries", Fields::from(vec![key, value]), false); + DataType::Map(FieldRef::new(entries), map_dtype.keys_sorted()) + } DType::Struct(struct_dtype, _) => { let mut fields = Vec::with_capacity(struct_dtype.names().len()); for (field_name, field_dt) in struct_dtype.names().iter().zip(struct_dtype.fields()) { @@ -433,9 +430,111 @@ fn variant_storage_fields_minimal() -> Fields { ]) } +/// Impls of the deprecated conversion traits, delegating to the naive free functions above. +#[allow(deprecated)] +mod deprecated_impls { + use super::*; + + impl TryFromArrowType<&DataType> for PType { + fn try_from_arrow(value: &DataType) -> VortexResult { + ptype_from_arrow(value) + } + } + + impl TryFromArrowType<&DataType> for DecimalDType { + fn try_from_arrow(value: &DataType) -> VortexResult { + decimal_dtype_from_arrow(value) + } + } + + impl FromArrowType<&ArrowTimeUnit> for TimeUnit { + fn from_arrow(value: &ArrowTimeUnit) -> Self { + from_arrow_time_unit(*value) + } + } + + impl FromArrowType for TimeUnit { + fn from_arrow(value: ArrowTimeUnit) -> Self { + from_arrow_time_unit(value) + } + } + + impl FromArrowType for DType { + fn from_arrow(value: SchemaRef) -> Self { + Self::from_arrow(value.as_ref()) + } + } + + impl TryFromArrowType for DType { + fn try_from_arrow(value: SchemaRef) -> VortexResult { + from_arrow_schema_naive(value.as_ref()) + } + } + + impl FromArrowType<&Schema> for DType { + fn from_arrow(value: &Schema) -> Self { + from_arrow_schema_naive(value).vortex_expect("arrow schema to dtype") + } + } + + impl TryFromArrowType<&Schema> for DType { + fn try_from_arrow(value: &Schema) -> VortexResult { + from_arrow_schema_naive(value) + } + } + + impl FromArrowType<&Fields> for StructFields { + fn from_arrow(value: &Fields) -> Self { + from_arrow_fields_naive(value).vortex_expect("arrow fields to struct fields") + } + } + + impl TryFromArrowType<&Fields> for StructFields { + fn try_from_arrow(value: &Fields) -> VortexResult { + from_arrow_fields_naive(value) + } + } + + impl FromArrowType<(&DataType, Nullability)> for DType { + fn from_arrow(value: (&DataType, Nullability)) -> Self { + from_arrow_data_type(value.0, value.1).vortex_expect("arrow data type to dtype") + } + } + + impl TryFromArrowType<(&DataType, Nullability)> for DType { + fn try_from_arrow( + (data_type, nullability): (&DataType, Nullability), + ) -> VortexResult { + from_arrow_data_type(data_type, nullability) + } + } + + impl FromArrowType<&Field> for DType { + fn from_arrow(field: &Field) -> Self { + from_arrow_field_naive(field).vortex_expect("arrow field to dtype") + } + } + + impl TryFromArrowType<&Field> for DType { + fn try_from_arrow(field: &Field) -> VortexResult { + from_arrow_field_naive(field) + } + } + + impl ToArrowType for DType { + fn to_arrow_schema(&self) -> VortexResult { + to_arrow_schema_naive(self) + } + + fn to_arrow_dtype(&self) -> VortexResult { + to_data_type_naive(self) + } + } +} + #[cfg(test)] mod test { - #![expect(deprecated, reason = "tests for deprecated methods")] + #![expect(deprecated, reason = "tests for deprecated traits and methods")] use arrow_schema::DataType; use arrow_schema::Field; use arrow_schema::FieldRef; @@ -646,6 +745,92 @@ mod test { assert_eq!(original_dtype, roundtripped_dtype); } + #[test] + fn map_dtype_roundtrip_uses_conventional_export_names_and_positional_import() -> VortexResult<()> + { + let dtype = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + Nullability::Nullable, + )?; + + let arrow = dtype.to_arrow_dtype()?; + let DataType::Map(entries, keys_sorted) = &arrow else { + panic!("expected Map, got {arrow:?}"); + }; + assert!(*keys_sorted); + assert_eq!(entries.name(), "entries"); + assert!(!entries.is_nullable()); + let DataType::Struct(fields) = entries.data_type() else { + panic!("expected map entries to be a struct"); + }; + assert_eq!(fields[0].name(), "key"); + assert!(!fields[0].is_nullable()); + assert_eq!(fields[1].name(), "value"); + assert!(fields[1].is_nullable()); + assert_eq!( + DType::try_from_arrow((&arrow, Nullability::Nullable))?, + dtype + ); + + let positional = DataType::Map( + Arc::new(Field::new_struct( + "anything", + Fields::from(vec![ + Field::new("first", DataType::Int32, false), + Field::new("second", DataType::Utf8, true), + ]), + false, + )), + false, + ); + assert_eq!( + DType::try_from_arrow((&positional, Nullability::NonNullable))?, + DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + false, + Nullability::NonNullable, + )? + ); + + Ok(()) + } + + #[test] + fn map_dtype_import_rejects_invalid_arrow_shape() { + let invalid_entries = [ + Field::new_struct( + "entries", + Fields::from(vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Utf8, true), + ]), + true, + ), + Field::new("entries", DataType::Int32, false), + Field::new_struct( + "entries", + Fields::from(vec![Field::new("key", DataType::Int32, false)]), + false, + ), + Field::new_struct( + "entries", + Fields::from(vec![ + Field::new("key", DataType::Int32, true), + Field::new("value", DataType::Utf8, true), + ]), + false, + ), + ]; + + for entries in invalid_entries { + let data_type = DataType::Map(Arc::new(entries), false); + assert!(DType::try_from_arrow((&data_type, Nullability::NonNullable)).is_err()); + } + } + // Regression test for https://github.com/vortex-data/vortex/issues/8346: unsupported Arrow // types must return an error instead of panicking with `unimplemented!`. #[rstest] diff --git a/vortex-arrow/src/executor/bool.rs b/vortex-arrow/src/executor/bool.rs index 64bffa804f2..f6815884b33 100644 --- a/vortex-arrow/src/executor/bool.rs +++ b/vortex-arrow/src/executor/bool.rs @@ -9,7 +9,9 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::arrays::BoolArray; use vortex_array::arrays::bool::BoolArrayExt; +use vortex_array::dtype::DType; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use crate::null_buffer::to_null_buffer; @@ -33,6 +35,11 @@ pub(super) fn to_arrow_bool( array: ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult { + vortex_ensure!( + matches!(array.dtype(), DType::Bool(_)), + "Cannot convert Vortex array with dtype {} to an Arrow Boolean array", + array.dtype() + ); let bool_array = array.execute::(ctx)?; canonical_bool_to_arrow(&bool_array, ctx) } diff --git a/vortex-arrow/src/executor/byte.rs b/vortex-arrow/src/executor/byte.rs index ae40ee01c26..43c71d24ba9 100644 --- a/vortex-arrow/src/executor/byte.rs +++ b/vortex-arrow/src/executor/byte.rs @@ -5,53 +5,98 @@ use std::sync::Arc; use arrow_array::ArrayRef as ArrowArrayRef; use arrow_array::GenericByteArray; -use arrow_array::types::BinaryViewType; use arrow_array::types::ByteArrayType; -use arrow_array::types::StringViewType; +use arrow_buffer::ArrowNativeType; +use arrow_buffer::Buffer as ArrowBuffer; +use arrow_buffer::NullBuffer; +use arrow_buffer::OffsetBuffer; +use arrow_schema::DataType; +use num_traits::AsPrimitive; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::Canonical; use vortex_array::ExecutionCtx; +use vortex_array::arrays::Chunked; +use vortex_array::arrays::Constant; use vortex_array::arrays::VarBin; -use vortex_array::arrays::VarBinViewArray; -use vortex_array::arrays::varbin::VarBinArrayExt; +use vortex_array::arrays::varbin::VarBinArraySlotsExt; +use vortex_array::builders::VarBinBuilder; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; -use vortex_error::VortexError; +use vortex_array::dtype::OffsetBuilderPType; +use vortex_array::matcher::Matcher; use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; -use crate::byte_view::execute_varbinview_to_arrow; use crate::executor::validity::to_arrow_null_buffer; +/// Matches the encodings [`to_arrow_byte_array`] requires for export. +/// +/// `Chunked` and `Constant` are matched to stop execution before it destroys them: they have +/// specialized `append_to_builder` impls (chunk-wise append, scalar repeat) that the builder +/// fallback exploits. +struct ArrowByteExportable; + +impl Matcher for ArrowByteExportable { + type Match<'a> = &'a ArrayRef; + + fn try_match(array: &ArrayRef) -> Option> { + (array.is::() || array.is::() || array.is::()).then_some(array) + } +} + /// Convert a Vortex array into an Arrow GenericBinaryArray. pub(super) fn to_arrow_byte_array( array: ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult where - T::Offset: NativePType, + T::Offset: OffsetBuilderPType, + usize: AsPrimitive, { - // If the Vortex array is already in VarBin format, we can directly convert it. + if !matches!(array.dtype(), DType::Utf8(_) | DType::Binary(_)) { + vortex_bail!( + "Cannot convert Vortex array with dtype {} to Arrow byte array type {}", + array.dtype(), + T::DATA_TYPE + ); + } + + // A logical dtype mismatch changes nothing about the physical export except that a `Binary` + // source exported to `Utf8` has to have its bytes validated. + let source_is_utf8 = matches!(array.dtype(), DType::Utf8(_)); + let target_is_utf8 = matches!(T::DATA_TYPE, DataType::Utf8 | DataType::LargeUtf8); + let validate_utf8 = target_is_utf8 && !source_is_utf8; + + let array = array.execute_until::(ctx)?; + + // If the Vortex array is in VarBin format, we can directly convert it. if let Some(array) = array.as_opt::() { - return varbin_to_byte_array::(array, ctx); - } - - // Otherwise, we execute the array to a VarBinViewArray and convert to Arrow ByteView, - // then cast to the target byte array type. - let varbinview = array.execute::(ctx)?; - let binary_view = match varbinview.dtype() { - DType::Utf8(_) => execute_varbinview_to_arrow::(&varbinview, ctx), - DType::Binary(_) => execute_varbinview_to_arrow::(&varbinview, ctx), - _ => unreachable!("VarBinViewArray must have Utf8 or Binary dtype"), - }?; - arrow_cast::cast(&binary_view, &T::DATA_TYPE).map_err(VortexError::from) + return varbin_to_byte_array::(array, validate_utf8, ctx); + } + + // The builder's offset type matches the Arrow target, so `varbin_to_byte_array` hands the + // offsets buffer straight to Arrow without a cast. + let mut builder = VarBinBuilder::::with_capacity_in( + array.dtype().clone(), + array.len(), + ctx.allocator(), + ); + array.append_to_builder(&mut builder, ctx)?; + varbin_to_byte_array::(builder.finish_into_varbin().as_view(), validate_utf8, ctx) } /// Convert a Vortex VarBinArray into an Arrow GenericBinaryArray. +/// +/// `validate_utf8` must be set when the Vortex dtype does not already guarantee the bytes are +/// valid for `T`, i.e. a `Binary` source exported to a `Utf8` target. fn varbin_to_byte_array( array: ArrayView<'_, VarBin>, + validate_utf8: bool, ctx: &mut ExecutionCtx, ) -> VortexResult where @@ -69,11 +114,78 @@ where let data = array.bytes().clone().into_arrow_buffer(); let null_buffer = to_arrow_null_buffer(array.validity()?, array.len(), ctx)?; + if validate_utf8 { + validate_live_values_utf8::(&offsets, &data, null_buffer.as_ref())?; + } + // SAFETY: `T::DATA_TYPE` is a `Utf8` variant only when the source dtype already guarantees the + // bytes are UTF-8, or when `validate_utf8` made us prove it above. Ok(Arc::new(unsafe { GenericByteArray::::new_unchecked(offsets, data, null_buffer) })) } +/// Checks that every non-null value spanned by `offsets` is valid UTF-8. +/// +/// [`GenericByteArray::try_new`] would instead validate the whole byte range the offsets span, +/// which also covers the extents of null slots. Those bytes are never exposed as a value, and +/// Vortex only requires non-null values to be UTF-8 (see `VarBinArray::validate_utf8`), so +/// checking them would reject perfectly well-formed exports. +fn validate_live_values_utf8( + offsets: &OffsetBuffer, + values: &ArrowBuffer, + nulls: Option<&NullBuffer>, +) -> VortexResult<()> { + let value_at = |index: usize, start: usize, end: usize| -> VortexResult<&[u8]> { + values + .get(start..end) + .ok_or_else(|| vortex_err!("Offsets {start}..{end} at index {index} are out of bounds")) + }; + + let Some(nulls) = nulls.filter(|nulls| nulls.null_count() > 0) else { + // With no nulls every byte the offsets span belongs to a value, so a single pass validates + // all of them at once. The offsets then only have to land on character boundaries, which + // is what stops `value(i)` from slicing a multi-byte character in half. + let (Some(first), Some(last)) = (offsets.first(), offsets.last()) else { + return Ok(()); + }; + let start = first.as_usize(); + let bytes = value_at(0, start, last.as_usize())?; + let validated = utf8_from_bytes(bytes) + .map_err(|err| vortex_err!("Encountered non UTF-8 data: {err}"))?; + for (index, offset) in offsets.iter().enumerate() { + let boundary = offset + .as_usize() + .checked_sub(start) + .filter(|boundary| validated.is_char_boundary(*boundary)); + vortex_ensure!( + boundary.is_some(), + "Offset {} at index {index} does not fall on a UTF-8 character boundary", + offset.as_usize() + ); + } + return Ok(()); + }; + + for (index, window) in offsets.windows(2).enumerate() { + if nulls.is_null(index) { + continue; + } + let bytes = value_at(index, window[0].as_usize(), window[1].as_usize())?; + utf8_from_bytes(bytes) + .map_err(|err| vortex_err!("Encountered non UTF-8 data at index {index}: {err}"))?; + } + Ok(()) +} + +/// Validates `bytes` as UTF-8, re-running the slower checker only to describe a failure. +fn utf8_from_bytes(bytes: &[u8]) -> Result<&str, simdutf8::compat::Utf8Error> { + simdutf8::basic::from_utf8(bytes).map_err(|_| { + #[expect(clippy::unwrap_used)] + // Run validation again using the `compat` module to get a more detailed error message. + simdutf8::compat::from_utf8(bytes).unwrap_err() + }) +} + #[cfg(test)] mod tests { use arrow_array::Array; @@ -84,13 +196,46 @@ mod tests { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::VarBinArray; + use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; + use vortex_array::scalar_fn::fns::mask::Mask as MaskFn; + use vortex_array::validity::Validity; + use vortex_buffer::ByteBuffer; + use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_mask::Mask; use crate::ArrowSessionExt; - use crate::executor::byte::VarBinViewArray; + + #[test] + fn mask_wrapped_varbin_exports() -> VortexResult<()> { + let session = array_session(); + let mut ctx = session.create_execution_ctx(); + + let varbin = VarBinArray::from_vec( + vec!["hello", "world", "vortex"], + DType::Utf8(Nullability::NonNullable), + ); + let mask = BoolArray::from_iter([true, false, true]); + let masked = MaskFn::try_new(varbin.into_array(), mask.into_array())?.into_array(); + + let field = Field::new("s", DataType::Utf8, true); + let arrow = session + .arrow() + .execute_arrow(masked, Some(&field), &mut ctx)?; + + let strings = arrow.as_string::(); + assert_eq!(strings.len(), 3); + assert!(!strings.is_null(0)); + assert!(strings.is_null(1)); + assert_eq!(strings.value(0), "hello"); + assert_eq!(strings.value(2), "vortex"); + Ok(()) + } fn make_utf8_array() -> VarBinViewArray { VarBinViewArray::from_iter_str(["hello", "world", "this is a longer string for testing"]) @@ -105,13 +250,13 @@ mod tests { } #[rstest] - // Utf8 source can convert to all string types and binary types (via arrow_cast) + // Utf8 source can convert to all string types and binary types #[case::utf8_to_binary(make_utf8_array(), DataType::Binary)] #[case::utf8_to_large_binary(make_utf8_array(), DataType::LargeBinary)] #[case::utf8_to_utf8(make_utf8_array(), DataType::Utf8)] #[case::utf8_to_large_utf8(make_utf8_array(), DataType::LargeUtf8)] #[case::utf8_to_utf8_view(make_utf8_array(), DataType::Utf8View)] - // Binary source can convert to all binary types and string types (via arrow_cast) + // Binary source can convert to all binary types and string types #[case::binary_to_binary(make_binary_array(), DataType::Binary)] #[case::binary_to_large_binary(make_binary_array(), DataType::LargeBinary)] #[case::binary_to_utf8(make_binary_array(), DataType::Utf8)] @@ -187,6 +332,125 @@ mod tests { assert!(!arrow.is_null(2)); } + #[rstest] + #[case(DataType::Utf8)] + #[case(DataType::LargeUtf8)] + fn binary_with_invalid_utf8_to_string_returns_error(#[case] target_dtype: DataType) { + let vortex_array = VarBinViewArray::from_iter_bin([ + b"hello".as_slice(), + b"\xff\xfe invalid utf8".as_slice(), + ]); + + let session = array_session(); + let mut ctx = session.create_execution_ctx(); + let field = Field::new("test_field", target_dtype, true); + let result = + session + .arrow() + .execute_arrow(vortex_array.into_array(), Some(&field), &mut ctx); + + assert!(result.is_err()); + } + + /// The whole values buffer is valid UTF-8, but each individual value slices a two-byte + /// character in half. Arrow's `value(i)` reinterprets the bytes unchecked, so the export has + /// to reject this rather than hand out a malformed `&str`. + #[rstest] + #[case(DataType::Utf8)] + #[case(DataType::LargeUtf8)] + fn binary_to_string_rejects_offsets_splitting_a_character(#[case] target_dtype: DataType) { + let array = VarBinArray::new( + buffer![0i32, 1, 2].into_array(), + ByteBuffer::copy_from("é".as_bytes()), + DType::Binary(Nullability::NonNullable), + Validity::NonNullable, + ); + + let session = array_session(); + let mut ctx = session.create_execution_ctx(); + let field = Field::new("test_field", target_dtype, false); + let result = session + .arrow() + .execute_arrow(array.into_array(), Some(&field), &mut ctx); + + assert!(result.is_err()); + } + + /// Multi-byte characters that do line up with the offsets must still export. + #[rstest] + #[case(DataType::Utf8)] + #[case(DataType::LargeUtf8)] + fn binary_to_string_accepts_multibyte_characters( + #[case] target_dtype: DataType, + ) -> VortexResult<()> { + let array = VarBinArray::new( + buffer![0i32, 2, 5].into_array(), + ByteBuffer::copy_from("é☃".as_bytes()), + DType::Binary(Nullability::NonNullable), + Validity::NonNullable, + ); + + let session = array_session(); + let mut ctx = session.create_execution_ctx(); + let field = Field::new("test_field", target_dtype.clone(), false); + let arrow = session + .arrow() + .execute_arrow(array.into_array(), Some(&field), &mut ctx)?; + + let values: Vec<&str> = match target_dtype { + DataType::Utf8 => (0..2).map(|i| arrow.as_string::().value(i)).collect(), + _ => (0..2).map(|i| arrow.as_string::().value(i)).collect(), + }; + assert_eq!(values, ["é", "☃"]); + Ok(()) + } + + /// The `VarBin` fast path hands Arrow the untrimmed bytes buffer, so a `Binary` to `Utf8` + /// export must not validate bytes that no live value points at. + #[rstest] + #[case(DataType::Utf8)] + #[case(DataType::LargeUtf8)] + fn varbin_binary_to_string_ignores_bytes_outside_values( + #[case] target_dtype: DataType, + ) -> VortexResult<()> { + // Slot 1 is null and its offsets still span two bytes that are not valid UTF-8. + let array = VarBinArray::new( + buffer![0i32, 5, 7, 12].into_array(), + ByteBuffer::copy_from(b"hello\xff\xfeworld".as_slice()), + DType::Binary(Nullability::Nullable), + Validity::from_mask(Mask::from_iter([true, false, true]), Nullability::Nullable), + ); + + let session = array_session(); + let mut ctx = session.create_execution_ctx(); + let field = Field::new("test_field", target_dtype, true); + let arrow = session + .arrow() + .execute_arrow(array.into_array(), Some(&field), &mut ctx)?; + + assert_eq!(arrow.len(), 3); + assert!(arrow.is_null(1)); + Ok(()) + } + + #[rstest] + #[case(DataType::Utf8)] + #[case(DataType::LargeUtf8)] + #[case(DataType::Binary)] + #[case(DataType::LargeBinary)] + fn incompatible_vortex_dtype_returns_error(#[case] target_dtype: DataType) { + let session = array_session(); + let mut ctx = session.create_execution_ctx(); + let field = Field::new("test_field", target_dtype, true); + let result = session.arrow().execute_arrow( + PrimitiveArray::from_iter([1i32, 2, 3]).into_array(), + Some(&field), + &mut ctx, + ); + + assert!(result.is_err()); + } + #[test] fn filtered_utf8_view_export_does_not_retain_unselected_buffers() -> VortexResult<()> { let unselected = "x".repeat(1 << 20); diff --git a/vortex-arrow/src/executor/byte_view.rs b/vortex-arrow/src/executor/byte_view.rs index ffddc49a4dd..3026317b511 100644 --- a/vortex-arrow/src/executor/byte_view.rs +++ b/vortex-arrow/src/executor/byte_view.rs @@ -6,16 +6,15 @@ use std::sync::Arc; use arrow_array::ArrayRef as ArrowArrayRef; use arrow_array::GenericByteViewArray; use arrow_array::types::ByteViewType; -use arrow_buffer::ScalarBuffer; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::arrays::VarBinViewArray; use vortex_array::builtins::ArrayBuiltins; -use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; +use vortex_buffer::Buffer; use vortex_error::VortexResult; -use crate::dtype::FromArrowType; +use crate::dtype::from_arrow_data_type; use crate::null_buffer::to_null_buffer; /// Convert a canonical VarBinViewArray directly to Arrow. @@ -23,8 +22,8 @@ pub fn canonical_varbinview_to_arrow( array: &VarBinViewArray, ctx: &mut ExecutionCtx, ) -> VortexResult { - let views = - ScalarBuffer::::from(array.views_handle().as_host().clone().into_arrow_buffer()); + let views = Buffer::::from_byte_buffer(array.views_handle().as_host().clone()) + .into_arrow_scalar_buffer(); let buffers: Vec<_> = array .data_buffers() .iter() @@ -59,9 +58,29 @@ pub(super) fn to_arrow_byte_view( // We do this in case the vortex array is Utf8, and we want Binary or vice versa. By casting // first, we may push this down through the Vortex array tree. We choose nullable to be most // flexible since there's no prescribed nullability in Arrow types. - let array = array.cast(DType::from_arrow((&T::DATA_TYPE, Nullability::Nullable)))?; + let array = array.cast(from_arrow_data_type(&T::DATA_TYPE, Nullability::Nullable)?)?; let array = array.execute::(ctx)?; let varbinview = array.execute::(ctx)?; execute_varbinview_to_arrow::(&varbinview, ctx) } + +#[cfg(test)] +mod tests { + use arrow_array::types::StringViewType; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + + use super::*; + + #[test] + fn empty_views_are_aligned() -> VortexResult<()> { + let array = VarBinViewArray::from_iter_str(std::iter::empty::<&str>()); + let mut ctx = array_session().create_execution_ctx(); + + let arrow = canonical_varbinview_to_arrow::(&array, &mut ctx)?; + + assert!(arrow.is_empty()); + Ok(()) + } +} diff --git a/vortex-arrow/src/executor/decimal.rs b/vortex-arrow/src/executor/decimal.rs index d5eabc77145..41ffa93630f 100644 --- a/vortex-arrow/src/executor/decimal.rs +++ b/vortex-arrow/src/executor/decimal.rs @@ -16,9 +16,11 @@ use num_traits::ToPrimitive; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::arrays::DecimalArray; +use vortex_array::dtype::DType; use vortex_array::dtype::DecimalType; use vortex_buffer::Buffer; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_error::vortex_err; use crate::null_buffer::to_null_buffer; @@ -28,6 +30,13 @@ pub(super) fn to_arrow_decimal( data_type: &DataType, ctx: &mut ExecutionCtx, ) -> VortexResult { + vortex_ensure!( + matches!(array.dtype(), DType::Decimal(..)), + "Cannot convert Vortex array with dtype {} to an Arrow {} array", + array.dtype(), + data_type + ); + // Execute the array as a DecimalArray. let decimal_array = array.execute::(ctx)?; @@ -265,7 +274,11 @@ mod tests { #[case] _decimal_type: T, ) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - let mut decimal = DecimalBuilder::new::(DecimalDType::new(2, 1), false.into()); + let mut decimal = DecimalBuilder::new_in::( + DecimalDType::new(2, 1), + false.into(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); decimal.append_value(10); decimal.append_value(11); decimal.append_value(12); @@ -293,7 +306,11 @@ mod tests { use arrow_array::Decimal32Array; let mut ctx = array_session().create_execution_ctx(); - let mut decimal = DecimalBuilder::new::(DecimalDType::new(2, 1), false.into()); + let mut decimal = DecimalBuilder::new_in::( + DecimalDType::new(2, 1), + false.into(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); decimal.append_value(10); decimal.append_value(11); decimal.append_value(12); @@ -321,7 +338,11 @@ mod tests { use arrow_array::Decimal64Array; let mut ctx = array_session().create_execution_ctx(); - let mut decimal = DecimalBuilder::new::(DecimalDType::new(2, 1), false.into()); + let mut decimal = DecimalBuilder::new_in::( + DecimalDType::new(2, 1), + false.into(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); decimal.append_value(10); decimal.append_value(11); decimal.append_value(12); @@ -349,7 +370,11 @@ mod tests { #[case] _decimal_type: T, ) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - let mut decimal = DecimalBuilder::new::(DecimalDType::new(2, 1), false.into()); + let mut decimal = DecimalBuilder::new_in::( + DecimalDType::new(2, 1), + false.into(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); decimal.append_value(10); decimal.append_value(11); decimal.append_value(12); diff --git a/vortex-arrow/src/executor/dictionary.rs b/vortex-arrow/src/executor/dictionary.rs index f5be1076a4b..1a17d9de1dd 100644 --- a/vortex-arrow/src/executor/dictionary.rs +++ b/vortex-arrow/src/executor/dictionary.rs @@ -18,18 +18,32 @@ use vortex_array::arrays::ConstantArray; use vortex_array::arrays::Dict; use vortex_array::arrays::DictArray; use vortex_array::arrays::dict::DictArraySlotsExt; +use vortex_array::matcher::Matcher; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; use crate::ArrowArrayExecutor; +/// Matches the encodings [`to_arrow_dictionary`] requires for export. +struct ArrowDictExportable; + +impl Matcher for ArrowDictExportable { + type Match<'a> = &'a ArrayRef; + + fn try_match(array: &ArrayRef) -> Option> { + (array.is::() || array.is::()).then_some(array) + } +} + pub(super) fn to_arrow_dictionary( array: ArrayRef, codes_type: &DataType, values_type: &DataType, ctx: &mut ExecutionCtx, ) -> VortexResult { + let array = array.execute_until::(ctx)?; + let array = match array.try_downcast::() { Ok(dict) => return dict_to_dict(dict, codes_type, values_type, ctx), Err(array) => array, @@ -141,6 +155,7 @@ mod tests { use std::sync::Arc; use arrow_array::DictionaryArray as ArrowDictArray; + use arrow_array::StringArray; use arrow_array::types::UInt8Type; use arrow_array::types::UInt32Type; use arrow_schema::DataType; @@ -148,11 +163,13 @@ mod tests { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; + use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability::Nullable; use vortex_array::scalar::Scalar; + use vortex_array::scalar_fn::fns::mask::Mask; use vortex_buffer::buffer; use vortex_error::VortexResult; @@ -224,4 +241,23 @@ mod tests { assert_eq!(expected.as_ref(), actual.as_ref()); Ok(()) } + + #[test] + fn mask_wrapped_dict_exports() -> VortexResult<()> { + // Dictionary behind a lazy `mask` scalar-fn — the shape a scan produces when a row + // mask is applied to a dict-encoded column. + let dict = DictArray::try_new( + buffer![0u8, 1, 0].into_array(), + VarBinViewArray::from_iter_str(["a", "b"]).into_array(), + )?; + let mask = BoolArray::from_iter([true, false, true]); + let masked = Mask::try_new(dict.into_array(), mask.into_array())?.into_array(); + + let actual = execute(masked, &dict_type(DataType::UInt8, DataType::Utf8))?; + let flat = arrow_cast::cast(&actual, &DataType::Utf8)?; + + let expected = StringArray::from(vec![Some("a"), None, Some("a")]); + assert_eq!(flat.as_ref(), &expected as &dyn arrow_array::Array); + Ok(()) + } } diff --git a/vortex-arrow/src/executor/fixed_size_list.rs b/vortex-arrow/src/executor/fixed_size_list.rs index 72ce7c57ff0..7220bbe13f7 100644 --- a/vortex-arrow/src/executor/fixed_size_list.rs +++ b/vortex-arrow/src/executor/fixed_size_list.rs @@ -9,6 +9,8 @@ use vortex_array::ExecutionCtx; use vortex_array::arrays::FixedSizeList; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; +use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use vortex_array::dtype::DType; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -21,6 +23,12 @@ pub(super) fn to_arrow_fixed_list( elements_field: &FieldRef, ctx: &mut ExecutionCtx, ) -> VortexResult { + vortex_ensure!( + matches!(array.dtype(), DType::FixedSizeList(..)), + "Cannot convert Vortex array with dtype {} to an Arrow FixedSizeList array", + array.dtype() + ); + // Check for Vortex FixedSizeListArray and convert directly. if let Some(array) = array.as_opt::() { return list_to_list(&array.into_owned(), elements_field, list_size, ctx); diff --git a/vortex-arrow/src/executor/list.rs b/vortex-arrow/src/executor/list.rs index 0b0d7f625b5..f995661a0c1 100644 --- a/vortex-arrow/src/executor/list.rs +++ b/vortex-arrow/src/executor/list.rs @@ -17,14 +17,15 @@ use vortex_array::arrays::ListArray; use vortex_array::arrays::ListView; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::chunked::ChunkedArrayExt; -use vortex_array::arrays::list::ListArrayExt; -use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::list::ListArraySlotsExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::arrays::listview::ListViewDataParts; use vortex_array::arrays::listview::ListViewRebuildMode; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; +use vortex_array::matcher::Matcher; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -33,6 +34,17 @@ use vortex_error::vortex_ensure; use crate::executor::validity::to_arrow_null_buffer; use crate::session::ArrowSessionExt; +/// Matches the encodings [`to_arrow_list`] requires for export. +struct ArrowListExportable; + +impl Matcher for ArrowListExportable { + type Match<'a> = &'a ArrayRef; + + fn try_match(array: &ArrayRef) -> Option> { + (array.is::() || array.is::() || array.is::()).then_some(array) + } +} + #[allow(rustdoc::broken_intra_doc_links)] /// Convert a Vortex VarBinArray into an Arrow [`GenericListArray`](arrow_array:array::GenericListArray). pub(super) fn to_arrow_list( @@ -40,9 +52,17 @@ pub(super) fn to_arrow_list( elements_field: &FieldRef, ctx: &mut ExecutionCtx, ) -> VortexResult { + vortex_ensure!( + matches!(array.dtype(), DType::List(..)), + "Cannot convert Vortex array with dtype {} to an Arrow list array", + array.dtype() + ); + + let array = array.execute_until::(ctx)?; + // If the Vortex array is already in List format, we can directly convert it. - if let Some(array) = array.as_opt::() { - return list_to_list::(&array.into_owned(), elements_field, ctx); + if let Some(list) = array.as_opt::() { + return list_to_list::(&list.into_owned(), elements_field, ctx); } // Converting each chunk individually, then using the fast concat logic from arrow @@ -56,24 +76,15 @@ pub(super) fn to_arrow_list( return Ok(arrow_select::concat::concat(&refs)?); } - // If the Vortex array is a ListViewArray, rebuild to ZCTL if needed and convert. - let array = match array.try_downcast::() { - Ok(array) => { - let zctl = if array.is_zero_copy_to_list() { - array - } else { - array.rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx)? - }; - return list_view_zctl::(zctl, elements_field, ctx); - } - Err(a) => a, - }; - - // Otherwise, we execute the array to become a ListViewArray, then rebuild to ZCTL. + // Otherwise the array is canonical: a ListViewArray, which we rebuild to ZCTL if needed. // Note: arrow_cast::cast supports ListView → List (apache/arrow-rs#8735), but it // unconditionally uses take. Our rebuild uses a heuristic that picks list-by-list // for large lists, which avoids materializing a large index buffer. - let list_view = array.execute::(ctx)?; + let list_view = array + .as_opt::() + .vortex_expect("Must be ListView from Matcher") + .into_owned(); + let zctl = if list_view.is_zero_copy_to_list() { list_view } else { @@ -213,15 +224,19 @@ mod tests { use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::SliceArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability::NonNullable; + use vortex_array::scalar_fn::fns::mask::Mask; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::ArrowArrayExecutor; + use crate::executor::list::ListArray; use crate::executor::list::ListViewArray; /// A shared session for these list-executor tests, used to create execution contexts. @@ -364,6 +379,92 @@ mod tests { Ok(()) } + #[test] + fn slice_wrapped_list_exports() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + // Lists [[1, 2], [3], [4, 5, 6]] behind a lazy Slice wrapper selecting rows 1..3. + let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6], Validity::NonNullable); + let elements_ptr = elements.as_slice::().as_ptr(); + let offsets = PrimitiveArray::new(buffer![0i32, 2, 3, 6], Validity::NonNullable); + let list = ListArray::try_new( + elements.into_array(), + offsets.into_array(), + Validity::NonNullable, + )?; + let sliced = SliceArray::new(list.into_array(), 1..3).into_array(); + + let field = Field::new("item", DataType::Int32, false); + let arrow_dt = DataType::List(field.into()); + let arrow_array = sliced.execute_arrow(Some(&arrow_dt), &mut ctx)?; + + let arrow_list = arrow_array + .as_any() + .downcast_ref::>() + .unwrap(); + assert_eq!(arrow_list.len(), 2); + let first = arrow_list.value(0); + let first_vals = first.as_any().downcast_ref::().unwrap(); + assert_eq!(first_vals.values(), &[3]); + let second = arrow_list.value(1); + let second_vals = second.as_any().downcast_ref::().unwrap(); + assert_eq!(second_vals.values(), &[4, 5, 6]); + + // The conversion shares the elements buffer regardless of which path resolves the + // Slice (revealed List or zero-copy-to-list ListView). + let values = arrow_list + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.values().as_ptr(), elements_ptr); + Ok(()) + } + + #[test] + fn mask_wrapped_list_exports() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + // Lists [[1, 2], [3], [4, 5, 6]] behind a lazy `mask` scalar-fn nulling out row 1 — + // the shape a scan produces when a row mask is applied to a List-encoded column. + let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6], Validity::NonNullable); + let elements_ptr = elements.as_slice::().as_ptr(); + let offsets = PrimitiveArray::new(buffer![0i32, 2, 3, 6], Validity::NonNullable); + let list = ListArray::try_new( + elements.into_array(), + offsets.into_array(), + Validity::NonNullable, + )?; + let mask = BoolArray::from_iter([true, false, true]); + let masked = Mask::try_new(list.into_array(), mask.into_array())?.into_array(); + + let field = Field::new("item", DataType::Int32, false); + let arrow_dt = DataType::List(field.into()); + let arrow_array = masked.execute_arrow(Some(&arrow_dt), &mut ctx)?; + + let arrow_list = arrow_array + .as_any() + .downcast_ref::>() + .unwrap(); + assert_eq!(arrow_list.len(), 3); + assert!(!arrow_list.is_null(0)); + assert!(arrow_list.is_null(1)); + assert!(!arrow_list.is_null(2)); + let first = arrow_list.value(0); + let first_vals = first.as_any().downcast_ref::().unwrap(); + assert_eq!(first_vals.values(), &[1, 2]); + let third = arrow_list.value(2); + let third_vals = third.as_any().downcast_ref::().unwrap(); + assert_eq!(third_vals.values(), &[4, 5, 6]); + + // Masking only touches validity, so the conversion still shares the elements buffer. + let values = arrow_list + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.values().as_ptr(), elements_ptr); + Ok(()) + } + #[test] fn test_to_arrow_list_empty_zctl() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); diff --git a/vortex-arrow/src/executor/list_view.rs b/vortex-arrow/src/executor/list_view.rs index da242eb9719..53d0071f5ff 100644 --- a/vortex-arrow/src/executor/list_view.rs +++ b/vortex-arrow/src/executor/list_view.rs @@ -13,6 +13,7 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::listview::DEFAULT_REBUILD_DENSITY_THRESHOLD; use vortex_array::arrays::listview::DEFAULT_TRIM_ELEMENTS_THRESHOLD; use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::arrays::listview::ListViewDataParts; use vortex_array::arrays::listview::ListViewRebuildMode; use vortex_array::builtins::ArrayBuiltins; @@ -30,6 +31,12 @@ pub(super) fn to_arrow_list_view( elements_field: &FieldRef, ctx: &mut ExecutionCtx, ) -> VortexResult { + vortex_ensure!( + matches!(array.dtype(), DType::List(..)), + "Cannot convert Vortex array with dtype {} to an Arrow list-view array", + array.dtype() + ); + let array = array.execute::(ctx)?; // Reclaim unreferenced elements before handing the array to Arrow. Otherwise downstream @@ -75,6 +82,7 @@ fn list_view_to_list_view( .. } = array.into_data_parts(); + let n_elements = elements.len(); let elements = ctx.session().clone().arrow().execute_arrow( elements, Some(elements_field.as_ref()), @@ -84,6 +92,13 @@ fn list_view_to_list_view( elements_field.is_nullable() || elements.null_count() == 0, "Elements field is non-nullable but elements array contains nulls" ); + // The unchecked construction below needs the views in bounds of the *exported* elements, so + // confirm the export preserved the length the Vortex invariant was checked against. + vortex_ensure!( + elements.len() == n_elements, + "Arrow export changed the elements length: {n_elements} became {}", + elements.len() + ); let offsets = offsets .cast(DType::Primitive(O::PTYPE, NonNullable))? @@ -98,13 +113,17 @@ fn list_view_to_list_view( let null_buffer = to_arrow_null_buffer(validity, offsets.len(), ctx)?; - Ok(Arc::new(GenericListViewArray::::new( - Arc::clone(elements_field), - offsets, - sizes, - elements, - null_buffer, - ))) + // SAFETY: `ListViewData::validate` guarantees `offsets[i] + sizes[i] <= elements.len()`, and + // the length check above carries that across the export. Skips Arrow's unvectorizable re-check. + Ok(Arc::new(unsafe { + GenericListViewArray::::new_unchecked( + Arc::clone(elements_field), + offsets, + sizes, + elements, + null_buffer, + ) + })) } #[cfg(test)] diff --git a/vortex-arrow/src/executor/map.rs b/vortex-arrow/src/executor/map.rs index 350daa75fd3..016edeb5e1e 100644 --- a/vortex-arrow/src/executor/map.rs +++ b/vortex-arrow/src/executor/map.rs @@ -3,49 +3,521 @@ use std::sync::Arc; +use arrow_array::Array as ArrowArray; use arrow_array::ArrayRef as ArrowArrayRef; -use arrow_array::GenericListArray; use arrow_array::MapArray as ArrowMapArray; -use arrow_array::StructArray as ArrowStructArray; +use arrow_array::cast::AsArray; +use arrow_schema::DataType; use arrow_schema::FieldRef; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; +use vortex_array::arrays::Map; +use vortex_array::arrays::MapArray; +use vortex_array::arrays::map::MapArrayExt; +use vortex_array::arrays::map::MapArraySlotsExt; use vortex_error::VortexResult; -use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; use crate::executor::list::to_arrow_list; -/// Convert a Vortex `List>` array into an Arrow [`MapArray`](ArrowMapArray). +/// Converts a Vortex Map array into an Arrow [`MapArray`](ArrowMapArray). /// -/// Arrow's `Map` has no Vortex `DType` of its own; it is aliased to `List>` -/// on the way in (see [`crate::dtype`]), so restoring the `Map` on the way out is a matter of -/// re-attaching the entries field and the `ordered` flag to the same offsets and entries. +/// The Map entries are exported through the ListView-to-Arrow List path and then repackaged with +/// Arrow's Map field metadata. When the requested Arrow field asserts sorted keys, the Vortex Map +/// dtype must already make the same assertion. pub(super) fn to_arrow_map( array: ArrayRef, entries_field: &FieldRef, - ordered: bool, + keys_sorted: bool, ctx: &mut ExecutionCtx, ) -> VortexResult { - // Map offsets are always i32. - let list_array = to_arrow_list::(array, entries_field, ctx)?; - - let Some(list_array) = list_array.as_any().downcast_ref::>() else { - vortex_bail!("to_arrow_list returned a non-ListArray when building a MapArray"); + let array = match array.try_downcast::() { + Ok(map) => map, + Err(array) => array.execute::(ctx)?, }; - let (_list_field, offsets, entries, nulls) = list_array.clone().into_parts(); - let Some(entries_struct) = entries.as_any().downcast_ref::() else { - vortex_bail!( - "Map entries must be a StructArray, got {}", - entries.data_type() - ); - }; + vortex_ensure!( + !keys_sorted || array.keys_sorted(), + "Cannot convert unsorted Vortex map to Arrow MapArray with keys_sorted=true" + ); + + let entries = array.entries().clone(); + let entries_list_type = DataType::List(Arc::clone(entries_field)); + let entries_list = to_arrow_list::(entries, entries_field, ctx)?; + vortex_ensure!( + entries_list.data_type() == &entries_list_type, + "Arrow Map entries converted to {}, expected {entries_list_type}", + entries_list.data_type() + ); - Ok(Arc::new(ArrowMapArray::try_new( + let entries_list = entries_list.as_list::(); + let entries = entries_list.values().as_struct().clone(); + let map = ArrowMapArray::try_new( Arc::clone(entries_field), - offsets, - entries_struct.clone(), - nulls, - ordered, - )?)) + entries_list.offsets().clone(), + entries, + entries_list.nulls().cloned(), + keys_sorted, + )?; + + Ok(Arc::new(map)) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_array::Array as ArrowArray; + use arrow_array::ArrayRef as ArrowArrayRef; + use arrow_array::FixedSizeBinaryArray; + use arrow_array::Int32Array; + use arrow_array::MapArray as ArrowMapArray; + use arrow_array::StringArray; + use arrow_array::StructArray as ArrowStructArray; + use arrow_array::builder::Int32Builder; + use arrow_array::builder::MapBuilder as ArrowMapBuilder; + use arrow_array::builder::StringBuilder; + use arrow_array::cast::AsArray; + use arrow_buffer::NullBuffer; + use arrow_buffer::OffsetBuffer; + use arrow_buffer::ScalarBuffer; + use arrow_schema::DataType; + use arrow_schema::Field; + use arrow_schema::Fields; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::ListView; + use vortex_array::arrays::Map as VortexMap; + use vortex_array::arrays::listview::ListViewArrayExt; + use vortex_array::arrays::map::MapArraySlotsExt; + use vortex_array::builders::ArrayBuilder; + use vortex_array::builders::MapBuilder; + use vortex_array::dtype::DType; + use vortex_array::dtype::MapDType; + use vortex_array::dtype::Nullability::NonNullable; + use vortex_array::dtype::Nullability::Nullable; + use vortex_array::dtype::PType; + use vortex_array::scalar::Scalar; + use vortex_error::VortexResult; + use vortex_mask::Mask; + + use crate::FromArrowArray as _; + use crate::session::ArrowSessionExt as _; + + fn i32_utf8_map_field(keys_sorted: bool, nullable: bool) -> Field { + let fields = Fields::from(vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Utf8, true), + ]); + Field::new( + "maps", + DataType::Map( + Arc::new(Field::new_struct("entries", fields, false)), + keys_sorted, + ), + nullable, + ) + } + + fn i32_utf8_map_array( + offsets: Vec, + keys: Vec, + values: Vec>, + nulls: Option, + keys_sorted: bool, + ) -> VortexResult { + let fields = Fields::from(vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Utf8, true), + ]); + let entries = ArrowStructArray::try_new( + fields.clone(), + vec![ + Arc::new(Int32Array::from(keys)), + Arc::new(StringArray::from(values)), + ], + None, + )?; + ArrowMapArray::try_new( + Arc::new(Field::new_struct("entries", fields, false)), + OffsetBuffer::new(ScalarBuffer::from(offsets)), + entries, + nulls, + keys_sorted, + ) + .map_err(Into::into) + } + + #[test] + fn map_roundtrip_preserves_null_empty_duplicate_and_sorted_rows() -> VortexResult<()> { + let vortex_session = array_session(); + let mut ctx = vortex_session.create_execution_ctx(); + let session = vortex_session.arrow(); + let field = i32_utf8_map_field(true, true); + let arrow = i32_utf8_map_array( + vec![0, 2, 2, 2, 4], + vec![1, 2, 1, 1], + vec![Some("one"), None, Some("dup-old"), Some("dup-new")], + Some(NullBuffer::from_iter([true, false, true, true])), + true, + )?; + + let vortex = session.from_arrow_array(Arc::new(arrow), &field)?; + assert_eq!( + vortex.dtype(), + &DType::map( + DType::Primitive(PType::I32, NonNullable), + DType::Utf8(Nullable), + true, + Nullable, + )? + ); + + let exported = session.execute_arrow(vortex, Some(&field), &mut ctx)?; + let map = exported.as_map(); + assert_eq!(map.value_offsets(), &[0, 2, 2, 2, 4]); + assert!(map.is_valid(0)); + assert!(map.is_null(1)); + assert!(map.is_valid(2)); + assert!(map.is_valid(3)); + assert_eq!( + map.keys() + .as_primitive::() + .values(), + &[1, 2, 1, 1] + ); + let values = map.values().as_string::(); + assert_eq!(values.value(0), "one"); + assert!(values.is_null(1)); + assert_eq!(values.value(2), "dup-old"); + assert_eq!(values.value(3), "dup-new"); + + Ok(()) + } + + #[test] + fn legacy_from_arrow_map_imports_canonical_map() -> VortexResult<()> { + let arrow = i32_utf8_map_array( + vec![0, 1, 2], + vec![1, 2], + vec![Some("one"), Some("two")], + None, + false, + )?; + let vortex = ArrayRef::from_arrow(&arrow, false)?; + + assert!(vortex.is::()); + assert_eq!( + vortex.dtype(), + &DType::map( + DType::Primitive(PType::I32, NonNullable), + DType::Utf8(Nullable), + false, + NonNullable, + )? + ); + + Ok(()) + } + + #[test] + fn arrow_rs_map_builder_default_field_names_import() -> VortexResult<()> { + let vortex_session = array_session(); + let mut ctx = vortex_session.create_execution_ctx(); + let session = vortex_session.arrow(); + let mut builder = ArrowMapBuilder::new(None, StringBuilder::new(), Int32Builder::new()); + + builder.keys().append_value("joe"); + builder.values().append_value(1); + builder.append(true)?; + builder.keys().append_value("blogs"); + builder.values().append_value(2); + builder.keys().append_value("foo"); + builder.values().append_value(4); + builder.append(true)?; + builder.append(true)?; + builder.append(false)?; + + let arrow = builder.finish(); + let DataType::Map(entries_field, keys_sorted) = arrow.data_type() else { + panic!("expected Arrow map dtype, got {}", arrow.data_type()); + }; + assert!(!*keys_sorted); + assert_eq!(entries_field.name(), "entries"); + let DataType::Struct(fields) = entries_field.data_type() else { + panic!("expected Arrow map entries struct, got {entries_field:?}"); + }; + assert_eq!(fields[0].name(), "keys"); + assert_eq!(fields[1].name(), "values"); + + let expected_dtype = DType::map( + DType::Utf8(NonNullable), + DType::Primitive(PType::I32, Nullable), + false, + Nullable, + )?; + let legacy = ArrayRef::from_arrow(&arrow, true)?; + assert_eq!(legacy.dtype(), &expected_dtype); + + let field = Field::new("maps", arrow.data_type().clone(), true); + let vortex = session.from_arrow_array(Arc::new(arrow), &field)?; + assert_eq!(vortex.dtype(), &expected_dtype); + let row0 = vortex.execute_scalar(0, &mut ctx)?; + let row0_entries = row0.as_map().entries().collect::>(); + assert_eq!(row0_entries.len(), 1); + assert_eq!(row0_entries[0].0, Scalar::utf8("joe", NonNullable)); + assert_eq!(row0_entries[0].1, Scalar::primitive(1_i32, Nullable)); + let row1 = vortex.execute_scalar(1, &mut ctx)?; + let row1_entries = row1.as_map().entries().collect::>(); + assert_eq!(row1_entries.len(), 2); + assert_eq!(row1_entries[0].0, Scalar::utf8("blogs", NonNullable)); + assert_eq!(row1_entries[0].1, Scalar::primitive(2_i32, Nullable)); + assert!(vortex.execute_scalar(2, &mut ctx)?.as_map().is_empty()); + assert!(vortex.execute_scalar(3, &mut ctx)?.is_null()); + + Ok(()) + } + + #[test] + fn arrow_map_import_sets_zero_copy_to_list() -> VortexResult<()> { + let vortex_session = array_session(); + let session = vortex_session.arrow(); + let field = i32_utf8_map_field(false, false); + let arrow = i32_utf8_map_array( + vec![0, 2, 3], + vec![1, 2, 3], + vec![Some("one"), Some("two"), Some("three")], + None, + false, + )?; + + let vortex = session.from_arrow_array(Arc::new(arrow), &field)?; + let map = vortex.as_::(); + assert!(map.entries().as_::().is_zero_copy_to_list()); + + let sliced = i32_utf8_map_array( + vec![0, 1, 3, 4], + vec![9, 10, 11, 12], + vec![Some("nine"), Some("ten"), Some("eleven"), Some("twelve")], + None, + false, + )?; + let sliced: ArrowArrayRef = Arc::new(sliced.slice(1, 2)); + + let vortex = session.from_arrow_array(sliced, &field)?; + let map = vortex.as_::(); + let entries = map.entries().as_::(); + assert!(entries.is_zero_copy_to_list()); + assert_eq!(entries.offset_at(0), 1); + assert_eq!(entries.offset_at(1), 3); + + Ok(()) + } + + #[test] + fn sliced_arrow_map_import_preserves_nonzero_offsets() -> VortexResult<()> { + let vortex_session = array_session(); + let mut ctx = vortex_session.create_execution_ctx(); + let session = vortex_session.arrow(); + let field = i32_utf8_map_field(true, false); + let arrow = i32_utf8_map_array( + vec![0, 1, 3, 4], + vec![9, 10, 11, 12], + vec![Some("nine"), Some("ten"), Some("eleven"), Some("twelve")], + None, + true, + )?; + let sliced: ArrowArrayRef = Arc::new(arrow.slice(1, 2)); + + let vortex = session.from_arrow_array(sliced, &field)?; + let map = vortex.as_::(); + let entries = map.entries().as_::(); + assert_eq!(entries.offset_at(0), 1); + assert_eq!(entries.offset_at(1), 3); + assert_eq!(entries.size_at(0), 2); + assert_eq!(entries.size_at(1), 1); + + let exported = session.execute_arrow(vortex, Some(&field), &mut ctx)?; + let map = exported.as_map(); + + assert_eq!(map.value_offsets(), &[1, 3, 4]); + assert_eq!( + map.value(0) + .column(0) + .as_primitive::() + .values(), + &[10, 11] + ); + assert_eq!( + map.value(1) + .column(0) + .as_primitive::() + .values(), + &[12] + ); + + Ok(()) + } + + #[test] + fn filtered_map_with_entry_gap_exports_correct_rows() -> VortexResult<()> { + let vortex_session = array_session(); + let mut ctx = vortex_session.create_execution_ctx(); + let session = vortex_session.arrow(); + let map_dtype = MapDType::try_new( + DType::Primitive(PType::I32, NonNullable), + DType::Utf8(Nullable), + false, + )?; + let dtype = DType::Map(map_dtype.clone(), Nullable); + let mut builder = MapBuilder::::with_capacity_in( + map_dtype, + Nullable, + 3, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + let rows: [&[(i32, &str)]; 3] = [ + &[(1, "a"), (2, "b")], + &[(3, "c"), (4, "d"), (5, "e")], + &[(6, "f")], + ]; + for row in rows { + let entries = row + .iter() + .map(|(key, value)| { + ( + Scalar::primitive(*key, NonNullable), + Scalar::utf8(*value, Nullable), + ) + }) + .collect::>(); + builder.append_scalar(&Scalar::try_map(dtype.clone(), entries)?)?; + } + let source = builder.finish_into_map().into_array(); + + // Dropping the non-empty middle row leaves a gap in the entry elements, so the export + // must rebuild the entries list instead of reusing the builder's contiguous layout. + let filtered = source.filter(Mask::from_iter([true, false, true]))?; + + let field = i32_utf8_map_field(false, true); + let exported = session.execute_arrow(filtered, Some(&field), &mut ctx)?; + let map = exported.as_map(); + + assert_eq!(map.value_offsets(), &[0, 2, 3]); + assert_eq!( + map.keys() + .as_primitive::() + .values(), + &[1, 2, 6] + ); + let values = map.values().as_string::(); + assert_eq!(values.value(0), "a"); + assert_eq!(values.value(1), "b"); + assert_eq!(values.value(2), "f"); + + Ok(()) + } + + #[test] + fn unsorted_map_rejects_sorted_arrow_target() -> VortexResult<()> { + let vortex_session = array_session(); + let mut ctx = vortex_session.create_execution_ctx(); + let session = vortex_session.arrow(); + let source_field = i32_utf8_map_field(false, false); + let target_field = i32_utf8_map_field(true, false); + let arrow = i32_utf8_map_array( + vec![0, 2], + vec![2, 1], + vec![Some("two"), Some("one")], + None, + false, + )?; + let vortex = session.from_arrow_array(Arc::new(arrow), &source_field)?; + + let error = session + .execute_arrow(vortex, Some(&target_field), &mut ctx) + .unwrap_err(); + assert!(error.to_string().contains("keys_sorted=true")); + + Ok(()) + } + + #[test] + fn malformed_map_field_errors_without_panic() -> VortexResult<()> { + let session = array_session(); + let arrow_session = session.arrow(); + let arrow = i32_utf8_map_array(vec![0, 1], vec![1], vec![Some("one")], None, false)?; + let fields = Fields::from(vec![ + Field::new("key", DataType::Int32, true), + Field::new("value", DataType::Utf8, true), + ]); + let bad_field = Field::new( + "maps", + DataType::Map(Arc::new(Field::new_struct("entries", fields, false)), false), + false, + ); + + let error = arrow_session + .from_arrow_array(Arc::new(arrow), &bad_field) + .unwrap_err(); + assert!(error.to_string().contains("key field must be non-nullable")); + + Ok(()) + } + + #[test] + fn map_roundtrip_preserves_nested_uuid_fields() -> VortexResult<()> { + use arrow_schema::extension::Uuid as ArrowUuid; + + let vortex_session = array_session(); + let mut ctx = vortex_session.create_execution_ctx(); + let session = vortex_session.arrow(); + + let mut key_field = Field::new("key", DataType::FixedSizeBinary(16), false); + key_field.try_with_extension_type(ArrowUuid)?; + let mut value_field = Field::new("value", DataType::FixedSizeBinary(16), true); + value_field.try_with_extension_type(ArrowUuid)?; + let fields = Fields::from(vec![key_field, value_field]); + let entries_field = Arc::new(Field::new_struct("entries", fields.clone(), false)); + let field = Field::new( + "ids", + DataType::Map(Arc::clone(&entries_field), true), + false, + ); + let keys = FixedSizeBinaryArray::try_from_iter( + [ + b"0123456789abcdef".as_slice(), + b"fedcba9876543210".as_slice(), + ] + .into_iter(), + )?; + let values = FixedSizeBinaryArray::try_from_sparse_iter_with_size( + [Some(b"aaaaaaaaaaaaaaaa".as_slice()), None].into_iter(), + 16, + )?; + let entries = + ArrowStructArray::try_new(fields, vec![Arc::new(keys), Arc::new(values)], None)?; + let arrow = ArrowMapArray::try_new( + entries_field, + OffsetBuffer::new(ScalarBuffer::from(vec![0, 2])), + entries, + None, + true, + )?; + + let vortex = session.from_arrow_array(Arc::new(arrow), &field)?; + let DType::Map(map_dtype, NonNullable) = vortex.dtype() else { + panic!("expected map dtype, got {}", vortex.dtype()); + }; + assert!(map_dtype.key_dtype().is_extension()); + assert!(map_dtype.value_dtype().is_extension()); + + let exported = session.execute_arrow(vortex, Some(&field), &mut ctx)?; + assert_eq!(exported.data_type(), field.data_type()); + assert_eq!(exported.as_map().value_offsets(), &[0, 2]); + + Ok(()) + } } diff --git a/vortex-arrow/src/executor/mod.rs b/vortex-arrow/src/executor/mod.rs index 5b5fa479f29..036954ffe66 100644 --- a/vortex-arrow/src/executor/mod.rs +++ b/vortex-arrow/src/executor/mod.rs @@ -35,15 +35,14 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::arrays::List; use vortex_array::arrays::VarBin; -use vortex_array::arrays::list::ListArrayExt; -use vortex_array::arrays::varbin::VarBinArrayExt; +use vortex_array::arrays::list::ListArraySlotsExt; +use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_array::dtype::DType; use vortex_array::dtype::PType; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; -use crate::dtype::to_data_type_naive; use crate::executor::bool::to_arrow_bool; use crate::executor::byte::to_arrow_byte_array; use crate::executor::byte_view::to_arrow_byte_view; @@ -133,7 +132,7 @@ pub(crate) fn execute_arrow_naive( let resolved_type: DataType = match data_type { Some(dt) => dt.clone(), - None => infer_nearest_arrow_type(&array)?, + None => infer_nearest_arrow_type(&array, ctx)?, }; let arrow = match &resolved_type { @@ -193,7 +192,9 @@ pub(crate) fn execute_arrow_naive( dt @ (DataType::Date32 | DataType::Date64) => to_arrow_date(array, dt, ctx), dt @ (DataType::Time32(_) | DataType::Time64(_)) => to_arrow_time(array, dt, ctx), dt @ DataType::Timestamp(..) => to_arrow_timestamp(array, dt, ctx), - DataType::Map(entries_field, ordered) => to_arrow_map(array, entries_field, *ordered, ctx), + DataType::Map(entries_field, keys_sorted) => { + to_arrow_map(array, entries_field, *keys_sorted, ctx) + } DataType::FixedSizeBinary(_) | DataType::Duration(_) | DataType::Interval(_) @@ -213,11 +214,11 @@ pub(crate) fn execute_arrow_naive( /// Determine the preferred (cheapest) Arrow type for an array. /// -/// For most arrays, this returns the canonical Arrow type from `dtype.to_arrow_dtype()`. +/// For most arrays, this returns the canonical Arrow type for the array's dtype. /// However, some encodings have cheaper Arrow representations: /// - `VarBinArray`: Uses `Utf8`/`Binary` (offset-based) instead of `Utf8View`/`BinaryView` /// - `ListArray`: Uses `List` instead of `ListView` -fn infer_nearest_arrow_type(array: &ArrayRef) -> VortexResult { +fn infer_nearest_arrow_type(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { // VarBinArray: use offset-based Binary/Utf8 instead of View types if let Some(varbin) = array.as_opt::() { let offsets_ptype = PType::try_from(varbin.offsets().dtype())?; @@ -236,12 +237,13 @@ fn infer_nearest_arrow_type(array: &ArrayRef) -> VortexResult { if let Some(list) = array.as_opt::() { let offsets_ptype = PType::try_from(list.offsets().dtype())?; let use_large = matches!(offsets_ptype, PType::I64 | PType::U64); - // Recursively get the preferred type for elements - let elem_dtype = infer_nearest_arrow_type(list.elements())?; - let field = FieldRef::new(Field::new_list_field( - elem_dtype, - list.elements().dtype().is_nullable(), - )); + // Recursively get the preferred field for elements, so extension elements keep the + // `ARROW:extension:name` metadata their export plugin assigns. + let field = FieldRef::new(infer_nearest_arrow_field( + list.elements(), + Field::LIST_FIELD_DEFAULT_NAME, + ctx, + )?); return Ok(if use_large { DataType::LargeList(field) @@ -250,6 +252,121 @@ fn infer_nearest_arrow_type(array: &ArrayRef) -> VortexResult { }); } - // Everything else: use canonical dtype conversion - to_data_type_naive(array.dtype()) + // Everything else: defer to the session's canonical conversion, which additionally resolves + // extension dtypes (including ones nested inside containers) through their export plugins. + Ok(infer_nearest_arrow_field(array, "", ctx)? + .data_type() + .clone()) +} + +/// Determine the preferred (cheapest) Arrow [`Field`] for an array. +/// +/// Unlike [`infer_nearest_arrow_type`] this preserves the Field-level `ARROW:extension:name` +/// metadata that export plugins assign, which a bare [`DataType`] cannot carry. +pub(crate) fn infer_nearest_arrow_field( + array: &ArrayRef, + name: &str, + ctx: &mut ExecutionCtx, +) -> VortexResult { + // Only the encodings with a cheaper-than-canonical Arrow type need the inference; every other + // dtype goes through the session, which is the only thing that knows how to map extensions. + if array.is::() || array.is::() { + let data_type = infer_nearest_arrow_type(array, ctx)?; + return Ok(Field::new(name, data_type, array.dtype().is_nullable())); + } + ctx.session() + .clone() + .arrow() + .to_arrow_field(name, array.dtype()) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_schema::DataType; + use arrow_schema::Field; + use arrow_schema::TimeUnit; + use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::VarBinViewArray; + + use crate::ArrowSessionExt; + + fn utf8() -> ArrayRef { + VarBinViewArray::from_iter_str(["a", "bb"]).into_array() + } + + fn primitive() -> ArrayRef { + PrimitiveArray::from_iter([1i32, 2]).into_array() + } + + fn boolean() -> ArrayRef { + BoolArray::from_iter([true, false]).into_array() + } + + fn list_target() -> DataType { + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))) + } + + /// Must error rather than panic inside `execute::`. + #[rstest] + #[case::bool_from_utf8(utf8(), DataType::Boolean)] + #[case::bool_from_primitive(primitive(), DataType::Boolean)] + #[case::null_from_utf8(utf8(), DataType::Null)] + #[case::null_from_primitive(primitive(), DataType::Null)] + #[case::decimal_from_utf8(utf8(), DataType::Decimal128(10, 2))] + #[case::decimal_from_primitive(primitive(), DataType::Decimal128(10, 2))] + #[case::list_from_utf8(utf8(), list_target())] + #[case::list_from_primitive(primitive(), list_target())] + #[case::list_view_from_primitive( + primitive(), + DataType::ListView(Arc::new(Field::new("item", DataType::Int32, true))) + )] + #[case::fixed_size_list_from_utf8( + utf8(), + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, true)), 2) + )] + #[case::byte_from_primitive(primitive(), DataType::Utf8)] + #[case::byte_from_bool(boolean(), DataType::Binary)] + fn incompatible_target_returns_error(#[case] array: ArrayRef, #[case] target: DataType) { + let session = array_session(); + let mut ctx = session.create_execution_ctx(); + let field = Field::new("f", target.clone(), array.dtype().is_nullable()); + + let result = session.arrow().execute_arrow(array, Some(&field), &mut ctx); + + assert!( + result.is_err(), + "expected an error exporting to {target:?}, got Ok" + ); + } + + /// Cross-class conversions that are genuinely supported must keep working. + #[rstest] + #[case::bool_to_int32(boolean(), DataType::Int32)] + #[case::bool_to_float64(boolean(), DataType::Float64)] + #[case::primitive_to_int64(primitive(), DataType::Int64)] + #[case::primitive_to_date32(primitive(), DataType::Date32)] + #[case::primitive_to_timestamp(primitive(), DataType::Timestamp(TimeUnit::Microsecond, None))] + #[case::utf8_to_binary(utf8(), DataType::Binary)] + #[case::utf8_to_large_utf8(utf8(), DataType::LargeUtf8)] + fn supported_cross_class_target_still_works(#[case] array: ArrayRef, #[case] target: DataType) { + let session = array_session(); + let mut ctx = session.create_execution_ctx(); + let field = Field::new("f", target.clone(), array.dtype().is_nullable()); + + let result = session.arrow().execute_arrow(array, Some(&field), &mut ctx); + + assert!( + result.is_ok(), + "expected {target:?} export to succeed, got {:?}", + result.err() + ); + } } diff --git a/vortex-arrow/src/executor/null.rs b/vortex-arrow/src/executor/null.rs index c5ffc991006..7e5dac3e9c4 100644 --- a/vortex-arrow/src/executor/null.rs +++ b/vortex-arrow/src/executor/null.rs @@ -8,7 +8,9 @@ use arrow_array::NullArray as ArrowNullArray; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::arrays::NullArray; +use vortex_array::dtype::DType; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; /// Convert a canonical NullArray directly to Arrow. pub fn canonical_null_to_arrow(array: &NullArray) -> ArrowArrayRef { @@ -19,6 +21,11 @@ pub(super) fn to_arrow_null( array: ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult { + vortex_ensure!( + matches!(array.dtype(), DType::Null), + "Cannot convert Vortex array with dtype {} to an Arrow Null array", + array.dtype() + ); let null_array = array.execute::(ctx)?; Ok(canonical_null_to_arrow(&null_array)) } diff --git a/vortex-arrow/src/executor/run_end.rs b/vortex-arrow/src/executor/run_end.rs index b6913c1ea1c..7dc16b303cc 100644 --- a/vortex-arrow/src/executor/run_end.rs +++ b/vortex-arrow/src/executor/run_end.rs @@ -16,6 +16,7 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; +use vortex_array::matcher::Matcher; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -23,8 +24,21 @@ use vortex_error::vortex_err; use vortex_runend::RunEnd; use vortex_runend::RunEndArray; use vortex_runend::RunEndArrayExt; +use vortex_runend::RunEndArraySlotsExt; use crate::ArrowArrayExecutor; +use crate::session::ArrowSessionExt; + +/// Matches the encodings [`to_arrow_run_end`] requires for export. +struct ArrowRunEndExportable; + +impl Matcher for ArrowRunEndExportable { + type Match<'a> = &'a ArrayRef; + + fn try_match(array: &ArrayRef) -> Option> { + (array.is::() || array.is::()).then_some(array) + } +} pub(super) fn to_arrow_run_end( array: ArrayRef, @@ -32,28 +46,37 @@ pub(super) fn to_arrow_run_end( values_type: &Field, ctx: &mut ExecutionCtx, ) -> VortexResult { + let array = array.execute_until::(ctx)?; + let array = match array.try_downcast::() { - Ok(constant) => { - return constant_to_run_end(constant, ends_type, values_type, ctx); - } + Ok(constant) => return constant_to_run_end(constant, ends_type, values_type, ctx), + Err(array) => array, + }; + let array = match array.try_downcast::() { + Ok(run_end) => return run_end_to_arrow(run_end, ends_type, values_type, ctx), Err(array) => array, }; - // Execute to unwrap any wrapper VTables (Slice, Filter, etc.) which may - // reveal a RunEndArray. - let array = array.execute::(ctx)?; - match array.try_downcast::() { - Ok(run_end) => run_end_to_arrow(run_end, ends_type, values_type, ctx), - Err(array) => { - // Fallback: canonicalize to flat Arrow, then cast to REE. - let flat = array.execute_arrow(Some(values_type.data_type()), ctx)?; - let ree_type = DataType::RunEndEncoded( - Arc::new(Field::new("run_ends", ends_type.clone(), false)), - Arc::new(values_type.clone()), - ); - arrow_cast::cast(&flat, &ree_type).map_err(VortexError::from) - } - } + // Fallback: canonicalize to flat Arrow, then cast to REE. + let flat = export_values(array, values_type, ctx)?; + let ree_type = DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", ends_type.clone(), false)), + Arc::new(values_type.clone()), + ); + arrow_cast::cast(&flat, &ree_type).map_err(VortexError::from) +} + +/// Export the values of a run-end array through the session, so extension metadata on the target +/// values field dispatches to its export plugin instead of the naive Arrow mapping. +fn export_values( + values: ArrayRef, + values_type: &Field, + ctx: &mut ExecutionCtx, +) -> VortexResult { + ctx.session() + .clone() + .arrow() + .execute_arrow(values, Some(values_type), ctx) } fn run_end_to_arrow( @@ -66,10 +89,7 @@ fn run_end_to_arrow( let offset = array.offset(); let arrow_ends = array.ends().clone().execute_arrow(Some(ends_type), ctx)?; - let arrow_values = array - .values() - .clone() - .execute_arrow(Some(values_type.data_type()), ctx)?; + let arrow_values = export_values(array.values().clone(), values_type, ctx)?; match ends_type { DataType::Int16 => build_run_array::(&arrow_ends, &arrow_values, offset, length), @@ -137,9 +157,11 @@ fn constant_to_run_end( return Ok(new_null_array(&ree_type, len)); } - let values = ConstantArray::new(scalar.clone(), 1) - .into_array() - .execute_arrow(Some(values_type.data_type()), ctx)?; + let values = export_values( + ConstantArray::new(scalar.clone(), 1).into_array(), + values_type, + ctx, + )?; match ends_type { DataType::Int16 => build_constant_run_array::(len, &values), @@ -177,18 +199,26 @@ mod tests { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::SliceArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability::Nullable; use vortex_array::dtype::PType; use vortex_array::scalar::Scalar; + use vortex_array::validity::Validity; + use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_session::VortexSession; use crate::ArrowArrayExecutor; use crate::executor::run_end::ConstantArray; + use crate::executor::run_end::RunEnd; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_runend::initialize(&session); + session + }); fn ree_type(ends: DataType, values_dtype: DataType) -> DataType { DataType::RunEndEncoded( @@ -261,6 +291,38 @@ mod tests { Ok(()) } + #[test] + fn nested_wrappers_reveal_run_end_fast_path() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + // Runs [7, 7, 8, 8] behind two nested lazy Slice wrappers. A single execution step + // only merges the slices, so the exporter must keep stepping to reveal the + // RunEndArray underneath. + let ends = PrimitiveArray::new(buffer![2u32, 4], Validity::NonNullable); + let values = PrimitiveArray::new(buffer![7i64, 8], Validity::NonNullable); + let values_ptr = values.as_slice::().as_ptr(); + let ree = RunEnd::try_new(ends.into_array(), values.into_array(), &mut ctx)?; + let wrapped = SliceArray::new(SliceArray::new(ree.into_array(), 0..4).into_array(), 0..4) + .into_array(); + + let result = execute(wrapped, &ree_type(DataType::Int32, DataType::Int64))?; + + let ree_arrow = result + .as_any() + .downcast_ref::>() + .ok_or_else(|| vortex_err!("expected Int32 run-end array"))?; + assert_eq!(ree_arrow.run_ends().values(), &[2, 4]); + + // The RunEnd fast path exports the values buffer zero-copy; the flat-then-recode + // fallback would materialize new buffers. + let arrow_values = ree_arrow + .values() + .as_any() + .downcast_ref::() + .ok_or_else(|| vortex_err!("expected Int64 values"))?; + assert_eq!(arrow_values.values().as_ptr(), values_ptr); + Ok(()) + } + #[test] fn primitive_to_ree() -> VortexResult<()> { let array = PrimitiveArray::from_iter(vec![10i32, 10, 20, 20, 20]).into_array(); diff --git a/vortex-arrow/src/executor/struct_.rs b/vortex-arrow/src/executor/struct_.rs index e2ce2700d4a..a7553dde78b 100644 --- a/vortex-arrow/src/executor/struct_.rs +++ b/vortex-arrow/src/executor/struct_.rs @@ -21,16 +21,32 @@ use vortex_array::arrays::struct_::StructDataParts; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::FieldNames; -use vortex_array::dtype::StructFields; +use vortex_array::matcher::Matcher; use vortex_array::scalar_fn::fns::pack::Pack; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use crate::ArrowArrayExecutor; -use crate::dtype::FromArrowType; +use crate::executor::infer_nearest_arrow_field; use crate::executor::validity::to_arrow_null_buffer; use crate::session::ArrowSessionExt; +/// Matches the encodings [`to_arrow_struct`] requires for export. +struct ArrowStructExportable; + +impl Matcher for ArrowStructExportable { + type Match<'a> = &'a ArrayRef; + + fn try_match(array: &ArrayRef) -> Option> { + (array.is::() + || array.is::() + || array + .as_opt::() + .is_some_and(|scalar_fn| scalar_fn.scalar_fn().as_opt::().is_some())) + .then_some(array) + } +} + pub(super) fn to_arrow_struct( array: ArrayRef, target_fields: Option<&Fields>, @@ -38,6 +54,8 @@ pub(super) fn to_arrow_struct( ) -> VortexResult { let len = array.len(); + let array = array.execute_until::(ctx)?; + // If the array is chunked, then we invert the chunk-of-struct to struct-of-chunk. let array = match array.try_downcast::() { Ok(array) => { @@ -87,7 +105,7 @@ pub(super) fn to_arrow_struct( // Otherwise, we fall back to executing to a StructArray. let array = if let Some(fields) = target_fields { - let vx_fields = StructFields::from_arrow(fields); + let vx_fields = ctx.session().arrow().from_arrow_fields(fields)?; // We apply a cast to ensure we push down casting where possible into the struct fields. array.cast(DType::Struct( vx_fields, @@ -158,26 +176,27 @@ fn create_from_fields( })) } Err(names) => { - // No target fields specified - use preferred types for each child + // No target fields specified - use preferred types for each child. The inferred field + // is what carries any `ARROW:extension:name` metadata, which the executed array's bare + // `DataType` cannot, so build the child fields from it rather than from the array. let mut arrow_arrays = Vec::with_capacity(vortex_fields.len()); - for vx_field in vortex_fields.iter() { + let mut arrow_fields = Vec::with_capacity(vortex_fields.len()); + for (name, vx_field) in names.iter().zip_eq(vortex_fields.iter()) { + let inferred = infer_nearest_arrow_field(vx_field, name.as_ref(), ctx)?; let arrow_array = vx_field.clone().execute_arrow(None, ctx)?; + // The executed array is authoritative for the physical type; only the metadata is + // taken from the inferred field. + arrow_fields.push(Arc::new( + Field::new( + name.as_ref(), + arrow_array.data_type().clone(), + vx_field.dtype().is_nullable(), + ) + .with_metadata(inferred.metadata().clone()), + )); arrow_arrays.push(arrow_array); } - - // Build the Arrow fields from the resulting arrays - let arrow_fields: Fields = names - .iter() - .zip_eq(arrow_arrays.iter()) - .zip_eq(vortex_fields.iter().map(|f| f.dtype().is_nullable())) - .map(|((name, arr), vx_nullable)| { - Arc::new(Field::new( - name.as_ref(), - arr.data_type().clone(), - vx_nullable, - )) - }) - .collect(); + let arrow_fields = Fields::from(arrow_fields); Ok(Arc::new(unsafe { ArrowStructArray::new_unchecked_with_length( @@ -196,6 +215,7 @@ mod tests { use std::sync::Arc; use arrays::varbinview::VarBinViewArray; + use arrow_array::Array; use arrow_array::ArrayRef; use arrow_array::PrimitiveArray as ArrowPrimitiveArray; use arrow_array::StringViewArray; @@ -204,20 +224,21 @@ mod tests { use arrow_buffer::NullBuffer; use arrow_schema::DataType; use arrow_schema::Field; - use vortex_array as array; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; use vortex_array::arrays; + use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::dtype::FieldNames; + use vortex_array::scalar_fn::fns::mask::Mask; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use crate::ArrowArrayExecutor; - use crate::FromArrowArray; + use crate::convert::from_arrow_dyn; use crate::dtype::to_data_type_naive; #[test] @@ -339,6 +360,40 @@ mod tests { Ok(()) } + #[test] + fn mask_wrapped_struct_exports_via_struct_fast_path() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // A struct behind a lazy `mask` scalar-fn nulling out row 1 — the shape a scan + // produces when a row mask is applied to a top-level struct batch. + let xs = PrimitiveArray::new(buffer![1i64, 2, 3], Validity::NonNullable); + let struct_array = StructArray::try_new( + FieldNames::from(["xs"]), + vec![xs.into_array()], + 3, + Validity::NonNullable, + )?; + let mask = BoolArray::from_iter([true, false, true]); + let masked = Mask::try_new(struct_array.into_array(), mask.into_array())?.into_array(); + + let arrow = masked.execute_arrow(None, &mut ctx)?; + + let arrow_struct = arrow + .as_any() + .downcast_ref::() + .expect("struct array"); + assert_eq!(arrow_struct.len(), 3); + assert!(!arrow_struct.is_null(0)); + assert!(arrow_struct.is_null(1)); + assert!(!arrow_struct.is_null(2)); + let xs_col = arrow_struct + .column(0) + .as_any() + .downcast_ref::>() + .expect("int64 column"); + assert_eq!(xs_col.values(), &[1, 2, 3]); + Ok(()) + } + #[test] fn to_arrow_with_non_nullable_fields() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); @@ -357,7 +412,7 @@ mod tests { )?; let orig_dtype = array.dtype().clone(); let arrow_array = array.into_array().execute_arrow(None, &mut ctx)?; - let from_arrow = array::ArrayRef::from_arrow(arrow_array.as_ref(), false)?; + let from_arrow = from_arrow_dyn(arrow_array.as_ref(), false)?; assert_eq!(&orig_dtype, from_arrow.dtype()); Ok(()) } diff --git a/vortex-arrow/src/iter.rs b/vortex-arrow/src/iter.rs index 0d7ebe5c82a..730585ea0e5 100644 --- a/vortex-arrow/src/iter.rs +++ b/vortex-arrow/src/iter.rs @@ -1,25 +1,50 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use arrow_array::RecordBatchReader; use arrow_array::ffi_stream; +use arrow_schema::SchemaRef; use vortex_array::ArrayRef; use vortex_array::dtype::DType; use vortex_array::iter::ArrayIterator; use vortex_error::VortexError; use vortex_error::VortexResult; -use crate::FromArrowArray; -use crate::dtype::FromArrowType; +use crate::ArrowSession; -/// An adapter for converting an `ArrowArrayStreamReader` into a Vortex `ArrayStream`. +/// An adapter for converting an `ArrowArrayStreamReader` into a Vortex `ArrayIterator`. +/// +/// The stream's schema and batches are imported through the [`ArrowSession`], so Arrow extension +/// types are routed to their registered import plugins instead of being flattened into their +/// storage types. pub struct ArrowArrayStreamAdapter { stream: ffi_stream::ArrowArrayStreamReader, + session: ArrowSession, + schema: SchemaRef, dtype: DType, } impl ArrowArrayStreamAdapter { - pub fn new(stream: ffi_stream::ArrowArrayStreamReader, dtype: DType) -> Self { - Self { stream, dtype } + /// Adapt `stream`, importing its schema and each of its batches through `session`. + /// + /// The adapter holds a clone of `session`, which shares the plugin registries with it, so + /// plugins registered after construction are still observed. + /// + /// The schema declared by the stream is the authoritative schema for the import: every batch + /// is converted against it, so extension metadata survives even if an individual batch has + /// lost it. + pub fn try_new( + session: &ArrowSession, + stream: ffi_stream::ArrowArrayStreamReader, + ) -> VortexResult { + let schema = stream.schema(); + let dtype = session.from_arrow_schema(schema.as_ref())?; + Ok(Self { + stream, + session: session.clone(), + schema, + dtype, + }) } } @@ -36,8 +61,75 @@ impl Iterator for ArrowArrayStreamAdapter { let batch = self.stream.next()?; Some(batch.map_err(VortexError::from).and_then(|b| { - debug_assert_eq!(&self.dtype, &DType::from_arrow(b.schema())); - ArrayRef::from_arrow(b, false) + self.session + .from_arrow_record_batch(b, self.schema.as_ref()) })) } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_array::FixedSizeBinaryArray; + use arrow_array::RecordBatch; + use arrow_array::RecordBatchIterator; + use arrow_schema::DataType; + use arrow_schema::Field; + use arrow_schema::Schema; + use arrow_schema::extension::Uuid as ArrowUuid; + use vortex_array::array_session; + use vortex_array::arrays::Struct; + use vortex_array::arrays::struct_::StructArrayExt; + use vortex_array::extension::uuid::Uuid; + use vortex_error::VortexExpect; + + use super::*; + use crate::ArrowSessionExt; + + /// The adapter imports through the [`ArrowSession`], so a UUID column arrives as a Vortex + /// extension array rather than its `FixedSizeList` storage. + #[test] + fn stream_preserves_extension_types() -> VortexResult<()> { + let mut field = Field::new("id", DataType::FixedSizeBinary(16), false); + field.try_with_extension_type(ArrowUuid)?; + let schema = Arc::new(Schema::new(vec![field])); + let ids = FixedSizeBinaryArray::try_from_iter( + [*b"0123456789abcdef", *b"fedcba9876543210"].into_iter(), + )?; + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(ids)])?; + + let reader = RecordBatchIterator::new([Ok(batch)], schema); + let stream = ffi_stream::ArrowArrayStreamReader::try_new( + ffi_stream::FFI_ArrowArrayStream::new(Box::new(reader)), + )?; + + let vortex_session = array_session(); + let mut adapter = ArrowArrayStreamAdapter::try_new(&vortex_session.arrow(), stream)?; + + let DType::Struct(fields, _) = adapter.dtype().clone() else { + panic!("expected a struct dtype, got {}", adapter.dtype()); + }; + assert!( + fields + .field_by_index(0) + .vortex_expect("one field") + .as_extension() + .is::() + ); + + let array = adapter.next().vortex_expect("one batch")?; + assert_eq!(array.dtype(), adapter.dtype()); + assert!( + array + .as_::() + .unmasked_field(0) + .dtype() + .as_extension() + .is::() + ); + assert!(adapter.next().is_none()); + + Ok(()) + } +} diff --git a/vortex-arrow/src/lib.rs b/vortex-arrow/src/lib.rs index 73fa4c724ce..c8536e10edf 100644 --- a/vortex-arrow/src/lib.rs +++ b/vortex-arrow/src/lib.rs @@ -3,10 +3,15 @@ //! Apache Arrow interoperability for Vortex. //! -//! This crate owns every conversion between Vortex and Arrow: importing Arrow arrays, record -//! batches, schemas, and fields into Vortex ([`FromArrowArray`], [`FromArrowType`], -//! [`TryFromArrowType`]), and executing Vortex arrays into Arrow ([`ArrowSession::execute_arrow`] -//! and the [`ArrowArrayExecutor`] convenience trait). +//! This crate owns every conversion between Vortex and Arrow. The authoritative entry point is +//! the [`ArrowSession`]: importing Arrow schemas, fields, and data types into Vortex +//! ([`ArrowSession::from_arrow_schema`], [`ArrowSession::from_arrow_field`], +//! [`ArrowSession::from_arrow_datatype`]), importing Arrow arrays and record batches +//! ([`ArrowSession::from_arrow_array`], [`ArrowSession::from_arrow_record_batch`]), +//! exporting Vortex dtypes to Arrow +//! ([`ArrowSession::to_arrow_schema`], [`ArrowSession::to_arrow_field`], +//! [`ArrowSession::to_arrow_datatype`]), and executing Vortex arrays into Arrow +//! ([`ArrowSession::execute_arrow`] and the [`ArrowArrayExecutor`] convenience trait). //! //! The [`ArrowSession`] session variable is created lazily on first use of //! [`ArrowSessionExt::arrow`], so no explicit registration is required for plain conversions. @@ -19,7 +24,7 @@ use vortex_array::legacy_session; use vortex_error::VortexResult; use vortex_session::VortexSession; -mod convert; +pub mod convert; mod datum; pub mod dtype; mod executor; @@ -33,8 +38,11 @@ mod uuid; pub use convert::IntoVortexArray; pub(crate) use convert::nulls; pub use datum::*; +#[allow(deprecated)] pub use dtype::FromArrowType; +#[allow(deprecated)] pub use dtype::ToArrowType; +#[allow(deprecated)] pub use dtype::TryFromArrowType; pub use executor::*; pub use iter::*; @@ -54,6 +62,7 @@ pub fn initialize(session: &VortexSession) { /// /// Implementations reuse the underlying Arrow buffers without copying wherever the Arrow and /// Vortex memory layouts allow it. +#[deprecated(note = "Use `ArrowSession` (`from_arrow_array`, `from_arrow_record_batch`) instead")] pub trait FromArrowArray { /// Convert `array` into a Vortex array whose [`DType`](vortex_array::dtype::DType) has the requested /// `nullable` [`Nullability`](vortex_array::dtype::Nullability). @@ -72,6 +81,9 @@ pub trait FromArrowArray { /// Returns an error if `nullable` is `false` but `array` physically contains one or more nulls /// (including an Arrow `NullArray`, which is entirely null), or if the Arrow data type is not /// supported. + #[deprecated( + note = "Use `ArrowSession` (`from_arrow_array`, `from_arrow_record_batch`) instead" + )] fn from_arrow(array: A, nullable: bool) -> VortexResult where Self: Sized; diff --git a/vortex-arrow/src/run_end_import.rs b/vortex-arrow/src/run_end_import.rs index 69deba71fdf..6a301298907 100644 --- a/vortex-arrow/src/run_end_import.rs +++ b/vortex-arrow/src/run_end_import.rs @@ -1,62 +1,69 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use arrow_array::Array as _; use arrow_array::RunArray; use arrow_array::types::RunEndIndexType; +use arrow_buffer::ArrowNativeType; use vortex_array::ArrayRef; use vortex_array::IntoArray; -use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::dtype::NativePType; -use vortex_array::legacy_session; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexResult; -use vortex_runend::RunEndData; -use vortex_runend::ops::find_physical_index; -use vortex_runend::ops::find_slice_end_index; - -use crate::FromArrowArray; - -impl FromArrowArray<&RunArray> for RunEndData +use vortex_error::vortex_ensure; +use vortex_runend::RunEnd; + +/// Build a Vortex run-end array from an Arrow [`RunArray`] and its already-converted `values`. +/// +/// The caller converts `values` separately (rather than this function calling +/// [`crate::FromArrowArray`] on `array.values()`) so that extension metadata on the Arrow values +/// field dispatches through the session's registered import plugins. +pub(crate) fn run_end_from_arrow( + array: &RunArray, + values: ArrayRef, +) -> VortexResult where R::Native: NativePType, { - #[allow(clippy::disallowed_methods)] - fn from_arrow(array: &RunArray, nullable: bool) -> VortexResult { + let ends_buf = Buffer::::from_arrow_scalar_buffer(array.run_ends().inner().clone()); + let ends = PrimitiveArray::new(ends_buf, Validity::NonNullable) + .reinterpret_cast(R::Native::PTYPE.to_unsigned()) + .into_array(); + + vortex_ensure!( + ends.len() == values.len(), + "Arrow run-end array has {} run ends but {} values", + ends.len(), + values.len() + ); + + // Arrow slices a RunArray by adjusting the logical offset/length while keeping the full + // children, so trim the runs outside the logical range. The run ends are sorted, so this is + // a binary search over the Arrow buffer, avoiding the `ExecutionCtx` the equivalent Vortex + // search would need. + let (ends, values, offset) = if array.is_empty() { + (ends.slice(0..0)?, values.slice(0..0)?, 0) + } else { let offset = array.run_ends().offset(); - let len = array.run_ends().len(); - let ends_buf = - Buffer::::from_arrow_scalar_buffer(array.run_ends().inner().clone()); - let ends = PrimitiveArray::new(ends_buf, Validity::NonNullable) - .reinterpret_cast(R::Native::PTYPE.to_unsigned()); - let values = ArrayRef::from_arrow(array.values().as_ref(), nullable)?; - - let ends_array = PrimitiveArray::from_buffer_handle( - ends.buffer_handle().clone(), - ends.ptype(), - ends.validity()?, + let run_ends = array.run_ends().values(); + let first = run_ends.partition_point(|end| end.as_usize() <= offset); + let last = run_ends.partition_point(|end| end.as_usize() < offset + array.len()); + ( + ends.slice(first..last + 1)?, + values.slice(first..last + 1)?, + offset, ) - .into_array(); - let mut ctx = legacy_session().create_execution_ctx(); - - let (ends_slice, values_slice) = if offset == 0 && len == array.run_ends().max_value() { - (ends_array, values) - } else { - let slice_begin = find_physical_index(&ends_array, offset, &mut ctx)?; - let slice_end = find_slice_end_index(&ends_array, offset + len, &mut ctx)?; - - ( - ends_array.slice(slice_begin..slice_end)?, - values.slice(slice_begin..slice_end)?, - ) - }; - - // SAFETY: arrow-rs enforces the RunEndArray invariants, we inherit their guarantees. - RunEndData::validate_parts(&ends_slice, &values_slice, offset, len, &mut ctx)?; - Ok(unsafe { RunEndData::new_unchecked(offset) }) - } + }; + + // SAFETY: `RunEndData::validate_parts` requires unsigned, strictly increasing run ends of the + // same length as the values, covering `offset..offset + length`. arrow-rs guarantees strictly + // increasing positive run ends (so the reinterpret_cast above preserves both value and order) + // covering the logical range, the length equality is checked above, and the trim keeps only + // runs that intersect that range. + Ok(unsafe { RunEnd::new_unchecked(ends, values, offset, array.len()) }.into_array()) } #[cfg(test)] @@ -79,24 +86,17 @@ mod tests { use vortex_array::IntoArray as _; use vortex_array::VortexSessionExecute as _; use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::validity::Validity; - use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_runend::RunEnd; - use vortex_runend::RunEndArray; - use vortex_runend::ops::find_physical_index; - use vortex_runend::ops::find_slice_end_index; use vortex_session::VortexSession; use crate::ArrowSessionExt; - use crate::FromArrowArray; static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session(); @@ -104,41 +104,17 @@ mod tests { session }); + /// Decode through the session import path, which is what production callers use. fn decode_run_array( array: &RunArray, nullable: bool, - ) -> VortexResult + ) -> VortexResult where R::Native: NativePType, { - let offset = array.run_ends().offset(); - let len = array.run_ends().len(); - let ends_buf = - Buffer::::from_arrow_scalar_buffer(array.run_ends().inner().clone()); - let ends = PrimitiveArray::new(ends_buf, Validity::NonNullable) - .reinterpret_cast(R::Native::PTYPE.to_unsigned()); - let values = ArrayRef::from_arrow(array.values().as_ref(), nullable)?; - - let ends_array = PrimitiveArray::from_buffer_handle( - ends.buffer_handle().clone(), - ends.ptype(), - ends.validity()?, - ) - .into_array(); - let mut ctx = SESSION.create_execution_ctx(); - let (ends_slice, values_slice) = if offset == 0 && len == array.run_ends().max_value() { - (ends_array, values) - } else { - let slice_begin = find_physical_index(&ends_array, offset, &mut ctx)?; - let slice_end = find_slice_end_index(&ends_array, offset + len, &mut ctx)?; - - ( - ends_array.slice(slice_begin..slice_end)?, - values.slice(slice_begin..slice_end)?, - ) - }; - - RunEnd::try_new_offset_length(ends_slice, values_slice, offset, array.len(), &mut ctx) + SESSION + .arrow() + .from_arrow_array(Arc::new(array.clone()), nullable) } #[test] diff --git a/vortex-arrow/src/scalar.rs b/vortex-arrow/src/scalar.rs index aeac56f43d3..09fe7cfcdee 100644 --- a/vortex-arrow/src/scalar.rs +++ b/vortex-arrow/src/scalar.rs @@ -7,16 +7,21 @@ use std::sync::Arc; use arrow_array::Scalar as ArrowScalar; use arrow_array::*; +use arrow_buffer::NullBuffer; +use arrow_buffer::OffsetBuffer; +use arrow_schema::Field; +use arrow_schema::Fields; use vortex_array::dtype::DType; use vortex_array::dtype::PType; +use vortex_array::dtype::i256; use vortex_array::extension::datetime::AnyTemporal; use vortex_array::extension::datetime::TemporalMetadata; use vortex_array::extension::datetime::TimeUnit; use vortex_array::scalar::BinaryScalar; use vortex_array::scalar::BoolScalar; use vortex_array::scalar::DecimalScalar; -use vortex_array::scalar::DecimalValue; use vortex_array::scalar::ExtScalar; +use vortex_array::scalar::MapScalar; use vortex_array::scalar::PrimitiveScalar; use vortex_array::scalar::Scalar; use vortex_array::scalar::Utf8Scalar; @@ -24,6 +29,8 @@ use vortex_error::VortexError; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use crate::dtype::to_data_type_naive; + /// Arrow represents scalars as single-element arrays. This constant is the length of those arrays. const SCALAR_ARRAY_LEN: usize = 1; @@ -68,11 +75,14 @@ impl ToArrowDatum for Scalar { DType::Decimal(..) => decimal_to_arrow(value.as_decimal()), DType::Utf8(_) => utf8_to_arrow(value.as_utf8()), DType::Binary(_) => binary_to_arrow(value.as_binary()), - DType::List(..) => unimplemented!("list scalar conversion"), - DType::FixedSizeList(..) => unimplemented!("fixed-size list scalar conversion"), - DType::Struct(..) => unimplemented!("struct scalar conversion"), - DType::Union(..) => unimplemented!("union scalar conversion"), - DType::Variant(_) => unimplemented!("Variant scalar conversion"), + DType::List(..) => vortex_bail!("list scalar conversion is not supported"), + DType::FixedSizeList(..) => { + vortex_bail!("fixed-size list scalar conversion is not supported") + } + DType::Map(..) => map_to_arrow(value.as_map()), + DType::Struct(..) => vortex_bail!("struct scalar conversion is not supported"), + DType::Union(..) => vortex_bail!("union scalar conversion is not supported"), + DType::Variant(_) => vortex_bail!("Variant scalar conversion is not supported"), DType::Extension(..) => extension_to_arrow(value.as_extension()), } } @@ -102,20 +112,50 @@ fn primitive_to_arrow(scalar: PrimitiveScalar<'_>) -> Result, Vor /// Convert a [`DecimalScalar`] to an Arrow [`Datum`]. fn decimal_to_arrow(scalar: DecimalScalar<'_>) -> Result, VortexError> { + let DType::Decimal(decimal_dtype, _) = scalar.dtype() else { + vortex_bail!("Expected decimal scalar, got {}", scalar.dtype()); + }; + let precision = decimal_dtype.precision(); + let scale = decimal_dtype.scale(); // TODO(joe): Replace with decimal32, etc. once Arrow supports them. match scalar.decimal_value() { - Some(DecimalValue::I8(v)) => Ok(Arc::new(Decimal128Array::new_scalar(v as i128))), - Some(DecimalValue::I16(v)) => Ok(Arc::new(Decimal128Array::new_scalar(v as i128))), - Some(DecimalValue::I32(v)) => Ok(Arc::new(Decimal128Array::new_scalar(v as i128))), - Some(DecimalValue::I64(v)) => Ok(Arc::new(Decimal128Array::new_scalar(v as i128))), - Some(DecimalValue::I128(v128)) => Ok(Arc::new(Decimal128Array::new_scalar(v128))), - Some(DecimalValue::I256(v256)) => Ok(Arc::new(Decimal256Array::new_scalar(v256.into()))), - None => Ok(Arc::new(arrow_array::Scalar::new( - Decimal128Array::new_null(SCALAR_ARRAY_LEN), - ))), + Some(value) => { + let value = value.as_i256(); + if precision <= 38 { + let value = value.maybe_i128().ok_or_else(|| { + vortex_err!( + "Decimal value {value} cannot fit in Arrow Decimal128 for precision {precision}" + ) + })?; + decimal128_scalar(value, precision, scale) + } else { + decimal256_scalar(value, precision, scale) + } + } + None => { + let data_type = to_data_type_naive(scalar.dtype())?; + Ok(Arc::new(ArrowScalar::new(new_null_array( + &data_type, + SCALAR_ARRAY_LEN, + )))) + } } } +fn decimal128_scalar(value: i128, precision: u8, scale: i8) -> Result, VortexError> { + let array = Decimal128Array::new_scalar(value) + .into_inner() + .with_precision_and_scale(precision, scale)?; + Ok(Arc::new(ArrowScalar::new(array))) +} + +fn decimal256_scalar(value: i256, precision: u8, scale: i8) -> Result, VortexError> { + let array = Decimal256Array::new_scalar(value.into()) + .into_inner() + .with_precision_and_scale(precision, scale)?; + Ok(Arc::new(ArrowScalar::new(array))) +} + /// Convert a [`Utf8Scalar`] to an Arrow [`Datum`]. fn utf8_to_arrow(scalar: Utf8Scalar<'_>) -> Result, VortexError> { value_to_arrow_scalar!(scalar.value(), StringViewArray) @@ -126,6 +166,74 @@ fn binary_to_arrow(scalar: BinaryScalar<'_>) -> Result, VortexErr value_to_arrow_scalar!(scalar.value(), BinaryViewArray) } +/// Convert a [`MapScalar`] to an Arrow [`Datum`]. +fn map_to_arrow(scalar: MapScalar<'_>) -> Result, VortexError> { + let map_dtype = scalar.map_dtype(); + let key_dtype = map_dtype.key_dtype(); + let value_dtype = map_dtype.value_dtype(); + let key_field = Field::new("key", to_data_type_naive(&key_dtype)?, false); + let value_field = Field::new( + "value", + to_data_type_naive(&value_dtype)?, + value_dtype.is_nullable(), + ); + let fields = Fields::from(vec![key_field, value_field]); + + let entries = scalar.entries().collect::>(); + let keys = entries + .iter() + .map(|(key, _)| key.to_arrow_datum()) + .collect::, _>>()?; + let values = entries + .iter() + .map(|(_, value)| value.to_arrow_datum()) + .collect::, _>>()?; + + let key_array = concat_scalar_arrays(&keys, &key_dtype)?; + let value_array = concat_scalar_arrays(&values, &value_dtype)?; + let entries = StructArray::try_new_with_length( + fields.clone(), + vec![key_array, value_array], + None, + entries.len(), + )?; + + let entries_len = entries.len(); + let entries_len = i32::try_from(entries_len).map_err(|_| { + vortex_err!( + "Cannot convert map scalar with {entries_len} entries to Arrow: MapArray offsets are i32" + ) + })?; + let offsets = OffsetBuffer::new(vec![0_i32, entries_len].into()); + let entries_field = Arc::new(Field::new_struct("entries", fields, false)); + let nulls = scalar + .is_null() + .then(|| NullBuffer::new_null(SCALAR_ARRAY_LEN)); + let map = MapArray::try_new( + entries_field, + offsets, + entries, + nulls, + map_dtype.keys_sorted(), + )?; + Ok(Arc::new(ArrowScalar::new(map))) +} + +fn concat_scalar_arrays( + scalars: &[Arc], + dtype: &DType, +) -> Result { + if scalars.is_empty() { + return Ok(new_empty_array(&to_data_type_naive(dtype)?)); + } + + let arrays = scalars + .iter() + .map(|scalar| scalar.get().0) + .collect::>(); + Ok(arrow_select::concat::concat(&arrays)?) +} + /// Convert an [`ExtScalar`] to an Arrow [`Datum`]. /// /// Currently only temporal extension types (timestamps, dates, and times) are supported. @@ -199,6 +307,12 @@ fn extension_to_arrow(scalar: ExtScalar<'_>) -> Result, VortexErr mod tests { use std::sync::Arc; + use arrow_array::Array; + use arrow_array::Decimal128Array; + use arrow_array::Int32Array; + use arrow_array::MapArray; + use arrow_array::StringViewArray; + use arrow_schema::DataType; use rstest::rstest; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; @@ -360,8 +474,16 @@ mod tests { assert!(result.is_ok()); } + fn assert_arrow_scalar_data_type(scalar: &Scalar, expected: DataType) -> VortexResult<()> { + let datum = scalar.to_arrow_datum()?; + let (array, is_scalar) = datum.get(); + assert!(is_scalar); + assert_eq!(array.data_type(), &expected); + Ok(()) + } + #[test] - fn test_decimal_scalars_to_arrow() { + fn test_decimal_scalars_to_arrow() -> VortexResult<()> { // Test various decimal value types let decimal_dtype = DecimalDType::new(5, 2); @@ -370,37 +492,35 @@ mod tests { decimal_dtype, Nullability::NonNullable, ); - assert!(scalar_i8.to_arrow_datum().is_ok()); + assert_arrow_scalar_data_type(&scalar_i8, DataType::Decimal128(5, 2))?; let scalar_i16 = Scalar::decimal( DecimalValue::I16(10000), decimal_dtype, Nullability::NonNullable, ); - assert!(scalar_i16.to_arrow_datum().is_ok()); + assert_arrow_scalar_data_type(&scalar_i16, DataType::Decimal128(5, 2))?; let scalar_i32 = Scalar::decimal( DecimalValue::I32(99999), decimal_dtype, Nullability::NonNullable, ); - assert!(scalar_i32.to_arrow_datum().is_ok()); + assert_arrow_scalar_data_type(&scalar_i32, DataType::Decimal128(5, 2))?; let scalar_i64 = Scalar::decimal( DecimalValue::I64(99999), decimal_dtype, Nullability::NonNullable, ); - assert!(scalar_i64.to_arrow_datum().is_ok()); + assert_arrow_scalar_data_type(&scalar_i64, DataType::Decimal128(5, 2))?; let scalar_i128 = Scalar::decimal( DecimalValue::I128(99999), decimal_dtype, Nullability::NonNullable, ); - assert!(scalar_i128.to_arrow_datum().is_ok()); - - // Test i256 + assert_arrow_scalar_data_type(&scalar_i128, DataType::Decimal128(5, 2))?; let value_i256 = i256::from_i128(99999); let scalar_i256 = Scalar::decimal( @@ -408,19 +528,214 @@ mod tests { decimal_dtype, Nullability::NonNullable, ); - assert!(scalar_i256.to_arrow_datum().is_ok()); + assert_arrow_scalar_data_type(&scalar_i256, DataType::Decimal128(5, 2))?; + + Ok(()) + } + + #[test] + fn decimal_i64_with_wide_precision_exports_decimal256() -> VortexResult<()> { + let scalar = Scalar::decimal( + DecimalValue::I64(1), + DecimalDType::new(39, 0), + Nullability::NonNullable, + ); + + assert_arrow_scalar_data_type(&scalar, DataType::Decimal256(39, 0)) + } + + #[test] + fn decimal_i256_with_narrow_precision_exports_decimal128() -> VortexResult<()> { + let scalar = Scalar::decimal( + DecimalValue::I256(i256::from_i128(1234)), + DecimalDType::new(4, 2), + Nullability::NonNullable, + ); + + assert_arrow_scalar_data_type(&scalar, DataType::Decimal128(4, 2)) } #[test] - fn test_null_decimal_to_arrow() { + fn test_null_decimal_to_arrow() -> VortexResult<()> { let decimal_dtype = DecimalDType::new(10, 2); let scalar = Scalar::null(DType::Decimal(decimal_dtype, Nullability::Nullable)); - let result = scalar.to_arrow_datum(); - assert!(result.is_ok()); + assert_arrow_scalar_data_type(&scalar, DataType::Decimal128(10, 2))?; + + let decimal_dtype = DecimalDType::new(39, 2); + let scalar = Scalar::null(DType::Decimal(decimal_dtype, Nullability::Nullable)); + assert_arrow_scalar_data_type(&scalar, DataType::Decimal256(39, 2)) + } + + #[test] + fn decimal_scalar_to_arrow_preserves_precision_and_scale() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(12, 3); + let scalar = Scalar::decimal( + DecimalValue::I128(12345), + decimal_dtype, + Nullability::NonNullable, + ); + + let datum = scalar.to_arrow_datum()?; + let (array, is_scalar) = datum.get(); + assert!(is_scalar); + let decimal = array + .as_any() + .downcast_ref::() + .expect("decimal scalar should convert to Decimal128"); + assert_eq!(decimal.precision(), 12); + assert_eq!(decimal.scale(), 3); + + Ok(()) + } + + #[test] + fn test_map_scalar_to_arrow() -> VortexResult<()> { + let dtype = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + Nullability::Nullable, + )?; + let scalar = Scalar::try_map( + dtype, + [ + ( + Scalar::primitive(1i32, Nullability::NonNullable), + Scalar::utf8("one", Nullability::Nullable), + ), + ( + Scalar::primitive(2i32, Nullability::NonNullable), + Scalar::null(DType::Utf8(Nullability::Nullable)), + ), + ], + )?; + + let datum = scalar.to_arrow_datum()?; + let (array, is_scalar) = datum.get(); + assert!(is_scalar); + let map = array + .as_any() + .downcast_ref::() + .expect("map scalar should convert to MapArray"); + assert_eq!(map.len(), 1); + assert_eq!(map.value_offsets(), &[0, 2]); + assert!(map.is_valid(0)); + assert_eq!( + map.keys() + .as_any() + .downcast_ref::() + .expect("map key array should be Int32") + .values(), + &[1, 2] + ); + let values = map + .values() + .as_any() + .downcast_ref::() + .expect("map value array should be StringView"); + assert_eq!(values.value(0), "one"); + assert!(values.is_null(1)); + + Ok(()) + } + + #[test] + fn map_decimal_scalar_to_arrow_preserves_decimal_type() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(9, 2); + let dtype = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Decimal(decimal_dtype, Nullability::Nullable), + false, + Nullability::NonNullable, + )?; + let scalar = Scalar::try_map( + dtype, + [( + Scalar::primitive(1i32, Nullability::NonNullable), + Scalar::decimal( + DecimalValue::I256(i256::from_i128(12345)), + decimal_dtype, + Nullability::Nullable, + ), + )], + )?; + + let datum = scalar.to_arrow_datum()?; + let (array, is_scalar) = datum.get(); + assert!(is_scalar); + let map = array + .as_any() + .downcast_ref::() + .expect("map scalar should convert to MapArray"); + assert_eq!(map.values().data_type(), &DataType::Decimal128(9, 2)); + + let wide_decimal_dtype = DecimalDType::new(39, 2); + let dtype = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Decimal(wide_decimal_dtype, Nullability::Nullable), + false, + Nullability::NonNullable, + )?; + let scalar = Scalar::try_map( + dtype, + [( + Scalar::primitive(1i32, Nullability::NonNullable), + Scalar::decimal( + DecimalValue::I64(12345), + wide_decimal_dtype, + Nullability::Nullable, + ), + )], + )?; + + let datum = scalar.to_arrow_datum()?; + let (array, is_scalar) = datum.get(); + assert!(is_scalar); + let map = array + .as_any() + .downcast_ref::() + .expect("map scalar should convert to MapArray"); + assert_eq!(map.values().data_type(), &DataType::Decimal256(39, 2)); + + Ok(()) + } + + #[test] + fn map_scalar_with_unsupported_nested_value_errors_without_panic() -> VortexResult<()> { + let struct_dtype = DType::Struct( + StructFields::from_iter([( + "field1", + FieldDType::from(DType::Primitive(PType::I32, Nullability::NonNullable)), + )]), + Nullability::NonNullable, + ); + let dtype = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + struct_dtype.clone(), + false, + Nullability::NonNullable, + )?; + let scalar = Scalar::try_map( + dtype, + [( + Scalar::primitive(1i32, Nullability::NonNullable), + Scalar::struct_( + struct_dtype, + vec![Scalar::primitive(42i32, Nullability::NonNullable)], + ), + )], + )?; + + let error = scalar + .to_arrow_datum() + .err() + .expect("unsupported nested map value should error"); + assert!(error.to_string().contains("struct scalar conversion")); + + Ok(()) } #[test] - #[should_panic(expected = "struct scalar conversion")] fn test_struct_scalar_to_arrow_todo() { let struct_dtype = DType::Struct( StructFields::from_iter([( @@ -434,11 +749,14 @@ mod tests { struct_dtype, vec![Scalar::primitive(42i32, Nullability::NonNullable)], ); - struct_scalar.to_arrow_datum().unwrap(); + let error = struct_scalar + .to_arrow_datum() + .err() + .expect("struct scalar should error"); + assert!(error.to_string().contains("struct scalar conversion")); } #[test] - #[should_panic(expected = "list scalar conversion")] fn test_list_scalar_to_arrow_todo() { let element_dtype = Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)); let list_scalar = Scalar::list( @@ -450,11 +768,14 @@ mod tests { Nullability::NonNullable, ); - list_scalar.to_arrow_datum().unwrap(); + let error = list_scalar + .to_arrow_datum() + .err() + .expect("list scalar should error"); + assert!(error.to_string().contains("list scalar conversion")); } #[test] - #[should_panic(expected = "Cannot convert extension scalar")] fn test_non_temporal_extension_to_arrow_todo() { #[derive(Clone, Debug, Default, PartialEq, Eq, Hash)] struct SomeExt; @@ -492,7 +813,15 @@ mod tests { Scalar::primitive(42i32, Nullability::NonNullable), ); - scalar.to_arrow_datum().unwrap(); + let error = scalar + .to_arrow_datum() + .err() + .expect("non-temporal extension scalar should error"); + assert!( + error + .to_string() + .contains("Cannot convert extension scalar") + ); } #[rstest] diff --git a/vortex-arrow/src/session.rs b/vortex-arrow/src/session.rs index ea53df4f724..d3a6a532233 100644 --- a/vortex-arrow/src/session.rs +++ b/vortex-arrow/src/session.rs @@ -21,15 +21,22 @@ //! the next. use std::any::Any; +use std::borrow::Cow; use std::fmt::Debug; use std::sync::Arc; -use arrow_array::Array as _; +use arrow_array::Array as ArrowArray; use arrow_array::ArrayRef as ArrowArrayRef; use arrow_array::RecordBatch; +use arrow_array::RunArray; use arrow_array::make_array; +use arrow_array::types::Int16Type; +use arrow_array::types::Int32Type; +use arrow_array::types::Int64Type; +use arrow_array::types::RunEndIndexType; use arrow_schema::DataType; use arrow_schema::Field; +use arrow_schema::FieldRef; use arrow_schema::Fields; use arrow_schema::Schema; use arrow_schema::extension::EXTENSION_TYPE_NAME_KEY; @@ -39,7 +46,7 @@ use tracing::trace; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::arc_swap_map::ArcSwapMap; +use vortex_array::arrays::DictArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::ListViewArray; @@ -47,6 +54,7 @@ use vortex_array::arrays::StructArray; use vortex_array::dtype::DType; use vortex_array::dtype::FieldName; use vortex_array::dtype::FieldNames; +use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::dtype::StructFields; use vortex_array::dtype::extension::ExtId; @@ -56,18 +64,22 @@ use vortex_array::validity::Validity; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::ArcSwapMap; use vortex_session::SessionExt; use vortex_session::SessionGuard; use vortex_session::SessionVar; use vortex_session::registry::Id; -use crate::FromArrowArray; use crate::IntoVortexArray; +use crate::convert::from_arrow_dyn; +use crate::convert::map_from_arrow_parts; use crate::convert::nulls; use crate::convert::remove_nulls; -use crate::dtype::TryFromArrowType; +use crate::dtype::from_arrow_data_type; use crate::dtype::to_data_type_naive; use crate::executor::execute_arrow_naive; +use crate::run_end_import::run_end_from_arrow; /// Outcome of a successful call to [`ArrowExportVTable::execute_arrow`]. /// @@ -93,6 +105,47 @@ pub enum ArrowImport { Imported(ArrayRef), } +/// The Arrow type description of an array being imported by [`ArrowSession::from_arrow_array`]. +/// +/// Callers holding an Arrow [`Field`] (or [`FieldRef`]) should pass it: its +/// `ARROW:extension:name` metadata is what dispatches the array to a registered +/// [`ArrowImportVTable`]. +/// +/// An Arrow array can carry a validity (null) buffer regardless of whether its schema declares +/// the field nullable, so when no [`Field`] is in hand the caller passes the desired nullability +/// instead, as a [`bool`] or a [`Nullability`]. An anonymous field is then synthesized from the +/// array's own data type, which means no extension plugin is dispatched for the array itself; +/// fields nested inside container data types still carry their metadata and are routed through +/// their importers. +pub trait IntoArrowField<'a> { + /// Resolve to the Arrow [`Field`] describing an array of `data_type`. + fn into_arrow_field(self, data_type: &DataType) -> Cow<'a, Field>; +} + +impl<'a> IntoArrowField<'a> for &'a Field { + fn into_arrow_field(self, _data_type: &DataType) -> Cow<'a, Field> { + Cow::Borrowed(self) + } +} + +impl<'a> IntoArrowField<'a> for &'a FieldRef { + fn into_arrow_field(self, _data_type: &DataType) -> Cow<'a, Field> { + Cow::Borrowed(self.as_ref()) + } +} + +impl<'a> IntoArrowField<'a> for bool { + fn into_arrow_field(self, data_type: &DataType) -> Cow<'a, Field> { + Cow::Owned(Field::new("", data_type.clone(), self)) + } +} + +impl<'a> IntoArrowField<'a> for Nullability { + fn into_arrow_field(self, data_type: &DataType) -> Cow<'a, Field> { + self.is_nullable().into_arrow_field(data_type) + } +} + /// Plugin layer for exporting a Vortex array to an Arrow extension type. /// /// This is purely an implementation trait, its methods should not be called directly. Instead, @@ -139,25 +192,46 @@ pub trait ArrowImportVTable: 'static + Send + Sync + Debug { /// Build the Vortex [`DType`] that corresponds to `field` (which carries this plugin's /// Arrow extension metadata). + /// + /// `session` is provided so plugins can resolve nested or storage fields through the + /// plugin-aware conversion (e.g. [`ArrowSession::from_arrow_datatype`]) instead of the + /// naive Arrow → Vortex mapping. #[allow(clippy::wrong_self_convention)] - fn from_arrow_field(&self, field: &Field) -> VortexResult>; + fn from_arrow_field( + &self, + field: &Field, + session: &ArrowSession, + ) -> VortexResult>; /// Convert an Arrow array into a Vortex array of `dtype`. /// /// Returns ownership of `array` via [`ArrowImport::Unsupported`] when the plugin cannot /// handle the input. + /// + /// `session` is provided so plugins can convert storage or nested arrays through the + /// session (e.g. [`ArrowSession::from_arrow_array`], passing a nullability rather than a + /// [`Field`] so the plugin is not dispatched again) instead of the deprecated + /// `FromArrowArray` trait. #[allow(clippy::wrong_self_convention)] fn from_arrow_array( &self, array: ArrowArrayRef, field: &Field, dtype: &DType, + session: &ArrowSession, ) -> VortexResult; } pub type ArrowExportVTableRef = Arc; pub type ArrowImportVTableRef = Arc; +/// Registry of Arrow exporters, keyed by target Arrow extension [`Id`]. +type ArrowExporterRegistry = ArcSwapMap>; +/// Registry of Arrow exporters, keyed by source Vortex extension [`ExtId`]. +type VortexExporterRegistry = ArcSwapMap>; +/// Registry of Arrow importers, keyed by source Arrow extension [`Id`]. +type ArrowImporterRegistry = ArcSwapMap>; + /// Session-scoped registry of Arrow extension plugins. /// /// Exporters are stored in two indices: one keyed by Arrow extension Id (used for @@ -169,17 +243,17 @@ pub type ArrowImportVTableRef = Arc; /// need plugins. #[derive(Clone, Debug)] pub struct ArrowSession { - exporters: ArcSwapMap>, - exporters_by_vortex: ArcSwapMap>, - importers: ArcSwapMap>, + exporters: ArrowExporterRegistry, + exporters_by_vortex: VortexExporterRegistry, + importers: ArrowImporterRegistry, } impl Default for ArrowSession { fn default() -> Self { let session = Self { - exporters: ArcSwapMap::default(), - exporters_by_vortex: ArcSwapMap::default(), - importers: ArcSwapMap::default(), + exporters: ArrowExporterRegistry::default(), + exporters_by_vortex: VortexExporterRegistry::default(), + importers: ArrowImporterRegistry::default(), }; session.register_exporter(Arc::new(Uuid)); @@ -240,6 +314,16 @@ impl ArrowSession { nullability.is_nullable(), )) } + DType::Map(map_dtype, nullability) => { + let key = self.to_arrow_field("key", &map_dtype.key_dtype())?; + let value = self.to_arrow_field("value", &map_dtype.value_dtype())?; + let entries = Field::new_struct("entries", Fields::from(vec![key, value]), false); + Ok(Field::new( + name, + DataType::Map(Arc::new(entries), map_dtype.keys_sorted()), + nullability.is_nullable(), + )) + } DType::Struct(fields, nullability) => { let arrow_fields = Fields::from_iter( fields @@ -309,25 +393,55 @@ impl ArrowSession { Ok(Schema::new(fields)) } + /// Returns the Arrow [`DataType`] that best corresponds to the given Vortex [`DType`], + /// dispatching [`DType::Extension`]s through registered export plugins. + /// + /// Note that a bare [`DataType`] cannot carry `ARROW:extension:name` metadata; use + /// [`Self::to_arrow_field`] when extension identity must survive the roundtrip. + pub fn to_arrow_datatype(&self, dtype: &DType) -> VortexResult { + Ok(self.to_arrow_field("", dtype)?.data_type().clone()) + } + /// Build the Vortex [`DType`] for an Arrow [`Field`]. /// /// Plugins registered against the field's Arrow extension name are tried in /// registration order; the first plugin to return `Some(dtype)` wins. If none - /// match (or all return `None`), recurses into container types ([`DataType::List`] - /// family, [`DataType::FixedSizeList`], [`DataType::Struct`]) so extension metadata - /// on nested element/struct fields is preserved. Leaf types use the canonical - /// Arrow → Vortex mapping via [`DType::try_from_arrow`]. + /// match (or all return `None`), the builtin `arrow.parquet.variant` extension maps + /// to [`DType::Variant`], and any other field converts through + /// [`Self::from_arrow_datatype`] so extension metadata on nested element/struct + /// fields is preserved. #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] pub fn from_arrow_field(&self, field: &Field) -> VortexResult { if let Some(name) = field.metadata().get(EXTENSION_TYPE_NAME_KEY) { for plugin in self.importers(&Id::new(name)).iter() { - if let Some(dtype) = plugin.from_arrow_field(field)? { + if let Some(dtype) = plugin.from_arrow_field(field, self)? { return Ok(dtype); } } + // Parquet Variant is understood even without a registered importer plugin. + if name == "arrow.parquet.variant" { + return Ok(DType::Variant(field.is_nullable().into())); + } } - let nullability: Nullability = field.is_nullable().into(); - Ok(match field.data_type() { + self.from_arrow_datatype(field.data_type(), field.is_nullable().into()) + } + + /// Build the Vortex [`DType`] for an Arrow [`DataType`]. + /// + /// Recurses into container types ([`DataType::List`] family, [`DataType::FixedSizeList`], + /// [`DataType::Struct`], [`DataType::Map`], [`DataType::RunEndEncoded`]) via + /// [`Self::from_arrow_field`] so extension metadata on nested fields dispatches through + /// registered import plugins. Leaf types use the canonical Arrow → Vortex mapping. + /// + /// [`DataType::Dictionary`] is a partial exception: Arrow models its values as a bare + /// [`DataType`] rather than a [`Field`], so the values themselves cannot carry an extension + /// name. Fields nested *inside* that data type still do, and are dispatched normally. + pub fn from_arrow_datatype( + &self, + data_type: &DataType, + nullability: Nullability, + ) -> VortexResult { + Ok(match data_type { DataType::List(elem) | DataType::LargeList(elem) | DataType::ListView(elem) @@ -339,41 +453,61 @@ impl ArrowSession { *size as u32, nullability, ), - // A Map is stored as its entries list, so it recurses like one: extension - // metadata on a key or value field must reach its importer exactly as it does - // for the `List>` the map is aliased to. - DataType::Map(entries, _ordered) => DType::List( - Arc::new(self.from_arrow_field(entries.as_ref())?), - nullability, - ), - DataType::Struct(fields) => { - let entries = fields - .iter() - .map(|f| { - self.from_arrow_field(f) - .map(|dt| (FieldName::from(f.name().as_str()), dt)) - }) - .collect::>>()?; - DType::Struct(StructFields::from_iter(entries), nullability) + DataType::Map(entries, keys_sorted) => { + vortex_ensure!( + !entries.is_nullable(), + "Arrow map entries field must be non-nullable" + ); + let DataType::Struct(fields) = entries.data_type() else { + vortex_bail!( + "Arrow map entries field must have Struct type, got {:?}", + entries.data_type() + ); + }; + vortex_ensure!( + fields.len() == 2, + "Arrow map entries struct must contain exactly two fields" + ); + vortex_ensure!( + !fields[0].is_nullable(), + "Arrow map key field must be non-nullable" + ); + DType::map( + self.from_arrow_field(fields[0].as_ref())?, + self.from_arrow_field(fields[1].as_ref())?, + *keys_sorted, + nullability, + )? + } + DataType::Struct(fields) => DType::Struct(self.from_arrow_fields(fields)?, nullability), + DataType::Dictionary(_, value_type) => { + self.from_arrow_datatype(value_type.as_ref(), nullability)? + } + DataType::RunEndEncoded(_, value_field) => { + self.from_arrow_field(&run_end_values_field(value_field, nullability))? } - _ => DType::try_from_arrow(field)?, + _ => from_arrow_data_type(data_type, nullability)?, }) } - /// Build the Vortex [`DType`] for an Arrow [`Schema`], dispatching extension fields - /// through registered import plugins. The result is a top-level non-nullable struct - /// matching the schema's fields. - pub fn from_arrow_schema(&self, schema: &Schema) -> VortexResult { - let entries = schema - .fields() + /// Build Vortex [`StructFields`] for Arrow [`Fields`], dispatching each field through + /// [`Self::from_arrow_field`]. + pub fn from_arrow_fields(&self, fields: &Fields) -> VortexResult { + fields .iter() .map(|f| { self.from_arrow_field(f) .map(|dt| (FieldName::from(f.name().as_str()), dt)) }) - .collect::>>()?; + .collect::>() + } + + /// Build the Vortex [`DType`] for an Arrow [`Schema`], dispatching extension fields + /// through registered import plugins. The result is a top-level non-nullable struct + /// matching the schema's fields. + pub fn from_arrow_schema(&self, schema: &Schema) -> VortexResult { Ok(DType::Struct( - StructFields::from_iter(entries), + self.from_arrow_fields(schema.fields())?, Nullability::NonNullable, )) } @@ -405,7 +539,7 @@ impl ArrowSession { ); let mut columns = Vec::with_capacity(schema.fields().len()); for (col, field) in batch.columns().iter().zip(schema.fields().iter()) { - columns.push(self.from_arrow_array(ArrowArrayRef::clone(col), field)?); + columns.push(self.from_arrow_array_inner(ArrowArrayRef::clone(col), field)?); } Ok(StructArray::try_new(names, columns, length, Validity::NonNullable)?.into_array()) } @@ -477,6 +611,11 @@ impl ArrowSession { /// Decode an Arrow array into a Vortex array. /// + /// `field` describes the Arrow type the array is imported from: pass an Arrow [`Field`] when + /// one is available, or just the desired nullability (`true` / `false` / + /// [`Nullability`]) to synthesize an anonymous field from the array's own data type. See + /// [`IntoArrowField`] for the trade-off between the two. + /// /// Routes through the registered import plugin if `field` carries an Arrow extension /// name we recognize, probing each plugin in registration order until one handles the /// input or all return [`ArrowImport::Unsupported`]. Otherwise recurses into container @@ -484,7 +623,29 @@ impl ArrowSession { /// [`arrow_array::FixedSizeListArray`], [`arrow_array::GenericListViewArray`]) so /// extension fields nested inside containers reach their importers; leaf types fall /// through to the canonical Arrow → Vortex array conversion. - pub fn from_arrow_array(&self, array: ArrowArrayRef, field: &Field) -> VortexResult { + /// + /// # Errors + /// + /// Returns an error if the field (or requested nullability) is non-nullable but the array + /// physically contains nulls, or if the Arrow data type is unsupported. + pub fn from_arrow_array<'a>( + &self, + array: ArrowArrayRef, + field: impl IntoArrowField<'a>, + ) -> VortexResult { + let field = field.into_arrow_field(array.data_type()); + self.from_arrow_array_inner(array, field.as_ref()) + } + + /// [`Self::from_arrow_array`] with the Arrow [`Field`] already resolved: probe the import + /// plugins registered for the field's extension name, then fall back to the canonical + /// conversion. Also the recursion point for nested fields, which already have a [`Field`]. + #[allow(clippy::wrong_self_convention)] + fn from_arrow_array_inner( + &self, + array: ArrowArrayRef, + field: &Field, + ) -> VortexResult { if let Some(extension_name) = field.metadata().get(EXTENSION_TYPE_NAME_KEY) { #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] let importers = self.importers(&Id::new(extension_name)); @@ -492,23 +653,23 @@ impl ArrowSession { let dtype = self.from_arrow_field(field)?; let mut current = array; for plugin in importers.iter() { - match plugin.from_arrow_array(current, field, &dtype)? { + match plugin.from_arrow_array(current, field, &dtype, self)? { ArrowImport::Imported(arr) => return Ok(arr), ArrowImport::Unsupported(arr) => current = arr, } } - return ArrayRef::from_arrow(current.as_ref(), field.is_nullable()); + return self.from_arrow_array_canonical(current.as_ref(), field); } } - self.from_arrow_array_canonical(array, field) + self.from_arrow_array_canonical(array.as_ref(), field) } /// Recurse into Arrow container arrays so nested fields with extension metadata reach - /// their importers, falling through to [`ArrayRef::from_arrow`] for leaf types. + /// their importers, falling through to the canonical conversion for leaf types. #[allow(clippy::wrong_self_convention)] fn from_arrow_array_canonical( &self, - array: ArrowArrayRef, + array: &dyn ArrowArray, field: &Field, ) -> VortexResult { use arrow_array::cast::AsArray; @@ -531,7 +692,7 @@ impl ArrowSession { } else { ArrowArrayRef::clone(col) }; - self.from_arrow_array(inner, child_field.as_ref()) + self.from_arrow_array_inner(inner, child_field.as_ref()) }) .collect::>>()?; let validity = nulls(arrow_struct.nulls(), field.is_nullable())?; @@ -558,8 +719,10 @@ impl ArrowSession { } DataType::FixedSizeList(elem_field, list_size) => { let fsl = array.as_fixed_size_list(); - let elements = - self.from_arrow_array(ArrowArrayRef::clone(fsl.values()), elem_field.as_ref())?; + let elements = self.from_arrow_array_inner( + ArrowArrayRef::clone(fsl.values()), + elem_field.as_ref(), + )?; let validity = nulls(fsl.nulls(), field.is_nullable())?; Ok( FixedSizeListArray::try_new(elements, *list_size as u32, validity, fsl.len())? @@ -584,22 +747,76 @@ impl ArrowSession { let validity = nulls(list.nulls(), field.is_nullable())?; Ok(ListViewArray::try_new(elements, offsets, sizes, validity)?.into_array()) } - // Arrow requires a map's entries struct to be non-nullable, so only the map - // itself carries validity; the entries recurse so nested extension fields reach - // their importers. - DataType::Map(entries_field, _ordered) => { + DataType::Map(entries_field, keys_sorted) => { let map = array.as_map(); - let entries = self.from_arrow_array( - Arc::new(map.entries().clone()) as ArrowArrayRef, - entries_field.as_ref(), - )?; - let offsets = map.offsets().clone().into_array(); - let validity = nulls(map.nulls(), field.is_nullable())?; - Ok(ListArray::try_new(entries, offsets, validity)?.into_array()) + let entries_array: ArrowArrayRef = Arc::new(map.entries().clone()); + let entries = self.from_arrow_array_inner(entries_array, entries_field.as_ref())?; + map_from_arrow_parts( + entries, + map.offsets(), + map.nulls(), + *keys_sorted, + field.is_nullable(), + ) } - _ => ArrayRef::from_arrow(array.as_ref(), field.is_nullable()), + DataType::RunEndEncoded(ends_field, values_field) => { + let values_field = run_end_values_field(values_field, field.is_nullable().into()); + match ends_field.data_type() { + DataType::Int16 => self.run_end_from_arrow::(array, &values_field), + DataType::Int32 => self.run_end_from_arrow::(array, &values_field), + DataType::Int64 => self.run_end_from_arrow::(array, &values_field), + ends_dt => vortex_bail!( + "Arrow run-end array run ends must be Int16, Int32 or Int64, got {ends_dt}" + ), + } + } + DataType::Dictionary(..) => { + let dict = array.as_any_dictionary(); + // Arrow models dictionary values as a bare `DataType`, so there is no field + // metadata to carry an extension name for the values themselves. Fields *nested + // inside* that data type (list elements, struct fields, map entries) do keep + // their metadata, so importing the values by nullability alone still routes them + // back through the plugin-aware conversion. + let values = self + .from_arrow_array(ArrowArrayRef::clone(dict.values()), field.is_nullable())?; + let codes = dict.keys(); + let codes = from_arrow_dyn(codes, codes.is_nullable())?; + // SAFETY: arrow-rs enforces the dictionary invariants on construction, so the + // codes are in-bounds for the values. + Ok(unsafe { DictArray::new_unchecked(codes, values) }.into_array()) + } + _ => from_arrow_dyn(array, field.is_nullable()), } } + + /// Decode an Arrow run-end array, recursing into its values so extension metadata on the + /// values field reaches its importer. + #[allow(clippy::wrong_self_convention)] + fn run_end_from_arrow( + &self, + array: &dyn ArrowArray, + values_field: &Field, + ) -> VortexResult + where + R::Native: NativePType, + { + let run_array = array + .as_any() + .downcast_ref::>() + .ok_or_else(|| vortex_err!("expected an Arrow RunArray, got {}", array.data_type()))?; + let values = + self.from_arrow_array_inner(ArrowArrayRef::clone(run_array.values()), values_field)?; + run_end_from_arrow(run_array, values) + } +} + +/// The values field of an Arrow [`DataType::RunEndEncoded`], re-stamped with the run-end array's +/// own nullability. +fn run_end_values_field(values_field: &FieldRef, nullability: Nullability) -> Field { + values_field + .as_ref() + .clone() + .with_nullable(nullability.into()) } // NOTE(aduffy): We should remove this once we bump Arrow to 0.59.0. This is replicating the @@ -639,13 +856,24 @@ impl ArrowSessionExt for S { mod tests { use std::sync::Arc; + use arrow_array::DictionaryArray; use arrow_array::FixedSizeBinaryArray; + use arrow_array::Int32Array; + use arrow_array::ListArray as ArrowListArray; + use arrow_array::StringArray; + use arrow_array::StructArray as ArrowStructArray; use arrow_array::cast::AsArray; + use arrow_buffer::OffsetBuffer; use arrow_schema::DataType; use arrow_schema::Field; use arrow_schema::extension::Uuid as ArrowUuid; + use rstest::rstest; + use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; + use vortex_array::arrays::Dict; + use vortex_array::arrays::ListArray; + use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::DType; use vortex_array::dtype::FieldName; use vortex_array::dtype::Nullability; @@ -655,6 +883,7 @@ mod tests { use vortex_array::dtype::extension::ExtVTable; use vortex_array::extension::uuid::Uuid; use vortex_array::extension::uuid::UuidMetadata; + use vortex_error::VortexExpect; use vortex_error::VortexResult; use super::*; @@ -722,6 +951,40 @@ mod tests { Ok(()) } + #[test] + fn schema_roundtrip_preserves_map_uuid_fields() -> VortexResult<()> { + let session = ArrowSession::default(); + let map = DType::map( + uuid_dtype(false), + uuid_dtype(true), + true, + Nullability::Nullable, + )?; + let dtype = DType::Struct( + StructFields::from_iter([(FieldName::from("ids"), map)]), + Nullability::NonNullable, + ); + + let schema = session.to_arrow_schema(&dtype)?; + let field = schema.field(0); + let DataType::Map(entries, keys_sorted) = field.data_type() else { + panic!("expected Map, got {:?}", field.data_type()); + }; + assert!(*keys_sorted); + assert_eq!(entries.name(), "entries"); + assert!(!entries.is_nullable()); + let DataType::Struct(fields) = entries.data_type() else { + panic!("expected map entries struct, got {:?}", entries.data_type()); + }; + assert!(has_valid_extension_type::(&fields[0])); + assert!(has_valid_extension_type::(&fields[1])); + assert!(!fields[0].is_nullable()); + assert!(fields[1].is_nullable()); + + assert_eq!(session.from_arrow_schema(&schema)?, dtype); + Ok(()) + } + #[test] fn to_arrow_schema_struct_of_struct_uuid() -> VortexResult<()> { let session = ArrowSession::default(); @@ -779,6 +1042,78 @@ mod tests { Ok(()) } + #[test] + fn to_arrow_datatype_dispatches_plugins() -> VortexResult<()> { + let session = ArrowSession::default(); + assert_eq!( + session.to_arrow_datatype(&uuid_dtype(false))?, + DataType::FixedSizeBinary(16) + ); + assert_eq!( + session.to_arrow_datatype(&DType::Utf8(Nullability::Nullable))?, + DataType::Utf8View + ); + Ok(()) + } + + #[test] + fn from_arrow_datatype_recurses_into_nested_extension_fields() -> VortexResult<()> { + let session = ArrowSession::default(); + let mut elem = Field::new("item", DataType::FixedSizeBinary(16), false); + elem.try_with_extension_type(ArrowUuid)?; + let data_type = DataType::List(Arc::new(elem)); + + let dtype = session.from_arrow_datatype(&data_type, Nullability::Nullable)?; + let DType::List(inner_dt, Nullability::Nullable) = dtype else { + panic!("expected nullable List dtype, got {dtype}"); + }; + assert!( + matches!(inner_dt.as_ref(), DType::Extension(ext) if ext.id() == Uuid.id()), + "expected Uuid extension element, got {inner_dt}", + ); + Ok(()) + } + + #[test] + fn from_arrow_fields_matches_schema_conversion() -> VortexResult<()> { + let session = ArrowSession::default(); + let fields = Fields::from(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Utf8View, true), + ]); + let struct_fields = session.from_arrow_fields(&fields)?; + let schema_dtype = session.from_arrow_schema(&Schema::new(fields))?; + assert_eq!( + schema_dtype, + DType::Struct(struct_fields, Nullability::NonNullable) + ); + Ok(()) + } + + #[test] + fn from_arrow_field_maps_variant_without_importer() -> VortexResult<()> { + let session = ArrowSession::default(); + let storage = DataType::Struct( + vec![ + Field::new("metadata", DataType::BinaryView, false), + Field::new("value", DataType::BinaryView, true), + ] + .into(), + ); + let field = Field::new("v", storage, true).with_metadata( + [( + "ARROW:extension:name".to_string(), + "arrow.parquet.variant".to_string(), + )] + .into(), + ); + assert_eq!( + session.from_arrow_field(&field)?, + DType::Variant(Nullability::Nullable) + ); + Ok(()) + } + #[test] fn execute_arrow_target_none_preserves_top_level_uuid_metadata() -> VortexResult<()> { let vortex_session = array_session(); @@ -805,79 +1140,241 @@ mod tests { Ok(()) } - /// A map's key and value fields are ordinary fields and may carry extension metadata, so - /// importing a map through the session has to recurse into its entries exactly as the - /// `List>` it is aliased to does. Without that recursion the whole map - /// delegates to the canonical conversion, which is plugin-blind: a `Map` - /// would be read as raw `FixedSizeBinary(16)` and fail to import at all. + /// Import an Arrow FixedSizeBinary UUID column as a Vortex extension array. + fn uuid_array(session: &ArrowSession) -> VortexResult { + let mut field = Field::new("id", DataType::FixedSizeBinary(16), false); + field.try_with_extension_type(ArrowUuid)?; + let arrow_array: ArrowArrayRef = Arc::new(FixedSizeBinaryArray::try_from_iter( + [*b"0123456789abcdef", *b"fedcba9876543210"].into_iter(), + )?); + session.from_arrow_array(arrow_array, &field) + } + + /// Exporting a struct that contains an extension column with no target field must still route + /// the column through its export plugin *and* re-attach the Arrow extension metadata to the + /// inferred child field. #[test] - fn map_import_recurses_into_entries_so_nested_extensions_reach_their_importer() - -> VortexResult<()> { - use arrow_array::MapArray; - use arrow_array::StringArray; - use arrow_array::StructArray; - use arrow_buffer::OffsetBuffer; + fn execute_arrow_target_none_preserves_nested_uuid_metadata() -> VortexResult<()> { + let vortex_session = array_session(); + let mut ctx = vortex_session.create_execution_ctx(); + let session = vortex_session.arrow(); - let session = ArrowSession::default(); + let uuids = uuid_array(&session)?; + let struct_array = StructArray::try_new( + FieldNames::from(["id"]), + vec![uuids], + 2, + Validity::NonNullable, + )? + .into_array(); + + let exported = session.execute_arrow(struct_array, None, &mut ctx)?; + let DataType::Struct(fields) = exported.data_type() else { + panic!("expected Struct, got {:?}", exported.data_type()); + }; + assert_eq!(fields[0].data_type(), &DataType::FixedSizeBinary(16)); + assert!(has_valid_extension_type::(&fields[0])); - let mut value_field = Field::new("values", DataType::FixedSizeBinary(16), false); - value_field.try_with_extension_type(ArrowUuid)?; - let value_field = Arc::new(value_field); - let key_field = Arc::new(Field::new("keys", DataType::Utf8, false)); - let entries_field = Arc::new(Field::new( - "entries", - DataType::Struct(vec![Arc::clone(&key_field), Arc::clone(&value_field)].into()), - false, - )); - let map_field = Field::new("m", DataType::Map(Arc::clone(&entries_field), false), true); + let uuids = exported.as_struct().column(0).as_fixed_size_binary(); + assert_eq!(uuids.value(0), b"0123456789abcdef"); + assert_eq!(uuids.value(1), b"fedcba9876543210"); + Ok(()) + } - // The dtype has to see through the map to the UUID extension. - let dtype = session.from_arrow_field(&map_field)?; - let DType::List(entries_dtype, _) = &dtype else { - panic!("expected a List dtype for a map, got {dtype}"); - }; - let DType::Struct(entry_fields, _) = entries_dtype.as_ref() else { - panic!("expected Struct entries, got {entries_dtype}"); + /// Exporting a list of extension elements with no target field must infer an element field that + /// still carries the Arrow extension metadata. + #[test] + fn execute_arrow_target_none_preserves_list_element_uuid_metadata() -> VortexResult<()> { + let vortex_session = array_session(); + let mut ctx = vortex_session.create_execution_ctx(); + let session = vortex_session.arrow(); + + let list = ListArray::try_new( + uuid_array(&session)?, + PrimitiveArray::from_iter([0i32, 1, 2]).into_array(), + Validity::NonNullable, + )? + .into_array(); + + let exported = session.execute_arrow(list, None, &mut ctx)?; + let DataType::List(elem) = exported.data_type() else { + panic!("expected List, got {:?}", exported.data_type()); }; - let values_dtype = entry_fields - .field("values") - .expect("entries should carry a `values` field"); + assert_eq!(elem.data_type(), &DataType::FixedSizeBinary(16)); + assert!(has_valid_extension_type::(elem)); + + let uuids = exported.as_list::().values().as_fixed_size_binary(); + assert_eq!(uuids.value(0), b"0123456789abcdef"); + assert_eq!(uuids.value(1), b"fedcba9876543210"); + Ok(()) + } + + /// An Arrow run-end array whose values field carries extension metadata must import as that + /// extension, through both the dtype and the array conversion. + #[test] + fn run_end_recurses_into_extension_values() -> VortexResult<()> { + let vortex_session = array_session(); + let mut ctx = vortex_session.create_execution_ctx(); + let session = vortex_session.arrow(); + + let mut values_field = Field::new("values", DataType::FixedSizeBinary(16), false); + values_field.try_with_extension_type(ArrowUuid)?; + let field = Field::new( + "id", + DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int32, false)), + Arc::new(values_field), + ), + false, + ); + + let dtype = session.from_arrow_field(&field)?; assert!( - matches!(&values_dtype, DType::Extension(ext) if ext.id() == Uuid.id()), - "expected the map's values to import as the Uuid extension, got {values_dtype}", + dtype.as_extension().is::(), + "expected a Uuid extension dtype, got {dtype}" ); - // And so does the array import. - let entries = StructArray::try_new( - vec![Arc::clone(&key_field), Arc::clone(&value_field)].into(), - vec![ - Arc::new(StringArray::from(vec!["a", "b"])) as ArrowArrayRef, - Arc::new(FixedSizeBinaryArray::try_from_iter( - [*b"0123456789abcdef", *b"fedcba9876543210"].into_iter(), - )?) as ArrowArrayRef, - ], - None, + let values = FixedSizeBinaryArray::try_from_iter( + [*b"0123456789abcdef", *b"fedcba9876543210"].into_iter(), )?; - let map: ArrowArrayRef = Arc::new(MapArray::try_new( - entries_field, + let run_array: ArrowArrayRef = Arc::new(RunArray::::try_new( + &Int32Array::from(vec![2i32, 5]), + &values, + )?); + + let vortex_array = session.from_arrow_array(run_array, &field)?; + assert_eq!(vortex_array.len(), 5); + // The array conversion must agree with the dtype conversion. + assert_eq!(vortex_array.dtype(), &dtype); + + // And the values must round-trip back out through the export plugin. + let exported = session.execute_arrow(vortex_array, Some(&field), &mut ctx)?; + assert_eq!(exported.len(), 5); + let ree = exported + .as_any() + .downcast_ref::>() + .ok_or_else(|| { + vortex_err!( + "expected an Int32 run-end array, got {}", + exported.data_type() + ) + })?; + let values = ree.values().as_fixed_size_binary(); + assert_eq!(values.value(0), b"0123456789abcdef"); + assert_eq!(values.value(1), b"fedcba9876543210"); + Ok(()) + } + + /// An Arrow dictionary array cannot carry extension metadata on the values themselves, but + /// fields nested inside the values data type can. Both the dtype and the array conversion + /// must dispatch those through their importer, and must agree with each other. + #[test] + fn dictionary_recurses_into_nested_extension_values() -> VortexResult<()> { + let session = ArrowSession::default(); + + let mut elem = Field::new("item", DataType::FixedSizeBinary(16), false); + elem.try_with_extension_type(ArrowUuid)?; + + let uuids = FixedSizeBinaryArray::try_from_iter( + [*b"0123456789abcdef", *b"fedcba9876543210"].into_iter(), + )?; + let values = ArrowListArray::try_new( + Arc::new(elem), OffsetBuffer::new(vec![0, 1, 2].into()), - entries, + Arc::new(uuids), None, - false, + )?; + let dict: ArrowArrayRef = Arc::new(DictionaryArray::::try_new( + Int32Array::from(vec![0, 1, 0]), + Arc::new(values), )?); + let field = Field::new("ids", dict.data_type().clone(), false); - let imported = session.from_arrow_array(Arc::clone(&map), &map_field)?; - assert_eq!(imported.dtype(), &dtype); + let dtype = session.from_arrow_field(&field)?; + let DType::List(elem_dt, _) = &dtype else { + panic!("expected a List dtype, got {dtype}"); + }; + assert!( + elem_dt.as_extension().is::(), + "expected a Uuid extension element, got {elem_dt}" + ); - // The export half already recurses - `to_arrow_map` hands the entries field to the - // list executor, which dispatches through the session - so the map round-trips whole. - let vortex_session = array_session(); - let exported = vortex_session.arrow().execute_arrow( - imported, - Some(&map_field), - &mut vortex_session.create_execution_ctx(), - )?; - assert_eq!(exported.as_ref(), map.as_ref()); + let array = session.from_arrow_array(dict, &field)?; + assert_eq!(array.len(), 3); + // The array conversion must agree with the dtype conversion. + assert_eq!(array.dtype(), &dtype); + + // Arrow's dictionary values data type does carry the nested element field, so the + // extension survives the export as well. + let mut ctx = array_session().create_execution_ctx(); + let exported = session.execute_arrow(array, Some(&field), &mut ctx)?; + assert_eq!(exported.data_type(), field.data_type()); + let values = exported.as_any_dictionary().values().as_list::(); + let uuids = values.values().as_fixed_size_binary(); + assert_eq!(uuids.value(0), b"0123456789abcdef"); + assert_eq!(uuids.value(1), b"fedcba9876543210"); + Ok(()) + } + + /// Importing by nullability instead of by [`Field`] synthesizes an anonymous field from the + /// array's own data type, so extension metadata on *nested* fields still reaches its importer. + /// A [`bool`] and the equivalent [`Nullability`] must agree. + #[rstest] + #[case(true)] + #[case(false)] + fn from_arrow_array_by_nullability(#[case] nullable: bool) -> VortexResult<()> { + let session = ArrowSession::default(); + + let mut uuid_field = Field::new("id", DataType::FixedSizeBinary(16), false); + uuid_field.try_with_extension_type(ArrowUuid)?; + let uuids: ArrowArrayRef = Arc::new(FixedSizeBinaryArray::try_from_iter( + [*b"0123456789abcdef", *b"fedcba9876543210"].into_iter(), + )?); + let arrow_struct: ArrowArrayRef = Arc::new(ArrowStructArray::try_new( + Fields::from(vec![uuid_field]), + vec![uuids], + None, + )?); + + let array = session.from_arrow_array(ArrowArrayRef::clone(&arrow_struct), nullable)?; + assert_eq!(array.dtype().nullability(), nullable.into()); + let DType::Struct(fields, _) = array.dtype() else { + panic!("expected a Struct dtype, got {}", array.dtype()); + }; + assert!( + fields + .field_by_index(0) + .vortex_expect("struct dtype has one field") + .as_extension() + .is::(), + "expected the nested Uuid extension to survive, got {}", + array.dtype() + ); + + // The `Nullability` form is equivalent to the `bool` form. + let by_nullability = session.from_arrow_array(arrow_struct, Nullability::from(nullable))?; + assert_eq!(by_nullability.dtype(), array.dtype()); + Ok(()) + } + + /// A plain Arrow dictionary imports as a Vortex `Dict` array over the dictionary values, + /// matching the dtype the schema conversion reports. + #[test] + fn dictionary_imports_as_dict_encoding() -> VortexResult<()> { + let session = ArrowSession::default(); + let dict: ArrowArrayRef = Arc::new(DictionaryArray::::try_new( + Int32Array::from(vec![0, 1, 0, 1]), + Arc::new(StringArray::from(vec!["a", "b"])), + )?); + let field = Field::new("s", dict.data_type().clone(), false); + + let dtype = session.from_arrow_field(&field)?; + assert_eq!(dtype, DType::Utf8(Nullability::NonNullable)); + + let array = session.from_arrow_array(dict, &field)?; + assert!(array.is::(), "expected a Dict encoding, got {array}"); + assert_eq!(array.dtype(), &dtype); + assert_eq!(array.len(), 4); Ok(()) } } diff --git a/vortex-arrow/src/uuid.rs b/vortex-arrow/src/uuid.rs index 5901f56389b..33228816283 100644 --- a/vortex-arrow/src/uuid.rs +++ b/vortex-arrow/src/uuid.rs @@ -104,7 +104,11 @@ impl ArrowImportVTable for Uuid { *ARROW_UUID } - fn from_arrow_field(&self, field: &Field) -> VortexResult> { + fn from_arrow_field( + &self, + field: &Field, + _session: &ArrowSession, + ) -> VortexResult> { if !has_valid_extension_type::(field) { return Ok(None); } @@ -125,6 +129,7 @@ impl ArrowImportVTable for Uuid { array: ArrowArrayRef, _field: &Field, dtype: &DType, + _session: &ArrowSession, ) -> VortexResult { let DType::Extension(dtype) = dtype else { return Ok(ArrowImport::Unsupported(array)); diff --git a/vortex-arrow/tests/canonical.rs b/vortex-arrow/tests/canonical.rs index f67e0b73115..5b1b70e29ac 100644 --- a/vortex-arrow/tests/canonical.rs +++ b/vortex-arrow/tests/canonical.rs @@ -16,7 +16,6 @@ use arrow_array::PrimitiveArray as ArrowPrimitiveArray; use arrow_array::StringArray; use arrow_array::StringViewArray; use arrow_array::StructArray as ArrowStructArray; -use arrow_array::cast::AsArray; use arrow_array::types::Int32Type; use arrow_array::types::Int64Type; use arrow_array::types::UInt64Type; @@ -24,13 +23,11 @@ use arrow_buffer::NullBufferBuilder; use arrow_buffer::OffsetBuffer; use arrow_schema::DataType; use arrow_schema::Field; -use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::StructArray; use vortex_arrow::ArrowSessionExt; -use vortex_arrow::FromArrowArray; use vortex_buffer::buffer; use vortex_session::VortexSession; @@ -118,7 +115,7 @@ fn roundtrip_struct() { None, ])); - let arrow_struct = ArrowStructArray::new( + let arrow_struct = Arc::new(ArrowStructArray::new( vec![ Arc::new(Field::new("name", DataType::Utf8View, true)), Arc::new(Field::new("age", DataType::Int32, true)), @@ -126,14 +123,17 @@ fn roundtrip_struct() { .into(), vec![names, ages], nulls.finish(), - ); + )) as ArrowArrayRef; - let vortex_struct = ArrayRef::from_arrow(&arrow_struct, true).unwrap(); + let vortex_struct = SESSION + .arrow() + .from_arrow_array(Arc::clone(&arrow_struct), true) + .unwrap(); let vortex_struct = SESSION .arrow() .execute_arrow(vortex_struct, None, &mut ctx) .unwrap(); - assert_eq!(&arrow_struct, vortex_struct.as_struct()); + assert_eq!(&arrow_struct, &vortex_struct); } #[test] @@ -145,24 +145,24 @@ fn roundtrip_list() { Some("Mikhail"), ])); - let arrow_list = ArrowListArray::new( + let arrow_list = Arc::new(ArrowListArray::new( Arc::new(Field::new_list_field(DataType::Utf8, true)), OffsetBuffer::from_lengths(vec![0, 2, 1]), names, None, - ); + )) as ArrowArrayRef; let list_data_type = arrow_list.data_type(); let list_field = Field::new(String::new(), list_data_type.clone(), true); - let vortex_list = ArrayRef::from_arrow(&arrow_list, true).unwrap(); + let vortex_list = SESSION + .arrow() + .from_arrow_array(Arc::clone(&arrow_list), true) + .unwrap(); let rt_arrow_list = SESSION .arrow() .execute_arrow(vortex_list, Some(&list_field), &mut ctx) .unwrap(); - assert_eq!( - (Arc::new(arrow_list.clone()) as ArrowArrayRef).as_ref(), - rt_arrow_list.as_ref() - ); + assert_eq!(&arrow_list, &rt_arrow_list); } diff --git a/vortex-bench/Cargo.toml b/vortex-bench/Cargo.toml index f228e1571d3..a63f954a637 100644 --- a/vortex-bench/Cargo.toml +++ b/vortex-bench/Cargo.toml @@ -13,6 +13,10 @@ repository = { workspace = true } rust-version = { workspace = true } version = { workspace = true } +[[bin]] +name = "data-gen" +test = false + [lints] workspace = true @@ -24,14 +28,16 @@ vortex = { workspace = true, features = [ "zstd", ] } vortex-arrow = { workspace = true } -vortex-geo = { workspace = true } +vortex-spatial = { workspace = true } vortex-tensor = { workspace = true } # TODO(connor): In the future, this might be inside vortex. anyhow = { workspace = true } arrow-array = { workspace = true } +arrow-ipc = { workspace = true } arrow-schema = { workspace = true } arrow-select = { workspace = true } async-trait = { workspace = true } +bytes = { workspace = true } bzip2 = { workspace = true } clap = { workspace = true, features = ["derive"] } futures = { workspace = true } @@ -54,12 +60,14 @@ regex = { workspace = true } reqwest = { workspace = true, features = ["stream"] } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } +sha2 = { workspace = true } spatialbench = { workspace = true } spatialbench-arrow = { workspace = true } spatialbench-parquet = { workspace = true } sysinfo = { workspace = true } tabled = { workspace = true, features = ["std"] } target-lexicon = { workspace = true } +tempfile = { workspace = true } tokio = { workspace = true, features = ["full"] } tokio-stream = { workspace = true } tokio-util = { workspace = true } @@ -78,6 +86,4 @@ wkb = { workspace = true } [dev-dependencies] insta = { workspace = true } - -[features] -unstable_encodings = ["vortex/unstable_encodings"] +rstest = { workspace = true } diff --git a/vortex-bench/REUSE.toml b/vortex-bench/REUSE.toml index 5fb9bd33466..a39e10c12b5 100644 --- a/vortex-bench/REUSE.toml +++ b/vortex-bench/REUSE.toml @@ -10,6 +10,12 @@ path = "**/*.csv" SPDX-FileCopyrightText = "Copyright the Vortex contributors" SPDX-License-Identifier = "Apache-2.0" +# Benchmark suite explainer docs linked from CI PR comments. +[[annotations]] +path = "sql/**/*.md" +SPDX-FileCopyrightText = "Copyright the Vortex contributors" +SPDX-License-Identifier = "CC-BY-4.0" + # `insta` snapshot files do not allow leading comment lines. [[annotations]] path = "src/snapshots/**.snap" diff --git a/vortex-bench/sql/appian/README.md b/vortex-bench/sql/appian/README.md new file mode 100644 index 00000000000..9304bcee2a0 --- /dev/null +++ b/vortex-bench/sql/appian/README.md @@ -0,0 +1,22 @@ +# Appian benchmark + +Mirrors the queries from DuckDB's in-tree +[`appian_benchmarks`](https://github.com/duckdb/duckdb/tree/main/benchmark/appian_benchmarks) +suite — a real-world business-application workload (customers, orders, addresses) with +selective filters and joins over camelCase-named columns. + +The eight queries live in this directory (`q1.sql` ... `q8.sql`). Upstream ships the data +as a single `.duckdb` blob (~593 MB); the harness downloads it once and projects each +table into Parquet with lowercased column names so that DataFusion and DuckDB resolve the +verbatim queries identically — see [`src/appian`](../../src/appian) for the details. + +## CI variant + +CI runs this suite from local NVMe as the `Appian on NVME` PR comment, comparing +DataFusion and DuckDB over Parquet and Vortex files (plus a native DuckDB baseline). + +## Running locally + +```bash +vx-bench run appian --engine datafusion,duckdb --format parquet,vortex +``` diff --git a/vortex-bench/sql/clickbench.md b/vortex-bench/sql/clickbench.md new file mode 100644 index 00000000000..eb3db810649 --- /dev/null +++ b/vortex-bench/sql/clickbench.md @@ -0,0 +1,30 @@ +# ClickBench benchmark + +[ClickBench](https://github.com/ClickHouse/ClickBench) is ClickHouse's web-analytics +benchmark: 43 queries over a single wide `hits` table (~100M rows of real-ish traffic +data). It is heavy on aggregations, `GROUP BY`s over high-cardinality columns, and +selective string filters, and is the main "wide table scan" workload in CI. + +The queries live in [`clickbench_queries.sql`](./clickbench_queries.sql) (one query per +line, numbered from Q0 in file order). The harness lives in +[`src/clickbench`](../src/clickbench). + +## CI variant + +CI runs this suite from local NVMe as the `Clickbench on NVME` PR comment, comparing +DataFusion and DuckDB over Parquet, Vortex, and vortex-compact files (plus a native +DuckDB baseline). + +## Sorted variant + +`Clickbench Sorted on NVME` runs the same table sorted by event time, split into 100 +shards whose filenames are shuffled so engines cannot rely on file order, exercising sort +pushdown and zone-map-style pruning on a subset of the queries. See +[`ClickBenchSortedBenchmark`](../src/clickbench/benchmark.rs). + +## Running locally + +```bash +vx-bench run clickbench --engine datafusion,duckdb --format parquet,vortex +vx-bench run clickbench-sorted --engine datafusion,duckdb --format parquet,vortex +``` diff --git a/vortex-bench/sql/fineweb.md b/vortex-bench/sql/fineweb.md new file mode 100644 index 00000000000..13421cc5456 --- /dev/null +++ b/vortex-bench/sql/fineweb.md @@ -0,0 +1,21 @@ +# FineWeb benchmark + +A string-heavy benchmark over a sample shard of the HuggingFace +[FineWeb](https://huggingface.co/datasets/HuggingFaceFW/fineweb) web-crawl dataset +(URLs, dates, and full page text). The dataset exercises dictionary and FSST string +compression heavily, and the hand-crafted queries in [`fineweb.sql`](./fineweb.sql) +(numbered from Q0 in file order) focus on string predicates: equality filters, `LIKE` +prefix and containment patterns, and aggregations over string columns. + +The harness lives in [`src/fineweb`](../src/fineweb). + +## CI variants + +CI runs this suite from local NVMe (`FineWeb NVMe`) and from S3 (`FineWeb S3`), comparing +DataFusion and DuckDB over Parquet, Vortex, and vortex-compact files. + +## Running locally + +```bash +vx-bench run fineweb --engine datafusion,duckdb --format parquet,vortex +``` diff --git a/vortex-bench/sql/gharchive.md b/vortex-bench/sql/gharchive.md new file mode 100644 index 00000000000..16cebd2d1a9 --- /dev/null +++ b/vortex-bench/sql/gharchive.md @@ -0,0 +1,14 @@ +# GitHub Archive benchmark + +A deeply nested real-world dataset from the [GitHub Archive](https://www.gharchive.org/): +GitHub event records whose `payload`, `repo`, `actor`, and `org` fields are nested structs. +The queries in [`gharchive.sql`](./gharchive.sql) (numbered from Q0 in file order) filter +and aggregate on those nested fields, making this the main struct-field-pushdown workload. + +The harness lives in [`src/realnest/gharchive.rs`](../src/realnest/gharchive.rs). + +## Running locally + +```bash +cargo run --bin datafusion-bench --profile release_debug -- --benchmark gharchive +``` diff --git a/vortex-bench/sql/polarsignals.md b/vortex-bench/sql/polarsignals.md new file mode 100644 index 00000000000..499466593e7 --- /dev/null +++ b/vortex-bench/sql/polarsignals.md @@ -0,0 +1,23 @@ +# PolarSignals Profiling benchmark + +A benchmark over continuous-profiling stacktrace data from +[Polar Signals](https://www.polarsignals.com/). The schema features a sparse struct (ten +nullable label fields covering five fill-rate tiers), deeply nested locations +(`List>>>`), and several low-cardinality string columns — +making it the main deeply-nested / sparse-data workload in CI. + +The queries in [`polarsignals.sql`](./polarsignals.sql) select stacktrace samples over +increasingly wide time ranges with equality filters on the low-cardinality metadata +columns. The harness lives in [`src/polarsignals`](../src/polarsignals). + +## CI variant + +CI runs this suite from local NVMe as the `PolarSignals Profiling` PR comment. It runs +DataFusion over Vortex only (there is no Parquet control for this suite, so its geomean is +reported without a drift-corrected verdict). + +## Running locally + +```bash +vx-bench run polarsignals --engine datafusion --format vortex +``` diff --git a/vortex-bench/sql/public-bi.md b/vortex-bench/sql/public-bi.md new file mode 100644 index 00000000000..60859677741 --- /dev/null +++ b/vortex-bench/sql/public-bi.md @@ -0,0 +1,15 @@ +# Public BI benchmark + +Real-world BI datasets from the +[Public BI benchmark](https://github.com/cwida/public_bi_benchmark) — anonymized Tableau +workbooks with messy, wide, string-heavy tables. A subset of the datasets (Arade, Bimbo, +CMSprovider, Euro2016, Food, HashTags) also feeds the +[Compression benchmark](../../benchmarks/compress-bench/README.md). + +The datasets, schemas, and queries are defined in [`src/public_bi.rs`](../src/public_bi.rs). + +## Running locally + +```bash +vx-bench run public-bi +``` diff --git a/vortex-bench/sql/spatialbench.md b/vortex-bench/sql/spatialbench.md new file mode 100644 index 00000000000..4b67405aaeb --- /dev/null +++ b/vortex-bench/sql/spatialbench.md @@ -0,0 +1,28 @@ +# SpatialBench benchmark + +The [Apache Sedona SpatialBench](https://sedona.apache.org/spatialbench/) spatial +analytics benchmark: twelve queries (Q1 ... Q12 in [`spatialbench.sql`](./spatialbench.sql), +DuckDB dialect) over a trips/zones schema, exercising spatial predicates and functions such +as `ST_DWithin`, `ST_Intersects`, and `ST_Distance`. The query logic matches upstream +`sedona-spatialbench`; only formatting differs. + +The harness lives in [`src/spatialbench`](../src/spatialbench). + +## Running locally + +```bash +vx-bench run spatialbench +``` + +The default command compares the Parquet and Vortex WKB representations with DuckDB. To run the +native Vortex spatial representation explicitly: + +```bash +vx-bench run spatialbench --engine duckdb --format vortex-spatial-native +``` + +To compare all three representations in one run: + +```bash +vx-bench run spatialbench --engine duckdb --format parquet,vortex,vortex-spatial-native +``` diff --git a/vortex-bench/sql/statpopgen.md b/vortex-bench/sql/statpopgen.md new file mode 100644 index 00000000000..c0b90a54f54 --- /dev/null +++ b/vortex-bench/sql/statpopgen.md @@ -0,0 +1,25 @@ +# Statistical and Population Genetics benchmark + +A genomics analytics benchmark using the gnomAD v3.1.2 release of the jointly called One +Thousand Genomes + Human Genome Diversity Project (1kG+HGDP) dataset — a prefix of +chromosome 21. The data source is . + +The distinguishing feature is the `GT` genotype column: a large variable-length list of +small nullable integers per row, which stresses list encodings and list-aware compute. +The queries in [`statpopgen.sql`](./statpopgen.sql) cover common genomics patterns: +allele-frequency calculations, Hardy-Weinberg equilibrium statistics, variant filtering by +frequency and by locus interval, and random access to a specific variant. + +The harness lives in [`src/statpopgen`](../src/statpopgen). + +## CI variant + +CI runs this suite from local NVMe as the `Statistical and Population Genetics` PR +comment. Only DuckDB is exercised (over Parquet, Vortex, and vortex-compact files), +because the queries rely on DuckDB list lambdas. + +## Running locally + +```bash +vx-bench run statpopgen --engine duckdb --format parquet,vortex +``` diff --git a/vortex-bench/sql/statpopgen.sql b/vortex-bench/sql/statpopgen.sql index 893cfa26f01..9e519163809 100644 --- a/vortex-bench/sql/statpopgen.sql +++ b/vortex-bench/sql/statpopgen.sql @@ -4,7 +4,7 @@ SELECT COUNT(*) FROM statpopgen; -- -- Count the number of samples (e.g. human beings) in this dataset. The GT arrays are all the same -- length, even though they're not stored as a fixed-list type. -SELECT array_length("GT", 1) FROM statpopgen LIMIT 1; +SELECT array_length("GT") FROM statpopgen LIMIT 1; -- 2. -- -- Extract just the genotypes as a dense list of lists. Usually this is sent to a linear algebra diff --git a/vortex-bench/sql/tpcds/README.md b/vortex-bench/sql/tpcds/README.md new file mode 100644 index 00000000000..3e4258ce251 --- /dev/null +++ b/vortex-bench/sql/tpcds/README.md @@ -0,0 +1,20 @@ +# TPC-DS benchmark + +The industry-standard [TPC-DS](https://www.tpc.org/tpcds/) decision-support benchmark: 99 +analytical queries (`01.sql` ... `99.sql`) over a retail snowflake schema. Compared to +TPC-H it has many more queries, a wider schema, and heavier use of joins, subqueries, and +window functions, making it a broad regression net for scan and pushdown behavior. + +The benchmark harness lives in [`src/tpcds`](../../src/tpcds). + +## CI variant + +CI runs this suite at scale factor 1 from local NVMe as the `TPC-DS SF=1 on NVME` PR +comment, comparing DataFusion and DuckDB over Parquet, Vortex, and vortex-compact files +(plus a native DuckDB baseline). + +## Running locally + +```bash +vx-bench run tpcds --engine datafusion,duckdb --format parquet,vortex --opt scale-factor=1.0 +``` diff --git a/vortex-bench/sql/tpch/README.md b/vortex-bench/sql/tpch/README.md new file mode 100644 index 00000000000..aa5a16fbab2 --- /dev/null +++ b/vortex-bench/sql/tpch/README.md @@ -0,0 +1,23 @@ +# TPC-H benchmark + +The industry-standard [TPC-H](https://www.tpc.org/tpch/) decision-support benchmark: 22 +analytical queries over an 8-table warehouse schema (`lineitem`, `orders`, `customer`, ...). +It stresses scan throughput, filter pushdown, joins, and aggregations over mostly numeric +and date columns. + +The queries in this directory (`q1.sql` ... `q22.sql`) are the standard TPC-H queries. +Data is generated deterministically by [`tpchgen`](../../src/tpch/tpchgen.rs); the benchmark +harness lives in [`src/tpch`](../../src/tpch). + +## CI variants + +CI runs this suite at scale factor 1 and 10, reading either from local NVMe or from S3 +(the data is uploaded to a bucket first), as the `TPC-H SF={1,10} on {NVME,S3}` PR +comments. Each variant compares DataFusion and DuckDB over Parquet, Vortex, and +vortex-compact files (plus in-memory Arrow and native DuckDB baselines on NVMe). + +## Running locally + +```bash +vx-bench run tpch --engine datafusion,duckdb --format parquet,vortex --opt scale-factor=1.0 +``` diff --git a/vortex-bench/sql/vortex/1_sum.sql b/vortex-bench/sql/vortex/1_sum.sql index 7640c0d3394..5d45f504612 100644 --- a/vortex-bench/sql/vortex/1_sum.sql +++ b/vortex-bench/sql/vortex/1_sum.sql @@ -1,3 +1,3 @@ --- When Footer changes land, vortex-duckdb should populate statistics from --- Footer without loading and decoding the data. +-- All zone maps are included since we don't have a filter. Vortex doesn't load +-- data but instead populates Sum() stat from zone maps SELECT sum(col) FROM test; diff --git a/vortex-bench/sql/vortex/2_sum-range.sql b/vortex-bench/sql/vortex/2_sum-range.sql new file mode 100644 index 00000000000..91c7f3b6261 --- /dev/null +++ b/vortex-bench/sql/vortex/2_sum-range.sql @@ -0,0 +1,3 @@ +-- Some zone maps are pruned, two zones are included partially, and all others +-- are included. Only two zone maps should be decoded fully. +SELECT sum(col) FROM test WHERE id >= 400000123 AND id < 2100000456; diff --git a/vortex-bench/sql/vortex/README.md b/vortex-bench/sql/vortex/README.md new file mode 100644 index 00000000000..e215394534d --- /dev/null +++ b/vortex-bench/sql/vortex/README.md @@ -0,0 +1,28 @@ +# Vortex queries benchmark + +A small suite of microbenchmark queries targeting Vortex-specific scan paths, run on +**every PR commit** (unlike the label-gated suites) with a high iteration count, so it is +the most sensitive — and most frequently seen — benchmark comment. + +[`init.sql`](./init.sql) generates a 2.5B-row three-column table (`id`, `col`, `col2`); the +numbered queries each pin down one scan behavior: + +- [`0_sum-with-filter.sql`](./0_sum-with-filter.sql): a filtered aggregation that forces a + linear scan today; once statistics are propagated to arrays it should use zone maps + instead of decoding every row. +- [`1_sum.sql`](./1_sum.sql): an unfiltered aggregation that should eventually be answered + from footer statistics without decoding data. +- [`2_sum-range.sql`](./2_sum-range.sql): a range filter on `id` that should prune most + zone maps and only fully decode the two zones that overlap the bounds. + +## CI variant + +CI runs this as the `Vortex queries` PR comment (see +[`.github/workflows/sql-vortex-pr.yml`](../../../.github/workflows/sql-vortex-pr.yml)), +comparing DataFusion and DuckDB over Parquet and Vortex files with 100 iterations. + +## Running locally + +```bash +vx-bench run vortex --engine datafusion,duckdb --format parquet,vortex +``` diff --git a/vortex-bench/sql/vortex/init.sql b/vortex-bench/sql/vortex/init.sql index 90002573995..aee587a2d2a 100644 --- a/vortex-bench/sql/vortex/init.sql +++ b/vortex-bench/sql/vortex/init.sql @@ -2,7 +2,8 @@ COPY ( SELECT - i % 1000 AS col, - (i * 2654435761) % 100000 AS col2 - FROM range(25000000) t(i) + i AS id, + (i % 1000)::INTEGER AS col, + ((i * 2654435761) % 100000)::INTEGER AS col2 + FROM range(2500000000) t(i) ) TO 'test.parquet' (FORMAT parquet); diff --git a/vortex-bench/src/appian/mod.rs b/vortex-bench/src/appian/mod.rs index cdc82ac200a..20dabb5676a 100644 --- a/vortex-bench/src/appian/mod.rs +++ b/vortex-bench/src/appian/mod.rs @@ -121,6 +121,10 @@ impl AppianBenchmark { #[async_trait::async_trait] impl Benchmark for AppianBenchmark { + fn doc_path(&self) -> &'static str { + "vortex-bench/sql/appian/README.md" + } + fn queries(&self) -> anyhow::Result> { Ok(appian_queries().collect()) } diff --git a/vortex-bench/src/benchmark.rs b/vortex-bench/src/benchmark.rs index d59708886a6..47de30a5faf 100644 --- a/vortex-bench/src/benchmark.rs +++ b/vortex-bench/src/benchmark.rs @@ -45,8 +45,8 @@ pub trait Benchmark: Send + Sync { /// This is the canonical source data that can be converted to other formats. /// This should be idempotent - safe to call multiple times. /// - /// Format-specific benchmark binaries (like lance-bench, datafusion-bench, duckdb-bench) should - /// call this method to ensure base data exists, then perform their own format conversion. + /// The `data-gen` binary calls this method before creating requested derived formats. + /// Engine-specific benchmark binaries should only run against prepared data. async fn generate_base_data(&self) -> anyhow::Result<()>; /// Prepare benchmark- and format-specific data beyond the Parquet base that @@ -64,6 +64,10 @@ pub trait Benchmark: Send + Sync { fn dataset(&self) -> BenchmarkDataset; + /// Repo-relative path of the markdown explainer for this benchmark suite, linked from the + /// title of CI benchmark PR comments. Required: every suite must ship a doc. + fn doc_path(&self) -> &'static str; + /// Get the name of the benchmark dataset fn dataset_name(&self) -> &str; diff --git a/vortex-bench/src/clickbench/benchmark.rs b/vortex-bench/src/clickbench/benchmark.rs index 346496dbda0..fcba1b9f9c1 100644 --- a/vortex-bench/src/clickbench/benchmark.rs +++ b/vortex-bench/src/clickbench/benchmark.rs @@ -76,6 +76,10 @@ fn read_clickbench_queries(queries_file: Option<&str>) -> Result &'static str { + "vortex-bench/sql/clickbench.md" + } + fn queries(&self) -> Result> { read_clickbench_queries(self.queries_file.as_deref()) } @@ -120,6 +124,10 @@ impl Benchmark for ClickBenchBenchmark { #[async_trait::async_trait] impl Benchmark for ClickBenchSortedBenchmark { + fn doc_path(&self) -> &'static str { + "vortex-bench/sql/clickbench.md#sorted-variant" + } + fn queries(&self) -> Result> { Ok(read_clickbench_queries(self.queries_file.as_deref())? .into_iter() diff --git a/vortex-bench/src/compress/mod.rs b/vortex-bench/src/compress/mod.rs index 6e261b871d2..d98bd6a795d 100644 --- a/vortex-bench/src/compress/mod.rs +++ b/vortex-bench/src/compress/mod.rs @@ -1,15 +1,28 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::Any; use std::borrow::Cow; use std::fmt; +use std::fs::File; use std::path::Path; +use std::sync::Arc; use std::time::Duration; +use anyhow::Context; use anyhow::Result; +use anyhow::bail; +use arrow_array::RecordBatch; +use arrow_schema::Schema; use async_trait::async_trait; +use bytes::Bytes; use clap::ValueEnum; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use serde::Serialize; +use tempfile::NamedTempFile; +use tempfile::TempDir; +use vortex::array::ArrayRef; +use vortex::expr::stats::Stat; use vortex::utils::aliases::hash_map::HashMap; use crate::Format; @@ -84,10 +97,104 @@ impl fmt::Display for CompressOp { } } +/// Source data in the form a format compresses, read from Parquet once per format. +pub enum Uncompressed { + /// Arrow record batches sharing one schema. + Arrow { + schema: Arc, + batches: Vec, + }, + /// A Vortex array. + Vortex(ArrayRef), + /// Format-private input, for backends whose Arrow crate version differs from the + /// workspace's and so cannot share [`Self::Arrow`]. + Opaque(Box), +} + +impl Uncompressed { + /// Read a Parquet file into Arrow record batches. + pub fn read_arrow(parquet_path: &Path) -> Result { + let builder = ParquetRecordBatchReaderBuilder::try_new(File::open(parquet_path)?)?; + let schema = Arc::clone(builder.schema()); + let batches = builder.build()?.collect::, _>>()?; + + Ok(Self::Arrow { schema, batches }) + } + + /// The Arrow schema and batches, or an error if this is not Arrow data. + pub fn arrow(&self) -> Result<(&Arc, &[RecordBatch])> { + match self { + Self::Arrow { schema, batches } => Ok((schema, batches)), + _ => bail!("expected Arrow record batches"), + } + } + + /// The Vortex array, or an error if this is not Vortex data. + pub fn vortex(&self) -> Result<&ArrayRef> { + match self { + Self::Vortex(array) => Ok(array), + _ => bail!("expected a Vortex array"), + } + } + + /// The format-private input as `T`, or an error if this holds something else. + pub fn opaque(&self) -> Result<&T> { + match self { + Self::Opaque(input) => input + .downcast_ref() + .ok_or_else(|| anyhow::anyhow!("format-private input has an unexpected type")), + _ => bail!("expected format-private input"), + } + } + + /// Return the input to the state [`Self::read_arrow`] or a fresh conversion would produce. + /// + /// The Vortex writer computes statistics inside the timed region and caches them on the + /// array, so reusing one array across iterations would let every run after the first skip + /// that work. Clearing the cache keeps each iteration's measurement comparable. + pub fn reset(&self) { + if let Self::Vortex(array) = self { + clear_stats(array); + } + } +} + +/// Clear cached statistics on `array` and every array beneath it. +fn clear_stats(array: &ArrayRef) { + for stat in Stat::all() { + array.statistics().clear(stat); + } + for child in array.children_iter() { + clear_stats(child); + } +} + +/// Where a format keeps its compressed output. +/// +/// Temporary files and directories are removed on drop, so the output lives exactly as long as +/// the value does. +pub enum CompressedData { + /// In-memory file image. + Bytes(Bytes), + /// Single file on disk. + File(NamedTempFile), + /// Directory of files on disk. + Dir(TempDir), +} + +/// Output of one compression run: the compressed data, its size and the compression time. +pub struct Compressed { + pub data: CompressedData, + pub size: u64, + pub elapsed: Duration, +} + /// Result of a compression benchmark run. pub struct CompressResult { pub time: Duration, pub compressed_size: u64, + /// Output of the last iteration, kept so decompression need not compress again. + pub compressed: Compressed, pub timing: CompressionTimingMeasurement, /// Per-iteration encode wall times. Captured for v3 emission. pub all_runs: Vec, @@ -108,51 +215,57 @@ pub struct DecompressResult { /// (e.g., Vortex, Parquet, Lance). The benchmark functions use this trait /// to run timing measurements. /// -/// The input data is provided as a path to a Parquet file, which implementations -/// read and convert as needed for their target format. +/// The input data is provided as a path to a Parquet file. [`Self::load`] reads it once into +/// the form the format compresses; the timed operations never touch the file again. #[async_trait] pub trait Compressor: Send + Sync { /// The format this compressor handles. fn format(&self) -> Format; - /// Compress data from a Parquet file, returning the compressed size in bytes and elapsed time. + /// Read a Parquet file into the form [`Self::compress`] consumes. Not timed. + async fn load(&self, parquet_path: &Path) -> Result; + + /// Compress output previously produced by [`Self::load`]. /// - /// The implementation should read the Parquet file and compress it - /// to the target format. - async fn compress(&self, parquet_path: &Path) -> Result<(u64, Duration)>; + /// Only the compression itself should be timed. + async fn compress(&self, input: &Uncompressed) -> Result; - /// Decompress data from the Parquet file (after compressing), returning the decompressed size. + /// Decompress output previously produced by [`Self::compress`]. /// - /// This method first compresses the data to the target format, then decompresses it. /// The timing returned should only measure the decompression phase. /// /// Format implementations apply the fixed wide-table read projection when the input schema /// matches the projection benchmark. - async fn decompress(&self, parquet_path: &Path) -> Result; + async fn decompress(&self, compressed: &Compressed) -> Result; } /// Run a compression benchmark for the given compressor. /// -/// Executes compression `iterations` times and returns timing statistics. +/// Compresses the same `input` `iterations` times and returns timing statistics. The input is +/// [reset](Uncompressed::reset) before every iteration so none of them starts warm. pub async fn benchmark_compress( compressor: &dyn Compressor, - parquet_path: &Path, + input: &Uncompressed, iterations: usize, bench_name: &str, ) -> Result { let format = compressor.format(); let mut fastest = Duration::MAX; - let mut compressed_size = 0u64; let mut all_runs = Vec::with_capacity(iterations); + let mut compressed = None; for _ in 0..iterations { - let (size, elapsed) = compressor.compress(parquet_path).await?; + input.reset(); + let result = compressor.compress(input).await?; - compressed_size = size; - fastest = fastest.min(elapsed); - all_runs.push(elapsed); + fastest = fastest.min(result.elapsed); + all_runs.push(result.elapsed); + compressed = Some(result); } + let compressed = compressed.context("--iterations must be at least 1")?; + let compressed_size = compressed.size; + let ratios = vec![CustomUnitMeasurement { name: format!("{} size/{bench_name}", format.name()), format, @@ -169,6 +282,7 @@ pub async fn benchmark_compress( Ok(CompressResult { time: fastest, compressed_size, + compressed, timing, all_runs, ratios, @@ -177,10 +291,10 @@ pub async fn benchmark_compress( /// Run a decompression benchmark for the given compressor. /// -/// Benchmarks decompression `iterations` times. +/// Decompresses the same `compressed` output `iterations` times. pub async fn benchmark_decompress( compressor: &dyn Compressor, - parquet_path: &Path, + compressed: &Compressed, iterations: usize, bench_name: &str, ) -> Result { @@ -189,7 +303,7 @@ pub async fn benchmark_decompress( let mut all_runs = Vec::with_capacity(iterations); for _ in 0..iterations { - let elapsed = compressor.decompress(parquet_path).await?; + let elapsed = compressor.decompress(compressed).await?; fastest = fastest.min(elapsed); all_runs.push(elapsed); diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index 92064c25785..6b4ed871f2e 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -6,6 +6,8 @@ use std::path::Path; use std::path::PathBuf; use std::sync::Arc; +use arrow_array::RecordBatch; +use arrow_select::concat::concat_batches; use futures::StreamExt; use futures::TryStreamExt; use parquet::arrow::AsyncArrowWriter; @@ -32,7 +34,7 @@ use vortex::array::arrays::Struct; use vortex::array::arrays::StructArray; use vortex::array::arrays::VarBinViewArray; use vortex::array::arrays::struct_::StructArrayExt; -use vortex::array::builders::builder_with_capacity; +use vortex::array::builders::builder_with_capacity_in; use vortex::array::stream::ArrayStreamAdapter; use vortex::array::stream::ArrayStreamExt; use vortex::compressor::BtrBlocksCompressorBuilder; @@ -52,10 +54,10 @@ use vortex::layout::layouts::compressed::CompressingStrategy; use vortex::layout::layouts::flat::writer::FlatLayoutStrategy; use vortex::session::VortexSession; use vortex::utils::aliases::hash_set::HashSet; -use vortex_arrow::FromArrowArray; -use vortex_arrow::FromArrowType; -use vortex_geo::extension::GeoMetadata; -use vortex_geo::extension::WellKnownBinary; +use vortex::utils::parallelism::get_available_parallelism; +use vortex_arrow::ArrowSessionExt; +use vortex_spatial::extension::SpatialMetadata; +use vortex_spatial::extension::WellKnownBinary; use wkb::Endianness; use wkb::reader::read_wkb; use wkb::writer::WriteOptions; @@ -64,6 +66,7 @@ use wkb::writer::write_geometry; use crate::CompactionStrategy; use crate::Format; use crate::SESSION; +use crate::retain_edition_encodings; use crate::utils::file::idempotent_async; /// Memory budget per concurrent conversion stream in GB. This is somewhat arbitary. @@ -72,9 +75,6 @@ const MEMORY_PER_STREAM_GB: u64 = 4; /// Minimum number of concurrent conversion streams. const MIN_CONCURRENCY: u64 = 1; -/// Maximum number of concurrent conversion streams. This is somewhat arbitary. -const MAX_CONCURRENCY: u64 = 16; - /// Returns the available system memory in bytes. fn available_memory_bytes() -> u64 { System::new_all().available_memory() @@ -83,14 +83,15 @@ fn available_memory_bytes() -> u64 { /// Calculate appropriate concurrency based on available memory. fn calculate_concurrency() -> usize { let available_gb = available_memory_bytes() / (1024 * 1024 * 1024); - let concurrency = (available_gb / MEMORY_PER_STREAM_GB).clamp(MIN_CONCURRENCY, MAX_CONCURRENCY); + let max_concurrency = get_available_parallelism().unwrap_or(1) as u64; + let concurrency = (available_gb / MEMORY_PER_STREAM_GB).clamp(MIN_CONCURRENCY, max_concurrency); info!( "Available memory: {}GB, maximum concurrency is: {}", available_gb, concurrency ); - concurrency as usize + usize::try_from(concurrency).unwrap_or(1) } /// Read a Parquet file and return it as a Vortex [`ChunkedArray`]. @@ -98,18 +99,71 @@ fn calculate_concurrency() -> usize { /// Note: This loads the entire file into memory. For large files, use the streaming conversion like /// in [`parquet_to_vortex_stream`] instead. pub async fn parquet_to_vortex_chunks(parquet_path: PathBuf) -> anyhow::Result { + parquet_to_vortex_chunks_with_batch_size(parquet_path, None).await +} + +/// Read a Parquet file as a Vortex [`ChunkedArray`] with chunks of exactly `batch_size` rows. +/// +/// With `batch_size` set, the source batches are concatenated and re-sliced on exact boundaries, +/// so every chunk but the last has the requested length. Setting the Arrow reader's batch size +/// is not enough on its own: the reader also breaks at the source file's row group boundaries, +/// so a file whose row groups are not a multiple of the batch size still yields short batches. +/// +/// This matters when comparing against a format whose physical partitioning is explicit. Chunk +/// size becomes the Vortex file's partition size, and small chunks mean many small compressed +/// blocks — and, on the GPU, many small kernel launches. +/// +/// `None` keeps whatever batches the Parquet reader produces. +pub async fn parquet_to_vortex_chunks_with_batch_size( + parquet_path: PathBuf, + batch_size: Option, +) -> anyhow::Result { let file = File::open(parquet_path).await?; let builder = ParquetRecordBatchStreamBuilder::new(file).await?; - let reader = builder.build()?; - let chunks: Vec = parquet_to_vortex_stream(reader) - .map(|r| r.map_err(anyhow::Error::from)) + let Some(batch_size) = batch_size.filter(|size| *size > 0) else { + let chunks: Vec = parquet_to_vortex_stream(builder.build()?) + .map(|r| r.map_err(anyhow::Error::from)) + .try_collect() + .await?; + return Ok(ChunkedArray::from_iter(chunks)); + }; + + let batches: Vec = builder + .with_batch_size(batch_size) + .build()? + .map_err(anyhow::Error::from) .try_collect() .await?; + let schema = batches + .first() + .map(RecordBatch::schema) + .ok_or_else(|| anyhow::anyhow!("cannot convert an empty Parquet file"))?; + let combined = concat_batches(&schema, &batches)?; + + let mut chunks = Vec::with_capacity(combined.num_rows().div_ceil(batch_size)); + for start in (0..combined.num_rows()).step_by(batch_size) { + let len = batch_size.min(combined.num_rows() - start); + chunks.push(record_batch_to_vortex(combined.slice(start, len))?); + } + Ok(ChunkedArray::from_iter(chunks)) } +/// Convert one Arrow [`RecordBatch`] into a canonical Vortex array. +fn record_batch_to_vortex(batch: RecordBatch) -> VortexResult { + let schema = batch.schema(); + let chunk = SESSION.arrow().from_arrow_record_batch(batch, &schema)?; + let mut ctx = VortexSession::default().create_execution_ctx(); + let mut builder = builder_with_capacity_in(chunk.dtype(), chunk.len(), ctx.allocator()); + + // Canonicalize the chunk. + chunk.append_to_builder(builder.as_mut(), &mut ctx)?; + + Ok(builder.finish()) +} + /// Create a streaming Vortex array from a Parquet reader. /// /// Streams record batches and converts them to Vortex arrays on-the-fly, avoiding loading the @@ -118,18 +172,9 @@ pub fn parquet_to_vortex_stream( reader: ParquetRecordBatchStream, ) -> impl futures::Stream> { reader.map(move |result| { - result.map_err(|e| vortex_err!(External: e)).and_then(|rb| { - let chunk = ArrayRef::from_arrow(rb, false)?; - let mut builder = builder_with_capacity(chunk.dtype(), chunk.len()); - - // Canonicalize the chunk. - chunk.append_to_builder( - builder.as_mut(), - &mut VortexSession::default().create_execution_ctx(), - )?; - - Ok(builder.finish()) - }) + result + .map_err(|e| vortex_err!(External: e)) + .and_then(record_batch_to_vortex) }) } @@ -145,10 +190,15 @@ pub async fn convert_parquet_file_to_vortex( let builder = ParquetRecordBatchStreamBuilder::new(file).await?; // GeoParquet geometry tagging. - let geo_columns = geoparquet_columns(builder.metadata()); - let dtype = tag_geo_dtype(DType::from_arrow(builder.schema().as_ref()), &geo_columns)?; + let spatial_columns = geoparquet_columns(builder.metadata()); + let dtype = tag_spatial_dtype( + SESSION + .arrow() + .from_arrow_schema(builder.schema().as_ref())?, + &spatial_columns, + )?; let stream = parquet_to_vortex_stream(builder.build()?) - .map(move |chunk| chunk.and_then(|chunk| tag_geo_array(chunk, &geo_columns))); + .map(move |chunk| chunk.and_then(|chunk| tag_spatial_array(chunk, &spatial_columns))); let mut output_file = OpenOptions::new() .write(true) @@ -198,8 +248,10 @@ fn write_options_for( let mut builder = WriteStrategyBuilder::default(); if matches!(compaction, CompactionStrategy::Compact) { - builder = - builder.with_btrblocks_builder(BtrBlocksCompressorBuilder::default().with_compact()); + builder = builder.with_btrblocks_builder(retain_edition_encodings( + &SESSION, + BtrBlocksCompressorBuilder::default().with_compact(), + )); } for name in binary_fields { builder = builder.with_field_writer(FieldPath::from_name(name), no_dict_layout()); @@ -211,7 +263,7 @@ fn write_options_for( fn no_dict_layout() -> Arc { Arc::new(CompressingStrategy::new( ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), - BtrBlocksCompressorBuilder::default().build(), + retain_edition_encodings(&SESSION, BtrBlocksCompressorBuilder::default()).build(), )) } @@ -302,7 +354,10 @@ pub async fn write_parquet_as_vortex( } /// Add GeoParquet `geo` file metadata to externally-sourced parquet we don't generate (e.g. SpatialBench `zone`). -pub async fn add_geoparquet_metadata(parquet_path: &Path, geo_json: &str) -> anyhow::Result<()> { +pub async fn add_geoparquet_metadata( + parquet_path: &Path, + geoparquet_json: &str, +) -> anyhow::Result<()> { let builder = ParquetRecordBatchStreamBuilder::new(File::open(parquet_path).await?).await?; let already_tagged = builder .metadata() @@ -321,7 +376,10 @@ pub async fn add_geoparquet_metadata(parquet_path: &Path, geo_json: &str) -> any while let Some(batch) = reader.try_next().await? { writer.write(&batch).await?; } - writer.append_key_value_metadata(KeyValue::new("geo".to_string(), Some(geo_json.to_string()))); + writer.append_key_value_metadata(KeyValue::new( + "geo".to_string(), + Some(geoparquet_json.to_string()), + )); writer.close().await?; tokio::fs::rename(&tmp_path, parquet_path).await?; Ok(()) @@ -334,7 +392,7 @@ fn geoparquet_columns(metadata: &ParquetMetaData) -> HashSet { .key_value_metadata() .and_then(|kvs| kvs.iter().find(|kv| kv.key == "geo")) .and_then(|kv| kv.value.as_deref()) - .and_then(|geo| serde_json::from_str::(geo).ok()) + .and_then(|metadata| serde_json::from_str::(metadata).ok()) .and_then(|value| { value .get("columns") @@ -344,15 +402,18 @@ fn geoparquet_columns(metadata: &ParquetMetaData) -> HashSet { .unwrap_or_default() } -/// The erased `vortex.geo.wkb` extension dtype over a binary `storage` dtype. +/// The erased `vortex.st.wkb` extension dtype over a binary `storage` dtype. fn wkb_ext_dtype(storage: &DType) -> VortexResult { - Ok(ExtDType::::try_new(GeoMetadata { crs: None }, storage.clone())?.erased()) + Ok( + ExtDType::::try_new(SpatialMetadata { crs: None }, storage.clone())? + .erased(), + ) } -/// Re-type the named binary columns of a struct `dtype` as `vortex.geo.wkb`, so the column +/// Re-type the named binary columns of a struct `dtype` as `vortex.st.wkb`, so the column /// self-describes as geometry. -fn tag_geo_dtype(dtype: DType, geo_columns: &HashSet) -> VortexResult { - if geo_columns.is_empty() { +fn tag_spatial_dtype(dtype: DType, spatial_columns: &HashSet) -> VortexResult { + if spatial_columns.is_empty() { return Ok(dtype); } let DType::Struct(fields, nullability) = &dtype else { @@ -363,7 +424,7 @@ fn tag_geo_dtype(dtype: DType, geo_columns: &HashSet) -> VortexResult

) -> VortexResult
) -> VortexResult { - if geo_columns.is_empty() { +fn tag_spatial_array(chunk: ArrayRef, spatial_columns: &HashSet) -> VortexResult { + if spatial_columns.is_empty() { return Ok(chunk); } let Some(struct_array) = chunk.as_opt::() else { @@ -391,7 +452,7 @@ fn tag_geo_array(chunk: ArrayRef, geo_columns: &HashSet) -> VortexResult let mut ctx = SESSION.create_execution_ctx(); let mut tagged = Vec::with_capacity(names.len()); for (name, field) in names.iter().zip(struct_array.iter_unmasked_fields()) { - if geo_columns.contains(name.as_ref()) && field.dtype().is_binary() { + if spatial_columns.contains(name.as_ref()) && field.dtype().is_binary() { let ext = wkb_ext_dtype(field.dtype())?; let little_endian = wkb_field_to_little_endian(field, &mut ctx)?; tagged.push(ExtensionArray::try_new(ext, little_endian)?.into_array()); diff --git a/vortex-bench/src/datasets/feature_vectors.rs b/vortex-bench/src/datasets/feature_vectors.rs index e09f29f49ec..173989d6e83 100644 --- a/vortex-bench/src/datasets/feature_vectors.rs +++ b/vortex-bench/src/datasets/feature_vectors.rs @@ -25,6 +25,7 @@ use crate::conversions::write_parquet_as_vortex; use crate::idempotent_async; use crate::random_access::BenchDataset; use crate::random_access::data_path; +use crate::random_access::parquet_to_arrow_file; /// Dataset identifier used for data path generation. pub const DATASET: &str = "feature_vectors"; @@ -46,6 +47,10 @@ impl BenchDataset for FeatureVectorsData { async fn path(&self, format: Format) -> Result { match format { + Format::ArrowIpc => { + let parquet_path = feature_vectors_parquet().await?; + parquet_to_arrow_file(parquet_path, data_path(DATASET, Format::ArrowIpc)) + } Format::OnDiskVortex => feature_vectors_vortex().await, Format::VortexCompact => feature_vectors_vortex_compact().await, Format::Parquet => feature_vectors_parquet().await, diff --git a/vortex-bench/src/datasets/nested_lists.rs b/vortex-bench/src/datasets/nested_lists.rs index 3c89ed8eb2b..92e8ae3c3c3 100644 --- a/vortex-bench/src/datasets/nested_lists.rs +++ b/vortex-bench/src/datasets/nested_lists.rs @@ -24,6 +24,7 @@ use crate::conversions::write_parquet_as_vortex; use crate::idempotent_async; use crate::random_access::BenchDataset; use crate::random_access::data_path; +use crate::random_access::parquet_to_arrow_file; /// Dataset identifier used for data path generation. pub const DATASET: &str = "nested_lists"; @@ -45,6 +46,10 @@ impl BenchDataset for NestedListsData { async fn path(&self, format: Format) -> Result { match format { + Format::ArrowIpc => { + let parquet_path = nested_lists_parquet().await?; + parquet_to_arrow_file(parquet_path, data_path(DATASET, Format::ArrowIpc)) + } Format::OnDiskVortex => nested_lists_vortex().await, Format::VortexCompact => nested_lists_vortex_compact().await, Format::Parquet => nested_lists_parquet().await, diff --git a/vortex-bench/src/datasets/nested_structs.rs b/vortex-bench/src/datasets/nested_structs.rs index c4a522184d3..4f0903661c7 100644 --- a/vortex-bench/src/datasets/nested_structs.rs +++ b/vortex-bench/src/datasets/nested_structs.rs @@ -25,6 +25,7 @@ use crate::conversions::write_parquet_as_vortex; use crate::idempotent_async; use crate::random_access::BenchDataset; use crate::random_access::data_path; +use crate::random_access::parquet_to_arrow_file; /// Dataset identifier used for data path generation. pub const DATASET: &str = "nested_structs"; @@ -46,6 +47,10 @@ impl BenchDataset for NestedStructsData { async fn path(&self, format: Format) -> Result { match format { + Format::ArrowIpc => { + let parquet_path = nested_structs_parquet().await?; + parquet_to_arrow_file(parquet_path, data_path(DATASET, Format::ArrowIpc)) + } Format::OnDiskVortex => nested_structs_vortex().await, Format::VortexCompact => nested_structs_vortex_compact().await, Format::Parquet => nested_structs_parquet().await, diff --git a/vortex-bench/src/datasets/struct_list_of_ints.rs b/vortex-bench/src/datasets/struct_list_of_ints.rs index 164ad4534fc..a315712e74a 100644 --- a/vortex-bench/src/datasets/struct_list_of_ints.rs +++ b/vortex-bench/src/datasets/struct_list_of_ints.rs @@ -27,7 +27,6 @@ use vortex::array::arrays::listview::recursive_list_from_list_view; use vortex::array::validity::Validity; use vortex::dtype::FieldNames; use vortex_arrow::ArrowSessionExt; -use vortex_arrow::ToArrowType; use crate::IdempotentPath; use crate::SESSION; @@ -147,7 +146,7 @@ impl Dataset for StructListOfInts { for chunk in chunked.iter_chunks() { let converted = recursive_list_from_list_view(chunk.clone(), &mut ctx)?; - let schema = converted.dtype().to_arrow_schema()?; + let schema = SESSION.arrow().to_arrow_schema(converted.dtype())?; let schema = Field::new_struct("", schema.fields, false); let batch = SESSION .arrow() diff --git a/vortex-bench/src/datasets/taxi_data.rs b/vortex-bench/src/datasets/taxi_data.rs index 919a0fa4dff..e592d84aec9 100644 --- a/vortex-bench/src/datasets/taxi_data.rs +++ b/vortex-bench/src/datasets/taxi_data.rs @@ -23,6 +23,8 @@ use crate::datasets::Dataset; use crate::datasets::data_downloads::download_data; use crate::idempotent_async; use crate::random_access::BenchDataset; +use crate::random_access::data_path; +use crate::random_access::parquet_to_arrow_file; /// Dataset identifier used for data path generation. pub const DATASET: &str = "taxi"; @@ -59,6 +61,10 @@ impl BenchDataset for TaxiData { async fn path(&self, format: Format) -> Result { match format { + Format::ArrowIpc => { + let parquet_path = taxi_data_parquet().await?; + parquet_to_arrow_file(parquet_path, data_path(DATASET, Format::ArrowIpc)) + } Format::OnDiskVortex => taxi_data_vortex().await, Format::VortexCompact => taxi_data_vortex_compact().await, Format::Parquet => taxi_data_parquet().await, diff --git a/vortex-bench/src/datasets/tpch_l_comment.rs b/vortex-bench/src/datasets/tpch_l_comment.rs index ebe6ff2e96a..c57bc91a65d 100644 --- a/vortex-bench/src/datasets/tpch_l_comment.rs +++ b/vortex-bench/src/datasets/tpch_l_comment.rs @@ -6,7 +6,6 @@ use std::path::PathBuf; use anyhow::Result; use async_trait::async_trait; use futures::TryStreamExt; -use glob::glob; use vortex::array::ArrayRef; use vortex::array::Canonical; use vortex::array::ExecutionCtx; @@ -39,12 +38,10 @@ async fn ensure_tpch_parquet() -> Result { generate_tpch_tables(options).await?; } - // Return the first lineitem parquet file - let pattern = data_dir.join("lineitem_*.parquet"); - glob(pattern.to_string_lossy().as_ref())? - .next() - .ok_or_else(|| anyhow::anyhow!("No lineitem parquet files found"))? - .map_err(Into::into) + let path = data_dir.join("lineitem.parquet"); + path.exists() + .then_some(path) + .ok_or_else(|| anyhow::anyhow!("No lineitem parquet file found")) } #[async_trait] @@ -58,7 +55,7 @@ impl Dataset for TPCHLCommentChunked { let scale_factor_dir = base_path.join("1.0"); let data_dir = scale_factor_dir.join(Format::OnDiskVortex.name()); - // Generate TPC-H CSV data if it doesn't exist + // Generate TPC-H Vortex data if it doesn't exist if !data_dir.exists() { // Use blocking call like TPC-H benchmark does let options = TpchGenOptions::new("1.0".to_string(), scale_factor_dir) @@ -67,30 +64,25 @@ impl Dataset for TPCHLCommentChunked { futures::executor::block_on(generate_tpch_tables(options))?; } - let mut chunks: Vec = vec![]; - for path in glob( - data_dir - .join("lineitem_*.vortex") - .to_string_lossy() - .as_ref(), - )? { - let file = SESSION.open_options().open_path(path?).await?; - let file_chunks: Vec<_> = file - .scan()? - .with_projection(pack(vec![("l_comment", col("l_comment"))], NonNullable)) - .map({ - let ctx = ctx.clone(); - move |a| { - let mut ctx = ctx.clone(); - let canonical = a.execute::(&mut ctx)?; - Ok(canonical.into_array()) - } - }) - .into_array_stream()? - .try_collect() - .await?; - chunks.extend(file_chunks); - } + let path = data_dir.join("lineitem.vortex"); + let file = SESSION.open_options().open_path(path).await?; + let projection = pack(vec![("l_comment", col("l_comment"))], NonNullable) + .optimize_recursive(file.dtype())? + .bind(file.dtype())?; + let chunks: Vec<_> = file + .scan()? + .with_projection(projection) + .map({ + let ctx = ctx.clone(); + move |a| { + let mut ctx = ctx.clone(); + let canonical = a.execute::(&mut ctx)?; + Ok(canonical.into_array()) + } + }) + .into_array_stream()? + .try_collect() + .await?; Ok(ChunkedArray::from_iter(chunks).into_array()) } diff --git a/vortex-bench/src/display.rs b/vortex-bench/src/display.rs index a8e903df25a..97203abdb4d 100644 --- a/vortex-bench/src/display.rs +++ b/vortex-bench/src/display.rs @@ -30,6 +30,16 @@ pub fn render_table( all_measurements: Vec, targets: &[Target], ) -> anyhow::Result<()> { + // `all_measurements` is empty for decompress-only runs such as `compress-bench + // --gpu-decompress`: every ratio compares vortex against parquet or lance, neither of + // which is benchmarked there, so no ratio is recorded even though a baseline target is + // still passed. `targets` is instead empty when the caller has no baseline format to + // report against. Either way there is no table to draw, and the indexing below would + // panic looking up a baseline that has no measurements. + if all_measurements.is_empty() || targets.is_empty() { + return Ok(()); + } + let mut measurements: HashMap> = HashMap::with_capacity(all_measurements.len().div_ceil(targets.len())); @@ -105,9 +115,14 @@ pub fn render_table( pub fn print_measurements_json( writer: &mut dyn Write, all_measurements: Vec, + doc: &str, ) -> anyhow::Result<()> { for measurement in all_measurements { - writeln!(writer, "{}", measurement.to_json())?; + let mut json = measurement.to_json(); + if let Some(obj) = json.as_object_mut() { + obj.insert("doc".to_string(), doc.into()); + } + writeln!(writer, "{json}")?; } Ok(()) @@ -122,3 +137,30 @@ fn color(baseline: MeasurementValue, value: MeasurementValue) -> Color { Color::BG_BRIGHT_GREEN | Color::FG_BLACK } } + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + use crate::Engine; + use crate::Format; + use crate::measurements::CompressionTimingMeasurement; + + /// Decompress-only runs (`compress-bench --gpu-decompress`) collect no compression ratios, + /// and callers pass an empty target list when the baseline format is absent. Both used to + /// panic: the first on the missing baseline key, the second on a divide-by-zero capacity. + #[rstest] + #[case::no_measurements(&[Target::new(Engine::Vortex, Format::OnDiskVortex)])] + #[case::no_targets(&[])] + fn render_table_without_a_baseline_is_a_noop(#[case] targets: &[Target]) -> anyhow::Result<()> { + let mut out = Vec::new(); + render_table( + &mut out, + Vec::::new(), + targets, + )?; + assert!(out.is_empty()); + Ok(()) + } +} diff --git a/vortex-bench/src/fineweb/mod.rs b/vortex-bench/src/fineweb/mod.rs index b8c86925905..aefa796af3a 100644 --- a/vortex-bench/src/fineweb/mod.rs +++ b/vortex-bench/src/fineweb/mod.rs @@ -53,6 +53,10 @@ impl FinewebBenchmark { #[async_trait::async_trait] impl Benchmark for FinewebBenchmark { + fn doc_path(&self) -> &'static str { + "vortex-bench/sql/fineweb.md" + } + /// Some basic string-focused queries, numbered from Q0 in `sql/fineweb.sql` file order. fn queries(&self) -> anyhow::Result> { // `;`-separated; a `;` must not appear in a comment, or it would split a statement in two. diff --git a/vortex-bench/src/lib.rs b/vortex-bench/src/lib.rs index abf4e1ccea3..569dfc74a4d 100644 --- a/vortex-bench/src/lib.rs +++ b/vortex-bench/src/lib.rs @@ -70,6 +70,8 @@ pub use datasets::BenchmarkDataset; pub use output::BenchmarkOutput; pub use output::create_output_writer; use vortex::VortexSessionDefault; +use vortex::editions::ComponentKind; +use vortex::editions::EditionSessionExt; pub use vortex::error::vortex_panic; use vortex::io::session::RuntimeSessionExt; use vortex::session::VortexSession; @@ -80,7 +82,7 @@ static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; pub static SESSION: LazyLock = LazyLock::new(|| { let session = VortexSession::default().with_tokio(); - vortex_geo::initialize(&session); + vortex_spatial::initialize(&session); session }); @@ -139,10 +141,10 @@ impl Display for Target { #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, ValueEnum, Serialize, Deserialize)] #[serde(rename_all = "kebab-case")] pub enum Format { + #[clap(name = "arrow-ipc")] + ArrowIpc, #[clap(name = "csv")] Csv, - #[clap(name = "arrow")] - Arrow, #[clap(name = "parquet")] Parquet, #[clap(name = "vortex")] @@ -151,9 +153,9 @@ pub enum Format { #[clap(name = "vortex-compact")] #[serde(rename = "vortex-compact")] VortexCompact, - #[clap(name = "vortex-geo-native")] - #[serde(rename = "vortex-geo-native")] - VortexNative, + #[clap(name = "vortex-spatial-native")] + #[serde(rename = "vortex-spatial-native")] + VortexSpatialNative, #[clap(name = "duckdb")] #[serde(rename = "duckdb")] OnDiskDuckDB, @@ -169,10 +171,15 @@ impl Display for Format { } /// Allowed formats for benchmark CLI arguments. -pub const ALLOWED_FORMATS: &[Format] = &[Format::Parquet, Format::OnDiskVortex, Format::Lance]; +pub const ALLOWED_FORMATS: &[Format] = &[ + Format::ArrowIpc, + Format::Parquet, + Format::OnDiskVortex, + Format::Lance, +]; impl Format { - /// Clap value parser that only accepts parquet, vortex, and lance. + /// Clap value parser that only accepts formats supported by random-access benchmarks. pub fn parse_allowed(s: &str) -> Result { let format = Format::from_str(s, true)?; if ALLOWED_FORMATS.contains(&format) { @@ -188,12 +195,12 @@ impl Format { pub fn name(&self) -> &'static str { match self { + Format::ArrowIpc => "arrow-ipc", Format::Csv => "csv", - Format::Arrow => "arrow", Format::Parquet => "parquet", Format::OnDiskVortex => "vortex-file-compressed", Format::VortexCompact => "vortex-compact", - Format::VortexNative => "vortex-geo-native", + Format::VortexSpatialNative => "vortex-spatial-native", Format::OnDiskDuckDB => "duckdb", Format::Lance => "lance", } @@ -201,12 +208,12 @@ impl Format { pub fn ext(&self) -> &'static str { match self { + Format::ArrowIpc => "arrow", Format::Csv => "csv", - Format::Arrow => "arrow", Format::Parquet => "parquet", Format::OnDiskVortex => "vortex", Format::VortexCompact => "vortex", - Format::VortexNative => "vortex", + Format::VortexSpatialNative => "vortex", Format::OnDiskDuckDB => "duckdb", Format::Lance => "lance", } @@ -218,7 +225,6 @@ impl Format { pub enum Engine { #[default] Vortex, - Arrow, #[clap(name = "datafusion")] #[serde(rename = "datafusion")] DataFusion, @@ -233,7 +239,6 @@ impl Display for Engine { Engine::DataFusion => write!(f, "datafusion"), Engine::DuckDB => write!(f, "duckdb"), Engine::Vortex => write!(f, "vortex"), - Engine::Arrow => write!(f, "arrow"), } } } @@ -250,7 +255,10 @@ impl CompactionStrategy { match self { CompactionStrategy::Compact => options.with_strategy( WriteStrategyBuilder::default() - .with_btrblocks_builder(BtrBlocksCompressorBuilder::default().with_compact()) + .with_btrblocks_builder(retain_edition_encodings( + &SESSION, + BtrBlocksCompressorBuilder::default().with_compact(), + )) .build(), ), CompactionStrategy::Default => options, @@ -258,6 +266,81 @@ impl CompactionStrategy { } } +/// Restrict `builder` to the encodings permitted by the session's enabled editions. +/// +/// The default writer applies this filter itself. An explicit strategy bypasses it, so a +/// benchmark that builds its own compressor applies it here to stay within editions. +pub fn retain_edition_encodings( + session: &VortexSession, + builder: BtrBlocksCompressorBuilder, +) -> BtrBlocksCompressorBuilder { + let allowed = session + .enabled_component_ids(ComponentKind::Array) + .into_iter() + .collect(); + builder.retain_allowed_encodings(&allowed) +} + +/// Verify that local data has already been prepared for the requested benchmark formats. +/// +/// Engine-specific benchmark binaries call this before running queries. Data generation itself +/// belongs to the `data-gen` binary. +pub fn require_prepared_data(benchmark: &B, formats: &[Format]) -> anyhow::Result<()> +where + B: Benchmark + ?Sized, +{ + if benchmark.data_url().scheme() != "file" { + return Ok(()); + } + + let base_path = benchmark + .data_url() + .to_file_path() + .map_err(|_| anyhow::anyhow!("Invalid file URL: {}", benchmark.data_url()))?; + + let mut missing = Vec::new(); + for format in formats.iter().copied().unique() { + let required_path = match format { + Format::Parquet => base_path.join(Format::Parquet.name()), + Format::OnDiskVortex => base_path.join(Format::OnDiskVortex.name()), + Format::VortexCompact => base_path.join(Format::VortexCompact.name()), + Format::OnDiskDuckDB => base_path + .join(Format::OnDiskDuckDB.name()) + .join("duckdb.db"), + format => base_path.join(format.name()), + }; + + if !required_path.exists() { + missing.push((format, required_path)); + } + } + + if missing.is_empty() { + return Ok(()); + } + + let missing_data = missing + .iter() + .map(|(format, path)| format!("{format} ({})", path.display())) + .join(", "); + let requested_formats = formats + .iter() + .copied() + .unique() + .map(|format| format!("\"{format}\"")) + .join(","); + + anyhow::bail!( + "prepared data is missing for {}: {missing_data}. Generate it first with \ + `vx-bench prepare-data {} --formats-json '[{requested_formats}]'` or \ + `cargo run --bin data-gen -- {} --formats {}` using the same --opt values.", + benchmark.dataset_display(), + benchmark.dataset_name(), + benchmark.dataset_name(), + formats.iter().copied().unique().join(","), + ); +} + /// CLI argument for selecting which benchmark to run. #[derive(clap::ValueEnum, Clone, Copy)] pub enum BenchmarkArg { diff --git a/vortex-bench/src/measurements.rs b/vortex-bench/src/measurements.rs index 4f6abe65563..9833d2b8bba 100644 --- a/vortex-bench/src/measurements.rs +++ b/vortex-bench/src/measurements.rs @@ -351,11 +351,12 @@ pub struct CompressionTimingMeasurement { impl ToJson for CompressionTimingMeasurement { fn to_json(&self) -> serde_json::Value { let (name, engine) = match self.format { + Format::ArrowIpc => (format!("arrow-ipc {}", self.name), Engine::Vortex), Format::OnDiskVortex => (self.name.to_string(), Engine::Vortex), - Format::Parquet => (format!("parquet_rs-zstd {}", self.name), Engine::Arrow), - Format::Lance => (format!("lance {}", self.name), Engine::Arrow), + Format::Parquet => (format!("parquet_rs-zstd {}", self.name), Engine::Vortex), + Format::Lance => (format!("lance {}", self.name), Engine::Vortex), _ => vortex_panic!( - "CompressionTimingMeasurement only supports vortex, lance, and parquet formats" + "CompressionTimingMeasurement only supports arrow-ipc, vortex, lance, and parquet formats" ), }; @@ -397,8 +398,8 @@ impl ToJson for CustomUnitMeasurement { fn to_json(&self) -> serde_json::Value { let engine = match self.format { Format::OnDiskVortex | Format::VortexCompact => Engine::Vortex, - Format::Parquet => Engine::Arrow, - Format::Lance => Engine::Arrow, + Format::Parquet => Engine::Vortex, + Format::Lance => Engine::Vortex, _ => Engine::Vortex, // Default to Vortex for other formats. }; @@ -524,3 +525,23 @@ pub struct MemoryMeasurementJson { pub target: Target, pub env_triple: TripleJson, } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn compression_timing_json_supports_arrow_ipc() { + let measurement = CompressionTimingMeasurement { + name: "compress time/demo".to_string(), + format: Format::ArrowIpc, + time: Duration::from_nanos(42), + }; + + let json = measurement.to_json(); + + assert_eq!(json["name"], "arrow-ipc compress time/demo"); + assert_eq!(json["target"]["format"], "arrow-ipc"); + assert_eq!(json["value"], 42.0); + } +} diff --git a/vortex-bench/src/polarsignals/benchmark.rs b/vortex-bench/src/polarsignals/benchmark.rs index fce692ceb68..d1097e83bd1 100644 --- a/vortex-bench/src/polarsignals/benchmark.rs +++ b/vortex-bench/src/polarsignals/benchmark.rs @@ -59,6 +59,10 @@ impl PolarSignalsBenchmark { #[async_trait::async_trait] impl Benchmark for PolarSignalsBenchmark { + fn doc_path(&self) -> &'static str { + "vortex-bench/sql/polarsignals.md" + } + fn queries(&self) -> Result> { let queries_file = workspace_root() .join("vortex-bench") diff --git a/vortex-bench/src/public_bi.rs b/vortex-bench/src/public_bi.rs index f5c22bdea8f..3e651df5bb7 100644 --- a/vortex-bench/src/public_bi.rs +++ b/vortex-bench/src/public_bi.rs @@ -530,6 +530,10 @@ impl PublicBiBenchmark { #[async_trait] impl Benchmark for PublicBiBenchmark { + fn doc_path(&self) -> &'static str { + "vortex-bench/sql/public-bi.md" + } + fn queries(&self) -> anyhow::Result> { self.pbi_benchmark().queries() } diff --git a/vortex-bench/src/random_access/mod.rs b/vortex-bench/src/random_access/mod.rs index 3b9f95f5eaf..2c475bbb51d 100644 --- a/vortex-bench/src/random_access/mod.rs +++ b/vortex-bench/src/random_access/mod.rs @@ -1,21 +1,29 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::fs::File; use std::path::PathBuf; +use std::sync::Arc; use anyhow::Result; use arrow_array::RecordBatch; +use arrow_ipc::writer::FileWriter; use async_trait::async_trait; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use vortex::array::ArrayRef; use crate::Format; +use crate::idempotent; pub mod take; // Re-export implementations +pub use take::ArrowIpcRandomAccessor; pub use take::ParquetRandomAccessor; pub use take::VortexRandomAccessor; +pub const ARROW_ROW_OFFSETS_METADATA_KEY: &str = "vortex.row_offsets"; + /// Generate the data path for a random-access benchmark dataset file. /// /// Returns a path like `random_access/{dataset}/{dataset}.{ext}` @@ -28,6 +36,42 @@ pub fn data_path(dataset: &str, format: Format) -> String { } } +/// Convert a Parquet input into an uncompressed Arrow IPC file. +pub fn parquet_to_arrow_file(parquet_path: PathBuf, arrow_path: String) -> Result { + idempotent(&arrow_path, |output_path| { + let input = File::open(parquet_path)?; + let builder = ParquetRecordBatchReaderBuilder::try_new(input)?; + let schema = Arc::clone(builder.schema()); + let reader = builder.build()?; + let output = File::create(output_path)?; + let mut writer = FileWriter::try_new(output, schema.as_ref())?; + let mut row_offsets = vec![0_u64]; + + for batch in reader { + let batch = batch?; + let next_offset = row_offsets + .last() + .copied() + .unwrap_or_default() + .checked_add(u64::try_from(batch.num_rows())?) + .ok_or_else(|| anyhow::anyhow!("Arrow row count overflow"))?; + writer.write(&batch)?; + row_offsets.push(next_offset); + } + + writer.write_metadata( + ARROW_ROW_OFFSETS_METADATA_KEY, + row_offsets + .iter() + .map(u64::to_string) + .collect::>() + .join(","), + ); + writer.finish()?; + Ok(()) + }) +} + /// Trait for a benchmark dataset that knows how to prepare data files. #[async_trait] pub trait BenchDataset: Send + Sync { @@ -46,6 +90,7 @@ pub trait BenchDataset: Send + Sync { pub enum RandomAccessorRet { RecordBatch(RecordBatch), ArrayRef(ArrayRef), + Native(Box), } /// Trait for format-specific random access (take) operations. diff --git a/vortex-bench/src/random_access/take.rs b/vortex-bench/src/random_access/take.rs index e86e5d576d6..26cd93a9a78 100644 --- a/vortex-bench/src/random_access/take.rs +++ b/vortex-bench/src/random_access/take.rs @@ -1,17 +1,22 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::collections::BTreeMap; +use std::fs::File as StdFile; use std::iter::once; use std::path::PathBuf; use std::sync::Arc; +use anyhow::Context; use arrow_array::PrimitiveArray; use arrow_array::types::Int64Type; +use arrow_ipc::reader::FileReader; use arrow_select::concat::concat_batches; use arrow_select::take::take_record_batch; use async_trait::async_trait; use futures::stream; use itertools::Itertools; +use parking_lot::Mutex; use parquet::arrow::ParquetRecordBatchStreamBuilder; use parquet::arrow::arrow_reader::ArrowReaderMetadata; use parquet::arrow::arrow_reader::ArrowReaderOptions; @@ -25,13 +30,98 @@ use vortex::array::stream::ArrayStreamExt; use vortex::buffer::Buffer; use vortex::file::OpenOptionsSessionExt; use vortex::file::VortexFile; +use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; use vortex::utils::aliases::hash_map::HashMap; use crate::Format; use crate::SESSION; +use crate::random_access::ARROW_ROW_OFFSETS_METADATA_KEY; use crate::random_access::RandomAccessor; use crate::random_access::RandomAccessorRet; +/// Random accessor for uncompressed Arrow IPC files. +pub struct ArrowIpcRandomAccessor { + name: String, + reader: Mutex>, + row_offsets: Vec, + schema: arrow_schema::SchemaRef, +} + +impl ArrowIpcRandomAccessor { + pub fn open(path: PathBuf, name: impl Into) -> anyhow::Result { + let reader = FileReader::try_new(StdFile::open(path)?, None)?; + let row_offsets = reader + .custom_metadata() + .get(ARROW_ROW_OFFSETS_METADATA_KEY) + .context("Arrow IPC file is missing row-offset metadata")? + .split(',') + .map(str::parse) + .collect::, _>>()?; + anyhow::ensure!( + row_offsets.len() == reader.num_batches() + 1, + "Arrow IPC row-offset metadata does not match the record batches" + ); + anyhow::ensure!( + row_offsets.first() == Some(&0), + "Arrow IPC row offsets must start at zero" + ); + anyhow::ensure!( + row_offsets.windows(2).all(|window| window[0] <= window[1]), + "Arrow IPC row offsets must be sorted" + ); + let schema = reader.schema(); + Ok(Self { + name: name.into(), + reader: Mutex::new(reader), + row_offsets, + schema, + }) + } +} + +#[async_trait] +impl RandomAccessor for ArrowIpcRandomAccessor { + fn format(&self) -> Format { + Format::ArrowIpc + } + + fn name(&self) -> &str { + &self.name + } + + async fn take(&self, indices: &[u64]) -> anyhow::Result { + let mut by_batch = BTreeMap::>::new(); + for &index in indices { + let batch_idx = self.row_offsets.partition_point(|offset| *offset <= index); + anyhow::ensure!( + batch_idx > 0 && batch_idx < self.row_offsets.len(), + "Arrow row index {index} is out of bounds" + ); + let batch_idx = batch_idx - 1; + by_batch + .entry(batch_idx) + .or_default() + .push(i64::try_from(index - self.row_offsets[batch_idx])?); + } + + let mut reader = self.reader.lock(); + let mut batches = Vec::with_capacity(by_batch.len()); + for (batch_idx, local_indices) in by_batch { + reader.set_index(batch_idx)?; + let batch = reader + .next() + .context("Arrow IPC record batch is missing")??; + let indices = PrimitiveArray::::from(local_indices); + batches.push(take_record_batch(&batch, &indices)?); + } + + Ok(RandomAccessorRet::RecordBatch(concat_batches( + &self.schema, + &batches, + )?)) + } +} + /// Random accessor for Vortex format files. /// /// The file handle is opened at construction time and reused across `take()` calls. @@ -76,7 +166,7 @@ impl RandomAccessor for VortexRandomAccessor { let array = self .file .scan()? - .with_row_indices(indices_buf) + .with_row_indices(StrictSortedBuffer::try_new(indices_buf)?) .into_array_stream()? .read_all() .await?; @@ -189,3 +279,43 @@ impl RandomAccessor for ParquetRandomAccessor { Ok(RandomAccessorRet::RecordBatch(result)) } } + +#[cfg(test)] +mod tests { + use arrow_array::Int64Array; + use arrow_array::RecordBatch; + use arrow_ipc::writer::FileWriter; + use arrow_schema::DataType; + use arrow_schema::Field; + use arrow_schema::Schema; + + use super::*; + + #[tokio::test] + async fn arrow_ipc_random_accessor_takes_rows_across_record_batches() -> anyhow::Result<()> { + let file = tempfile::NamedTempFile::new()?; + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + { + let mut writer = FileWriter::try_new(file.reopen()?, schema.as_ref())?; + writer.write(&RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(vec![0, 1]))], + )?)?; + writer.write(&RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(vec![2, 3, 4]))], + )?)?; + writer.write_metadata(ARROW_ROW_OFFSETS_METADATA_KEY, "0,2,5"); + writer.finish()?; + } + + let accessor = ArrowIpcRandomAccessor::open(file.path().to_path_buf(), "arrow-ipc")?; + let RandomAccessorRet::RecordBatch(actual) = accessor.take(&[1, 3, 4]).await? else { + anyhow::bail!("Arrow accessor returned a Vortex array") + }; + let expected = + RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(vec![1, 3, 4]))])?; + assert_eq!(actual, expected); + Ok(()) + } +} diff --git a/vortex-bench/src/realnest/gharchive.rs b/vortex-bench/src/realnest/gharchive.rs index a84395f4a10..6596598d5ba 100644 --- a/vortex-bench/src/realnest/gharchive.rs +++ b/vortex-bench/src/realnest/gharchive.rs @@ -66,6 +66,10 @@ impl GithubArchiveBenchmark { #[async_trait::async_trait] impl Benchmark for GithubArchiveBenchmark { + fn doc_path(&self) -> &'static str { + "vortex-bench/sql/gharchive.md" + } + /// GitHub Archive queries, numbered from Q0 in `sql/gharchive.sql` file order. fn queries(&self) -> anyhow::Result> { // `;`-separated; a `;` must not appear in a comment, or it would split a statement in two. diff --git a/vortex-bench/src/runner.rs b/vortex-bench/src/runner.rs index 1983c63ba05..9e5c4db3d08 100644 --- a/vortex-bench/src/runner.rs +++ b/vortex-bench/src/runner.rs @@ -75,6 +75,7 @@ pub struct SqlBenchmarkRunner { formats: Vec, memory_tracker: Option, hide_progress_bar: bool, + doc: &'static str, query_measurements: Vec, memory_measurements: Vec, } @@ -106,6 +107,7 @@ impl SqlBenchmarkRunner { formats, memory_tracker, hide_progress_bar, + doc: benchmark.doc_path(), query_measurements: Vec::new(), memory_measurements: Vec::new(), }) @@ -223,6 +225,7 @@ impl SqlBenchmarkRunner { display_format, self.engine, &self.formats, + self.doc, f, ) } @@ -232,6 +235,7 @@ impl SqlBenchmarkRunner { display_format, self.engine, &self.formats, + self.doc, std::io::stdout().lock(), ), } @@ -249,6 +253,7 @@ impl SqlBenchmarkRunner { display_format, self.engine, &self.formats, + self.doc, output, ) } @@ -457,6 +462,7 @@ pub fn export_results( display_format: &DisplayFormat, engine: Engine, formats: &[Format], + doc: &str, mut output: W, ) -> anyhow::Result<()> { let targets = formats @@ -467,13 +473,13 @@ pub fn export_results( if !memory.is_empty() { match display_format { DisplayFormat::Table => render_table(&mut output, memory, &targets)?, - DisplayFormat::GhJson => print_measurements_json(&mut output, memory)?, + DisplayFormat::GhJson => print_measurements_json(&mut output, memory, doc)?, }; } match display_format { DisplayFormat::Table => render_table(&mut output, queries, &targets)?, - DisplayFormat::GhJson => print_measurements_json(&mut output, queries)?, + DisplayFormat::GhJson => print_measurements_json(&mut output, queries, doc)?, }; Ok(()) diff --git a/vortex-bench/src/snapshots/vortex_bench__v3__tests__snapshot_compression_size.snap b/vortex-bench/src/snapshots/vortex_bench__v3__tests__snapshot_compression_size.snap index 3a4ed1aa8a5..c026a0dd287 100644 --- a/vortex-bench/src/snapshots/vortex_bench__v3__tests__snapshot_compression_size.snap +++ b/vortex-bench/src/snapshots/vortex_bench__v3__tests__snapshot_compression_size.snap @@ -7,5 +7,6 @@ expression: render(&record) "commit_sha": "", "dataset": "taxi", "format": "lance", - "value_bytes": 12345678 + "value_bytes": 12345678, + "uncompressed_bytes": 98765432 } diff --git a/vortex-bench/src/snapshots/vortex_bench__v3__tests__snapshot_random_access_time.snap b/vortex-bench/src/snapshots/vortex_bench__v3__tests__snapshot_random_access_time.snap index 6fce08f5ee2..f0ff27df294 100644 --- a/vortex-bench/src/snapshots/vortex_bench__v3__tests__snapshot_random_access_time.snap +++ b/vortex-bench/src/snapshots/vortex_bench__v3__tests__snapshot_random_access_time.snap @@ -7,6 +7,7 @@ expression: render(&record) "commit_sha": "", "dataset": "taxi", "format": "parquet", + "open_mode": "cached", "value_ns": 850000, "all_runtimes_ns": [ 800000, diff --git a/vortex-bench/src/spatialbench/benchmark.rs b/vortex-bench/src/spatialbench/benchmark.rs index 48fa84f8ec1..693a174eae1 100644 --- a/vortex-bench/src/spatialbench/benchmark.rs +++ b/vortex-bench/src/spatialbench/benchmark.rs @@ -18,10 +18,10 @@ use crate::spatialbench::datagen::Table; use crate::utils::file::resolve_data_url; use crate::workspace_root; -/// Data-dir subfolder for the native-geometry Vortex files (the `vortex-geo-native` lane). -pub const NATIVE_DIR: &str = "vortex-geo-native"; +/// Data-dir subfolder for the native-geometry Vortex files (the `vortex-spatial-native` lane). +pub const SPATIAL_NATIVE_DIR: &str = "vortex-spatial-native"; -/// SpatialBench geospatial benchmark (Apache Sedona): a `trip` point table, `building` polygons, and +/// SpatialBench spatial benchmark (Apache Sedona): a `trip` point table, `building` polygons, and /// a `customer` attribute table, queried with spatial filters and joins. `zone` polygons are sourced /// externally and registered when present. See . pub struct SpatialBenchBenchmark { @@ -58,6 +58,10 @@ impl SpatialBenchBenchmark { #[async_trait::async_trait] impl Benchmark for SpatialBenchBenchmark { + fn doc_path(&self) -> &'static str { + "vortex-bench/sql/spatialbench.md" + } + /// All SpatialBench queries, numbered started at Q1 in `spatialbench.sql` file order. fn queries(&self) -> anyhow::Result> { // `;`-separated; a `;` must not appear in a comment, or it would split a statement in two. @@ -86,14 +90,15 @@ impl Benchmark for SpatialBenchBenchmark { .map_err(|_| anyhow::anyhow!("Invalid file URL: {}", self.data_url.as_str()))?; datagen::generate_tables(&self.scale_factor, base_data_dir.clone()).await?; - // `zone` is externally sourced (SpatialBench directly generate it), so - // re-tag its parquet with the geo metadata. - if let Some(geo) = datagen::wkb::geo_parquet_metadata(Table::Zone) { + // `zone` is externally sourced (SpatialBench directly generates it), so restore its + // GeoParquet metadata. + if let Some(geoparquet_metadata) = datagen::wkb::geoparquet_metadata(Table::Zone) { let zone_glob = base_data_dir .join(Format::Parquet.name()) .join("zone_*.parquet"); for zone_file in glob::glob(&zone_glob.to_string_lossy())?.flatten() { - crate::conversions::add_geoparquet_metadata(&zone_file, &geo).await?; + crate::conversions::add_geoparquet_metadata(&zone_file, &geoparquet_metadata) + .await?; } } @@ -102,19 +107,19 @@ impl Benchmark for SpatialBenchBenchmark { let derived_dirs = [ base_data_dir.join(Format::OnDiskVortex.name()), base_data_dir.join(Format::VortexCompact.name()), - base_data_dir.join(NATIVE_DIR), + base_data_dir.join(SPATIAL_NATIVE_DIR), ]; datagen::spatially_sort_tables(&base_data_dir.join(Format::Parquet.name()), &derived_dirs) .await?; Ok(()) } - /// The `vortex-geo-native` lane decodes each table's WKB geometry to native GeoArrow once, into the - /// `vortex-geo-native` dir, so its queries read DuckDB `GEOMETRY` directly. Idempotent. + /// The `vortex-spatial-native` lane decodes each table's WKB geometry to native GeoArrow once, into the + /// `vortex-spatial-native` dir, so its queries read DuckDB `GEOMETRY` directly. Idempotent. async fn prepare_format(&self, format: Format, base_path: &Path) -> anyhow::Result<()> { - if format == Format::VortexNative { + if format == Format::VortexSpatialNative { let parquet_dir = base_path.join(Format::Parquet.name()); - let native_dir = base_path.join(NATIVE_DIR); + let native_dir = base_path.join(SPATIAL_NATIVE_DIR); for table in self.base_tables() { datagen::write_native_vortex(table, &parquet_dir, &native_dir).await?; } @@ -126,11 +131,11 @@ impl Benchmark for SpatialBenchBenchmark { &self.data_url } - /// The `vortex-geo-native` lane reads the native-geometry Vortex dir; every other format reads its + /// The `vortex-spatial-native` lane reads the native-geometry Vortex dir; every other format reads its /// own `{format}` subfolder. fn format_path(&self, format: Format, base_url: &Url) -> anyhow::Result { let dir = match format { - Format::VortexNative => NATIVE_DIR, + Format::VortexSpatialNative => SPATIAL_NATIVE_DIR, other => other.name(), }; Ok(base_url.join(&format!("{dir}/"))?) @@ -144,7 +149,7 @@ impl Benchmark for SpatialBenchBenchmark { // shorter vec means the runner skips them). match self.scale_factor.as_str() { "1.0" => Some(vec![0, 94, 1, 22, 258, 316691, 3, 6000000, 369, 37]), - "10.0" => Some(vec![0, 994, 1, 79, 231, 3144328, 3, 60000000, 9357, 573]), + "10.0" => Some(vec![0, 994, 1, 79, 243, 3144328, 3, 60000000, 9357, 573]), _ => None, } } @@ -164,7 +169,7 @@ impl Benchmark for SpatialBenchBenchmark { } /// Both lanes register the same tables (WKB reads `parquet`/`vortex`, native reads - /// `vortex-geo-native`); `zone` is externally sourced and optional, registered only when present. + /// `vortex-spatial-native`); `zone` is externally sourced and optional, registered only when present. fn table_specs(&self) -> Vec { self.base_tables() .iter() diff --git a/vortex-bench/src/spatialbench/datagen/mod.rs b/vortex-bench/src/spatialbench/datagen/mod.rs index 4c0d4bd2f81..d144746712e 100644 --- a/vortex-bench/src/spatialbench/datagen/mod.rs +++ b/vortex-bench/src/spatialbench/datagen/mod.rs @@ -2,8 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors //! SpatialBench data preparation. [`wkb`] generates the canonical WKB base tables; [`native`] derives -//! native-geometry Vortex files from them for `points=native`. The [`table`] catalog is the single -//! source of truth for the base tables both stages share. +//! native-geometry Vortex files for the `vortex-spatial-native` format. The [`table`] catalog is +//! the single source of truth for the base tables both stages share. pub mod native; pub mod spatial_sort; diff --git a/vortex-bench/src/spatialbench/datagen/native.rs b/vortex-bench/src/spatialbench/datagen/native.rs index 2e03b61b39f..510682c22e7 100644 --- a/vortex-bench/src/spatialbench/datagen/native.rs +++ b/vortex-bench/src/spatialbench/datagen/native.rs @@ -1,8 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Native-geometry prep for `points=native`: decode a table's WKB geometry to native -//! `vortex.geo.{point,polygon,multipolygon}` via `geoarrow_cast` (so Vortex never decodes WKB), then +//! Native-geometry prep for `vortex-spatial-native`: decode a table's WKB geometry to native +//! `vortex.st.{point,polygon,multipolygon}` via `geoarrow_cast` (so Vortex never decodes WKB), then //! write a Vortex file. A one-time cost; queries then see DuckDB `GEOMETRY` directly. use std::path::Path; @@ -40,7 +40,7 @@ use super::table::Table; use crate::SESSION; use crate::utils::file::idempotent_async; -fn geo_metadata() -> Arc { +fn geoarrow_metadata() -> Arc { Arc::new(Metadata::new(Crs::default(), None)) } @@ -58,8 +58,7 @@ pub async fn write_native_vortex( let dtype = chunks[0].dtype().clone(); let chunked = ChunkedArray::try_new(chunks, dtype)?.into_array(); let mut file = TokioFile::create(&path).await?; - SESSION - .write_options() + SESSION.write_options() .write(&mut file, chunked.to_array_stream()) .await?; tracing::info!(path = %path.display(), table = table.name(), "wrote native geometry table"); @@ -103,7 +102,7 @@ fn native_record_batch(batch: RecordBatch, table: Table) -> anyhow::Result anyhow::Result cast( wkb.as_ref(), &GeoArrowType::Point( - PointType::new(Dimension::XY, geo_metadata()) + PointType::new(Dimension::XY, geoarrow_metadata()) .with_coord_type(CoordType::Separated), ), )?, GeometryKind::Polygon => cast( wkb.as_ref(), &GeoArrowType::Polygon( - PolygonType::new(Dimension::XY, geo_metadata()) + PolygonType::new(Dimension::XY, geoarrow_metadata()) .with_coord_type(CoordType::Separated), ), )?, @@ -138,7 +137,7 @@ fn native_record_batch(batch: RecordBatch, table: Table) -> anyhow::Result cast( wkb.as_ref(), &GeoArrowType::MultiPolygon( - MultiPolygonType::new(Dimension::XY, geo_metadata()) + MultiPolygonType::new(Dimension::XY, geoarrow_metadata()) .with_coord_type(CoordType::Separated), ), )?, diff --git a/vortex-bench/src/spatialbench/datagen/spatial_sort.rs b/vortex-bench/src/spatialbench/datagen/spatial_sort.rs index d77aea15846..e4382ffb4f3 100644 --- a/vortex-bench/src/spatialbench/datagen/spatial_sort.rs +++ b/vortex-bench/src/spatialbench/datagen/spatial_sort.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors //! Spatial clustering of the source parquet, in place, so every downstream lane (parquet, -//! vortex-WKB, `vortex-geo-native`) reads the same layout. +//! vortex-WKB, `vortex-spatial-native`) reads the same layout. //! //! The geometry zone-map prune skips a chunk only when its bounding box is disjoint from the query //! region. In generation order every chunk's box spans the whole map, so nothing prunes; ordering @@ -44,12 +44,12 @@ use tokio::fs::File as TokioFile; use tracing::info; use super::table::Table; -use super::wkb::geo_parquet_metadata; +use super::wkb::geoparquet_metadata; /// Parquet metadata marker: this file is already spatially sorted (makes the step idempotent). const SORTED_KEY: &str = "vortex_spatial_sorted"; -fn geo_metadata() -> Arc { +fn geoarrow_metadata() -> Arc { Arc::new(Metadata::new(Crs::default(), None)) } @@ -148,9 +148,11 @@ async fn sort_part(path: &Path, table: Table, geom_col: &str) -> anyhow::Result< let mut writer = AsyncArrowWriter::try_new(TokioFile::create(&tmp_path).await?, schema, Some(props))?; writer.write(&sorted).await?; - // A fresh write drops metadata: re-tag geo so DuckDB reads `GEOMETRY`, and add the sorted marker. - if let Some(geo) = geo_parquet_metadata(table) { - writer.append_key_value_metadata(KeyValue::new("geo".to_string(), Some(geo))); + // A fresh write drops metadata: restore the GeoParquet tag so DuckDB reads `GEOMETRY`, and add + // the sorted marker. + if let Some(geoparquet_metadata) = geoparquet_metadata(table) { + writer + .append_key_value_metadata(KeyValue::new("geo".to_string(), Some(geoparquet_metadata))); } writer.append_key_value_metadata(KeyValue::new( SORTED_KEY.to_string(), @@ -170,7 +172,7 @@ async fn sort_part(path: &Path, table: Table, geom_col: &str) -> anyhow::Result< /// The bounding-box center `(x, y)` of every geometry in a WKB column, whatever its geometry type. fn wkb_centers(column: &dyn Array) -> anyhow::Result<(Vec, Vec)> { - let wkb_type = WkbType::new(geo_metadata()); + let wkb_type = WkbType::new(geoarrow_metadata()); // Expanded per concrete WKB array type. macro_rules! centers { ($array:expr) => {{ diff --git a/vortex-bench/src/spatialbench/datagen/wkb.rs b/vortex-bench/src/spatialbench/datagen/wkb.rs index 51b34bf3016..417c95cdf62 100644 --- a/vortex-bench/src/spatialbench/datagen/wkb.rs +++ b/vortex-bench/src/spatialbench/datagen/wkb.rs @@ -101,8 +101,11 @@ pub async fn generate_tables(scale_factor: &str, output_dir: PathBuf) -> Result< } // Tag geometry columns with GeoParquet `geo` metadata so DuckDB's `read_parquet` // surfaces them as `GEOMETRY` directly. - if let Some(geo) = geo_parquet_metadata(table) { - writer.append_key_value_metadata(KeyValue::new("geo".to_string(), Some(geo))); + if let Some(geoparquet_metadata) = geoparquet_metadata(table) { + writer.append_key_value_metadata(KeyValue::new( + "geo".to_string(), + Some(geoparquet_metadata), + )); } writer.close().await?; @@ -162,7 +165,7 @@ async fn generate_zone(scale_factor: f64, parquet_dir: &Path) -> Result<()> { } /// GeoParquet metadata for WKB geometry columns, or `None` when it has none. -pub(crate) fn geo_parquet_metadata(table: Table) -> Option { +pub(crate) fn geoparquet_metadata(table: Table) -> Option { let geometry_columns = table.geometry_columns(); let primary = geometry_columns.first()?; let columns: serde_json::Map = geometry_columns diff --git a/vortex-bench/src/spatialbench/mod.rs b/vortex-bench/src/spatialbench/mod.rs index bba06bd7ef9..24de5fc3d64 100644 --- a/vortex-bench/src/spatialbench/mod.rs +++ b/vortex-bench/src/spatialbench/mod.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! SpatialBench geospatial analytics benchmark. +//! SpatialBench spatial analytics benchmark. //! //! See . diff --git a/vortex-bench/src/statpopgen/statpopgen_benchmark.rs b/vortex-bench/src/statpopgen/statpopgen_benchmark.rs index 0861f7bfc88..2c677fb79a3 100644 --- a/vortex-bench/src/statpopgen/statpopgen_benchmark.rs +++ b/vortex-bench/src/statpopgen/statpopgen_benchmark.rs @@ -103,6 +103,10 @@ impl StatPopGenBenchmark { #[async_trait::async_trait] impl Benchmark for StatPopGenBenchmark { + fn doc_path(&self) -> &'static str { + "vortex-bench/sql/statpopgen.md" + } + fn queries(&self) -> Result> { let queries_file = workspace_root() .join("vortex-bench") diff --git a/vortex-bench/src/tpcds/tpcds_benchmark.rs b/vortex-bench/src/tpcds/tpcds_benchmark.rs index 7f8f2d5281e..725dd204de7 100644 --- a/vortex-bench/src/tpcds/tpcds_benchmark.rs +++ b/vortex-bench/src/tpcds/tpcds_benchmark.rs @@ -39,7 +39,7 @@ impl TpcDsBenchmark { Url::from_directory_path(&data_dir_with_sf).map_err(|_| { anyhow!( "Failed to create URL from directory path: {:?}", - &data_dir_with_sf + data_dir_with_sf ) }) } @@ -56,6 +56,10 @@ impl TpcDsBenchmark { #[async_trait::async_trait] impl Benchmark for TpcDsBenchmark { + fn doc_path(&self) -> &'static str { + "vortex-bench/sql/tpcds/README.md" + } + fn queries(&self) -> Result> { Ok(tpcds_queries().collect()) } diff --git a/vortex-bench/src/tpch/benchmark.rs b/vortex-bench/src/tpch/benchmark.rs index 6851e5f5f06..5d768e459af 100644 --- a/vortex-bench/src/tpch/benchmark.rs +++ b/vortex-bench/src/tpch/benchmark.rs @@ -49,7 +49,7 @@ impl TpcHBenchmark { Url::from_directory_path(&data_dir_with_sf).map_err(|_| { anyhow::anyhow!( "Failed to create URL from directory path: {:?}", - &data_dir_with_sf + data_dir_with_sf ) }) } @@ -75,6 +75,10 @@ impl TpcHBenchmark { #[async_trait::async_trait] impl Benchmark for TpcHBenchmark { + fn doc_path(&self) -> &'static str { + "vortex-bench/sql/tpch/README.md" + } + fn queries(&self) -> anyhow::Result> { Ok(tpch_queries().collect()) } @@ -89,8 +93,7 @@ impl Benchmark for TpcHBenchmark { .to_file_path() .map_err(|_| anyhow::anyhow!("Invalid file URL: {}", self.data_url.as_str()))?; - let options = TpchGenOptions::new(self.scale_factor.clone(), base_data_dir) - .with_max_file_size_mb(Some(600)); + let options = TpchGenOptions::new(self.scale_factor.clone(), base_data_dir); tpchgen::generate_tpch_tables(options).await?; @@ -139,7 +142,7 @@ impl Benchmark for TpcHBenchmark { #[expect(clippy::expect_used)] fn pattern(&self, table_name: &str, format: Format) -> Option { Some( - format!("{}_*.{}", table_name, format.ext()) + format!("{}.{}", table_name, format.ext()) .parse() .expect("valid glob pattern"), ) diff --git a/vortex-bench/src/tpch/tpchgen.rs b/vortex-bench/src/tpch/tpchgen.rs index 79b91200e48..2c4a26fd527 100644 --- a/vortex-bench/src/tpch/tpchgen.rs +++ b/vortex-bench/src/tpch/tpchgen.rs @@ -33,11 +33,9 @@ use tpchgen_arrow::RecordBatchIterator; use tracing::info; use vortex::array::ArrayRef; use vortex::array::stream::ArrayStreamAdapter; -use vortex::dtype::DType; use vortex::error::VortexExpect; use vortex::file::WriteOptionsSessionExt; -use vortex_arrow::FromArrowArray; -use vortex_arrow::FromArrowType; +use vortex_arrow::ArrowSessionExt; use crate::CompactionStrategy; use crate::Format; @@ -56,8 +54,6 @@ pub struct TpchGenOptions { pub format: Format, /// Batch size for streaming pub batch_size: usize, - /// The max size of uncompressed file .tbl that we should generate - pub max_file_size_mb: Option, } impl Default for TpchGenOptions { @@ -67,7 +63,6 @@ impl Default for TpchGenOptions { output_dir: "tpch".to_data_path(), format: Format::Parquet, batch_size: 8192 * 64, - max_file_size_mb: None, } } } @@ -90,11 +85,6 @@ impl TpchGenOptions { self.batch_size = batch_size; self } - - pub fn with_max_file_size_mb(mut self, max_file_size_mb: Option) -> Self { - self.max_file_size_mb = max_file_size_mb; - self - } } /// Generate all TPC-H tables for a single scale factor @@ -125,11 +115,10 @@ pub async fn generate_tpch_tables(options: TpchGenOptions) -> Result<()> { ); let table_name = table_name.to_string(); - generate_table_files(table_name, *generator, options.clone()) + generate_table_file(table_name, *generator, options.clone()) }) .collect::>>()? .into_iter() - .flatten() .collect::>>>(); let limiter = Arc::new(Semaphore::new(MAX_CONCURRENT_FILES)); @@ -169,31 +158,14 @@ enum TableGenerator { LineItem, } -impl TableGenerator { - // Returns the size in MBs of the table at scale factor 1, if the table is constant size - // return None - fn uncompressed_data_size(&self) -> Option { - match self { - TableGenerator::Nation => None, - TableGenerator::Region => None, - TableGenerator::Part => Some(113), - TableGenerator::Supplier => Some(1), - TableGenerator::Customer => Some(23), - TableGenerator::PartSupp => Some(113), - TableGenerator::Orders => Some(164), - TableGenerator::LineItem => Some(725), - } - } -} - -/// Generate files for a specific table in streaming fashion -fn generate_table_files( +/// Generate one file for a specific table in streaming fashion +fn generate_table_file( table_name: String, generator: TableGenerator, options: TpchGenOptions, -) -> Result>>> { +) -> Result>> { let write_format = match options.format { - Format::Parquet | Format::Arrow | Format::OnDiskDuckDB => Format::Parquet, + Format::Parquet | Format::OnDiskDuckDB => Format::Parquet, Format::OnDiskVortex => Format::OnDiskVortex, Format::VortexCompact => Format::VortexCompact, f => { @@ -205,98 +177,60 @@ fn generate_table_files( fs::create_dir_all(&output_dir)?; - // Calculate number of partitions without creating expensive iterators - let num_parts = calculate_num_parts(generator, &options)?; - - let mut futures = Vec::new(); - - for partition_idx in 0..num_parts { - let output_file = output_dir.join(format!( - "{table_name}_{partition_idx}.{}", - write_format.ext() - )); - - // Clone necessary data for the async closure - let options_clone = options.clone(); - let table_name = table_name.to_string(); - - let future = async move { - let of = output_file.clone(); - idempotent_async(output_file.to_string_lossy().as_ref(), |path| async move { - info!( - "Generating {table_name} table as {write_format}, at {}", - of.to_string_lossy() - ); - - // Create the specific iterator for this partition only when we need to generate - let iter = create_single_batch_iterator(generator, &options_clone, partition_idx)?; - - // Create generator and process batches in streaming fashion - let schema = Arc::clone(iter.schema()); - - // Create writer based on format - let mut writer: Box = match write_format { - Format::Parquet => Box::new(ParquetWriter::new(path, schema).await?), - Format::OnDiskVortex => Box::new(VortexWriter::new( - path, - schema, - CompactionStrategy::Default, - )?), - Format::VortexCompact => Box::new(VortexWriter::new( - path, - schema, - CompactionStrategy::Compact, - )?), - _ => unreachable!(), - }; - - for batch in iter { - writer.write_batch(&batch).await?; - } - - writer.finalize().await?; - Ok::<(), anyhow::Error>(()) - }) - .await?; - Ok(()) - } - .boxed(); - - futures.push(future); - } - - Ok(futures) -} + let output_file = output_dir.join(format!("{table_name}.{}", write_format.ext())); -/// Calculate the number of partitions without creating expensive iterators -#[expect(clippy::cast_possible_truncation)] -fn calculate_num_parts(generator: TableGenerator, options: &TpchGenOptions) -> Result { - let scale_factor = options.scale_factor.parse::()?; + let future = async move { + let of = output_file.clone(); + idempotent_async(output_file.to_string_lossy().as_ref(), |path| async move { + info!( + "Generating {table_name} table as {write_format}, at {}", + of.to_string_lossy() + ); - let num_parts = if let Some((data_size, max_file_size)) = generator - .uncompressed_data_size() - .zip(options.max_file_size_mb) - { - #[expect(clippy::cast_precision_loss)] - let file_size = (data_size as f64 * scale_factor).ceil() as u64; - file_size.div_ceil(max_file_size) - } else { - 1 - }; + let iter = create_batch_iterator(generator, &options)?; + + // Create generator and process batches in streaming fashion + let schema = Arc::clone(iter.schema()); + + // Create writer based on format + let mut writer: Box = match write_format { + Format::Parquet => Box::new(ParquetWriter::new(path, schema).await?), + Format::OnDiskVortex => Box::new(VortexWriter::new( + path, + schema, + CompactionStrategy::Default, + )?), + Format::VortexCompact => Box::new(VortexWriter::new( + path, + schema, + CompactionStrategy::Compact, + )?), + _ => unreachable!(), + }; + + for batch in iter { + writer.write_batch(&batch).await?; + } + + writer.finalize().await?; + Ok::<(), anyhow::Error>(()) + }) + .await?; + Ok(()) + } + .boxed(); - Ok(num_parts as usize) + Ok(future) } -/// Create a single batch iterator for a specific partition -#[expect(clippy::cast_possible_truncation)] -fn create_single_batch_iterator( +/// Create a batch iterator for one full table +fn create_batch_iterator( generator: TableGenerator, options: &TpchGenOptions, - partition_idx: usize, ) -> Result> { let scale_factor = options.scale_factor.parse::()?; - let num_parts = calculate_num_parts(generator, options)? as i32; - let part = (partition_idx + 1) as i32; // 1-indexed + let part = 1; + let num_parts = 1; let batch_size = options.batch_size; let iterator: Box = match generator { @@ -370,7 +304,6 @@ impl ParquetWriter { let file = TokioFile::create(&path).await?; let props = WriterProperties::builder() .set_compression(Compression::SNAPPY) - .set_bloom_filter_enabled(true) .build(); let writer = AsyncArrowWriter::try_new(file, schema, Some(props))?; Ok(Self { writer }) @@ -403,7 +336,7 @@ impl VortexWriter { ) -> Result { // Increase buffer size to avoid backpressure issues let (sender, receiver) = mpsc::channel(2); - let dtype = DType::from_arrow(schema); + let dtype = SESSION.arrow().from_arrow_schema(schema.as_ref())?; let file_path = path; let write_task = Some(tokio::spawn(async move { let stream = ArrayStreamAdapter::new(dtype, ReceiverStream::new(receiver)); @@ -428,7 +361,10 @@ impl VortexWriter { #[async_trait::async_trait] impl FileWriter for VortexWriter { async fn write_batch(&mut self, batch: &RecordBatch) -> Result<()> { - let array = ArrayRef::from_arrow(batch, false)?; + let schema = batch.schema(); + let array = SESSION + .arrow() + .from_arrow_record_batch(batch.clone(), &schema)?; self.sender .as_ref() .vortex_expect("sender closed early") diff --git a/vortex-bench/src/utils/file.rs b/vortex-bench/src/utils/file.rs index 4af68c029ab..1abcf42d953 100644 --- a/vortex-bench/src/utils/file.rs +++ b/vortex-bench/src/utils/file.rs @@ -130,7 +130,7 @@ pub fn resolve_data_url(remote_data_dir: Option<&str>, local_subdir: &str) -> Re None => { let data_dir = data_dir().join(local_subdir); Url::from_directory_path(&data_dir).map_err(|_| { - anyhow::anyhow!("Failed to create URL from directory path: {:?}", &data_dir) + anyhow::anyhow!("Failed to create URL from directory path: {:?}", data_dir) }) } Some(remote_data_dir) => { diff --git a/vortex-bench/src/v3.rs b/vortex-bench/src/v3.rs index 4bdd67ed18f..51a2d71262b 100644 --- a/vortex-bench/src/v3.rs +++ b/vortex-bench/src/v3.rs @@ -1,19 +1,17 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! v3 wire-format records emitted by `--gh-json-v3`. +//! Benchmark wire-format records emitted by `--ingest-jsonl`. //! //! Each record on the wire is one of five `kind`-discriminated shapes that -//! map 1:1 to the v3 fact tables. The records here are **bare**: the -//! ingest envelope (`run_meta` + `commit`) is added by -//! `scripts/post-ingest.py` before POSTing to -//! `bench.vortex.dev/api/ingest` — keeps the Rust emitter dependency-light -//! and lets CI fill the commit fields from `${{ github.sha }}` plus -//! `git show`. +//! map 1:1 to the Postgres fact tables. The records here are **bare**; +//! `scripts/post-ingest.py` fills the commit metadata from `${{ github.sha }}` +//! plus `git show`, computes the internal `measurement_id`, and upserts them +//! directly into the database. //! -//! Wire-shape source of truth: [`vortex_bench_server::records`]. When -//! changing a shape, change both sides in the same commit and run the -//! server's snapshot tests. +//! When changing a shape, update `benchmarks-website/CONTRACT.md`, +//! `benchmarks-website/web/lib/schema-version.ts`, and +//! `scripts/post-ingest.py` in the same logical change. //! //! ## Producer mapping //! @@ -56,19 +54,6 @@ //! For SQL query suites (everything that flows through `query_measurements`), //! the dim columns are populated as documented on //! [`benchmark_dataset_dims`]. -//! -//! ## Historical-data side -//! -//! [`vortex_bench_migrate::classifier`] is the bug-for-bug port of v2's -//! `getGroup` that recovers the same `(kind, dim tuple)` triple from the -//! v2 S3 dump. It exists only for the one-shot migration; once cutover -//! lands and the historical archive is loaded, both the migrator and its -//! classifier go away. For new ingest, no classifier is needed — the -//! emitter writes v3-shape records directly. -//! -//! [`vortex_bench_server::records`]: ../../../benchmarks-website/server/src/records.rs -//! [`vortex_bench_migrate::classifier`]: ../../../benchmarks-website/migrate/src/classifier.rs - use std::io::Write; use std::sync::LazyLock; @@ -99,7 +84,7 @@ pub static ENV_TRIPLE: LazyLock = LazyLock::new(|| { /// Wire-format kind discriminator. One value per fact table. /// /// Each variant flattens its inner record next to a `"kind"` field, matching the -/// shape consumed by `/api/ingest`. +/// shape consumed by `scripts/post-ingest.py`. #[derive(Debug, Clone, Serialize)] #[serde(tag = "kind", rename_all = "snake_case")] pub enum V3Record { @@ -141,7 +126,7 @@ pub struct QueryMeasurementRecord { pub query_idx: u32, /// Storage backend the run targeted (`nvme` or `s3`). pub storage: String, - /// Query engine (`datafusion`, `duckdb`, `vortex`, `arrow`). + /// Query engine (`datafusion`, `duckdb`, `vortex`). pub engine: String, /// On-disk format (`parquet`, `vortex-file-compressed`, `lance`, ...). pub format: String, @@ -203,6 +188,8 @@ pub struct CompressionSizeRecord { pub format: String, /// Size in bytes. pub value_bytes: u64, + /// Arrow memory size after decoding the source Parquet file. + pub uncompressed_bytes: u64, } /// A single take-time timing from `random-access-bench`. @@ -214,6 +201,8 @@ pub struct RandomAccessTimeRecord { pub dataset: String, /// On-disk format the timing applies to. pub format: String, + /// File access mode: `cached` reuses an accessor and `reopen` opens one per take. + pub open_mode: String, /// Median per-iteration wall time in nanoseconds. pub value_ns: u64, /// Per-iteration wall times in nanoseconds. @@ -361,11 +350,7 @@ pub fn query_measurement_record( dataset, dataset_variant, scale_factor, - // Clamp at `i32::MAX`, not `u32::MAX`: the server's `query_idx` - // field is `i32`, so a saturated `u32::MAX` would fail serde - // deserialization. Real `query_idx` values are 0..200, so the - // saturation path is defensive rather than load-bearing. - query_idx: u32::try_from(qm.query_idx).unwrap_or(i32::MAX as u32), + query_idx: clamp_to_i32(qm.query_idx), storage: qm.storage.clone(), engine: engine_label(qm.target.engine).to_string(), format: qm.target.format.name().to_string(), @@ -416,6 +401,7 @@ pub fn compression_size_record( dataset_variant: Option<&str>, format: Format, value_bytes: u64, + uncompressed_bytes: u64, ) -> V3Record { V3Record::CompressionSize(CompressionSizeRecord { commit_sha: GIT_COMMIT_ID.clone(), @@ -423,17 +409,23 @@ pub fn compression_size_record( dataset_variant: dataset_variant.map(str::to_string), format: format.name().to_string(), value_bytes, + uncompressed_bytes, }) } /// Build a `random_access_time` record from a [`TimingMeasurement`]. -pub fn random_access_record(timing: &TimingMeasurement, dataset: &str) -> V3Record { +pub fn random_access_record( + timing: &TimingMeasurement, + dataset: &str, + open_mode: &str, +) -> V3Record { let value_ns = duration_as_ns(timing.median_time()); let all_runtimes_ns = timing.runs.iter().copied().map(duration_as_ns).collect(); V3Record::RandomAccessTime(RandomAccessTimeRecord { commit_sha: GIT_COMMIT_ID.clone(), dataset: dataset.to_string(), format: timing.target.format.name().to_string(), + open_mode: open_mode.to_string(), value_ns, all_runtimes_ns, env_triple: Some(ENV_TRIPLE.clone()), @@ -463,10 +455,7 @@ pub fn vector_search_record( rows_scanned: u64, bytes_scanned: u64, ) -> V3Record { - // Clamp at `i32::MAX`: server-side `iterations` is `i32` (see the - // same-named field on `vortex_bench_server::records::VectorSearchRun`), - // so saturating to `u32::MAX` would 400 the envelope. - let iterations = u32::try_from(all_runs_ns.len()).unwrap_or(i32::MAX as u32); + let iterations = clamp_to_i32(all_runs_ns.len()); V3Record::VectorSearchRun(VectorSearchRunRecord { commit_sha: GIT_COMMIT_ID.clone(), dataset: dims.dataset.to_string(), @@ -504,22 +493,20 @@ pub fn write_jsonl_to_path(path: &std::path::Path, records: &[V3Record]) -> std: write_jsonl(&mut file, records) } -/// Convert a `Duration` to `u64` nanoseconds, clamping at the largest -/// value the server-side `value_ns: i64` deserializer can accept. -/// -/// The wire field is `u64` (see `value_ns` on every record type below), -/// but the server's records mirrors them as `i64`. -/// Without the clamp, an overflowed measurement would land at `u64::MAX` -/// here, fail serde deserialization on the server, and 400 the whole -/// ingest envelope. +/// Clamp unsigned wire values to the Postgres `INTEGER` limit enforced by `scripts/post-ingest.py`. +fn clamp_to_i32(value: usize) -> u32 { + i32::try_from(value).unwrap_or(i32::MAX) as u32 +} + +/// Convert a duration to nanoseconds within the Postgres `BIGINT` limit enforced by +/// `scripts/post-ingest.py`, preserving the unsigned wire field. fn duration_as_ns(d: std::time::Duration) -> u64 { - u64::try_from(d.as_nanos()).unwrap_or(i64::MAX as u64) + i64::try_from(d.as_nanos()).unwrap_or(i64::MAX) as u64 } fn engine_label(engine: Engine) -> &'static str { match engine { Engine::Vortex => "vortex", - Engine::Arrow => "arrow", Engine::DataFusion => "datafusion", Engine::DuckDB => "duckdb", } @@ -538,6 +525,7 @@ mod tests { use insta::assert_snapshot; use insta::with_settings; + use rstest::rstest; use super::*; use crate::Target; @@ -554,6 +542,26 @@ mod tests { Ok(redact_env(&json)) } + #[rstest] + #[case::zero(0, 0)] + #[case::below_limit(i32::MAX as usize - 1, i32::MAX as u32 - 1)] + #[case::at_limit(i32::MAX as usize, i32::MAX as u32)] + #[case::above_limit(i32::MAX as usize + 1, i32::MAX as u32)] + #[case::overflow(usize::MAX, i32::MAX as u32)] + fn test_clamp_to_i32(#[case] value: usize, #[case] expected: u32) { + assert_eq!(clamp_to_i32(value), expected); + } + + #[rstest] + #[case::zero(Duration::ZERO, 0)] + #[case::below_limit(Duration::from_nanos(i64::MAX as u64 - 1), i64::MAX as u64 - 1)] + #[case::at_limit(Duration::from_nanos(i64::MAX as u64), i64::MAX as u64)] + #[case::above_limit(Duration::from_nanos(i64::MAX as u64 + 1), i64::MAX as u64)] + #[case::overflow(Duration::MAX, i64::MAX as u64)] + fn test_duration_as_ns(#[case] duration: Duration, #[case] expected: u64) { + assert_eq!(duration_as_ns(duration), expected); + } + #[test] fn snapshot_query_measurement_with_memory() -> anyhow::Result<()> { let qm = QueryMeasurement { @@ -631,7 +639,7 @@ mod tests { #[test] fn snapshot_compression_size() -> anyhow::Result<()> { - let record = compression_size_record("taxi", None, Format::Lance, 12_345_678); + let record = compression_size_record("taxi", None, Format::Lance, 12_345_678, 98_765_432); assert_snapshot!(render(&record)?); Ok(()) } @@ -663,7 +671,7 @@ mod tests { fn snapshot_random_access_time() -> anyhow::Result<()> { let timing = TimingMeasurement { name: "random-access/taxi/uniform/parquet-tokio-local-disk".to_string(), - target: Target::new(Engine::Arrow, Format::Parquet), + target: Target::new(Engine::Vortex, Format::Parquet), storage: "nvme".to_string(), runs: vec![ Duration::from_nanos(800_000), @@ -671,7 +679,7 @@ mod tests { Duration::from_nanos(850_000), ], }; - let record = random_access_record(&timing, "taxi"); + let record = random_access_record(&timing, "taxi", "cached"); assert_snapshot!(render(&record)?); Ok(()) } @@ -748,13 +756,13 @@ mod tests { #[test] fn compression_records_lowercase_dataset_for_v2_history_match() { - // The v2 → v3 migrate classifier stores `dataset = series.to_lowercase()` - // for compress-bench records (see `benchmarks-website/migrate/src/classifier.rs`). + // The historical v2 backfill stores `dataset = series.to_lowercase()` + // for compress-bench records. // Datasets whose `Dataset::name()` returns mixed case // (`TPC-H l_comment chunked`, every PBI name like `Arade`/`CMSprovider`) // would otherwise emit live records that do not merge with their - // migrated history. Lowercasing inside the v3 helpers keeps the trait - // API simple for non-v3 callers while still matching migrate's shape. + // migrated history. Lowercasing inside these wire helpers keeps the trait + // API simple for other callers while still matching the historical shape. let timing = CompressionTimingMeasurement { name: "compress time/TPC-H l_comment chunked".to_string(), format: Format::OnDiskVortex, @@ -772,7 +780,7 @@ mod tests { }; assert_eq!(time.dataset, "tpc-h l_comment chunked"); - let record = compression_size_record("CMSprovider", None, Format::OnDiskVortex, 42); + let record = compression_size_record("CMSprovider", None, Format::OnDiskVortex, 42, 420); let V3Record::CompressionSize(size) = &record else { panic!("expected CompressionSize variant, got {record:?}"); }; @@ -816,7 +824,7 @@ mod tests { #[test] fn jsonl_round_trips_one_record_per_line() -> anyhow::Result<()> { - let record = compression_size_record("taxi", None, Format::Parquet, 100); + let record = compression_size_record("taxi", None, Format::Parquet, 100, 1_000); let mut buf: Vec = Vec::new(); write_jsonl(&mut buf, &[record.clone(), record])?; let s = String::from_utf8(buf)?; diff --git a/vortex-bench/src/vector_dataset/download.rs b/vortex-bench/src/vector_dataset/download.rs index a7b7d8c6e89..a51102cea47 100644 --- a/vortex-bench/src/vector_dataset/download.rs +++ b/vortex-bench/src/vector_dataset/download.rs @@ -102,11 +102,7 @@ pub async fn download(ds: VectorDataset, layout: TrainLayout) -> Result = train_targets - .iter() - .cloned() - .zip(urls.into_iter()) - .collect(); + let train_downloads: Vec<(PathBuf, String)> = train_targets.iter().cloned().zip(urls).collect(); download_many(train_downloads) .await .with_context(|| format!("download train shards for {}", ds.name()))?; diff --git a/vortex-bench/src/vortex_queries.rs b/vortex-bench/src/vortex_queries.rs index f611238bb6b..63e9a658068 100644 --- a/vortex-bench/src/vortex_queries.rs +++ b/vortex-bench/src/vortex_queries.rs @@ -89,6 +89,10 @@ impl VortexBenchmark { #[async_trait::async_trait] impl Benchmark for VortexBenchmark { + fn doc_path(&self) -> &'static str { + "vortex-bench/sql/vortex/README.md" + } + fn queries(&self) -> Result> { self.query_files()? .iter() diff --git a/vortex-btrblocks/Cargo.toml b/vortex-btrblocks/Cargo.toml index 255d2fbd075..0aa7e9669c3 100644 --- a/vortex-btrblocks/Cargo.toml +++ b/vortex-btrblocks/Cargo.toml @@ -29,7 +29,7 @@ vortex-decimal-byte-parts = { workspace = true } vortex-error = { workspace = true } vortex-fastlanes = { workspace = true } vortex-fsst = { workspace = true } -vortex-onpair = { workspace = true, optional = true } +vortex-onpair = { workspace = true } vortex-pco = { workspace = true, optional = true } vortex-runend = { workspace = true } vortex-sequence = { workspace = true } @@ -39,16 +39,20 @@ vortex-zigzag = { workspace = true } vortex-zstd = { workspace = true, optional = true } [dev-dependencies] +arrow-array = { workspace = true } divan = { workspace = true } insta = { workspace = true } rstest = { workspace = true } test-with = { workspace = true } +tpchgen = { workspace = true } +tpchgen-arrow = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-arrow = { workspace = true } +vortex-edition = { workspace = true } +vortex-mask = { workspace = true } vortex-session = { workspace = true } [features] -# This feature enabled unstable encodings for which we don't guarantee stability. -unstable_encodings = ["dep:vortex-onpair", "vortex-zstd?/unstable_encodings"] pco = ["dep:pco", "dep:vortex-pco"] zstd = ["dep:vortex-zstd"] diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index e51fac9a9d5..c6faaba93d8 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -3,6 +3,7 @@ //! Builder for configuring `BtrBlocksCompressor` instances. +use vortex_array::ArrayId; use vortex_utils::aliases::hash_set::HashSet; use crate::BtrBlocksCompressor; @@ -36,7 +37,6 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ &integer::SequenceScheme, &integer::IntRLEScheme, // Prefer all other schemes above delta, for now (since its slower to decompress). - #[cfg(feature = "unstable_encodings")] &integer::DeltaScheme::new(1.25), //////////////////////////////////////////////////////////////////////////////////////////////// // Float schemes. @@ -53,13 +53,13 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ // Both string-fragmentation schemes are registered; the sample-based // selector keeps whichever is smaller per column. &string::FSSTScheme, - #[cfg(feature = "unstable_encodings")] &string::OnPairScheme, &string::NullDominatedSparseScheme, //////////////////////////////////////////////////////////////////////////////////////////////// // Binary schemes. //////////////////////////////////////////////////////////////////////////////////////////////// &binary::BinaryDictScheme, + &binary::VarBinScheme, // Decimal schemes. &decimal::DecimalScheme, // Temporal schemes. @@ -152,51 +152,43 @@ impl BtrBlocksCompressorBuilder { builder } - /// Excludes schemes without CUDA kernel support and adds Zstd for string and binary compression. + /// Excludes schemes without CUDA kernel support, keeps FSST for string compression, + /// and adds Zstd for binary compression. /// - /// With the `unstable_encodings` feature, buffer-level Zstd compression is used which - /// preserves the array buffer layout for zero-conversion GPU decompression. Without it, - /// interleaved Zstd compression is used. + /// Both the array-level and the buffer-level Zstd schemes are added. Buffer-level + /// compression preserves binary arrays' buffer layout for zero-conversion GPU decompression, + /// but belongs to the opt-in `zstd` edition, so callers filter the two through + /// [`retain_allowed_encodings`](Self::retain_allowed_encodings). /// /// This preset is intended for files that will be decoded by CUDA kernels. It may choose a /// larger encoded representation than the default compressor. pub fn only_cuda_compatible(self) -> Self { - // String fragmentation schemes (OnPair, FSST) require host-side - // dictionary expansion at decode time, which is incompatible with - // pure-GPU decompression paths. Strip whichever string-fragment - // scheme is enabled by feature. - #[cfg_attr( - not(any(feature = "pco", feature = "unstable_encodings")), - allow(unused_mut) - )] + // Keep FSST, which has a CUDA decoder and direct Arrow offset-based export. Other + // string fragmentation and dictionary schemes still require unsupported decode paths. + #[cfg_attr(not(any(feature = "pco", feature = "zstd")), allow(unused_mut))] let mut excluded: Vec = vec![ integer::SparseScheme.id(), integer::IntRLEScheme.id(), float::ALPRDScheme.id(), float::FloatRLEScheme.id(), float::NullDominatedSparseScheme.id(), + string::NullDominatedSparseScheme.id(), string::StringDictScheme.id(), - string::FSSTScheme.id(), binary::BinaryDictScheme.id(), ]; - #[cfg(feature = "unstable_encodings")] - excluded.push(string::OnPairScheme.id()); - // Delta has no GPU decode kernel and its prefix-sum decode is inherently sequential, so it - // is incompatible with pure-GPU decompression paths. - #[cfg(feature = "unstable_encodings")] + // Delta now has a CUDA decode kernel, so arrays that reach the GPU already encoded with + // it — the Delta children OnPair emits, for instance — decode there. It stays excluded + // from this preset until GPU delta decode is benchmarked against the schemes it would + // displace, since the preset picks encodings rather than merely decoding them. excluded.push(integer::DeltaScheme::default().id()); #[cfg(feature = "pco")] excluded.extend([integer::PcoScheme.id(), float::PcoScheme.id()]); let builder = self.exclude_schemes(excluded); - #[cfg(all(feature = "zstd", feature = "unstable_encodings"))] + #[cfg(feature = "zstd")] let builder = builder - .with_new_scheme(&string::ZstdBuffersScheme) + .with_new_scheme(&binary::ZstdScheme) .with_new_scheme(&binary::ZstdBuffersScheme); - #[cfg(all(feature = "zstd", not(feature = "unstable_encodings")))] - let builder = builder - .with_new_scheme(&string::ZstdScheme) - .with_new_scheme(&binary::ZstdScheme); builder } @@ -208,6 +200,16 @@ impl BtrBlocksCompressorBuilder { self } + /// Retains only schemes whose produced encodings all belong to `allowed`. + /// + /// The file writer uses this to restrict compression to the encodings of its configured + /// editions. + pub fn retain_allowed_encodings(mut self, allowed: &HashSet) -> Self { + self.schemes + .retain(|s| s.produced_encodings().iter().all(|id| allowed.contains(id))); + self + } + /// Builds the configured [`BtrBlocksCompressor`]. pub fn build(self) -> BtrBlocksCompressor { BtrBlocksCompressor(CascadingCompressor::new(self.schemes)) @@ -216,6 +218,9 @@ impl BtrBlocksCompressorBuilder { #[cfg(test)] mod tests { + use vortex_array::VTable; + use vortex_fastlanes::FoR; + use super::*; #[test] @@ -230,6 +235,27 @@ mod tests { assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); } + #[test] + fn retain_allowed_encodings_filters_schemes() { + let allowed: HashSet = [FoR.id()].into_iter().collect(); + let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed); + assert_eq!(builder.schemes.len(), 1); + assert_eq!(builder.schemes[0].id(), integer::FoRScheme.id()); + + let none = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&HashSet::new()); + assert!(none.schemes.is_empty()); + } + + #[test] + fn retaining_all_declared_outputs_keeps_every_scheme() { + let allowed: HashSet = ALL_SCHEMES + .iter() + .flat_map(|scheme| scheme.produced_encodings()) + .collect(); + let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed); + assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); + } + #[test] fn cuda_compatible_excludes_alprd() { let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); @@ -241,6 +267,40 @@ mod tests { ); } + /// `vortex.sparse` has no CUDA decode kernel, so no sparse scheme may survive this preset. + #[test] + fn cuda_compatible_excludes_every_sparse_scheme() { + let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); + for excluded in [ + integer::SparseScheme.id(), + float::NullDominatedSparseScheme.id(), + string::NullDominatedSparseScheme.id(), + ] { + assert!( + !builder.schemes.iter().any(|s| s.id() == excluded), + "{excluded} should be excluded" + ); + } + } + + #[test] + fn cuda_compatible_uses_fsst_for_strings() { + let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); + assert!( + builder + .schemes + .iter() + .any(|scheme| scheme.id() == string::FSSTScheme.id()) + ); + #[cfg(feature = "zstd")] + assert!( + !builder + .schemes + .iter() + .any(|scheme| scheme.id() == string::ZstdScheme.id()) + ); + } + #[test] #[cfg(feature = "pco")] fn cuda_compatible_excludes_pco() { diff --git a/vortex-btrblocks/src/canonical_compressor.rs b/vortex-btrblocks/src/canonical_compressor.rs index 9bcee31d990..ed528145d71 100644 --- a/vortex-btrblocks/src/canonical_compressor.rs +++ b/vortex-btrblocks/src/canonical_compressor.rs @@ -71,6 +71,10 @@ mod tests { use std::sync::LazyLock; use rstest::rstest; + #[cfg(feature = "zstd")] + use vortex_array::ArrayId; + #[cfg(feature = "zstd")] + use vortex_array::ArrayPlugin; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; @@ -88,6 +92,8 @@ mod tests { use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; + #[cfg(feature = "zstd")] + use vortex_utils::aliases::hash_set::HashSet; use crate::BtrBlocksCompressor; #[cfg(feature = "zstd")] @@ -272,9 +278,13 @@ mod tests { Ok(()) } - #[cfg(all(feature = "zstd", feature = "unstable_encodings"))] - #[test] - fn test_cuda_compatible_binary_zstd_buffers_compressed() -> VortexResult<()> { + #[cfg(feature = "zstd")] + #[rstest] + #[case::array_level(vortex_zstd::Zstd.id())] + #[case::buffer_level(vortex_zstd::ZstdBuffers.id())] + fn test_cuda_compatible_binary_zstd_follows_editions( + #[case] allowed: ArrayId, + ) -> VortexResult<()> { let values = (0..1024) .map(|idx| { let mut value = Vec::from(&b"common binary payload prefix "[..]); @@ -288,17 +298,16 @@ mod tests { DType::Binary(Nullability::NonNullable), ); + // The CUDA preset carries both Zstd schemes; the edition filter decides which one + // survives. let compressor = BtrBlocksCompressorBuilder::default() .only_cuda_compatible() + .retain_allowed_encodings(&HashSet::from([allowed])) .build(); let mut ctx = SESSION.create_execution_ctx(); let compressed = compressor.compress(&array.clone().into_array(), &mut ctx)?; - assert!( - compressed.is::(), - "expected ZstdBuffers, got {}", - compressed.encoding_id() - ); + assert_eq!(compressed.encoding_id(), allowed); assert_arrays_eq!(compressed, array, &mut ctx); Ok(()) } diff --git a/vortex-btrblocks/src/lib.rs b/vortex-btrblocks/src/lib.rs index 37915b427b6..1ca05c86b4e 100644 --- a/vortex-btrblocks/src/lib.rs +++ b/vortex-btrblocks/src/lib.rs @@ -70,6 +70,9 @@ mod builder; mod canonical_compressor; /// Compression scheme implementations. pub mod schemes; +#[cfg(test)] +#[cfg(not(codspeed))] +mod trace_tests; // Re-export framework types from vortex-compressor for backwards compatibility. // Btrblocks-specific exports. diff --git a/vortex-btrblocks/src/schemes/binary/mod.rs b/vortex-btrblocks/src/schemes/binary/mod.rs index af66d345cc6..d83b6eef919 100644 --- a/vortex-btrblocks/src/schemes/binary/mod.rs +++ b/vortex-btrblocks/src/schemes/binary/mod.rs @@ -3,14 +3,16 @@ //! Binary compression schemes. +mod varbin; #[cfg(feature = "zstd")] mod zstd; -#[cfg(all(feature = "zstd", feature = "unstable_encodings"))] +#[cfg(feature = "zstd")] mod zstd_buffers; // Re-export builtin schemes from vortex-compressor. +pub use varbin::VarBinScheme; pub use vortex_compressor::builtins::BinaryDictScheme; #[cfg(feature = "zstd")] pub use zstd::ZstdScheme; -#[cfg(all(feature = "zstd", feature = "unstable_encodings"))] +#[cfg(feature = "zstd")] pub use zstd_buffers::ZstdBuffersScheme; diff --git a/vortex-btrblocks/src/schemes/binary/varbin.rs b/vortex-btrblocks/src/schemes/binary/varbin.rs new file mode 100644 index 00000000000..849402a6493 --- /dev/null +++ b/vortex-btrblocks/src/schemes/binary/varbin.rs @@ -0,0 +1,100 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Offset-based storage for binary arrays. +//! +//! Canonical binary arrays are [`VarBinViewArray`](vortex_array::arrays::VarBinViewArray), which +//! spends a fixed 16 bytes per element on an opaque views buffer that no scheme can compress. +//! Re-encoding as [`VarBinArray`] replaces that buffer with an offsets child array, which the +//! cascading compressor can then compress with the ordinary integer schemes. For fixed-width +//! values the offsets are a constant-stride sequence and collapse to nothing. + +use vortex_array::ArrayId; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::VarBin; +use vortex_array::arrays::VarBinArray; +use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_array::arrays::varbin::VarBinArraySlotsExt; +use vortex_array::builders::VarBinBuilder; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::SchemeExt; +use vortex_error::VortexResult; + +use crate::ArrayAndStats; +use crate::CascadingCompressor; +use crate::CompressorContext; +use crate::Scheme; + +/// Offset-based (rather than view-based) storage for binary arrays. +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub struct VarBinScheme; + +impl Scheme for VarBinScheme { + fn scheme_name(&self) -> &'static str { + "vortex.binary.varbin" + } + + fn matches(&self, canonical: &Canonical) -> bool { + canonical.dtype().is_binary() + } + + fn produced_encodings(&self) -> Vec { + vec![VarBin.id()] + } + + fn num_children(&self) -> usize { + 1 + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Deferred(DeferredEstimate::Sample) + } + + fn compress( + &self, + compressor: &CascadingCompressor, + data: &ArrayAndStats, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + // `append_to_builder` resolves the views slice and data buffers once and appends + // borrowed slices into a single pre-sized allocation. Iterating the array per element + // instead would clone a buffer handle and allocate for every value. + let array = data.array(); + let mut builder = VarBinBuilder::::with_capacity_in( + array.dtype().clone(), + array.len(), + exec_ctx.allocator(), + ); + array.append_to_builder(&mut builder, exec_ctx)?; + let varbin = builder.finish_into_varbin(); + + let offsets = varbin + .offsets() + .clone() + .execute::(exec_ctx)? + .narrow(exec_ctx)? + .into_array(); + let compressed_offsets = + compressor.compress_child(&offsets, &compress_ctx, self.id(), 0, exec_ctx)?; + + Ok(VarBinArray::try_new( + compressed_offsets, + varbin.bytes().clone(), + varbin.dtype().clone(), + varbin.validity()?, + )? + .into_array()) + } +} diff --git a/vortex-btrblocks/src/schemes/binary/zstd.rs b/vortex-btrblocks/src/schemes/binary/zstd.rs index 83c09d24e94..d652e344db2 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd.rs @@ -3,10 +3,12 @@ //! Zstd compression for binary arrays. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; @@ -29,6 +31,10 @@ impl Scheme for ZstdScheme { canonical.dtype().is_binary() } + fn produced_encodings(&self) -> Vec { + vec![vortex_zstd::Zstd.id()] + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, diff --git a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs index a84672ef860..3f06d65b061 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs @@ -3,10 +3,12 @@ //! Zstd buffer-level binary compression preserving array layout for GPU decompression. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; @@ -29,6 +31,10 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_binary() } + fn produced_encodings(&self) -> Vec { + vec![vortex_zstd::ZstdBuffers.id()] + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, diff --git a/vortex-btrblocks/src/schemes/decimal.rs b/vortex-btrblocks/src/schemes/decimal.rs index 20c5cd80a3c..1dff2171f60 100644 --- a/vortex-btrblocks/src/schemes/decimal.rs +++ b/vortex-btrblocks/src/schemes/decimal.rs @@ -3,10 +3,12 @@ //! Decimal compression scheme using byte-part decomposition. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::decimal::narrowed_decimal; @@ -38,6 +40,10 @@ impl Scheme for DecimalScheme { matches!(canonical, Canonical::Decimal(_)) } + fn produced_encodings(&self) -> Vec { + vec![DecimalByteParts.id()] + } + /// Children: primitive=0. fn num_children(&self) -> usize { 1 diff --git a/vortex-btrblocks/src/schemes/float/alp.rs b/vortex-btrblocks/src/schemes/float/alp.rs index b8a26abf348..f9fc7066bf4 100644 --- a/vortex-btrblocks/src/schemes/float/alp.rs +++ b/vortex-btrblocks/src/schemes/float/alp.rs @@ -7,10 +7,12 @@ use vortex_alp::ALP; use vortex_alp::ALPArrayExt; use vortex_alp::ALPArraySlotsExt; use vortex_alp::alp_encode; +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::arrays::Patched; use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::arrays::primitive::PrimitiveArrayExt; @@ -40,6 +42,14 @@ impl Scheme for ALPScheme { canonical.dtype().is_float() } + fn produced_encodings(&self) -> Vec { + let mut encodings = vec![ALP.id()]; + if use_experimental_patches() { + encodings.push(Patched.id()); + } + encodings + } + /// Children: encoded_ints=0. fn num_children(&self) -> usize { 1 diff --git a/vortex-btrblocks/src/schemes/float/alprd.rs b/vortex-btrblocks/src/schemes/float/alprd.rs index 662a47a1d9c..09c0ee85b0d 100644 --- a/vortex-btrblocks/src/schemes/float/alprd.rs +++ b/vortex-btrblocks/src/schemes/float/alprd.rs @@ -6,10 +6,13 @@ use vortex_alp::ALPRDArrayExt; use vortex_alp::ALPRDArrayOwnedExt; use vortex_alp::RDEncoder; +use vortex_alp::RDEncoderExt; +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::dtype::PType; use vortex_compressor::scheme::CompressionEstimate; @@ -37,6 +40,10 @@ impl Scheme for ALPRDScheme { canonical.dtype().is_float() } + fn produced_encodings(&self) -> Vec { + vec![vortex_alp::ALPRD.id()] + } + fn expected_compression_ratio( &self, data: &ArrayAndStats, diff --git a/vortex-btrblocks/src/schemes/float/pco.rs b/vortex-btrblocks/src/schemes/float/pco.rs index dc9a96133d5..416668c2fd0 100644 --- a/vortex-btrblocks/src/schemes/float/pco.rs +++ b/vortex-btrblocks/src/schemes/float/pco.rs @@ -3,10 +3,12 @@ //! Pco (pcodec) float compression. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; @@ -29,6 +31,10 @@ impl Scheme for PcoScheme { canonical.dtype().is_float() } + fn produced_encodings(&self) -> Vec { + vec![vortex_pco::Pco.id()] + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, diff --git a/vortex-btrblocks/src/schemes/float/rle.rs b/vortex-btrblocks/src/schemes/float/rle.rs index c2683ed131c..71158b9dc3b 100644 --- a/vortex-btrblocks/src/schemes/float/rle.rs +++ b/vortex-btrblocks/src/schemes/float/rle.rs @@ -3,15 +3,18 @@ //! Run-length float encoding. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; +use vortex_array::VTable; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::DescendantExclusion; use vortex_compressor::scheme::EstimateVerdict; use vortex_error::VortexResult; +use vortex_fastlanes::RLE; use crate::ArrayAndStats; use crate::CascadingCompressor; @@ -35,6 +38,10 @@ impl Scheme for FloatRLEScheme { canonical.dtype().is_float() } + fn produced_encodings(&self) -> Vec { + vec![RLE.id()] + } + /// Children: values=0, indices=1, offsets=2. fn num_children(&self) -> usize { 3 diff --git a/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs index 4c7a5f85fa6..c6f54a9dafe 100644 --- a/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs @@ -59,7 +59,11 @@ fn test_dict_compressed() -> VortexResult<()> { #[test] fn test_null_dominated_compressed() -> VortexResult<()> { - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::Nullable, 100); + let mut builder = PrimitiveBuilder::::with_capacity_in( + Nullability::Nullable, + 100, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); for i in 0..5 { builder.append_value(i as f64); } diff --git a/vortex-btrblocks/src/schemes/float/sparse.rs b/vortex-btrblocks/src/schemes/float/sparse.rs index 42d09e323c5..3d9c25b18e4 100644 --- a/vortex-btrblocks/src/schemes/float/sparse.rs +++ b/vortex-btrblocks/src/schemes/float/sparse.rs @@ -3,10 +3,12 @@ //! Sparse encoding for null-dominated float arrays. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_compressor::scheme::ChildSelection; @@ -39,6 +41,10 @@ impl Scheme for NullDominatedSparseScheme { canonical.dtype().is_float() } + fn produced_encodings(&self) -> Vec { + vec![Sparse.id()] + } + /// Children: indices=0. fn num_children(&self) -> usize { 1 diff --git a/vortex-btrblocks/src/schemes/float/tests.rs b/vortex-btrblocks/src/schemes/float/tests.rs index 2d0a04542cc..bb5301b807b 100644 --- a/vortex-btrblocks/src/schemes/float/tests.rs +++ b/vortex-btrblocks/src/schemes/float/tests.rs @@ -50,7 +50,8 @@ fn test_compress() -> VortexResult<()> { .display_as(DisplayOptions::MetadataOnly) .to_string() .to_lowercase(); - assert_eq!(display, "vortex.dict(f32, len=1024)"); + assert_arrays_eq!(compressed, array, &mut SESSION.create_execution_ctx()); + assert_eq!(display, "vortex.alp(f32, len=1024)"); Ok(()) } @@ -77,7 +78,11 @@ fn test_rle_compression() -> VortexResult<()> { #[test] fn test_sparse_compression() -> VortexResult<()> { - let mut array = PrimitiveBuilder::::with_capacity(Nullability::Nullable, 100); + let mut array = PrimitiveBuilder::::with_capacity_in( + Nullability::Nullable, + 100, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); array.append_value(f32::NAN); array.append_value(-f32::NAN); array.append_value(f32::INFINITY); diff --git a/vortex-btrblocks/src/schemes/integer/bitpacking.rs b/vortex-btrblocks/src/schemes/integer/bitpacking.rs index 2424402649e..5ac7d0e4078 100644 --- a/vortex-btrblocks/src/schemes/integer/bitpacking.rs +++ b/vortex-btrblocks/src/schemes/integer/bitpacking.rs @@ -3,10 +3,12 @@ //! BitPacking integer encoding. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::arrays::Patched; use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::arrays::primitive::PrimitiveArrayExt; @@ -38,6 +40,14 @@ impl Scheme for BitPackingScheme { canonical.dtype().is_int() } + fn produced_encodings(&self) -> Vec { + let mut encodings = vec![BitPacked.id()]; + if use_experimental_patches() { + encodings.push(Patched.id()); + } + encodings + } + fn expected_compression_ratio( &self, data: &ArrayAndStats, diff --git a/vortex-btrblocks/src/schemes/integer/delta.rs b/vortex-btrblocks/src/schemes/integer/delta.rs index ed087f5d845..46b2f1e302e 100644 --- a/vortex-btrblocks/src/schemes/integer/delta.rs +++ b/vortex-btrblocks/src/schemes/integer/delta.rs @@ -3,10 +3,12 @@ //! FastLanes Delta integer encoding. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; @@ -21,6 +23,7 @@ use vortex_compressor::scheme::EstimateScore; use vortex_compressor::scheme::EstimateVerdict; use vortex_error::VortexResult; use vortex_fastlanes::Delta; +use vortex_fastlanes::FL_CHUNK_SIZE; use crate::ArrayAndStats; use crate::CascadingCompressor; @@ -69,6 +72,22 @@ const DELTA_PENALTY: f64 = 0.95; /// Minimum length before Delta is worth considering (one FastLanes chunk). const MIN_DELTA_LEN: usize = 1024; +/// Penalized compression ratio for `len` values of `full_width` bits whose transposed residuals +/// need `delta_bits` bits each. +/// +/// Delta writes two children and both are sized by the padded length, since `delta_compress` +/// rounds the input up to whole [`FL_CHUNK_SIZE`] chunks. The deltas child holds `padded_len` +/// residuals of `delta_bits`; the bases child holds one base per lane per chunk, and a chunk of +/// `W`-bit values has `FL_CHUNK_SIZE / W` lanes, so its bases occupy [`FL_CHUNK_SIZE`] bits per +/// chunk whatever `W` is: one further bit per residual. Charging the bases at full width is an +/// upper bound, as the cascade compresses that child like any other. +/// +/// `len` must be non-zero; callers gate on [`MIN_DELTA_LEN`]. +fn penalized_ratio(len: usize, full_width: f64, delta_bits: f64) -> f64 { + let padded_len = len.next_multiple_of(FL_CHUNK_SIZE) as f64; + (len as f64) * full_width / (padded_len * (delta_bits + 1.0)) * DELTA_PENALTY +} + impl Scheme for DeltaScheme { fn scheme_name(&self) -> &'static str { "vortex.int.delta" @@ -78,6 +97,10 @@ impl Scheme for DeltaScheme { canonical.dtype().is_int() } + fn produced_encodings(&self) -> Vec { + vec![Delta.id()] + } + fn num_children(&self) -> usize { 2 } @@ -136,11 +159,12 @@ impl Scheme for DeltaScheme { move |_compressor, data, best_so_far, _ctx, exec_ctx| { let primitive = data.array().clone().execute::(exec_ctx)?; let full_width = primitive.ptype().bit_width() as f64; + let len = primitive.len(); // Delta's best case is residuals collapsing to a single bit. If even that, after // the penalty, can't beat the incumbent, skip before doing the encode work. let threshold = best_so_far.and_then(EstimateScore::finite_ratio); - if threshold.is_some_and(|t| full_width * DELTA_PENALTY <= t) { + if threshold.is_some_and(|t| penalized_ratio(len, full_width, 1.0) <= t) { return Ok(EstimateVerdict::Skip); } @@ -159,7 +183,7 @@ impl Scheme for DeltaScheme { None => return Ok(EstimateVerdict::Skip), }; - let ratio = full_width / delta_bits * DELTA_PENALTY; + let ratio = penalized_ratio(len, full_width, delta_bits); if ratio <= min_ratio { return Ok(EstimateVerdict::Skip); } diff --git a/vortex-btrblocks/src/schemes/integer/for_.rs b/vortex-btrblocks/src/schemes/integer/for_.rs index 832ed4819fd..476a0dec282 100644 --- a/vortex-btrblocks/src/schemes/integer/for_.rs +++ b/vortex-btrblocks/src/schemes/integer/for_.rs @@ -3,10 +3,12 @@ //! Frame of Reference integer encoding. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; @@ -20,6 +22,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_fastlanes::FoR; use vortex_fastlanes::FoRArrayExt; +use vortex_fastlanes::FoRArraySlotsExt; use super::BitPackingScheme; use crate::ArrayAndStats; @@ -41,6 +44,10 @@ impl Scheme for FoRScheme { canonical.dtype().is_int() } + fn produced_encodings(&self) -> Vec { + vec![FoR.id()] + } + /// Dict codes always start at 0, so FoR (which subtracts the min) is a no-op. fn ancestor_exclusions(&self) -> Vec { vec![ diff --git a/vortex-btrblocks/src/schemes/integer/mod.rs b/vortex-btrblocks/src/schemes/integer/mod.rs index 3aae2ae5601..f9568cfcd1f 100644 --- a/vortex-btrblocks/src/schemes/integer/mod.rs +++ b/vortex-btrblocks/src/schemes/integer/mod.rs @@ -4,7 +4,6 @@ //! Integer compression schemes. mod bitpacking; -#[cfg(feature = "unstable_encodings")] mod delta; mod for_; mod rle; @@ -17,14 +16,12 @@ mod zigzag; mod pco; pub use bitpacking::BitPackingScheme; -#[cfg(feature = "unstable_encodings")] pub use delta::DeltaScheme; pub use for_::FoRScheme; #[cfg(feature = "pco")] pub use pco::PcoScheme; pub use rle::IntRLEScheme; pub(crate) use rle::rle_compress; -#[cfg(feature = "unstable_encodings")] pub(crate) use rle::try_compress_delta; pub use runend::RunEndScheme; pub use sequence::SequenceScheme; diff --git a/vortex-btrblocks/src/schemes/integer/pco.rs b/vortex-btrblocks/src/schemes/integer/pco.rs index d7f182f588e..675a112d44f 100644 --- a/vortex-btrblocks/src/schemes/integer/pco.rs +++ b/vortex-btrblocks/src/schemes/integer/pco.rs @@ -3,10 +3,12 @@ //! Pco (pcodec) integer compression. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::EstimateVerdict; @@ -30,6 +32,10 @@ impl Scheme for PcoScheme { canonical.dtype().is_int() } + fn produced_encodings(&self) -> Vec { + vec![vortex_pco::Pco.id()] + } + fn expected_compression_ratio( &self, data: &ArrayAndStats, diff --git a/vortex-btrblocks/src/schemes/integer/rle.rs b/vortex-btrblocks/src/schemes/integer/rle.rs index 5bc4e5296f1..0909e1428d4 100644 --- a/vortex-btrblocks/src/schemes/integer/rle.rs +++ b/vortex-btrblocks/src/schemes/integer/rle.rs @@ -3,10 +3,12 @@ //! Run-length integer encoding and shared RLE compression helpers. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_compressor::scheme::AncestorExclusion; @@ -14,14 +16,14 @@ use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::DescendantExclusion; use vortex_compressor::scheme::EstimateVerdict; -#[cfg(feature = "unstable_encodings")] use vortex_compressor::scheme::SchemeId; use vortex_error::VortexResult; -#[cfg(feature = "unstable_encodings")] use vortex_fastlanes::Delta; use vortex_fastlanes::RLE; use vortex_fastlanes::RLEArrayExt; +use vortex_fastlanes::RLEArraySlotsExt; +use super::DeltaScheme; use super::RUN_LENGTH_THRESHOLD; use crate::ArrayAndStats; use crate::CascadingCompressor; @@ -57,38 +59,25 @@ pub(crate) fn rle_compress( exec_ctx, )?; - // Delta is an unstable encoding, once we deem it stable we can switch over to this always. - #[cfg(feature = "unstable_encodings")] let compressed_indices = { let rle_indices_primitive = rle_array .indices() .clone() .execute::(exec_ctx)? .narrow(exec_ctx)?; - try_compress_delta( - compressor, - &rle_indices_primitive.into_array(), - &compress_ctx, - scheme.id(), - 1, - exec_ctx, - )? - }; - - #[cfg(not(feature = "unstable_encodings"))] - let compressed_indices = { - let rle_indices_primitive = rle_array - .indices() - .clone() - .execute::(exec_ctx)? - .narrow(exec_ctx)?; - compressor.compress_child( - &rle_indices_primitive.into_array(), - &compress_ctx, - scheme.id(), - 1, - exec_ctx, - )? + let rle_indices = rle_indices_primitive.into_array(); + if compressor.has_scheme(DeltaScheme::default().id()) { + try_compress_delta( + compressor, + &rle_indices, + &compress_ctx, + scheme.id(), + 1, + exec_ctx, + )? + } else { + compressor.compress_child(&rle_indices, &compress_ctx, scheme.id(), 1, exec_ctx)? + } }; let rle_offsets_primitive = rle_array @@ -117,7 +106,6 @@ pub(crate) fn rle_compress( } } -#[cfg(feature = "unstable_encodings")] pub(crate) fn try_compress_delta( compressor: &CascadingCompressor, child: &ArrayRef, @@ -156,6 +144,10 @@ impl Scheme for IntRLEScheme { canonical.dtype().is_int() } + fn produced_encodings(&self) -> Vec { + vec![RLE.id()] + } + /// Children: values=0, indices=1, offsets=2. fn num_children(&self) -> usize { 3 diff --git a/vortex-btrblocks/src/schemes/integer/runend.rs b/vortex-btrblocks/src/schemes/integer/runend.rs index 175f33e7406..6a97f7ec37d 100644 --- a/vortex-btrblocks/src/schemes/integer/runend.rs +++ b/vortex-btrblocks/src/schemes/integer/runend.rs @@ -3,10 +3,12 @@ //! Run-end integer encoding. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; @@ -46,6 +48,10 @@ impl Scheme for RunEndScheme { canonical.dtype().is_int() } + fn produced_encodings(&self) -> Vec { + vec![RunEnd.id()] + } + /// Children: values=0, ends=1. fn num_children(&self) -> usize { 2 diff --git a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs index e4227a472ec..ecd183b8820 100644 --- a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs @@ -158,7 +158,6 @@ fn test_rle_compressed() -> VortexResult<()> { /// (so Sequence skips), all-unique with no runs (so RunEnd/Dict skip), and a wide absolute range. /// Delta's residuals are far smaller than the FoR span, so Delta should win and round-trip, and /// it must appear at most once in the tree. -#[cfg(feature = "unstable_encodings")] #[test] fn test_delta_compressed() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); @@ -195,8 +194,69 @@ fn test_delta_compressed() -> VortexResult<()> { Ok(()) } +/// Same as [`test_delta_compressed`], but with a length that is not a multiple of 1024. +/// Zero-padding the trailing chunk used to inflate the delta span and cause DeltaScheme to skip. +#[test] +fn test_delta_compressed_unaligned_length() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + use vortex_array::assert_arrays_eq; + use vortex_fastlanes::Delta; + + let mut rng = StdRng::seed_from_u64(7u64); + let mut value = 500_000i32; + let values: Vec = (0..1025) + .map(|_| { + value += 1 + (rng.next_u32() % 6) as i32; + value + }) + .collect(); + let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); + + let btr = BtrBlocksCompressor::default(); + let compressed = btr.compress( + &array.clone().into_array(), + &mut SESSION.create_execution_ctx(), + )?; + assert!( + compressed.is::(), + "expected Delta for unaligned near-monotone column, got tree:\n{}", + compressed.display_tree() + ); + assert_arrays_eq!(compressed, array.into_array(), &mut ctx); + Ok(()) +} + +/// Nullable unaligned monotone must round-trip through Delta (and a cascaded residual). +/// +/// Mirrors `duckdb/aggregate_pushdown.slt`: `NULL` then `1..=100000` (length 100001). +#[test] +fn test_delta_nullable_unaligned_sum() -> VortexResult<()> { + use vortex_array::aggregate_fn::fns::sum::sum; + use vortex_array::assert_arrays_eq; + use vortex_fastlanes::Delta; + + let mut ctx = SESSION.create_execution_ctx(); + let array = + PrimitiveArray::from_option_iter(iter::once(None).chain((1i32..=100_000).map(Some))); + + let btr = BtrBlocksCompressor::default(); + let compressed = btr.compress(&array.clone().into_array(), &mut ctx)?; + assert!( + compressed.is::(), + "expected Delta, got tree:\n{}", + compressed.display_tree() + ); + assert_arrays_eq!(compressed, array.into_array(), &mut ctx); + + let expected_sum: i64 = (1i64..=100_000).sum(); + assert_eq!( + sum(&compressed, &mut ctx)?.as_primitive().as_::(), + Some(expected_sum), + ); + Ok(()) +} + /// Returns true if any `Delta` array appears below an ancestor `Delta` in the tree. -#[cfg(feature = "unstable_encodings")] fn has_nested_delta(array: &vortex_array::ArrayRef, under_delta: bool) -> bool { use vortex_fastlanes::Delta; diff --git a/vortex-btrblocks/src/schemes/integer/sequence.rs b/vortex-btrblocks/src/schemes/integer/sequence.rs index fda9995b510..edcefb99fc2 100644 --- a/vortex-btrblocks/src/schemes/integer/sequence.rs +++ b/vortex-btrblocks/src/schemes/integer/sequence.rs @@ -3,9 +3,11 @@ //! Sequence integer encoding for sequential patterns. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; +use vortex_array::VTable; use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; use vortex_compressor::builtins::IntDictScheme; @@ -19,6 +21,7 @@ use vortex_compressor::scheme::EstimateVerdict; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_sequence::Sequence; use vortex_sequence::sequence_encode; use crate::ArrayAndStats; @@ -40,6 +43,10 @@ impl Scheme for SequenceScheme { canonical.dtype().is_int() } + fn produced_encodings(&self) -> Vec { + vec![Sequence.id()] + } + /// Sequence encoding on dictionary codes just adds a layer of indirection without compressing /// the data. Dict codes are compact integers that benefit from BitPacking or FoR, not from /// sequence detection. diff --git a/vortex-btrblocks/src/schemes/integer/sparse.rs b/vortex-btrblocks/src/schemes/integer/sparse.rs index e843d9d7999..429ff5c1a31 100644 --- a/vortex-btrblocks/src/schemes/integer/sparse.rs +++ b/vortex-btrblocks/src/schemes/integer/sparse.rs @@ -3,10 +3,13 @@ //! Sparse integer encoding for single-value-dominated arrays. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; @@ -43,6 +46,10 @@ impl Scheme for SparseScheme { canonical.dtype().is_int() } + fn produced_encodings(&self) -> Vec { + vec![Sparse.id(), Constant.id()] + } + fn stats_options(&self) -> GenerateStatsOptions { GenerateStatsOptions { count_distinct_values: true, diff --git a/vortex-btrblocks/src/schemes/integer/zigzag.rs b/vortex-btrblocks/src/schemes/integer/zigzag.rs index 959889dabac..0e4be01845a 100644 --- a/vortex-btrblocks/src/schemes/integer/zigzag.rs +++ b/vortex-btrblocks/src/schemes/integer/zigzag.rs @@ -3,10 +3,12 @@ //! ZigZag integer encoding for signed integers. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; @@ -20,7 +22,7 @@ use vortex_compressor::scheme::DescendantExclusion; use vortex_compressor::scheme::EstimateVerdict; use vortex_error::VortexResult; use vortex_zigzag::ZigZag; -use vortex_zigzag::ZigZagArrayExt; +use vortex_zigzag::ZigZagArraySlotsExt; use vortex_zigzag::zigzag_encode; use super::RunEndScheme; @@ -44,6 +46,10 @@ impl Scheme for ZigZagScheme { canonical.dtype().is_int() } + fn produced_encodings(&self) -> Vec { + vec![ZigZag.id()] + } + /// Children: encoded=0. fn num_children(&self) -> usize { 1 diff --git a/vortex-btrblocks/src/schemes/string/fsst.rs b/vortex-btrblocks/src/schemes/string/fsst.rs index bd5bd010396..fd3fd28696a 100644 --- a/vortex-btrblocks/src/schemes/string/fsst.rs +++ b/vortex-btrblocks/src/schemes/string/fsst.rs @@ -3,19 +3,26 @@ //! FSST (Fast Static Symbol Table) string compression. +use std::sync::Arc; + +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::VarBin; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; -use vortex_array::arrays::varbin::VarBinArrayExt; +use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; use vortex_fsst::FSST; use vortex_fsst::FSSTArrayExt; +use vortex_fsst::FSSTArraySlotsExt; +use vortex_fsst::FSSTSymbolTable; use vortex_fsst::fsst_compress; use vortex_fsst::fsst_train_compressor; @@ -40,7 +47,11 @@ impl Scheme for FSSTScheme { } fn matches(&self, canonical: &Canonical) -> bool { - canonical.dtype().is_utf8() + canonical.dtype().is_utf8() || canonical.dtype().is_binary() + } + + fn produced_encodings(&self) -> Vec { + vec![FSST.id(), VarBin.id()] } /// Children: lengths=0, code_offsets=1. @@ -101,10 +112,14 @@ impl Scheme for FSSTScheme { fsst.codes().validity()?, )?; - let fsst = FSST::try_new( + // Reuse the padded symbol table as-is; only the codes and lengths change here. + let fsst = FSST::try_new_with_symbol_table( fsst.dtype().clone(), - fsst.symbols().clone(), - fsst.symbol_lengths().clone(), + Arc::new(FSSTSymbolTable::new_padded( + fsst.padded_symbols().clone(), + fsst.padded_symbol_lengths().clone(), + fsst.n_symbols(), + )?), compressed_codes, compressed_original_lengths, exec_ctx, diff --git a/vortex-btrblocks/src/schemes/string/mod.rs b/vortex-btrblocks/src/schemes/string/mod.rs index ac8e5b4b8df..5d04de4777f 100644 --- a/vortex-btrblocks/src/schemes/string/mod.rs +++ b/vortex-btrblocks/src/schemes/string/mod.rs @@ -8,14 +8,12 @@ mod sparse; #[cfg(feature = "zstd")] mod zstd; -#[cfg(all(feature = "zstd", feature = "unstable_encodings"))] +#[cfg(feature = "zstd")] mod zstd_buffers; -#[cfg(feature = "unstable_encodings")] mod onpair; pub use fsst::FSSTScheme; -#[cfg(feature = "unstable_encodings")] pub use onpair::OnPairScheme; pub use sparse::NullDominatedSparseScheme; // Re-export builtin schemes from vortex-compressor. @@ -23,7 +21,7 @@ pub use vortex_compressor::builtins::StringDictScheme; pub use vortex_compressor::stats::StringStats; #[cfg(feature = "zstd")] pub use zstd::ZstdScheme; -#[cfg(all(feature = "zstd", feature = "unstable_encodings"))] +#[cfg(feature = "zstd")] pub use zstd_buffers::ZstdBuffersScheme; #[cfg(test)] diff --git a/vortex-btrblocks/src/schemes/string/onpair.rs b/vortex-btrblocks/src/schemes/string/onpair.rs index f0a8966089c..a1bc8643775 100644 --- a/vortex-btrblocks/src/schemes/string/onpair.rs +++ b/vortex-btrblocks/src/schemes/string/onpair.rs @@ -3,17 +3,19 @@ //! OnPair short-string compression (dict-12). +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::SchemeId; use vortex_error::VortexResult; -use vortex_onpair::DEFAULT_DICT12_CONFIG; +use vortex_onpair::DEFAULT_CONFIG; use vortex_onpair::OnPair; use vortex_onpair::OnPairArrayExt; use vortex_onpair::OnPairArraySlotsExt; @@ -24,6 +26,7 @@ use crate::CascadingCompressor; use crate::CompressorContext; use crate::Scheme; use crate::SchemeExt; +use crate::schemes::integer::DeltaScheme; use crate::schemes::integer::try_compress_delta; /// OnPair short-string compression (dict-12). @@ -47,9 +50,13 @@ impl Scheme for OnPairScheme { canonical.dtype().is_utf8() } + fn produced_encodings(&self) -> Vec { + vec![OnPair.id()] + } + /// 4 primitive slot children flow through the cascading compressor: /// `dict_offsets` (u32 → typically `FoR`/`BitPacked`), `codes` (u16 → - /// `FastLanes::BitPacked` to exactly `bits` = 12 by default), + /// usually `FastLanes::BitPacked` after scheme selection), /// `codes_offsets` (u32 → `FoR`), `uncompressed_lengths` (i32 → narrow /// + `FoR`). Validity stays untouched. fn num_children(&self) -> usize { @@ -73,7 +80,10 @@ impl Scheme for OnPairScheme { exec_ctx: &mut ExecutionCtx, ) -> VortexResult { let utf8 = data.array_as_varbinview().into_owned(); - let onpair_array = onpair_compress(utf8.as_array(), DEFAULT_DICT12_CONFIG, exec_ctx)?; + let encoded = onpair_compress(utf8.as_array(), DEFAULT_CONFIG, exec_ctx)?; + let Some(onpair_array) = encoded.as_opt::() else { + return Ok(encoded); + }; let dict_offsets = compress_offsets_child( compressor, @@ -108,15 +118,14 @@ impl Scheme for OnPairScheme { exec_ctx, )?; - Ok(OnPair::try_new( + Ok(OnPair::try_new_with_data( onpair_array.dtype().clone(), - onpair_array.dict_bytes_handle().clone(), + onpair_array.data().clone(), dict_offsets, codes, codes_offsets, uncompressed_lengths, onpair_array.array_validity(), - onpair_array.bits(), )? .into_array()) } @@ -169,6 +178,9 @@ fn compress_offsets_child( if narrowed.len() < OFFSETS_DELTA_MIN_LEN { return Ok(plain); } + if !compressor.has_scheme(DeltaScheme::default().id()) { + return Ok(plain); + } let delta = try_compress_delta( compressor, &narrowed, diff --git a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs index c1473146607..aac0b4de4de 100644 --- a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs @@ -46,11 +46,10 @@ fn test_dict_compressed() -> VortexResult<()> { Ok(()) } -#[cfg(feature = "unstable_encodings")] #[test] -fn test_onpair_in_default_scheme_list() { +fn test_all_schemes_includes_onpair() { use crate::SchemeExt; - use crate::schemes::string::OnPairScheme; + use crate::schemes::string::onpair::OnPairScheme; let ids: Vec<_> = crate::ALL_SCHEMES.iter().map(|s| s.id()).collect(); assert!( @@ -59,12 +58,11 @@ fn test_onpair_in_default_scheme_list() { ); } -#[cfg(feature = "unstable_encodings")] #[test] -fn test_onpair_compressed() -> VortexResult<()> { +fn test_default_btrblocks_compressor_selects_onpair() -> VortexResult<()> { // Dictionary-style string corpus: high lexical overlap, short rows. // OnPair beats FSST on this corpus, so it wins the sample-based - // comparison even though both are registered by default. + // comparison even though both are registered. let mut strings = Vec::with_capacity(1000); for i in 0..1000 { strings.push(Some(format!( @@ -83,8 +81,8 @@ fn test_onpair_compressed() -> VortexResult<()> { Ok(()) } -/// FSST is registered in the default scheme list (alongside OnPair), and an -/// FSST-only builder still produces an FSST array. +/// FSST is registered in the default scheme list, and an FSST-only builder +/// still produces an FSST array. #[test] fn test_fsst_in_default_scheme_list() -> VortexResult<()> { use crate::BtrBlocksCompressorBuilder; diff --git a/vortex-btrblocks/src/schemes/string/sparse.rs b/vortex-btrblocks/src/schemes/string/sparse.rs index 750dd978030..8620c366f77 100644 --- a/vortex-btrblocks/src/schemes/string/sparse.rs +++ b/vortex-btrblocks/src/schemes/string/sparse.rs @@ -3,10 +3,12 @@ //! Sparse encoding for null-dominated string arrays. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_compressor::scheme::ChildSelection; @@ -40,6 +42,10 @@ impl Scheme for NullDominatedSparseScheme { canonical.dtype().is_utf8() } + fn produced_encodings(&self) -> Vec { + vec![Sparse.id()] + } + /// Children: indices=0. fn num_children(&self) -> usize { 1 diff --git a/vortex-btrblocks/src/schemes/string/tests.rs b/vortex-btrblocks/src/schemes/string/tests.rs index ec9c4060f64..1928f0065a2 100644 --- a/vortex-btrblocks/src/schemes/string/tests.rs +++ b/vortex-btrblocks/src/schemes/string/tests.rs @@ -45,7 +45,11 @@ fn test_strings() -> VortexResult<()> { #[test] fn test_sparse_nulls() -> VortexResult<()> { - let mut strings = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 100); + let mut strings = VarBinViewBuilder::with_capacity_in( + DType::Utf8(Nullability::Nullable), + 100, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); strings.append_nulls(99); strings.append_value("one little string"); diff --git a/vortex-btrblocks/src/schemes/string/zstd.rs b/vortex-btrblocks/src/schemes/string/zstd.rs index 177acfa1543..84e8860d626 100644 --- a/vortex-btrblocks/src/schemes/string/zstd.rs +++ b/vortex-btrblocks/src/schemes/string/zstd.rs @@ -3,10 +3,12 @@ //! Zstd string compression without dictionaries (nvCOMP compatible). +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; @@ -29,6 +31,10 @@ impl Scheme for ZstdScheme { canonical.dtype().is_utf8() } + fn produced_encodings(&self) -> Vec { + vec![vortex_zstd::Zstd.id()] + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, diff --git a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs index a8d16cc837f..cf691c70fcb 100644 --- a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs @@ -3,10 +3,12 @@ //! Zstd buffer-level string compression preserving array layout for GPU decompression. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; @@ -29,6 +31,10 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_utf8() } + fn produced_encodings(&self) -> Vec { + vec![vortex_zstd::ZstdBuffers.id()] + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, diff --git a/vortex-btrblocks/src/schemes/temporal.rs b/vortex-btrblocks/src/schemes/temporal.rs index 6989b6a2aba..79748b69450 100644 --- a/vortex-btrblocks/src/schemes/temporal.rs +++ b/vortex-btrblocks/src/schemes/temporal.rs @@ -3,10 +3,12 @@ //! Temporal compression scheme using datetime-part decomposition. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::TemporalArray; @@ -53,6 +55,10 @@ impl Scheme for TemporalScheme { ) } + fn produced_encodings(&self) -> Vec { + vec![DateTimeParts.id()] + } + /// Children: days=0, seconds=1, subseconds=2. fn num_children(&self) -> usize { 3 diff --git a/vortex-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs new file mode 100644 index 00000000000..21e440a7436 --- /dev/null +++ b/vortex-btrblocks/src/trace_tests.rs @@ -0,0 +1,455 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Snapshot tests tracing how TPC-H style scan operations reduce and execute over +//! BtrBlocks-compressed lineitem columns. +//! +//! The lineitem table (SF 0.001, deterministic) compresses to the same encodings a real +//! file scan would see: bitpacked integers, `decimal_byte_parts` decimals with dictionary +//! or bitpacked byte parts, `ext(date)` over FoR + bitpacking, dictionary-of-FSST for +//! low-cardinality strings, and FSST for comments. The tests below pin down which reduce +//! rules and execute kernels fire for the scan operations TPC-H queries perform over those +//! encodings. + +use std::sync::LazyLock; + +use arrow_array::RecordBatch; +use tpchgen::distribution::Distributions; +use tpchgen::generators::LineItemGenerator; +use tpchgen::text::TextPool; +use tpchgen_arrow::LineItemArrow; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::FilterArray; +use vortex_array::arrays::Patched; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::Struct; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::patched::use_experimental_patches; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::optimizer::ArrayOptimizer; +use vortex_array::scalar::DecimalValue; +use vortex_array::scalar::PValue; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::fns::binary::Binary; +use vortex_array::scalar_fn::fns::like::Like; +use vortex_array::scalar_fn::fns::like::LikeOptions; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::session::ArraySession; +use vortex_array::session::ArraySessionExt; +use vortex_array::test_harness::trace::Traced; +use vortex_array::test_harness::trace::trace_op; +use vortex_arrow::ArrowSessionExt; +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +use crate::BtrBlocksCompressor; + +/// A session with the default Vortex encodings registered. +/// +/// Reduce-parent and execute-parent kernels live in the session registry, so an encoding whose +/// `initialize` was never called still decodes correctly but silently contributes none of its +/// pushdown kernels — the traces below would degrade to plain canonicalization without failing +/// any value assertion. +/// +/// This mirrors `vortex_file::register_default_encodings`, copied rather than called so that +/// `vortex-btrblocks` does not depend on `vortex-file` (which depends on this crate). Keep the +/// two in step when encodings are added. `bytebool` and `tensor` are the only entries omitted: +/// this crate does not depend on them, so the compressor cannot emit them. +fn trace_session() -> VortexSession { + let session = VortexSession::empty().with::(); + + vortex_fsst::initialize(&session); + vortex_onpair::initialize(&session); + vortex_zigzag::initialize(&session); + #[cfg(feature = "zstd")] + vortex_zstd::initialize(&session); + + { + let arrays = session.arrays(); + #[cfg(feature = "pco")] + arrays.register(vortex_pco::Pco); + if use_experimental_patches() { + arrays.register(Patched); + } + } + + vortex_alp::initialize(&session); + vortex_datetime_parts::initialize(&session); + vortex_decimal_byte_parts::initialize(&session); + vortex_fastlanes::initialize(&session); + vortex_runend::initialize(&session); + vortex_sequence::initialize(&session); + vortex_sparse::initialize(&session); + + session +} + +fn execution_ctx() -> ExecutionCtx { + ExecutionCtx::new(trace_session()) +} + +/// A 1MiB TPC-H text pool. The spec-default 300MiB pool takes several seconds to initialize in +/// debug builds; a smaller pool keeps the same text distribution and is deterministic. +static TEXT_POOL: LazyLock = + LazyLock::new(|| TextPool::new(1 << 20, Distributions::static_default())); + +/// The first 4096 rows of TPC-H lineitem at scale factor 0.001. Both tpchgen and the +/// compressor (fixed sampling seed) are deterministic, so the resulting encodings are stable. +fn lineitem() -> VortexResult { + let generator = LineItemGenerator::new_with_distributions_and_text_pool( + 0.001, + 1, + 1, + Distributions::static_default(), + &TEXT_POOL, + ); + let batch: RecordBatch = LineItemArrow::new(generator) + .with_batch_size(1 << 12) + .next() + .expect("at least one batch"); + let schema = batch.schema(); + trace_session() + .arrow() + .from_arrow_record_batch(batch, &schema) +} + +fn compressed_lineitem() -> VortexResult { + BtrBlocksCompressor::default().compress(&lineitem()?, &mut execution_ctx()) +} + +fn field(array: &ArrayRef, name: &str) -> VortexResult { + Ok(array.as_::().unmasked_field_by_name(name)?.clone()) +} + +/// Trace the optimize pass and then the execution of `array`, asserting that the canonical +/// result matches running the same operation over the uncompressed column. +fn optimize_then_execute( + array: ArrayRef, + expected: &ArrayRef, +) -> VortexResult<[Traced; 2]> { + let optimized = trace_op(|| array.optimize())?; + let optimized_output = optimized.output.clone(); + let executed = trace_op(|| { + optimized_output + .execute::(&mut execution_ctx()) + .map(IntoArray::into_array) + })?; + let expected = expected + .clone() + .execute::(&mut execution_ctx())? + .into_array(); + assert_arrays_eq!(executed.output, expected, &mut execution_ctx()); + Ok([ + Traced { + output: optimized.output, + trace: optimized.trace, + }, + executed, + ]) +} + +/// Q6-style predicate over the shipdate column: `l_shipdate >= DATE '1994-01-01'`. +/// +/// The column compresses to `ext(date) -> for -> bitpacked`. +fn shipdate_predicate(column: ArrayRef, len: usize) -> VortexResult { + let DType::Extension(ext) = column.dtype().clone() else { + panic!("expected extension dtype for l_shipdate") + }; + // 8766 days since the epoch = 1994-01-01. + let cutoff = Scalar::extension_ref( + ext, + Scalar::primitive_value(PValue::I32(8766), PType::I32, Nullability::NonNullable), + ); + Binary::try_new( + column, + ConstantArray::new(cutoff, len).into_array(), + Operator::Gte, + ) + .map(IntoArray::into_array) +} + +#[test] +fn trace_scan_compare_on_compressed_shipdate() -> VortexResult<()> { + let compressed = compressed_lineitem()?; + let len = compressed.len(); + + let lazy = shipdate_predicate(field(&compressed, "l_shipdate")?, len)?; + let expected = shipdate_predicate(field(&lineitem()?, "l_shipdate")?, len)?; + let [optimized, executed] = optimize_then_execute(lazy, &expected)?; + + // No reduce rule rewrites a comparison over the extension array; the extension compare + // kernel handles it at execution time by comparing the underlying storage. + insta::assert_snapshot!(optimized.trace.to_string(), @""); + insta::assert_snapshot!(executed.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.binary(bool, len=4096) + iter 0 current=vortex.binary(bool, len=4096) builder_active=false + child_execute_parent session[0]:execute_parent_fn slot=0 parent=vortex.binary(bool, len=4096) child=vortex.ext(vortex.date[days](i32), len=4096) -> vortex.binary(bool, len=4096) + iter 1 current=vortex.binary(bool, len=4096) builder_active=false + execute_until target=AnyCanonical root=fastlanes.for(i32, len=4096) + iter 0 current=fastlanes.for(i32, len=4096) builder_active=false + execute_until target=AnyCanonical root=fastlanes.bitpacked(i32, len=4096) + iter 0 current=fastlanes.bitpacked(i32, len=4096) builder_active=false + Done array=vortex.primitive(i32, len=4096) + iter 1 current=vortex.primitive(i32, len=4096) builder_active=false + return output=vortex.primitive(i32, len=4096) + Done array=vortex.primitive(i32, len=4096) + iter 1 current=vortex.primitive(i32, len=4096) builder_active=false + return output=vortex.primitive(i32, len=4096) + Done array=vortex.bool(bool, len=4096) + iter 2 current=vortex.bool(bool, len=4096) builder_active=false + return output=vortex.bool(bool, len=4096) + "); + + Ok(()) +} + +/// Q6-style predicate over the quantity column: `l_quantity < 24`. +/// +/// The column compresses to `decimal_byte_parts -> dict -> bitpacked/sequence`. +fn quantity_predicate(column: ArrayRef, len: usize) -> VortexResult { + let cutoff = Scalar::decimal( + DecimalValue::I128(2400), + DecimalDType::new(15, 2), + Nullability::NonNullable, + ); + Binary::try_new( + column, + ConstantArray::new(cutoff, len).into_array(), + Operator::Lt, + ) + .map(IntoArray::into_array) +} + +#[test] +fn trace_scan_compare_on_compressed_quantity() -> VortexResult<()> { + let compressed = compressed_lineitem()?; + let len = compressed.len(); + + let lazy = quantity_predicate(field(&compressed, "l_quantity")?, len)?; + let expected = quantity_predicate(field(&lineitem()?, "l_quantity")?, len)?; + let [optimized, executed] = optimize_then_execute(lazy, &expected)?; + + // No reduce rule rewrites a comparison over decimal_byte_parts; its compare kernel fires + // at execution time and pushes the comparison into the byte-parts dictionary, whose values + // are then compared and decoded. + insta::assert_snapshot!(optimized.trace.to_string(), @""); + insta::assert_snapshot!(executed.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.binary(bool, len=4096) + iter 0 current=vortex.binary(bool, len=4096) builder_active=false + optimize root=vortex.binary(bool, len=4096) session=false + reduce_parent static:DictionaryScalarFnValuesPushDownRule slot=0 parent=vortex.binary(bool, len=4096) child=vortex.dict(i16, len=4096) -> vortex.dict(bool, len=4096) + done output=vortex.dict(bool, len=4096) + child_execute_parent session[0]:execute_parent_fn slot=0 parent=vortex.binary(bool, len=4096) child=vortex.decimal_byte_parts(decimal(15,2), len=4096) -> vortex.dict(bool, len=4096) + iter 1 current=vortex.dict(bool, len=4096) builder_active=false + ExecuteSlot slot=0 parent=vortex.dict(bool, len=4096) child=fastlanes.bitpacked(u8, len=4096) + iter 2 current=fastlanes.bitpacked(u8, len=4096) stack_parent=vortex.dict(bool, len=4096) slot=0 builder_active=false + Done array=vortex.primitive(u8, len=4096) + iter 3 current=vortex.primitive(u8, len=4096) stack_parent=vortex.dict(bool, len=4096) slot=0 builder_active=false + pop_frame slot=0 output=vortex.dict(bool, len=4096) + iter 4 current=vortex.dict(bool, len=4096) builder_active=false + ExecuteSlot slot=1 parent=vortex.dict(bool, len=4096) child=vortex.binary(bool, len=50) + iter 5 current=vortex.binary(bool, len=50) stack_parent=vortex.dict(bool, len=4096) slot=1 builder_active=false + execute_until target=AnyCanonical root=vortex.sequence(i16, len=50) + iter 0 current=vortex.sequence(i16, len=50) builder_active=false + Done array=vortex.primitive(i16, len=50) + iter 1 current=vortex.primitive(i16, len=50) builder_active=false + return output=vortex.primitive(i16, len=50) + Done array=vortex.bool(bool, len=50) + iter 6 current=vortex.bool(bool, len=50) stack_parent=vortex.dict(bool, len=4096) slot=1 builder_active=false + pop_frame slot=1 output=vortex.dict(bool, len=4096) + iter 7 current=vortex.dict(bool, len=4096) builder_active=false + child_execute_parent session[0]:execute_parent_fn slot=1 parent=vortex.dict(bool, len=4096) child=vortex.bool(bool, len=50) -> vortex.bool(bool, len=4096) + iter 8 current=vortex.bool(bool, len=4096) builder_active=false + return output=vortex.bool(bool, len=4096) + "); + + Ok(()) +} + +/// Q12-style predicate over the shipmode column: `l_shipmode = 'AIR'`. +/// +/// The column compresses to `dict -> {bitpacked codes, fsst values}`. +fn shipmode_predicate(column: ArrayRef, len: usize) -> VortexResult { + Binary::try_new( + column, + ConstantArray::new(Scalar::from("AIR"), len).into_array(), + Operator::Eq, + ) + .map(IntoArray::into_array) +} + +#[test] +fn trace_scan_compare_on_compressed_shipmode() -> VortexResult<()> { + let compressed = compressed_lineitem()?; + let len = compressed.len(); + + let lazy = shipmode_predicate(field(&compressed, "l_shipmode")?, len)?; + let expected = shipmode_predicate(field(&lineitem()?, "l_shipmode")?, len)?; + let [optimized, executed] = optimize_then_execute(lazy, &expected)?; + + insta::assert_snapshot!(optimized.trace.to_string(), @" + optimize root=vortex.binary(bool, len=4096) session=false + reduce_parent static:DictionaryScalarFnValuesPushDownRule slot=0 parent=vortex.binary(bool, len=4096) child=vortex.dict(utf8, len=4096) -> vortex.dict(bool, len=4096) + done output=vortex.dict(bool, len=4096) + "); + insta::assert_snapshot!(executed.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.dict(bool, len=4096) + iter 0 current=vortex.dict(bool, len=4096) builder_active=false + ExecuteSlot slot=0 parent=vortex.dict(bool, len=4096) child=fastlanes.bitpacked(u8, len=4096) + iter 1 current=fastlanes.bitpacked(u8, len=4096) stack_parent=vortex.dict(bool, len=4096) slot=0 builder_active=false + Done array=vortex.primitive(u8, len=4096) + iter 2 current=vortex.primitive(u8, len=4096) stack_parent=vortex.dict(bool, len=4096) slot=0 builder_active=false + pop_frame slot=0 output=vortex.dict(bool, len=4096) + iter 3 current=vortex.dict(bool, len=4096) builder_active=false + ExecuteSlot slot=1 parent=vortex.dict(bool, len=4096) child=vortex.binary(bool, len=7) + iter 4 current=vortex.binary(bool, len=7) stack_parent=vortex.dict(bool, len=4096) slot=1 builder_active=false + child_execute_parent session[0]:execute_parent_fn slot=0 parent=vortex.binary(bool, len=7) child=vortex.fsst(utf8, len=7) -> vortex.binary(bool, len=7) + iter 5 current=vortex.binary(bool, len=7) stack_parent=vortex.dict(bool, len=4096) slot=1 builder_active=false + child_execute_parent session[0]:execute_parent_fn slot=0 parent=vortex.binary(bool, len=7) child=vortex.varbin(binary, len=7) -> vortex.bool(bool, len=7) + iter 6 current=vortex.bool(bool, len=7) stack_parent=vortex.dict(bool, len=4096) slot=1 builder_active=false + pop_frame slot=1 output=vortex.dict(bool, len=4096) + iter 7 current=vortex.dict(bool, len=4096) builder_active=false + child_execute_parent session[0]:execute_parent_fn slot=1 parent=vortex.dict(bool, len=4096) child=vortex.bool(bool, len=7) -> vortex.bool(bool, len=4096) + iter 8 current=vortex.bool(bool, len=4096) builder_active=false + return output=vortex.bool(bool, len=4096) + "); + + Ok(()) +} + +/// Q13-style predicate over the comment column: `l_comment LIKE '%special%'`. +/// +/// The column compresses to `fsst -> delta offsets` with bitpacked residuals. +fn comment_predicate(column: ArrayRef, len: usize) -> VortexResult { + Like::try_new( + column, + ConstantArray::new(Scalar::from("%special%"), len).into_array(), + LikeOptions { + negated: false, + case_insensitive: false, + }, + ) + .map(IntoArray::into_array) +} + +#[test] +fn trace_scan_like_on_compressed_comment() -> VortexResult<()> { + let compressed = compressed_lineitem()?; + let len = compressed.len(); + + let lazy = comment_predicate(field(&compressed, "l_comment")?, len)?; + let expected = comment_predicate(field(&lineitem()?, "l_comment")?, len)?; + let [optimized, executed] = optimize_then_execute(lazy, &expected)?; + + // No reduce rule rewrites a like over FSST; the FSST like kernel compiles the pattern and + // matches in compressed space at execution time. + insta::assert_snapshot!(optimized.trace.to_string(), @""); + insta::assert_snapshot!(executed.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.like(bool, len=4096) + iter 0 current=vortex.like(bool, len=4096) builder_active=false + execute_until target=AnyCanonical root=fastlanes.delta(u16, len=4097) + iter 0 current=fastlanes.delta(u16, len=4097) builder_active=false + execute_until target=AnyCanonical root=fastlanes.bitpacked(u16, len=5120) + iter 0 current=fastlanes.bitpacked(u16, len=5120) builder_active=false + Done array=vortex.primitive(u16, len=5120) + iter 1 current=vortex.primitive(u16, len=5120) builder_active=false + return output=vortex.primitive(u16, len=5120) + Done array=vortex.primitive(u16, len=4097) + iter 1 current=vortex.primitive(u16, len=4097) builder_active=false + return output=vortex.primitive(u16, len=4097) + child_execute_parent session[0]:execute_parent_fn slot=0 parent=vortex.like(bool, len=4096) child=vortex.fsst(utf8, len=4096) -> vortex.bool(bool, len=4096) + iter 1 current=vortex.bool(bool, len=4096) builder_active=false + return output=vortex.bool(bool, len=4096) + "); + + Ok(()) +} + +/// A narrowed projection of lineitem with one column per interesting compressed encoding: +/// decimal_byte_parts, ext-over-FoR dates, and dict-of-FSST strings. +fn project(table: &ArrayRef) -> VortexResult { + Ok(StructArray::from_fields(&[ + ("l_quantity", field(table, "l_quantity")?), + ("l_shipdate", field(table, "l_shipdate")?), + ("l_shipmode", field(table, "l_shipmode")?), + ])? + .into_array()) +} + +#[test] +fn trace_scan_filter_on_compressed_table() -> VortexResult<()> { + let compressed = project(&compressed_lineitem()?)?; + let len = compressed.len(); + let mask = Mask::from_iter((0..len).map(|i| i % 97 == 0)); + + let lazy = FilterArray::try_new(compressed, mask.clone())?.into_array(); + let expected = FilterArray::try_new(project(&lineitem()?)?, mask)?.into_array(); + let [optimized, executed] = optimize_then_execute(lazy, &expected)?; + + insta::assert_snapshot!(optimized.trace.to_string(), @" + optimize root=vortex.filter({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=43) session=false + optimize root=vortex.filter(decimal(15,2), len=43) session=false + optimize root=vortex.filter(i16, len=43) session=false + reduce_parent static:FilterReduceAdaptor(Dict) slot=0 parent=vortex.filter(i16, len=43) child=vortex.dict(i16, len=4096) -> vortex.dict(i16, len=43) + done output=vortex.dict(i16, len=43) + reduce_parent static:DecimalBytePartsFilterPushDownRule slot=0 parent=vortex.filter(decimal(15,2), len=43) child=vortex.decimal_byte_parts(decimal(15,2), len=4096) -> vortex.decimal_byte_parts(decimal(15,2), len=43) + done output=vortex.decimal_byte_parts(decimal(15,2), len=43) + optimize root=vortex.filter(vortex.date[days](i32), len=43) session=false + optimize root=vortex.filter(i32, len=43) session=false + reduce_parent static:FoRFilterPushDownRule slot=0 parent=vortex.filter(i32, len=43) child=fastlanes.for(i32, len=4096) -> fastlanes.for(i32, len=43) + done output=fastlanes.for(i32, len=43) + reduce_parent static:ExtensionFilterPushDownRule slot=0 parent=vortex.filter(vortex.date[days](i32), len=43) child=vortex.ext(vortex.date[days](i32), len=4096) -> vortex.ext(vortex.date[days](i32), len=43) + done output=vortex.ext(vortex.date[days](i32), len=43) + optimize root=vortex.filter(utf8, len=43) session=false + reduce_parent static:FilterReduceAdaptor(Dict) slot=0 parent=vortex.filter(utf8, len=43) child=vortex.dict(utf8, len=4096) -> vortex.dict(utf8, len=43) + done output=vortex.dict(utf8, len=43) + reduce FilterStructRule: vortex.filter({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=43) -> vortex.struct({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=43) + done output=vortex.struct({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=43) + "); + insta::assert_snapshot!(executed.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.struct({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=43) + iter 0 current=vortex.struct({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=43) builder_active=false + return output=vortex.struct({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=43) + "); + + Ok(()) +} + +#[test] +fn trace_scan_take_on_compressed_table() -> VortexResult<()> { + let compressed = project(&compressed_lineitem()?)?; + let len = compressed.len() as u64; + // A take is expressed as a `DictArray` whose codes are the take indices. + let indices = PrimitiveArray::from_iter((0..64u64).map(|i| (i * 941) % len)).into_array(); + + let lazy = DictArray::try_new(indices.clone(), compressed)?.into_array(); + let expected = DictArray::try_new(indices, project(&lineitem()?)?)?.into_array(); + let [optimized, executed] = optimize_then_execute(lazy, &expected)?; + + insta::assert_snapshot!(optimized.trace.to_string(), @" + optimize root=vortex.dict({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=64) session=false + optimize root=vortex.dict(vortex.date[days](i32), len=64) session=false + reduce_parent static:TakeReduceAdaptor(Extension) slot=1 parent=vortex.dict(vortex.date[days](i32), len=64) child=vortex.ext(vortex.date[days](i32), len=4096) -> vortex.ext(vortex.date[days](i32), len=64) + done output=vortex.ext(vortex.date[days](i32), len=64) + reduce_parent static:TakeReduceAdaptor(Struct) slot=1 parent=vortex.dict({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=64) child=vortex.struct({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=4096) -> vortex.struct({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=64) + done output=vortex.struct({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=64) + "); + insta::assert_snapshot!(executed.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.struct({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=64) + iter 0 current=vortex.struct({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=64) builder_active=false + return output=vortex.struct({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=64) + "); + + Ok(()) +} diff --git a/vortex-btrblocks/tests/golden.rs b/vortex-btrblocks/tests/golden.rs index c60752b064a..fc636252c27 100644 --- a/vortex-btrblocks/tests/golden.rs +++ b/vortex-btrblocks/tests/golden.rs @@ -9,12 +9,12 @@ //! snapshot untouched, so snapshot churn in a later change is the reviewable signal of a //! behavior change. //! -//! Three variants cover the feature matrix: +//! Three variants cover the edition and feature matrix: //! -//! - `default`: the default feature set and [`BtrBlocksCompressor::default`]. -//! - `unstable`: `unstable_encodings` enabled, default builder — pins Delta / OnPair -//! selection (compiled out of `ALL_SCHEMES` otherwise). -//! - `compact`: `unstable_encodings` + `zstd` + `pco`, with +//! - `regular`: the schemes permitted by the default `core` edition, minus OnPair. +//! - `onpair`: the structured-string entry with OnPair enabled — pins OnPair selection. +//! - `compact`: the schemes permitted by the default `core` and opt-in `zstd` editions, with +//! the `zstd` + `pco` features and //! [`BtrBlocksCompressorBuilder::with_compact`] — pins Zstd / Pco selection. //! //! Every corpus entry is longer than 1024 values so the sampling-based estimation path is @@ -49,7 +49,15 @@ use vortex_array::dtype::Nullability; use vortex_array::extension::datetime::TimeUnit; use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressor; +use vortex_btrblocks::BtrBlocksCompressorBuilder; use vortex_buffer::Buffer; +use vortex_edition::ComponentKind; +use vortex_edition::EDITION_DECLARATIONS; +use vortex_edition::EDITION_FAMILIES; +use vortex_edition::EditionId; +use vortex_edition::EditionSession; +use vortex_edition::EditionSessionExt; +use vortex_edition::declarations::core::CORE_2026_08_3; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -65,7 +73,7 @@ const N: usize = 16_384; /// [`NbytesExtractor`]: vortex_array::display::NbytesExtractor struct ExactNbytesExtractor; -impl TreeExtractor for ExactNbytesExtractor { +impl TreeExtractor for ExactNbytesExtractor { fn write_header( &self, array: &ArrayRef, @@ -94,7 +102,17 @@ fn render(input: &ArrayRef, compressed: &ArrayRef) -> String { /// Compresses every corpus entry twice (direct determinism check) and snapshots the result /// under `{variant}__{entry}`. fn golden_corpus_snapshots(variant: &str, compressor: &BtrBlocksCompressor) -> VortexResult<()> { - for (name, array) in corpus()? { + golden_snapshots(variant, compressor, corpus()?) +} + +/// Compresses each entry twice (direct determinism check) and snapshots the result under +/// `{variant}__{entry}`. +fn golden_snapshots( + variant: &str, + compressor: &BtrBlocksCompressor, + entries: Vec<(&'static str, ArrayRef)>, +) -> VortexResult<()> { + for (name, array) in entries { let rendered = { let mut exec_ctx = SESSION.create_execution_ctx(); render(&array, &compressor.compress(&array, &mut exec_ctx)?) @@ -372,44 +390,83 @@ fn list_of_int_runs() -> VortexResult { Ok(ListArray::try_new(elements, offsets.into_array(), Validity::NonNullable)?.into_array()) } -/// Excludes OnPair from the golden compressors: its dictionary training (upstream `onpair` -/// crate) iterates randomly-seeded `hashbrown` maps, so its compressed output — and therefore -/// its sampled estimate — differs run-to-run. A nondeterministic scheme cannot serve as a -/// golden baseline; excluding it keeps the remaining unstable schemes pinned. -#[cfg(feature = "unstable_encodings")] -fn without_onpair( - builder: vortex_btrblocks::BtrBlocksCompressorBuilder, -) -> vortex_btrblocks::BtrBlocksCompressorBuilder { +/// Excludes OnPair from the `regular` and `compact` variants: it beats FSST on +/// `string_fsst_structured`, and those variants pin the FSST selection. OnPair's own decisions +/// are pinned by [`golden_onpair`]. +fn without_onpair(builder: BtrBlocksCompressorBuilder) -> BtrBlocksCompressorBuilder { use vortex_btrblocks::SchemeExt; use vortex_btrblocks::schemes::string::OnPairScheme; builder.exclude_schemes([OnPairScheme.id()]) } -#[cfg(not(feature = "unstable_encodings"))] -#[test] -fn golden_default() -> VortexResult<()> { - golden_corpus_snapshots("default", &BtrBlocksCompressor::default()) +fn edition_session(editions: &[EditionId]) -> VortexResult { + let session = vortex_array::array_session().with::(); + for family in EDITION_FAMILIES { + session.editions().declare_family(family)?; + } + for declaration in EDITION_DECLARATIONS { + session.register_edition(declaration)?; + } + for edition in editions { + session.enable_edition(*edition)?; + } + Ok(session) +} + +fn compressor_for_session( + session: &VortexSession, + builder: BtrBlocksCompressorBuilder, +) -> BtrBlocksCompressor { + let allowed = session + .enabled_component_ids(ComponentKind::Array) + .into_iter() + .collect(); + without_onpair(builder) + .retain_allowed_encodings(&allowed) + .build() +} + +/// Like [`compressor_for_session`] but keeps OnPair in the scheme pool. +fn compressor_with_onpair( + session: &VortexSession, + builder: BtrBlocksCompressorBuilder, +) -> BtrBlocksCompressor { + let allowed = session + .enabled_component_ids(ComponentKind::Array) + .into_iter() + .collect(); + builder.retain_allowed_encodings(&allowed).build() } -#[cfg(feature = "unstable_encodings")] #[test] -fn golden_unstable() -> VortexResult<()> { - use vortex_btrblocks::BtrBlocksCompressorBuilder; +fn golden_regular() -> VortexResult<()> { + let session = edition_session(&[CORE_2026_08_3])?; + let compressor = compressor_for_session(&session, BtrBlocksCompressorBuilder::default()); + golden_corpus_snapshots("regular", &compressor) +} - golden_corpus_snapshots( - "unstable", - &without_onpair(BtrBlocksCompressorBuilder::default()).build(), +/// Pins OnPair's selection over FSST on the structured-string entry. +#[test] +fn golden_onpair() -> VortexResult<()> { + let session = edition_session(&[CORE_2026_08_3])?; + let compressor = compressor_with_onpair(&session, BtrBlocksCompressorBuilder::default()); + golden_snapshots( + "onpair", + &compressor, + vec![("string_fsst_structured", string_fsst_structured())], ) } -#[cfg(all(feature = "unstable_encodings", feature = "zstd", feature = "pco"))] +#[cfg(all(feature = "zstd", feature = "pco"))] #[test] fn golden_compact() -> VortexResult<()> { - use vortex_btrblocks::BtrBlocksCompressorBuilder; - - golden_corpus_snapshots( - "compact", - &without_onpair(BtrBlocksCompressorBuilder::default().with_compact()).build(), - ) + let session = edition_session(&[CORE_2026_08_3])?; + vortex_zstd::initialize(&session); + session.enable_edition(vortex_zstd::editions::ZSTD_2026_02)?; + let compressor = compressor_for_session( + &session, + BtrBlocksCompressorBuilder::default().with_compact(), + ); + golden_corpus_snapshots("compact", &compressor) } diff --git a/vortex-btrblocks/tests/onpair_roundtrip.rs b/vortex-btrblocks/tests/onpair_roundtrip.rs index 272fa930bd0..31734d6a60e 100644 --- a/vortex-btrblocks/tests/onpair_roundtrip.rs +++ b/vortex-btrblocks/tests/onpair_roundtrip.rs @@ -2,11 +2,9 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors // //! End-to-end round-trip through the full Vortex compressor + decompressor -//! on string arrays. Lives in `vortex-btrblocks` (gated on `unstable_encodings`) -//! so it exercises the same code path the file writer uses, not just the -//! OnPair crate in isolation. +//! on string arrays. Lives in `vortex-btrblocks` so it exercises the same code path the file +//! writer uses, not just the OnPair crate in isolation. -#![cfg(feature = "unstable_encodings")] #![allow( clippy::cast_possible_truncation, clippy::tests_outside_test_module, @@ -188,7 +186,7 @@ fn empty_and_short_string_roundtrip() { } /// Regression for the Euro2016 compress-bench panic -/// (`onpair::decompress`: "dictionary offsets must be nondecreasing"). +/// (`onpair::decompress`: "dictionary offsets must be strictly increasing"). /// /// A large, high-cardinality corpus fills the OnPair dictionary toward its /// 4096-entry cap, so the cascading compressor narrows `dict_offsets` to `u16` diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__binary_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__compact__binary_low_cardinality.snap index d68474ce61c..788cf54ecd8 100644 --- a/vortex-btrblocks/tests/snapshots/golden__compact__binary_low_cardinality.snap +++ b/vortex-btrblocks/tests/snapshots/golden__compact__binary_low_cardinality.snap @@ -3,9 +3,11 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: binary, len=16384, nbytes=315856 -root: vortex.dict(binary, len=16384) nbytes=6199 +root: vortex.dict(binary, len=16384) nbytes=6196 metadata: all_values_referenced: true codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144 metadata: bit_width: 3, offset: 0 - values: vortex.zstd(binary, len=5) nbytes=55 - metadata: nrows: 5, slice: 0..5 + values: vortex.varbin(binary, len=5) nbytes=52 + metadata: + offsets: vortex.primitive(u8, len=6) nbytes=6 + metadata: ptype: u8 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__float_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__default__float_low_cardinality.snap deleted file mode 100644 index 38d952262d6..00000000000 --- a/vortex-btrblocks/tests/snapshots/golden__default__float_low_cardinality.snap +++ /dev/null @@ -1,19 +0,0 @@ ---- -source: vortex-btrblocks/tests/golden.rs -expression: rendered ---- -input: f64, len=16384, nbytes=131072 -root: vortex.alp(f64, len=16384) nbytes=10233 - metadata: exponents: e: 16, f: 12, patch_offset: 0 - encoded: vortex.dict(i64, len=16384) nbytes=6200 - metadata: all_values_referenced: true - codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144 - metadata: bit_width: 3, offset: 0 - values: vortex.primitive(i64, len=7) nbytes=56 - metadata: ptype: i64 - patch_indices: vortex.primitive(u16, len=1996) nbytes=3992 - metadata: ptype: u16 - patch_values: vortex.constant(f64, len=1996) nbytes=9 - metadata: scalar: 1000.03125f64 - patch_chunk_offsets: vortex.primitive(u16, len=16) nbytes=32 - metadata: ptype: u16 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__string_fsst_structured.snap b/vortex-btrblocks/tests/snapshots/golden__onpair__string_fsst_structured.snap similarity index 51% rename from vortex-btrblocks/tests/snapshots/golden__unstable__string_fsst_structured.snap rename to vortex-btrblocks/tests/snapshots/golden__onpair__string_fsst_structured.snap index 327f050b0d7..028dbe81dc9 100644 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__string_fsst_structured.snap +++ b/vortex-btrblocks/tests/snapshots/golden__onpair__string_fsst_structured.snap @@ -3,13 +3,17 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: utf8, len=16384, nbytes=653785 -root: vortex.fsst(utf8, len=16384) nbytes=151382 - metadata: len: 16384, nsymbols: 223 +root: vortex.onpair(utf8, len=16384) nbytes=139746 + metadata: dict_bytes_len: 11866 + dict_offsets: vortex.primitive(u16, len=1626) nbytes=3252 + metadata: ptype: u16 + codes: fastlanes.bitpacked(u16, len=63845) nbytes=88704 + metadata: bit_width: 11, offset: 0 + codes_offsets: vortex.primitive(u16, len=16385) nbytes=32770 + metadata: ptype: u16 uncompressed_lengths: vortex.sparse(u8, len=16384) nbytes=3154 metadata: fill_value: 24u8 patch_indices: vortex.primitive(u16, len=1575) nbytes=3150 metadata: ptype: u16 patch_values: vortex.constant(u8, len=1575) nbytes=2 metadata: scalar: 23u8 - codes_offsets: fastlanes.bitpacked(u32, len=16385) nbytes=36992 - metadata: bit_width: 17, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__binary_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__regular__binary_low_cardinality.snap similarity index 56% rename from vortex-btrblocks/tests/snapshots/golden__default__binary_low_cardinality.snap rename to vortex-btrblocks/tests/snapshots/golden__regular__binary_low_cardinality.snap index 56918c0d096..788cf54ecd8 100644 --- a/vortex-btrblocks/tests/snapshots/golden__default__binary_low_cardinality.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__binary_low_cardinality.snap @@ -3,9 +3,11 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: binary, len=16384, nbytes=315856 -root: vortex.dict(binary, len=16384) nbytes=6240 +root: vortex.dict(binary, len=16384) nbytes=6196 metadata: all_values_referenced: true codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144 metadata: bit_width: 3, offset: 0 - values: vortex.varbinview(binary, len=5) nbytes=96 - metadata: + values: vortex.varbin(binary, len=5) nbytes=52 + metadata: + offsets: vortex.primitive(u8, len=6) nbytes=6 + metadata: ptype: u8 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__bool_random.snap b/vortex-btrblocks/tests/snapshots/golden__regular__bool_random.snap similarity index 100% rename from vortex-btrblocks/tests/snapshots/golden__default__bool_random.snap rename to vortex-btrblocks/tests/snapshots/golden__regular__bool_random.snap diff --git a/vortex-btrblocks/tests/snapshots/golden__default__decimal_prices.snap b/vortex-btrblocks/tests/snapshots/golden__regular__decimal_prices.snap similarity index 100% rename from vortex-btrblocks/tests/snapshots/golden__default__decimal_prices.snap rename to vortex-btrblocks/tests/snapshots/golden__regular__decimal_prices.snap diff --git a/vortex-btrblocks/tests/snapshots/golden__default__float_alp_prices.snap b/vortex-btrblocks/tests/snapshots/golden__regular__float_alp_prices.snap similarity index 100% rename from vortex-btrblocks/tests/snapshots/golden__default__float_alp_prices.snap rename to vortex-btrblocks/tests/snapshots/golden__regular__float_alp_prices.snap diff --git a/vortex-btrblocks/tests/snapshots/golden__default__float_full_precision.snap b/vortex-btrblocks/tests/snapshots/golden__regular__float_full_precision.snap similarity index 100% rename from vortex-btrblocks/tests/snapshots/golden__default__float_full_precision.snap rename to vortex-btrblocks/tests/snapshots/golden__regular__float_full_precision.snap diff --git a/vortex-btrblocks/tests/snapshots/golden__regular__float_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__regular__float_low_cardinality.snap new file mode 100644 index 00000000000..9c2192c70d7 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__regular__float_low_cardinality.snap @@ -0,0 +1,13 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: f64, len=16384, nbytes=131072 +root: vortex.alp(f64, len=16384) nbytes=6208 + metadata: exponents: e: 16, f: 11 + encoded: vortex.dict(i64, len=16384) nbytes=6208 + metadata: all_values_referenced: true + codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144 + metadata: bit_width: 3, offset: 0 + values: vortex.primitive(i64, len=8) nbytes=64 + metadata: ptype: i64 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__float_mostly_null.snap b/vortex-btrblocks/tests/snapshots/golden__regular__float_mostly_null.snap similarity index 100% rename from vortex-btrblocks/tests/snapshots/golden__default__float_mostly_null.snap rename to vortex-btrblocks/tests/snapshots/golden__regular__float_mostly_null.snap diff --git a/vortex-btrblocks/tests/snapshots/golden__default__int_arithmetic_sequence.snap b/vortex-btrblocks/tests/snapshots/golden__regular__int_arithmetic_sequence.snap similarity index 100% rename from vortex-btrblocks/tests/snapshots/golden__default__int_arithmetic_sequence.snap rename to vortex-btrblocks/tests/snapshots/golden__regular__int_arithmetic_sequence.snap diff --git a/vortex-btrblocks/tests/snapshots/golden__default__int_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__regular__int_low_cardinality.snap similarity index 100% rename from vortex-btrblocks/tests/snapshots/golden__default__int_low_cardinality.snap rename to vortex-btrblocks/tests/snapshots/golden__regular__int_low_cardinality.snap diff --git a/vortex-btrblocks/tests/snapshots/golden__default__int_monotone_jitter.snap b/vortex-btrblocks/tests/snapshots/golden__regular__int_monotone_jitter.snap similarity index 100% rename from vortex-btrblocks/tests/snapshots/golden__default__int_monotone_jitter.snap rename to vortex-btrblocks/tests/snapshots/golden__regular__int_monotone_jitter.snap diff --git a/vortex-btrblocks/tests/snapshots/golden__default__int_mostly_null.snap b/vortex-btrblocks/tests/snapshots/golden__regular__int_mostly_null.snap similarity index 100% rename from vortex-btrblocks/tests/snapshots/golden__default__int_mostly_null.snap rename to vortex-btrblocks/tests/snapshots/golden__regular__int_mostly_null.snap diff --git a/vortex-btrblocks/tests/snapshots/golden__default__int_negatives.snap b/vortex-btrblocks/tests/snapshots/golden__regular__int_negatives.snap similarity index 100% rename from vortex-btrblocks/tests/snapshots/golden__default__int_negatives.snap rename to vortex-btrblocks/tests/snapshots/golden__regular__int_negatives.snap diff --git a/vortex-btrblocks/tests/snapshots/golden__default__int_runs.snap b/vortex-btrblocks/tests/snapshots/golden__regular__int_runs.snap similarity index 100% rename from vortex-btrblocks/tests/snapshots/golden__default__int_runs.snap rename to vortex-btrblocks/tests/snapshots/golden__regular__int_runs.snap diff --git a/vortex-btrblocks/tests/snapshots/golden__default__int_sparse_outliers.snap b/vortex-btrblocks/tests/snapshots/golden__regular__int_sparse_outliers.snap similarity index 100% rename from vortex-btrblocks/tests/snapshots/golden__default__int_sparse_outliers.snap rename to vortex-btrblocks/tests/snapshots/golden__regular__int_sparse_outliers.snap diff --git a/vortex-btrblocks/tests/snapshots/golden__default__int_wide_random.snap b/vortex-btrblocks/tests/snapshots/golden__regular__int_wide_random.snap similarity index 100% rename from vortex-btrblocks/tests/snapshots/golden__default__int_wide_random.snap rename to vortex-btrblocks/tests/snapshots/golden__regular__int_wide_random.snap diff --git a/vortex-btrblocks/tests/snapshots/golden__default__list_of_int_runs.snap b/vortex-btrblocks/tests/snapshots/golden__regular__list_of_int_runs.snap similarity index 100% rename from vortex-btrblocks/tests/snapshots/golden__default__list_of_int_runs.snap rename to vortex-btrblocks/tests/snapshots/golden__regular__list_of_int_runs.snap diff --git a/vortex-btrblocks/tests/snapshots/golden__default__string_fsst_structured.snap b/vortex-btrblocks/tests/snapshots/golden__regular__string_fsst_structured.snap similarity index 100% rename from vortex-btrblocks/tests/snapshots/golden__default__string_fsst_structured.snap rename to vortex-btrblocks/tests/snapshots/golden__regular__string_fsst_structured.snap diff --git a/vortex-btrblocks/tests/snapshots/golden__default__string_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__regular__string_low_cardinality.snap similarity index 100% rename from vortex-btrblocks/tests/snapshots/golden__default__string_low_cardinality.snap rename to vortex-btrblocks/tests/snapshots/golden__regular__string_low_cardinality.snap diff --git a/vortex-btrblocks/tests/snapshots/golden__default__struct_mixed.snap b/vortex-btrblocks/tests/snapshots/golden__regular__struct_mixed.snap similarity index 100% rename from vortex-btrblocks/tests/snapshots/golden__default__struct_mixed.snap rename to vortex-btrblocks/tests/snapshots/golden__regular__struct_mixed.snap diff --git a/vortex-btrblocks/tests/snapshots/golden__default__temporal_timestamp_micros.snap b/vortex-btrblocks/tests/snapshots/golden__regular__temporal_timestamp_micros.snap similarity index 100% rename from vortex-btrblocks/tests/snapshots/golden__default__temporal_timestamp_micros.snap rename to vortex-btrblocks/tests/snapshots/golden__regular__temporal_timestamp_micros.snap diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__binary_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__binary_low_cardinality.snap deleted file mode 100644 index 56918c0d096..00000000000 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__binary_low_cardinality.snap +++ /dev/null @@ -1,11 +0,0 @@ ---- -source: vortex-btrblocks/tests/golden.rs -expression: rendered ---- -input: binary, len=16384, nbytes=315856 -root: vortex.dict(binary, len=16384) nbytes=6240 - metadata: all_values_referenced: true - codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144 - metadata: bit_width: 3, offset: 0 - values: vortex.varbinview(binary, len=5) nbytes=96 - metadata: diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__bool_random.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__bool_random.snap deleted file mode 100644 index d80e2fa4464..00000000000 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__bool_random.snap +++ /dev/null @@ -1,7 +0,0 @@ ---- -source: vortex-btrblocks/tests/golden.rs -expression: rendered ---- -input: bool, len=16384, nbytes=2048 -root: vortex.bool(bool, len=16384) nbytes=2048 - metadata: offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__decimal_prices.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__decimal_prices.snap deleted file mode 100644 index f669755e4b1..00000000000 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__decimal_prices.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: vortex-btrblocks/tests/golden.rs -expression: rendered ---- -input: decimal(12,2), len=16384, nbytes=131072 -root: vortex.decimal_byte_parts(decimal(12,2), len=16384) nbytes=49152 - metadata: - msp: fastlanes.bitpacked(i32, len=16384) nbytes=49152 - metadata: bit_width: 24, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__float_alp_prices.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__float_alp_prices.snap deleted file mode 100644 index f70d90a301d..00000000000 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__float_alp_prices.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: vortex-btrblocks/tests/golden.rs -expression: rendered ---- -input: f64, len=16384, nbytes=131072 -root: vortex.alp(f64, len=16384) nbytes=49152 - metadata: exponents: e: 14, f: 12 - encoded: fastlanes.bitpacked(i64, len=16384) nbytes=49152 - metadata: bit_width: 24, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__float_full_precision.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__float_full_precision.snap deleted file mode 100644 index 487e7c7900b..00000000000 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__float_full_precision.snap +++ /dev/null @@ -1,15 +0,0 @@ ---- -source: vortex-btrblocks/tests/golden.rs -expression: rendered ---- -input: f64, len=16384, nbytes=131072 -root: vortex.alprd(f64, len=16384) nbytes=112880 - metadata: right_bit_width: 52, patch_offset: 0 - left_parts: fastlanes.bitpacked(u16, len=16384) nbytes=6144 - metadata: bit_width: 3, offset: 0 - right_parts: fastlanes.bitpacked(u64, len=16384) nbytes=106496 - metadata: bit_width: 52, offset: 0 - patch_indices: vortex.primitive(u16, len=60) nbytes=120 - metadata: ptype: u16 - patch_values: vortex.primitive(u16, len=60) nbytes=120 - metadata: ptype: u16 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__float_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__float_low_cardinality.snap deleted file mode 100644 index 38d952262d6..00000000000 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__float_low_cardinality.snap +++ /dev/null @@ -1,19 +0,0 @@ ---- -source: vortex-btrblocks/tests/golden.rs -expression: rendered ---- -input: f64, len=16384, nbytes=131072 -root: vortex.alp(f64, len=16384) nbytes=10233 - metadata: exponents: e: 16, f: 12, patch_offset: 0 - encoded: vortex.dict(i64, len=16384) nbytes=6200 - metadata: all_values_referenced: true - codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144 - metadata: bit_width: 3, offset: 0 - values: vortex.primitive(i64, len=7) nbytes=56 - metadata: ptype: i64 - patch_indices: vortex.primitive(u16, len=1996) nbytes=3992 - metadata: ptype: u16 - patch_values: vortex.constant(f64, len=1996) nbytes=9 - metadata: scalar: 1000.03125f64 - patch_chunk_offsets: vortex.primitive(u16, len=16) nbytes=32 - metadata: ptype: u16 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__float_mostly_null.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__float_mostly_null.snap deleted file mode 100644 index 7b88d819472..00000000000 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__float_mostly_null.snap +++ /dev/null @@ -1,13 +0,0 @@ ---- -source: vortex-btrblocks/tests/golden.rs -expression: rendered ---- -input: f64?, len=16384, nbytes=133120 -root: vortex.sparse(f64?, len=16384) nbytes=8609 - metadata: fill_value: null - patch_indices: vortex.primitive(u16, len=850) nbytes=1700 - metadata: ptype: u16 - patch_values: vortex.primitive(f64?, len=850) nbytes=6907 - metadata: ptype: f64 - validity: vortex.bool(bool, len=850) nbytes=107 - metadata: offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__int_arithmetic_sequence.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__int_arithmetic_sequence.snap deleted file mode 100644 index 873f3655a1f..00000000000 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__int_arithmetic_sequence.snap +++ /dev/null @@ -1,7 +0,0 @@ ---- -source: vortex-btrblocks/tests/golden.rs -expression: rendered ---- -input: i64, len=16384, nbytes=131072 -root: vortex.sequence(i64, len=16384) nbytes=0 - metadata: base: 10000i64, multiplier: 7i64 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__int_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__int_low_cardinality.snap deleted file mode 100644 index 8648200a138..00000000000 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__int_low_cardinality.snap +++ /dev/null @@ -1,11 +0,0 @@ ---- -source: vortex-btrblocks/tests/golden.rs -expression: rendered ---- -input: i64, len=16384, nbytes=131072 -root: vortex.dict(i64, len=16384) nbytes=6192 - metadata: all_values_referenced: true - codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144 - metadata: bit_width: 3, offset: 0 - values: vortex.primitive(i64, len=6) nbytes=48 - metadata: ptype: i64 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__int_monotone_jitter.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__int_monotone_jitter.snap deleted file mode 100644 index f1ded4c3e59..00000000000 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__int_monotone_jitter.snap +++ /dev/null @@ -1,15 +0,0 @@ ---- -source: vortex-btrblocks/tests/golden.rs -expression: rendered ---- -input: u64, len=16384, nbytes=131072 -root: fastlanes.delta(u64, len=16384) nbytes=19840 - metadata: offset: 0 - bases: vortex.primitive(u64, len=256) nbytes=2048 - metadata: ptype: u64 - deltas: vortex.dict(u64, len=16384) nbytes=17792 - metadata: all_values_referenced: true - codes: vortex.primitive(u8, len=16384) nbytes=16384 - metadata: ptype: u8 - values: fastlanes.bitpacked(u64, len=201) nbytes=1408 - metadata: bit_width: 11, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__int_mostly_null.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__int_mostly_null.snap deleted file mode 100644 index 73566369279..00000000000 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__int_mostly_null.snap +++ /dev/null @@ -1,13 +0,0 @@ ---- -source: vortex-btrblocks/tests/golden.rs -expression: rendered ---- -input: i32?, len=16384, nbytes=67584 -root: vortex.sparse(i32?, len=16384) nbytes=3031 - metadata: fill_value: null - patch_indices: vortex.primitive(u16, len=823) nbytes=1646 - metadata: ptype: u16 - patch_values: fastlanes.bitpacked(i32?, len=823) nbytes=1383 - metadata: bit_width: 10, offset: 0 - validity_child: vortex.bool(bool, len=823) nbytes=103 - metadata: offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__int_negatives.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__int_negatives.snap deleted file mode 100644 index 24949893190..00000000000 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__int_negatives.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: vortex-btrblocks/tests/golden.rs -expression: rendered ---- -input: i64, len=16384, nbytes=131072 -root: fastlanes.for(i64, len=16384) nbytes=16384 - metadata: reference: -128i64 - encoded: fastlanes.bitpacked(i64, len=16384) nbytes=16384 - metadata: bit_width: 8, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__int_runs.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__int_runs.snap deleted file mode 100644 index 271de8faf93..00000000000 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__int_runs.snap +++ /dev/null @@ -1,15 +0,0 @@ ---- -source: vortex-btrblocks/tests/golden.rs -expression: rendered ---- -input: i32, len=16384, nbytes=65536 -root: vortex.runend(i32, len=16384) nbytes=3968 - metadata: offset: 0 - ends: fastlanes.for(u16, len=1020) nbytes=1792 - metadata: reference: 13u16 - encoded: fastlanes.bitpacked(u16, len=1020) nbytes=1792 - metadata: bit_width: 14, offset: 0 - values: fastlanes.for(i32, len=1020) nbytes=2176 - metadata: reference: -49931i32 - encoded: fastlanes.bitpacked(i32, len=1020) nbytes=2176 - metadata: bit_width: 17, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__int_sparse_outliers.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__int_sparse_outliers.snap deleted file mode 100644 index 53a5d1e023c..00000000000 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__int_sparse_outliers.snap +++ /dev/null @@ -1,13 +0,0 @@ ---- -source: vortex-btrblocks/tests/golden.rs -expression: rendered ---- -input: i64, len=16384, nbytes=131072 -root: vortex.sparse(i64, len=16384) nbytes=5540 - metadata: fill_value: 1000000i64 - patch_indices: vortex.primitive(u16, len=848) nbytes=1696 - metadata: ptype: u16 - patch_values: fastlanes.for(i64, len=848) nbytes=3840 - metadata: reference: 1000830099i64 - encoded: fastlanes.bitpacked(i64, len=848) nbytes=3840 - metadata: bit_width: 30, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__int_wide_random.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__int_wide_random.snap deleted file mode 100644 index 90e1894f910..00000000000 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__int_wide_random.snap +++ /dev/null @@ -1,7 +0,0 @@ ---- -source: vortex-btrblocks/tests/golden.rs -expression: rendered ---- -input: u64, len=16384, nbytes=131072 -root: vortex.primitive(u64, len=16384) nbytes=131072 - metadata: ptype: u64 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__list_of_int_runs.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__list_of_int_runs.snap deleted file mode 100644 index c9554add05a..00000000000 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__list_of_int_runs.snap +++ /dev/null @@ -1,25 +0,0 @@ ---- -source: vortex-btrblocks/tests/golden.rs -expression: rendered ---- -input: list(i32), len=4066, nbytes=81804 -root: vortex.list(list(i32), len=4066) nbytes=11146 - metadata: - elements: vortex.runend(i32, len=16384) nbytes=3968 - metadata: offset: 0 - ends: fastlanes.for(u16, len=1020) nbytes=1792 - metadata: reference: 13u16 - encoded: fastlanes.bitpacked(u16, len=1020) nbytes=1792 - metadata: bit_width: 14, offset: 0 - values: fastlanes.for(i32, len=1020) nbytes=2176 - metadata: reference: -49931i32 - encoded: fastlanes.bitpacked(i32, len=1020) nbytes=2176 - metadata: bit_width: 17, offset: 0 - offsets: fastlanes.bitpacked(u16, len=4067) nbytes=7178 - metadata: bit_width: 14, offset: 0 - patch_indices: vortex.primitive(u16, len=1) nbytes=2 - metadata: ptype: u16 - patch_values: vortex.constant(u16, len=1) nbytes=4 - metadata: scalar: 16384u16 - patch_chunk_offsets: vortex.primitive(u8, len=4) nbytes=4 - metadata: ptype: u8 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__string_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__string_low_cardinality.snap deleted file mode 100644 index ec79989d46a..00000000000 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__string_low_cardinality.snap +++ /dev/null @@ -1,15 +0,0 @@ ---- -source: vortex-btrblocks/tests/golden.rs -expression: rendered ---- -input: utf8, len=16384, nbytes=262144 -root: vortex.dict(utf8, len=16384) nbytes=8373 - metadata: all_values_referenced: true - codes: fastlanes.bitpacked(u8, len=16384) nbytes=8192 - metadata: bit_width: 4, offset: 0 - values: vortex.fsst(utf8, len=12) nbytes=181 - metadata: len: 12, nsymbols: 9 - uncompressed_lengths: vortex.primitive(u8, len=12) nbytes=12 - metadata: ptype: u8 - codes_offsets: vortex.primitive(u8, len=13) nbytes=13 - metadata: ptype: u8 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__struct_mixed.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__struct_mixed.snap deleted file mode 100644 index be8f48b779b..00000000000 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__struct_mixed.snap +++ /dev/null @@ -1,23 +0,0 @@ ---- -source: vortex-btrblocks/tests/golden.rs -expression: rendered ---- -input: {id=i64, category=utf8, value=f64}, len=16384, nbytes=524288 -root: vortex.struct({id=i64, category=utf8, value=f64}, len=16384) nbytes=57525 - metadata: - id: vortex.sequence(i64, len=16384) nbytes=0 - metadata: base: 10000i64, multiplier: 7i64 - category: vortex.dict(utf8, len=16384) nbytes=8373 - metadata: all_values_referenced: true - codes: fastlanes.bitpacked(u8, len=16384) nbytes=8192 - metadata: bit_width: 4, offset: 0 - values: vortex.fsst(utf8, len=12) nbytes=181 - metadata: len: 12, nsymbols: 9 - uncompressed_lengths: vortex.primitive(u8, len=12) nbytes=12 - metadata: ptype: u8 - codes_offsets: vortex.primitive(u8, len=13) nbytes=13 - metadata: ptype: u8 - value: vortex.alp(f64, len=16384) nbytes=49152 - metadata: exponents: e: 14, f: 12 - encoded: fastlanes.bitpacked(i64, len=16384) nbytes=49152 - metadata: bit_width: 24, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__temporal_timestamp_micros.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__temporal_timestamp_micros.snap deleted file mode 100644 index a763fa28f4b..00000000000 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__temporal_timestamp_micros.snap +++ /dev/null @@ -1,13 +0,0 @@ ---- -source: vortex-btrblocks/tests/golden.rs -expression: rendered ---- -input: vortex.timestamp[µs, tz=UTC](i64), len=16384, nbytes=131072 -root: vortex.ext(vortex.timestamp[µs, tz=UTC](i64), len=16384) nbytes=43008 - metadata: - storage: fastlanes.delta(i64, len=16384) nbytes=43008 - metadata: offset: 0 - bases: vortex.primitive(i64, len=256) nbytes=2048 - metadata: ptype: i64 - deltas: fastlanes.bitpacked(i64, len=16384) nbytes=40960 - metadata: bit_width: 20, offset: 0 diff --git a/vortex-btrblocks/tests/varbin_scheme.rs b/vortex-btrblocks/tests/varbin_scheme.rs new file mode 100644 index 00000000000..d47c8280af1 --- /dev/null +++ b/vortex-btrblocks/tests/varbin_scheme.rs @@ -0,0 +1,167 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Measures `VarBinScheme` against the same compressor with the scheme excluded. + +#![allow(clippy::cast_possible_truncation, clippy::tests_outside_test_module)] + +use std::sync::LazyLock; + +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_btrblocks::BtrBlocksCompressorBuilder; +use vortex_btrblocks::SchemeExt; +use vortex_btrblocks::schemes::binary::VarBinScheme; +use vortex_btrblocks::schemes::string::OnPairScheme; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + +const N: usize = 100_000; + +fn lcg(state: &mut u64) -> u64 { + *state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + *state +} + +fn cases() -> Vec<(&'static str, ArrayRef)> { + let mut s = 42u64; + let random: Vec> = (0..N) + .map(|_| (0..16).map(|_| (lcg(&mut s) >> 33) as u8).collect()) + .collect(); + let prefixed: Vec> = (0..N) + .map(|i| format!("PREFIX_{i:09}").into_bytes()) + .collect(); + let mut s2 = 7u64; + let wide: Vec> = (0..N) + .map(|_| (0..256).map(|_| (lcg(&mut s2) >> 33) as u8).collect()) + .collect(); + + let nullable = VarBinViewArray::from_iter( + (0..N).map(|i| (i % 7 != 0).then(|| prefixed[i].as_slice())), + DType::Binary(Nullability::Nullable), + ) + .into_array(); + + vec![ + ("nulls every 7th", nullable), + ( + "random 16B (hash)", + VarBinViewArray::from_iter_bin(random.iter().map(|v| v.as_slice())).into_array(), + ), + ( + "shared prefix", + VarBinViewArray::from_iter_bin(prefixed.iter().map(|v| v.as_slice())).into_array(), + ), + ( + "random 256B", + VarBinViewArray::from_iter_bin(wide.iter().map(|v| v.as_slice())).into_array(), + ), + ] +} + +#[test] +fn varbin_scheme_shrinks_binary() -> VortexResult<()> { + let with = BtrBlocksCompressorBuilder::default().build(); + let without = BtrBlocksCompressorBuilder::default() + .exclude_schemes([VarBinScheme.id()]) + .build(); + + println!( + "{:<20}{:>12}{:>14}{:>14}{:>9}", + "case", "input", "without", "with", "ratio" + ); + for (name, array) in cases() { + let a = { + let mut ctx = SESSION.create_execution_ctx(); + without.compress(&array, &mut ctx)?.nbytes() + }; + let b = { + let mut ctx = SESSION.create_execution_ctx(); + with.compress(&array, &mut ctx)?.nbytes() + }; + println!( + "{:<20}{:>12}{:>14}{:>14}{:>9.2}", + name, + array.nbytes(), + a, + b, + b as f64 / a as f64 + ); + + assert!( + b <= a, + "{name}: enabling VarBinScheme grew the output, {a} -> {b}" + ); + + let mut ctx = SESSION.create_execution_ctx(); + let compressed = with.compress(&array, &mut ctx)?; + let decoded = compressed + .execute::(&mut ctx)? + .into_array(); + assert_arrays_eq!(&array, &decoded, &mut ctx); + } + Ok(()) +} + +/// Same bytes, two dtypes: FSST compresses bytes rather than codepoints, so the dtype must not +/// change the result. `OnPairScheme` only matches utf8 and would otherwise win the utf8 column. +#[test] +fn fsst_versus_varbin_on_identical_bytes() -> VortexResult<()> { + let builder = BtrBlocksCompressorBuilder::default(); + let builder = builder.exclude_schemes([OnPairScheme.id()]); + let compressor = builder.build(); + let mut seed = 99u64; + + let shared_prefix: Vec = (0..N).map(|i| format!("PREFIX_{i:09}")).collect(); + let hex_random: Vec = (0..N) + .map(|_| { + (0..16) + .map(|_| format!("{:02x}", (lcg(&mut seed) >> 33) as u8)) + .collect() + }) + .collect(); + + println!( + "{:<20}{:>14}{:>14}{:>9}", + "case", "binary", "utf8(fsst)", "fsst/bin" + ); + for (name, vals) in [ + ("shared prefix", &shared_prefix), + ("hex random", &hex_random), + ] { + let as_bin = VarBinViewArray::from_iter_bin(vals.iter().map(|v| v.as_bytes())).into_array(); + let as_str = VarBinViewArray::from_iter_str(vals.iter().map(|v| v.as_str())).into_array(); + + let bin_bytes = { + let mut ctx = SESSION.create_execution_ctx(); + compressor.compress(&as_bin, &mut ctx)?.nbytes() + }; + let utf8_bytes = { + let mut ctx = SESSION.create_execution_ctx(); + compressor.compress(&as_str, &mut ctx)?.nbytes() + }; + println!( + "{:<20}{:>14}{:>14}{:>9.2}", + name, + bin_bytes, + utf8_bytes, + utf8_bytes as f64 / bin_bytes as f64 + ); + + // FSST compresses bytes, not codepoints, so the dtype must not change the result. + assert_eq!( + bin_bytes, utf8_bytes, + "{name}: binary and utf8 must compress identically" + ); + } + Ok(()) +} diff --git a/vortex-buffer/Cargo.toml b/vortex-buffer/Cargo.toml index ed9df797535..02cf4ce7782 100644 --- a/vortex-buffer/Cargo.toml +++ b/vortex-buffer/Cargo.toml @@ -23,6 +23,7 @@ serde = ["dep:serde", "serde/serde_derive"] warn-copy = ["dep:tracing"] [dependencies] +allocator-api2 = { workspace = true } arrow-buffer = { workspace = true } bitvec = { workspace = true } bytes = { workspace = true } @@ -40,6 +41,7 @@ workspace = true divan = { workspace = true } num-traits = { workspace = true } rstest = { workspace = true } +vortex-bench-support = { workspace = true } [[bench]] name = "vortex_buffer" @@ -56,3 +58,7 @@ harness = false [[bench]] name = "collect_bool" harness = false + +[[bench]] +name = "allocation" +harness = false diff --git a/vortex-buffer/benches/allocation.rs b/vortex-buffer/benches/allocation.rs new file mode 100644 index 00000000000..6cd923b5dfb --- /dev/null +++ b/vortex-buffer/benches/allocation.rs @@ -0,0 +1,106 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use allocator_api2::alloc::Global; +use arrow_buffer::MutableBuffer; +use bytes::BytesMut; +use divan::Bencher; +use vortex_buffer::Alignment; +use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; +use vortex_buffer::BufferMut; + +const SIZES: &[usize] = &[0, 64, 256, 1024, 16_384, 65_536]; + +fn main() { + divan::main(); +} + +#[divan::bench(args = SIZES)] +fn allocate_drop_vortex(bencher: Bencher, size: usize) { + bencher.bench(|| drop(BufferMut::::with_capacity(size))); +} + +#[divan::bench(args = SIZES)] +fn allocate_drop_vortex_custom(bencher: Bencher, size: usize) { + bencher + .with_inputs(|| BufferAllocatorRef::new(Global)) + .bench_refs(|allocator| drop(allocator.with_capacity::(size))); +} + +#[divan::bench(args = SIZES)] +fn allocate_drop_vortex_minimal_alignment(bencher: Bencher, size: usize) { + bencher.bench(|| { + drop(BufferMut::::with_capacity_preferred_aligned( + size, + Alignment::of::(), + None, + )) + }); +} + +#[divan::bench(args = SIZES)] +fn allocate_drop_bytes(bencher: Bencher, size: usize) { + bencher.bench(|| drop(BytesMut::with_capacity(size))); +} + +#[divan::bench(args = SIZES)] +fn allocate_drop_arrow(bencher: Bencher, size: usize) { + bencher.bench(|| drop(MutableBuffer::with_capacity(size))); +} + +#[divan::bench(args = SIZES)] +fn allocate_freeze_drop_vortex(bencher: Bencher, size: usize) { + bencher.bench(|| drop(BufferMut::::with_capacity(size).freeze())); +} + +#[divan::bench(args = SIZES)] +fn allocate_freeze_drop_vortex_custom(bencher: Bencher, size: usize) { + bencher + .with_inputs(|| BufferAllocatorRef::new(Global)) + .bench_refs(|allocator| drop(allocator.with_capacity::(size).freeze())); +} + +#[divan::bench(args = SIZES)] +fn allocate_freeze_drop_vortex_minimal_alignment(bencher: Bencher, size: usize) { + bencher.bench(|| { + drop( + BufferMut::::with_capacity_preferred_aligned(size, Alignment::of::(), None) + .freeze(), + ) + }); +} + +#[divan::bench(args = SIZES)] +fn allocate_freeze_drop_bytes(bencher: Bencher, size: usize) { + bencher.bench(|| drop(BytesMut::with_capacity(size).freeze())); +} + +#[divan::bench(args = SIZES)] +fn allocate_freeze_drop_arrow(bencher: Bencher, size: usize) { + bencher.bench(|| { + let buffer: arrow_buffer::Buffer = MutableBuffer::with_capacity(size).into(); + drop(buffer) + }); +} + +#[divan::bench(args = SIZES)] +fn from_vec_drop_vortex(bencher: Bencher, size: usize) { + bencher + .with_inputs(|| vec![0u8; size]) + .bench_values(|values| drop(Buffer::from(values))); +} + +#[divan::bench(args = SIZES)] +fn from_vec_drop_bytes(bencher: Bencher, size: usize) { + bencher + .with_inputs(|| vec![0u8; size]) + .bench_values(|values| drop(bytes::Bytes::from(values))); +} + +#[divan::bench(args = SIZES)] +fn from_vec_drop_arrow(bencher: Bencher, size: usize) { + bencher + .with_inputs(|| vec![0u8; size]) + .bench_values(|values| drop(arrow_buffer::Buffer::from_vec(values))); +} diff --git a/vortex-buffer/benches/collect_bool.rs b/vortex-buffer/benches/collect_bool.rs index 3bfe9a55b69..864179471f8 100644 --- a/vortex-buffer/benches/collect_bool.rs +++ b/vortex-buffer/benches/collect_bool.rs @@ -4,9 +4,7 @@ //! Benchmarks for `collect_bool` bit packing. //! //! Three groups, each with a scalar baseline that replicates the previous bit-at-a-time -//! implementation. Under CodSpeed only the shipped entry points are tracked -//! (`from_bool_slice`, `collect_bool_u32_gt`); the baselines, the dispatch loop, -//! per-SIMD-level loops, and the portable SWAR kernel compile out and are for local A/B runs: +//! implementation: //! //! - `pack_words_*`: the portable SWAR kernel vs the scalar loop, word by word. The SIMD //! kernels are covered by `words_gather_*` instead, where they can inline. @@ -14,10 +12,21 @@ //! measuring the fused fill + pack pipeline per SIMD level. //! - `collect_bool_*` / `from_bool_slice`: the public entry points end to end, with a //! boolean-gather predicate and a `u32` comparison predicate. +//! +//! `words_gather_dispatch` and `words_gather_scalar` carry `#[cpu_features]`, so they are +//! measured on every walltime CPU-feature leg rather than in simulation. Both are written +//! once and compiled differently per leg: the shipped entry point picks its pack kernel +//! through `cfg(target_feature)`, and how well the scalar loop auto-vectorizes depends on +//! the build. Comparing them across legs is the point. +//! +//! The hand-written per-kernel benchmarks are not tagged. Each one needs an instruction set +//! extension the other legs do not build for, so they stay out of CodSpeed entirely and +//! remain local A/B tools, as do the historical bit-at-a-time baselines. +//! +//! A plain `cargo bench` ignores all of it and runs everything on the host. use divan::Bencher; use vortex_buffer::BitBuffer; -#[cfg(not(codspeed))] use vortex_buffer::collect_bool_word_scalar; #[cfg(not(codspeed))] use vortex_buffer::pack_bool_word_swar; @@ -86,7 +95,6 @@ fn pack_words_swar(bencher: Bencher, len: usize) { /// Benchmark a full multiversioned word loop with a bounds-check-free bool gather, measuring /// the packing pipeline itself (fill + pack fully inlined) rather than predicate evaluation. -#[cfg(not(codspeed))] fn bench_words_gather( bencher: Bencher, len: usize, @@ -98,7 +106,7 @@ fn bench_words_gather( .bench_refs(|words| collect(words, len, &bools)); } -#[cfg(not(codspeed))] +#[vortex_bench_support::cpu_features] #[divan::bench(args = INPUT_SIZE)] fn words_gather_dispatch(bencher: Bencher, len: usize) { bench_words_gather(bencher, len, |words, len, bools| { @@ -107,7 +115,7 @@ fn words_gather_dispatch(bencher: Bencher, len: usize) { }); } -#[cfg(not(codspeed))] +#[vortex_bench_support::cpu_features] #[divan::bench(args = INPUT_SIZE)] fn words_gather_scalar(bencher: Bencher, len: usize) { bench_words_gather(bencher, len, |words, len, bools| { @@ -164,7 +172,6 @@ fn words_gather_neon(bencher: Bencher, len: usize) { /// Faithful copy of the previous scalar-only `collect_bool_words` word loop, used as the /// baseline for the end-to-end comparison. -#[cfg(not(codspeed))] fn collect_bool_words_old(words: &mut [u64], len: usize, mut f: impl FnMut(usize) -> bool) { let full = len / 64; let remainder = len % 64; diff --git a/vortex-buffer/benches/vortex_bitbuffer.rs b/vortex-buffer/benches/vortex_bitbuffer.rs index b67b3f1bf39..6b7a2492a4d 100644 --- a/vortex-buffer/benches/vortex_bitbuffer.rs +++ b/vortex-buffer/benches/vortex_bitbuffer.rs @@ -114,7 +114,11 @@ fn append_buffer_vortex_buffer(bencher: Bencher, length: usize) { }); } -#[divan::bench(args = INPUT_SIZE)] +/// Arrow has no bulk append from a `BooleanBuffer`, so this appends bit-by-bit and runs an +/// order of magnitude slower than the vortex bulk path; cap the top size accordingly. +const APPEND_BUFFER_ARROW_INPUT_SIZE: &[usize] = &[128, 1024, 2048, 16_384, 32_768]; + +#[divan::bench(args = APPEND_BUFFER_ARROW_INPUT_SIZE)] fn append_buffer_arrow_buffer(bencher: Bencher, length: usize) { bencher .with_inputs(|| { @@ -153,30 +157,45 @@ fn value_arrow_buffer(bencher: Bencher, length: usize) { }); } -#[divan::bench(args = INPUT_SIZE)] -fn slice_vortex_buffer(bencher: Bencher, length: usize) { - let buffer = BitBuffer::from_iter((0..length).map(|i| i % 2 == 0)); - bencher - .with_inputs(|| (&buffer, length / 2)) - .bench_refs(|(buffer, mid)| { - let mid = *mid; - buffer.slice(mid / 2..mid + mid / 2) - }); +/// Slicing only adjusts an offset, a length and a refcount, so its cost is independent of buffer +/// length. Measure one length, with enough slices per iteration to stay above harness overhead. +const SLICE_ITERS: usize = 64; +const SLICE_INPUT_SIZE: usize = 65_536; + +#[divan::bench] +fn slice_vortex_buffer(bencher: Bencher) { + let buffer = BitBuffer::from_iter((0..SLICE_INPUT_SIZE).map(|i| i % 2 == 0)); + let mid = SLICE_INPUT_SIZE / 2; + bencher.with_inputs(|| &buffer).bench_refs(|buffer| { + let mut total = 0; + for offset in 0..SLICE_ITERS { + total += buffer.slice(mid / 2 + offset..mid + mid / 2).len(); + } + total + }); } #[cfg(not(codspeed))] -#[divan::bench(args = INPUT_SIZE)] -fn slice_arrow_buffer(bencher: Bencher, length: usize) { - let buffer = Arrow(BooleanBuffer::from_iter((0..length).map(|i| i % 2 == 0))); - bencher - .with_inputs(|| (&buffer, length / 2)) - .bench_refs(|(buffer, mid)| { - let mid = *mid; - buffer.0.slice(mid / 2, mid / 2) - }); +#[divan::bench] +fn slice_arrow_buffer(bencher: Bencher) { + let buffer = Arrow(BooleanBuffer::from_iter( + (0..SLICE_INPUT_SIZE).map(|i| i % 2 == 0), + )); + let mid = SLICE_INPUT_SIZE / 2; + bencher.with_inputs(|| &buffer).bench_refs(|buffer| { + let mut total = 0; + for offset in 0..SLICE_ITERS { + total += buffer.0.slice(mid / 2 + offset, mid / 2).len(); + } + total + }); } -#[divan::bench(args = INPUT_SIZE)] +/// Below a few thousand bits the measurement is mostly divan's fixed per-iteration overhead +/// rather than the popcount, so it moves when the count itself has not changed. +const TRUE_COUNT_INPUT_SIZE: &[usize] = &[16_384, 65_536]; + +#[divan::bench(args = TRUE_COUNT_INPUT_SIZE)] fn true_count_vortex_buffer(bencher: Bencher, length: usize) { let buffer = BitBuffer::from_iter((0..length).map(true_count_pattern)); @@ -185,7 +204,7 @@ fn true_count_vortex_buffer(bencher: Bencher, length: usize) { .bench_refs(|buffer| buffer.true_count()) } -#[divan::bench(args = INPUT_SIZE)] +#[divan::bench(args = TRUE_COUNT_INPUT_SIZE)] fn true_count_arrow_buffer(bencher: Bencher, length: usize) { let buffer = Arrow(BooleanBuffer::from_iter( (0..length).map(true_count_pattern), diff --git a/vortex-buffer/benches/vortex_buffer.rs b/vortex-buffer/benches/vortex_buffer.rs index e1e6683e505..66afcd0dd64 100644 --- a/vortex-buffer/benches/vortex_buffer.rs +++ b/vortex-buffer/benches/vortex_buffer.rs @@ -103,7 +103,11 @@ fn map_each + FromIterator>(bencher: Bencher, n: i32) .bench_values(|buffer| B::map_each(buffer, |i| (i as u32) + 1)); } -#[divan::bench(args = INPUT_SIZE)] +/// Element-at-a-time pushes cost tens of nanoseconds each, so cap the top size to keep the +/// benchmark runtime bounded. +const PUSH_INPUT_SIZE: &[i32] = &[128, 1024, 2048, 16_384, 32_768]; + +#[divan::bench(args = PUSH_INPUT_SIZE)] fn push_vortex_buffer(bencher: Bencher, length: i32) { bencher .with_inputs(|| BufferMut::::with_capacity(length as usize)) @@ -114,7 +118,7 @@ fn push_vortex_buffer(bencher: Bencher, length: i32) { }); } -#[divan::bench(args = INPUT_SIZE)] +#[divan::bench(args = PUSH_INPUT_SIZE)] fn push_arrow_buffer(bencher: Bencher, length: i32) { bencher .with_inputs(|| { diff --git a/vortex-buffer/src/alignment.rs b/vortex-buffer/src/alignment.rs index ad5c51761bd..f4f476c7691 100644 --- a/vortex-buffer/src/alignment.rs +++ b/vortex-buffer/src/alignment.rs @@ -2,19 +2,26 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::fmt::Display; -use std::ops::Deref; use vortex_error::VortexError; use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_err; /// The alignment of a buffer. /// -/// This type is a wrapper around `usize` that ensures the alignment is a non-zero power of 2. +/// This type stores the base-2 exponent of a non-zero power-of-two alignment. #[derive(Clone, Debug, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct Alignment(usize); +pub struct Alignment(u8); impl Alignment { + /// Largest alignment accepted from untrusted serialized input. + /// + /// This admits 64KiB page alignment, as used on some ARM systems, while bounding the extra + /// allocation required to satisfy an alignment from untrusted input. + pub const MAX_UNTRUSTED: Self = Alignment::new(64 * 1024); + /// Default alignment for device-to-host buffer copies. pub const HOST_COPY: Self = Alignment::new(256); @@ -33,10 +40,11 @@ impl Alignment { /// /// Panics if `align` is zero or is not a power of 2. #[inline] + #[expect(clippy::cast_possible_truncation, reason = "usize has at most 64 bits")] pub const fn new(align: usize) -> Self { assert!(align > 0, "Alignment must be greater than 0"); assert!(align.is_power_of_two(), "Alignment must be a power of 2"); - Self(align) + Self(align.trailing_zeros() as u8) } /// Create a new 1-byte alignment. @@ -96,7 +104,7 @@ impl Alignment { #[inline] pub const fn is_offset_aligned(&self, offset: usize) -> bool { // Alignment is always a power of 2, so a mask test is equivalent to `offset % self == 0`. - offset & (self.0 - 1) == 0 + offset & (self.as_usize() - 1) == 0 } /// Check if the given pointer is aligned to this alignment. @@ -107,33 +115,67 @@ impl Alignment { /// Returns the log2 of the alignment. pub fn exponent(&self) -> u8 { - u8::try_from(self.0.trailing_zeros()) - .vortex_expect("alignment is a power of 2 within usize, so its exponent fits in u8") + self.0 } /// Create from the log2 exponent of the alignment. /// /// ## Panics /// - /// Panics if `1 << exponent` overflows `usize`. + /// Panics if `1 << exponent` overflows `usize`. Use [`Self::try_from_exponent`] when parsing + /// untrusted input. #[inline] pub const fn from_exponent(exponent: u8) -> Self { - Self::new(1 << exponent) + assert!( + (exponent as u32) < usize::BITS, + "Alignment exponent must fit in usize" + ); + Self(exponent) } -} -impl Display for Alignment { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{}", self.0) + /// Create from the log2 exponent of the alignment, returning an error rather than panicking if + /// `1 << exponent` would overflow `usize`. + /// + /// Prefer this over [`from_exponent`](Self::from_exponent) when the exponent originates from + /// untrusted input such as a serialized file, where a too-large value must not panic. + #[inline] + pub fn try_from_exponent(exponent: u8) -> VortexResult { + if u32::from(exponent) >= usize::BITS { + vortex_bail!( + "Alignment exponent {exponent} is too large for a {}-bit usize", + usize::BITS + ); + } + Ok(Self::new(1 << exponent)) } -} -impl Deref for Alignment { - type Target = usize; + /// Create an alignment from an exponent in untrusted serialized input. + /// + /// In addition to rejecting exponents that do not fit in `usize`, this rejects alignments + /// large enough to cause an unreasonable allocation when a buffer needs to be copied to + /// satisfy the alignment. + #[inline] + pub fn try_from_untrusted_exponent(exponent: u8) -> VortexResult { + let alignment = Self::try_from_exponent(exponent)?; + if alignment > Self::MAX_UNTRUSTED { + vortex_bail!( + "Untrusted alignment {alignment} exceeds the {}-byte maximum", + Self::MAX_UNTRUSTED + ); + } + Ok(alignment) + } + /// Return the alignment in bytes. #[inline] - fn deref(&self) -> &Self::Target { - &self.0 + pub const fn as_usize(self) -> usize { + 1 << self.0 + } +} + +impl Display for Alignment { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.as_usize()) } } @@ -154,14 +196,14 @@ impl From for Alignment { impl From for usize { #[inline] fn from(value: Alignment) -> Self { - value.0 + value.as_usize() } } impl From for u32 { #[inline] fn from(value: Alignment) -> Self { - u32::try_from(value.0).vortex_expect("Alignment must fit into u32") + u32::try_from(value.as_usize()).vortex_expect("Alignment must fit into u32") } } @@ -179,7 +221,7 @@ impl TryFrom for Alignment { return Err(vortex_err!("Alignment must be a power of 2, got {value}")); } - Ok(Self(value)) + Ok(Self::new(value)) } } @@ -197,7 +239,7 @@ mod test { fn alignment_above_u16() { // 64KiB alignment (one past `u16::MAX`) is valid — common on ARM with 64K pages. let alignment = Alignment::new(u16::MAX as usize + 1); - assert_eq!(*alignment, 1 << 16); + assert_eq!(alignment.as_usize(), 1 << 16); assert_eq!(alignment, Alignment::from_exponent(16)); } @@ -236,6 +278,28 @@ mod test { assert!(Alignment::try_from(3u32).is_err()); } + #[test] + fn try_from_exponent() { + match Alignment::try_from_exponent(10) { + Ok(alignment) => assert_eq!(alignment, Alignment::new(1024)), + Err(err) => panic!("valid exponent should succeed: {err}"), + } + // Exponents whose `1 << exponent` would overflow a usize must error rather than panic. + // 64 is `>= usize::BITS` on both 32- and 64-bit targets. + assert!(Alignment::try_from_exponent(64).is_err()); + assert!(Alignment::try_from_exponent(u8::MAX).is_err()); + } + + #[test] + fn try_from_untrusted_exponent() { + assert_eq!( + Alignment::try_from_untrusted_exponent(16).unwrap(), + Alignment::new(64 * 1024) + ); + assert!(Alignment::try_from_untrusted_exponent(17).is_err()); + assert!(Alignment::try_from_untrusted_exponent(u8::MAX).is_err()); + } + #[test] fn into_u32() { let alignment = Alignment::new(64); diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs new file mode 100644 index 00000000000..22792408cca --- /dev/null +++ b/vortex-buffer/src/allocation.rs @@ -0,0 +1,538 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Allocator-backed storage for Vortex buffers. + +use std::alloc::Layout; +use std::any::Any; +use std::fmt; +use std::fmt::Debug; +use std::mem::ManuallyDrop; +use std::ptr::NonNull; +use std::sync::Arc; + +use allocator_api2::alloc::AllocError; +use allocator_api2::alloc::Allocator; +use allocator_api2::alloc::Global; +use allocator_api2::alloc::handle_alloc_error; +use vortex_error::VortexExpect; + +use crate::Alignment; +use crate::BufferMut; + +/// An allocator that can back a Vortex buffer. +/// +/// Vortex over-allocates raw storage and aligns the buffer within it. +pub trait BufferAllocator: Allocator + Debug + Send + Sync + 'static {} + +impl BufferAllocator for A where A: Allocator + Debug + Send + Sync + 'static {} + +/// A shared reference to a buffer allocator. +/// +/// The static allocator does not need shared ownership, so it is stored without an [`Arc`]. This +/// makes cloning the common static allocator a simple value copy. +#[derive(Clone)] +pub struct BufferAllocatorRef( + // `None` selects the static allocator without allocating or updating an Arc reference count. + // `Some` keeps a custom allocator alive for as long as its buffers need it. + Option>, +); + +impl BufferAllocatorRef { + /// Wrap an allocator in a shared reference. + pub fn new(allocator: impl BufferAllocator) -> Self { + Self(Some(Arc::new(allocator))) + } + + /// Return a shared reference to the static allocator. + pub fn statically_allocated() -> Self { + Self(None) + } + + /// Return a borrowed reference to the static allocator. + pub fn static_ref() -> &'static Self { + &STATIC_ALLOCATOR + } + + pub(crate) fn is_statically_allocated(&self) -> bool { + self.0.is_none() + } + + /// Returns true if both references point to the same allocator. + pub fn ptr_eq(&self, other: &Self) -> bool { + match (&self.0, &other.0) { + (None, None) => true, + (Some(lhs), Some(rhs)) => Arc::ptr_eq(lhs, rhs), + _ => false, + } + } + + /// Create a mutable buffer with this allocator. + pub fn with_capacity(&self, capacity: usize) -> BufferMut { + BufferMut::with_capacity_in(capacity, self.clone()) + } + + /// Create an aligned mutable buffer with this allocator. + pub fn with_capacity_aligned(&self, capacity: usize, alignment: Alignment) -> BufferMut { + BufferMut::with_capacity_aligned_in(capacity, alignment, self.clone()) + } + + /// Create a zeroed mutable buffer with this allocator. + pub fn zeroed(&self, len: usize) -> BufferMut { + BufferMut::zeroed_in(len, self.clone()) + } + + /// Copy values into a mutable buffer made by this allocator. + pub fn copy_from(&self, values: impl AsRef<[T]>) -> BufferMut { + BufferMut::copy_from_in(values, self.clone()) + } +} + +impl Debug for BufferAllocatorRef { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match &self.0 { + Some(allocator) => allocator.fmt(f), + None => StaticBufferAllocator.fmt(f), + } + } +} + +// SAFETY: all calls are forwarded to the same allocator value held by the Arc. +unsafe impl Allocator for BufferAllocatorRef { + fn allocate(&self, layout: Layout) -> Result, AllocError> { + match &self.0 { + Some(allocator) => allocator.allocate(layout), + None => Global.allocate(layout), + } + } + + fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { + match &self.0 { + Some(allocator) => allocator.allocate_zeroed(layout), + None => Global.allocate_zeroed(layout), + } + } + + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + // SAFETY: the caller upholds the Allocator contract. + match &self.0 { + Some(allocator) => unsafe { allocator.deallocate(ptr, layout) }, + None => unsafe { Global.deallocate(ptr, layout) }, + } + } + + unsafe fn grow( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + // SAFETY: the caller upholds the Allocator contract. + match &self.0 { + Some(allocator) => unsafe { allocator.grow(ptr, old_layout, new_layout) }, + None => unsafe { Global.grow(ptr, old_layout, new_layout) }, + } + } + + unsafe fn grow_zeroed( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + // SAFETY: the caller upholds the Allocator contract. + match &self.0 { + Some(allocator) => unsafe { allocator.grow_zeroed(ptr, old_layout, new_layout) }, + None => unsafe { Global.grow_zeroed(ptr, old_layout, new_layout) }, + } + } + + unsafe fn shrink( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + // SAFETY: the caller upholds the Allocator contract. + match &self.0 { + Some(allocator) => unsafe { allocator.shrink(ptr, old_layout, new_layout) }, + None => unsafe { Global.shrink(ptr, old_layout, new_layout) }, + } + } +} + +/// The allocator used by buffer APIs that do not take an allocator. +#[derive(Clone, Copy, Debug, Default)] +pub struct StaticBufferAllocator; + +impl StaticBufferAllocator { + /// Create a mutable buffer with the static allocator. + pub fn with_capacity(capacity: usize) -> BufferMut { + BufferMut::with_capacity(capacity) + } + + /// Create an aligned mutable buffer with the static allocator. + pub fn with_capacity_aligned(capacity: usize, alignment: Alignment) -> BufferMut { + BufferMut::with_capacity_aligned(capacity, alignment) + } + + /// Create a zeroed mutable buffer with the static allocator. + pub fn zeroed(len: usize) -> BufferMut { + BufferMut::zeroed(len) + } + + /// Copy values into a mutable buffer made by the static allocator. + pub fn copy_from(values: impl AsRef<[T]>) -> BufferMut { + BufferMut::copy_from(values) + } +} + +// SAFETY: Global satisfies the Allocator contract and this type only forwards to it. +unsafe impl Allocator for StaticBufferAllocator { + fn allocate(&self, layout: Layout) -> Result, AllocError> { + Global.allocate(layout) + } + + fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { + Global.allocate_zeroed(layout) + } + + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.deallocate(ptr, layout) } + } + + unsafe fn grow( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.grow(ptr, old_layout, new_layout) } + } + + unsafe fn grow_zeroed( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.grow_zeroed(ptr, old_layout, new_layout) } + } + + unsafe fn shrink( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.shrink(ptr, old_layout, new_layout) } + } +} + +static STATIC_ALLOCATOR: BufferAllocatorRef = BufferAllocatorRef(None); + +pub(crate) struct Allocation { + ptr: NonNull, + layout: Layout, + allocator: BufferAllocatorRef, +} + +// SAFETY: Allocation owns its memory, and its allocator is Send + Sync. +unsafe impl Send for Allocation {} +// SAFETY: shared access to Allocation never permits mutation of the allocation. +unsafe impl Sync for Allocation {} + +impl Allocation { + pub(crate) fn allocate(layout: Layout, allocator: BufferAllocatorRef) -> Self { + Self::allocate_impl(layout, allocator, false) + } + + pub(crate) fn allocate_zeroed(layout: Layout, allocator: BufferAllocatorRef) -> Self { + Self::allocate_impl(layout, allocator, true) + } + + pub(crate) fn from_vec(vec: Vec) -> Self { + assert!(!std::mem::needs_drop::()); + + let mut vec = ManuallyDrop::new(vec); + let layout = Layout::array::(vec.capacity()) + .unwrap_or_else(|_| unreachable!("a Vec capacity always has a valid layout")); + let ptr = NonNull::new(vec.as_mut_ptr().cast()) + .vortex_expect("a Vec always has a non-null pointer"); + + Self { + ptr, + layout, + allocator: BufferAllocatorRef::statically_allocated(), + } + } + + fn allocate_impl(layout: Layout, allocator: BufferAllocatorRef, zeroed: bool) -> Self { + if layout.size() == 0 { + return Self { + ptr: layout.dangling_ptr(), + layout, + allocator, + }; + } + + let allocation = if zeroed { + allocator.allocate_zeroed(layout) + } else { + allocator.allocate(layout) + } + .unwrap_or_else(|_| handle_alloc_error(layout)); + + Self { + ptr: allocation.cast(), + layout, + allocator, + } + } + + #[allow(clippy::inline_always)] + #[inline(always)] + pub(crate) fn ptr(&self) -> NonNull { + self.ptr + } + + #[allow(clippy::inline_always)] + #[inline(always)] + pub(crate) fn size(&self) -> usize { + self.layout.size() + } + + #[allow(clippy::inline_always)] + #[inline(always)] + pub(crate) fn alignment(&self) -> usize { + self.layout.align() + } + + #[allow(clippy::inline_always)] + #[inline(always)] + pub(crate) fn allocator(&self) -> &BufferAllocatorRef { + &self.allocator + } + + pub(crate) fn grow(&mut self, new_layout: Layout) { + let allocation = if self.layout.size() == 0 { + self.allocator.allocate(new_layout) + } else { + // SAFETY: ptr denotes a live block owned by allocator, old_layout fits the block, and + // the caller only grows the allocation. + unsafe { self.allocator.grow(self.ptr, self.layout, new_layout) } + } + .unwrap_or_else(|_| handle_alloc_error(new_layout)); + self.ptr = allocation.cast(); + self.layout = new_layout; + } +} + +impl Drop for Allocation { + fn drop(&mut self) { + if self.layout.size() == 0 { + return; + } + // SAFETY: ptr and layout describe a live block allocated by self.allocator. + unsafe { self.allocator.deallocate(self.ptr, self.layout) } + } +} + +pub(crate) enum BufferBacking { + Owned(Allocation), + Bytes(bytes::Bytes), + #[cfg(feature = "arrow")] + Arrow(arrow_buffer::Buffer), + External { + _owner: Box, + }, +} + +impl BufferBacking { + #[allow(clippy::inline_always)] + #[inline(always)] + pub(crate) fn allocator(&self) -> &BufferAllocatorRef { + match self { + Self::Owned(allocation) => allocation.allocator(), + Self::Bytes(_) | Self::External { .. } => &STATIC_ALLOCATOR, + #[cfg(feature = "arrow")] + Self::Arrow(_) => &STATIC_ALLOCATOR, + } + } +} + +#[cfg(test)] +mod tests { + use std::alloc::Layout; + use std::ptr::NonNull; + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + use allocator_api2::alloc::AllocError; + use allocator_api2::alloc::Allocator; + use allocator_api2::alloc::Global; + use rstest::rstest; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + + use crate::Alignment; + use crate::BufferAllocatorRef; + use crate::BufferMut; + + #[derive(Clone, Debug, Default)] + struct TrackingAllocator { + state: Arc, + } + + #[derive(Debug, Default)] + struct TrackingState { + allocations: AtomicUsize, + deallocations: AtomicUsize, + grows: AtomicUsize, + alignment: AtomicUsize, + } + + // SAFETY: this forwards all memory operations to Global and only records call metadata. + unsafe impl Allocator for TrackingAllocator { + fn allocate(&self, layout: Layout) -> Result, AllocError> { + self.state.allocations.fetch_add(1, Ordering::Relaxed); + self.state + .alignment + .store(layout.align(), Ordering::Relaxed); + Global.allocate(layout) + } + + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + self.state.deallocations.fetch_add(1, Ordering::Relaxed); + // SAFETY: the caller passes the pointer and layout returned by Global. + unsafe { Global.deallocate(ptr, layout) } + } + + unsafe fn grow( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + self.state.grows.fetch_add(1, Ordering::Relaxed); + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.grow(ptr, old_layout, new_layout) } + } + } + + #[test] + fn allocator_identity() { + let static_allocator = BufferAllocatorRef::statically_allocated(); + assert!(static_allocator.ptr_eq(&BufferAllocatorRef::statically_allocated())); + + let custom_allocator = BufferAllocatorRef::new(TrackingAllocator::default()); + assert!(custom_allocator.ptr_eq(&custom_allocator.clone())); + assert!(!custom_allocator.ptr_eq(&static_allocator)); + assert!(!custom_allocator.ptr_eq(&BufferAllocatorRef::new(TrackingAllocator::default()))); + } + + #[test] + fn allocation_lives_until_last_view() { + let allocator = TrackingAllocator::default(); + let state = Arc::clone(&allocator.state); + let buffer = BufferAllocatorRef::new(allocator) + .copy_from([1u32, 2, 3, 4]) + .freeze(); + let view = buffer.slice(0..2); + + assert_eq!(state.allocations.load(Ordering::Relaxed), 1); + assert_eq!( + state.alignment.load(Ordering::Relaxed), + Alignment::of::().as_usize() + ); + drop(buffer); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 0); + drop(view); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 1); + } + + #[rstest] + fn buffer_growth_uses_allocator_grow(#[values(4, 64, 4096)] alignment: usize) { + let allocator = TrackingAllocator::default(); + let state = Arc::clone(&allocator.state); + let alignment = Alignment::new(alignment); + let mut buffer = + BufferAllocatorRef::new(allocator).with_capacity_aligned::(1, alignment); + let initial_capacity = buffer.capacity(); + buffer.extend(std::iter::repeat_n(7, initial_capacity)); + + buffer.push(u32::MAX); + assert!(alignment.is_ptr_aligned(buffer.as_ptr())); + + assert_eq!(&buffer[..initial_capacity], vec![7; initial_capacity]); + assert_eq!(buffer[initial_capacity], u32::MAX); + assert_eq!(state.allocations.load(Ordering::Relaxed), 1); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 0); + assert_eq!(state.grows.load(Ordering::Relaxed), 1); + + drop(buffer); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 1); + } + + #[test] + fn zero_capacity_does_not_allocate() { + let allocator = TrackingAllocator::default(); + let state = Arc::clone(&allocator.state); + let mut buffer = BufferAllocatorRef::new(allocator).with_capacity::(0); + + assert_eq!(buffer.capacity(), 0); + assert!(Alignment::DEFAULT_ALIGNMENT.is_offset_aligned(buffer.as_ptr().addr())); + assert_eq!(state.allocations.load(Ordering::Relaxed), 0); + + buffer.push(42); + + assert_eq!(buffer.as_slice(), [42]); + assert_eq!(state.allocations.load(Ordering::Relaxed), 1); + assert_eq!(state.grows.load(Ordering::Relaxed), 0); + } + + #[test] + fn empty_buffers_preserve_allocator_without_allocating() -> VortexResult<()> { + let allocator = TrackingAllocator::default(); + let state = Arc::clone(&allocator.state); + let allocator = BufferAllocatorRef::new(allocator); + let buffer = BufferMut::::zeroed_in(0, allocator.clone()); + let buffer = buffer.freeze(); + let copy = buffer.clone().into_mut(); + assert!(copy.allocator().ptr_eq(&allocator)); + let mut buffer = buffer + .try_into_mut() + .map_err(|_| vortex_err!("unique buffer"))?; + buffer.reserve(0); + assert!(buffer.is_empty()); + assert!(buffer.allocator().ptr_eq(&allocator)); + drop((copy, buffer)); + assert_eq!(state.allocations.load(Ordering::Relaxed), 0); + assert_eq!(state.grows.load(Ordering::Relaxed), 0); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 0); + Ok(()) + } + + #[test] + fn shared_into_mut_preserves_allocator() { + let allocator = TrackingAllocator::default(); + let state = Arc::clone(&allocator.state); + let allocator = BufferAllocatorRef::new(allocator); + let original = allocator.copy_from([1u32, 2, 3]).freeze(); + let mut copy = original.clone().into_mut(); + assert!(copy.allocator().ptr_eq(&allocator)); + copy[0] = 42; + assert_eq!(original.as_slice(), [1, 2, 3]); + assert_eq!(copy.as_slice(), [42, 2, 3]); + assert_eq!(state.allocations.load(Ordering::Relaxed), 2); + drop(copy); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 1); + drop(original); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 2); + } +} diff --git a/vortex-buffer/src/arrow.rs b/vortex-buffer/src/arrow.rs index 63eec6deb60..424aa1878c0 100644 --- a/vortex-buffer/src/arrow.rs +++ b/vortex-buffer/src/arrow.rs @@ -3,7 +3,6 @@ use arrow_buffer::ArrowNativeType; use arrow_buffer::OffsetBuffer; -use bytes::Bytes; use vortex_error::vortex_panic; use crate::Alignment; @@ -13,7 +12,10 @@ use crate::ByteBuffer; impl Buffer { /// Converts the buffer zero-copy into a `arrow_buffer::Buffer`. pub fn into_arrow_scalar_buffer(self) -> arrow_buffer::ScalarBuffer { - let buffer = arrow_buffer::Buffer::from(self.into_inner()); + if self.is_empty() { + return Vec::new().into(); + } + let buffer = self.into_byte_buffer().into_arrow_buffer(); arrow_buffer::ScalarBuffer::from(buffer) } @@ -25,22 +27,18 @@ impl Buffer { /// alignment is not sufficient for type T. pub fn from_arrow_scalar_buffer(arrow: arrow_buffer::ScalarBuffer) -> Self { let length = arrow.len(); - let bytes = Bytes::from_owner(ArrowWrapper(arrow.into_inner())); + let arrow = arrow.into_inner(); let alignment = Alignment::of::(); - if bytes.as_ptr().align_offset(*alignment) != 0 { + if arrow.as_ptr().align_offset(alignment.as_usize()) != 0 { vortex_panic!( "Arrow buffer is not aligned to the requested alignment: {}", alignment ); } - Self { - bytes, - length, - alignment, - _marker: Default::default(), - } + debug_assert_eq!(length, arrow.len() / size_of::()); + Self::from_arrow_owner(arrow, length, alignment) } /// Converts the buffer zero-copy into a `arrow_buffer::OffsetBuffer`. @@ -55,7 +53,11 @@ impl Buffer { impl ByteBuffer { /// Converts the buffer zero-copy into a `arrow_buffer::Buffer`. pub fn into_arrow_buffer(self) -> arrow_buffer::Buffer { - arrow_buffer::Buffer::from(self.into_inner()) + if let Some(crate::BufferBacking::Arrow(arrow)) = self.backing.as_deref() { + let offset = self.ptr.addr().get() - arrow.as_ptr().addr(); + return arrow.slice_with_length(offset, self.length); + } + arrow_buffer::Buffer::from(self.into_bytes()) } /// Convert an Arrow scalar buffer into a Vortex scalar buffer. @@ -66,30 +68,14 @@ impl ByteBuffer { pub fn from_arrow_buffer(arrow: arrow_buffer::Buffer, alignment: Alignment) -> Self { let length = arrow.len(); - let bytes = Bytes::from_owner(ArrowWrapper(arrow)); - if bytes.as_ptr().align_offset(*alignment) != 0 { + if arrow.as_ptr().align_offset(alignment.as_usize()) != 0 { vortex_panic!( "Arrow buffer is not aligned to the requested alignment: {}", alignment ); } - Self { - bytes, - length, - alignment, - _marker: Default::default(), - } - } -} - -/// A wrapper struct to allow `arrow_buffer::Buffer` to implement `AsRef<[u8]>` for -/// `Bytes::from_owner`. -struct ArrowWrapper(arrow_buffer::Buffer); - -impl AsRef<[u8]> for ArrowWrapper { - fn as_ref(&self) -> &[u8] { - self.0.as_slice() + Self::from_arrow_owner(arrow, length, alignment) } } @@ -118,11 +104,22 @@ mod test { assert_eq!(scalar.as_ptr(), buf.as_ptr(), "Conversion not zero-copy") } + #[test] + fn empty_into_arrow_scalar_buffer() { + let scalar = Buffer::::empty().into_arrow_scalar_buffer(); + + assert!(scalar.is_empty()); + assert_eq!(scalar.as_ptr().align_offset(align_of::()), 0); + } + #[test] fn from_arrow_buffer() { let arrow = ArrowBuffer::from_vec(vec![0i32, 1, 2]); let buf = Buffer::from_arrow_buffer(arrow.clone(), Alignment::of::()); assert_eq!(arrow.as_ref(), buf.as_slice(), "Buffer values differ"); assert_eq!(arrow.as_ptr(), buf.as_ptr(), "Conversion not zero-copy"); + + let round_trip = buf.into_arrow_buffer(); + assert_eq!(round_trip.as_ptr(), arrow.as_ptr()); } } diff --git a/vortex-buffer/src/bit/buf.rs b/vortex-buffer/src/bit/buf.rs index 8dc5e38a42c..6bc2d3c3626 100644 --- a/vortex-buffer/src/bit/buf.rs +++ b/vortex-buffer/src/bit/buf.rs @@ -1,17 +1,18 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::convert::Infallible; use std::fmt::Display; use std::fmt::Formatter; use std::fmt::Result as FmtResult; use std::ops::BitAnd; use std::ops::BitOr; use std::ops::BitXor; -use std::ops::Bound; use std::ops::Not; use std::ops::RangeBounds; use crate::Alignment; +use crate::BitBufferMeta; use crate::BitBufferMut; use crate::Buffer; use crate::BufferMut; @@ -95,6 +96,7 @@ impl BitBuffer { /// Create a new `BoolBuffer` backed by a [`ByteBuffer`] with `len` bits in view. /// /// Panics if the buffer is not large enough to hold `len` bits. + #[inline] pub fn new(buffer: ByteBuffer, len: usize) -> Self { assert!( buffer.len() * 8 >= len, @@ -115,6 +117,7 @@ impl BitBuffer { /// the given `offset` (in bits). /// /// Panics if the buffer is not large enough to hold `len` bits after the offset. + #[inline] pub fn new_with_offset(buffer: ByteBuffer, len: usize, offset: usize) -> Self { assert!( len.saturating_add(offset) <= buffer.len().saturating_mul(8), @@ -142,6 +145,7 @@ impl BitBuffer { } /// Create a new `BoolBuffer` of length `len` where all bits are set (true). + #[inline] pub fn new_set(len: usize) -> Self { let words = len.div_ceil(8); let buffer = buffer![0xFF; words]; @@ -154,6 +158,7 @@ impl BitBuffer { } /// Create a new `BoolBuffer` of length `len` where all bits are unset (false). + #[inline] pub fn new_unset(len: usize) -> Self { let words = len.div_ceil(8); let buffer = Buffer::zeroed(words); @@ -171,11 +176,13 @@ impl BitBuffer { } /// Create a new empty `BitBuffer`. + #[inline] pub fn empty() -> Self { Self::new_set(0) } /// Create a new `BitBuffer` of length `len` where all bits are set to `value`. + #[inline] pub fn full(value: bool, len: usize) -> Self { if value { Self::new_set(len) @@ -184,39 +191,46 @@ impl BitBuffer { } } - /// Invokes `f` with indexes `0..len` collecting the boolean results into a new [`BitBuffer`]. + /// Collects `len` Boolean values from `f` into a packed [`BitBuffer`]. + /// + /// Calls `f` exactly once for each index in `0..len`, in order. + /// + /// # Code generation /// - /// `f` is invoked exactly once per index, in ascending order, and the results are packed - /// with the baseline SIMD byte→bit instruction of the target. + /// `collect_bool_words` calls `collect_bool_words_inline`, which selects a `pack_bool_word_*` + /// kernel at compile time and passes it to `collect_bool_words_with`. That shared word loop + /// materializes each full 64-value chunk as a byte-per-value `[bool; 64]`, then passes it to the + /// selected kernel. For simple predicates, LLVM vectorizes the loop and removes the physical + /// stack array. On AVX-512, it still combines the comparison masks, expands the result into 64 + /// `0` or `1` bytes with `vpbroadcastq` and `vmovdqu8`, then recreates the mask with `vptestmb`. + /// [LLVM issue #219235](https://github.com/llvm/llvm-project/issues/219235) tracks replacing + /// that round trip with a direct `kmovq` store. The conversion is per chunk. This method does + /// not create a full-column byte buffer. /// /// # Performance /// - /// The packing is a few instructions per 64 bits, so evaluating `f` is usually the - /// bottleneck. In particular, a bounds-checked slice access in `f` (`|i| values[i] > x`) - /// blocks vectorization of the gather and can cost ~10x the packing itself. Since `f` only - /// ever sees indices `0..len`, callers reading from a slice with `len <= values.len()` may - /// soundly use `|i| unsafe { *values.get_unchecked(i) }`. + /// `collect_bool_words_inline` and `collect_bool_words_with` can inline into the caller, so LLVM + /// sees `f`, the fill loop, and the packing kernel together. A retained bounds check inside `f` + /// can prevent vectorization. A caller that proves `len <= values.len()` can use + /// `unsafe { *values.get_unchecked(i) }` because this method only passes indices in `0..len`. /// - /// Prefer this entry point for every predicate. Only switch to - /// [`Self::collect_bool_multiversioned`] after carefully checking that your specific `f` - /// meets its contract (a trivially cheap, bounds-check-free gather or comparison) — - /// ideally with a benchmark. + /// Use this method for general predicates. Use [`Self::collect_bool_multiversioned`] only for + /// the specialized predicates described there. #[inline] pub fn collect_bool bool>(len: usize, f: F) -> Self { BitBufferMut::collect_bool(len, f).freeze() } - /// Like [`Self::collect_bool`], but compiles the packing loop — with `f` inside it — once - /// per CPU feature level (AVX-512BW/AVX2/baseline) and selects a clone by runtime feature - /// detection. + /// Collects Boolean values with a fill-and-pack loop selected for the current CPU. /// - /// Calling this asserts that `f` is small and simple enough (e.g. a bounds-check-free slice - /// gather or comparison) that duplicating it per feature level and paying a - /// `#[target_feature]` call boundary beats inlining it once into your function. For any - /// non-trivial `f` that assertion is false — the boundary deoptimizes the predicate — so - /// unless you have carefully checked (ideally benchmarked) that your specific `f` - /// qualifies, use [`Self::collect_bool`]. See - /// [`collect_bool_words_multiversioned`](crate::bit::collect_bool_words_multiversioned). + /// This has the same callback contract as [`Self::collect_bool`]. On x86-64, + /// `collect_bool_words_multiversioned` selects `collect_bool_words_avx512`, + /// `collect_bool_words_avx2`, or `collect_bool_words_inline` at runtime. Each wider version is a + /// `#[target_feature]` function, so Rust cannot inline it into a caller compiled without those + /// features. + /// + /// Use this method only for a small, bounds-check-free predicate whose wider loop has been + /// benchmarked. Use [`Self::collect_bool`] for general predicates. #[inline] pub fn collect_bool_multiversioned bool>(len: usize, f: F) -> Self { BitBufferMut::collect_bool_multiversioned(len, f).freeze() @@ -271,6 +285,7 @@ impl BitBuffer { } /// Clear all bits in the buffer, preserving existing capacity. + #[inline] pub fn clear(&mut self) { self.buffer.clear(); self.len = 0; @@ -293,12 +308,14 @@ impl BitBuffer { } /// Offset of the start of the buffer in bits. + #[allow(clippy::inline_always)] #[inline(always)] pub fn offset(&self) -> usize { self.offset } /// Get a reference to the underlying buffer. + #[allow(clippy::inline_always)] #[inline(always)] pub fn inner(&self) -> &ByteBuffer { &self.buffer @@ -345,26 +362,9 @@ impl BitBuffer { /// for `len` bits. /// /// Panics if the slice would extend beyond the end of the buffer. + #[inline] pub fn slice(&self, range: impl RangeBounds) -> Self { - let start = match range.start_bound() { - Bound::Included(&s) => s, - Bound::Excluded(&s) => s + 1, - Bound::Unbounded => 0, - }; - let end = match range.end_bound() { - Bound::Included(&e) => e + 1, - Bound::Excluded(&e) => e, - Bound::Unbounded => self.len, - }; - - assert!(start <= end); - assert!(start <= self.len); - assert!(end <= self.len); - let len = end - start; - - let offset = self.offset + start; - let byte_offset = offset / 8; - let bit_offset = offset % 8; + let (byte_offset, meta) = BitBufferMeta::new(self.offset, self.len).slice(range); // Trim whole bytes off the front directly rather than going through `new_with_offset`, // which would slice (and re-clone) the clone we'd have to pass it. @@ -376,8 +376,8 @@ impl BitBuffer { Self { buffer, - offset: bit_offset, - len, + offset: meta.offset(), + len: meta.len(), } } @@ -394,6 +394,7 @@ impl BitBuffer { } /// Access chunks of the buffer aligned to 8 byte boundary as [prefix, \, suffix] + #[inline] pub fn unaligned_chunks(&self) -> UnalignedBitChunk<'_> { UnalignedBitChunk::new(self.buffer.as_slice(), self.offset, self.len) } @@ -401,6 +402,7 @@ impl BitBuffer { /// Access chunks of the underlying buffer as 8 byte chunks with a final trailer /// /// If you're performing operations on a single buffer, prefer [BitBuffer::unaligned_chunks] + #[inline] pub fn chunks(&self) -> BitChunks<'_> { BitChunks::new(self.buffer.as_slice(), self.offset, self.len) } @@ -431,10 +433,41 @@ impl BitBuffer { /// which logical bit position holds that rank. /// /// Returns `None` if `nth` is greater than or equal to the number of set bits. + #[inline] pub fn select(&self, nth: usize) -> Option { bit_select(self.buffer.as_slice(), self.offset, self.len, nth) } + /// Returns the index of the last set bit, or `None` if every bit is unset. + /// + /// This scans from the end a word at a time, avoiding the full forward scan required by + /// [`Self::select`] when selecting the final set bit. + #[inline] + pub fn last_set_index(&self) -> Option { + let chunks = self.unaligned_chunks(); + let lead = chunks.lead_padding(); + let prefix_words = usize::from(chunks.prefix().is_some()); + + if let Some(word) = chunks.suffix() + && word != 0 + { + let word_index = prefix_words + chunks.chunks().len(); + return Some(word_index * 64 + 63 - word.leading_zeros() as usize - lead); + } + + for (index, &word) in chunks.chunks().iter().enumerate().rev() { + if word != 0 { + let word_index = prefix_words + index; + return Some(word_index * 64 + 63 - word.leading_zeros() as usize - lead); + } + } + + chunks.prefix().filter(|word| *word != 0).map(|word| { + debug_assert!(word.trailing_zeros() as usize >= lead); + 63 - word.leading_zeros() as usize - lead + }) + } + /// Get the number of unset bits in the buffer. #[inline] pub fn false_count(&self) -> usize { @@ -442,16 +475,19 @@ impl BitBuffer { } /// Iterator over bits in the buffer + #[inline] pub fn iter(&self) -> BitIterator<'_> { BitIterator::new(self.buffer.as_slice(), self.offset, self.len) } /// Iterator over set indices of the underlying buffer + #[inline] pub fn set_indices(&self) -> BitIndexIterator<'_> { BitIndexIterator::new(self.buffer.as_slice(), self.offset, self.len) } /// Iterator over set slices of the underlying buffer + #[inline] pub fn set_slices(&self) -> BitSliceIterator<'_> { BitSliceIterator::new(self.buffer.as_slice(), self.offset, self.len) } @@ -466,21 +502,37 @@ impl BitBuffer { /// (whose per-`next` iterator state does not inline as well). #[inline] pub fn for_each_set_index(&self, mut f: F) { + let Ok(()) = self.try_for_each_set_index(|index| { + f(index); + Ok::<_, Infallible>(()) + }); + } + + /// Fallible variant of [`for_each_set_index`](Self::for_each_set_index). + /// + /// Stops and returns the first error from `f`. + #[inline] + pub fn try_for_each_set_index(&self, mut f: F) -> Result<(), E> + where + F: FnMut(usize) -> Result<(), E>, + { let mut base = 0usize; for word in self.chunks().iter_padded() { if word == u64::MAX { for k in 0..64 { - f(base + k); + f(base + k)?; } } else { let mut w = word; while w != 0 { - f(base + w.trailing_zeros() as usize); + f(base + w.trailing_zeros() as usize)?; w &= w - 1; } } base += 64; } + + Ok(()) } /// Created a new BitBuffer with offset reset to 0 @@ -502,11 +554,13 @@ impl BitBuffer { impl BitBuffer { /// Returns the offset, len and underlying buffer. + #[inline] pub fn into_inner(self) -> (usize, usize, ByteBuffer) { (self.offset, self.len, self.buffer) } /// Attempt to convert this `BitBuffer` into a mutable version. + #[inline] pub fn try_into_mut(self) -> Result { match self.buffer.try_into_mut() { Ok(buffer) => Ok(BitBufferMut::from_buffer(buffer, self.offset, self.len)), @@ -528,6 +582,7 @@ impl From> for BitBuffer { } impl FromIterator for BitBuffer { + #[inline] fn from_iter>(iter: T) -> Self { BitBufferMut::from_iter(iter).freeze() } @@ -813,6 +868,30 @@ mod tests { BitBuffer::from_indices(5, [0, 5]); } + #[rstest] + #[case(0, 0, None)] + #[case(3, 7, None)] + #[case(0, 1, Some(0))] + #[case(8, 64, Some(63))] + #[case(13, 65, Some(0))] + #[case(13, 65, Some(64))] + #[case(67, 151, Some(97))] + #[case(67, 151, Some(150))] + fn last_set_index_handles_offsets_and_padding( + #[case] offset: usize, + #[case] len: usize, + #[case] expected: Option, + ) { + let backing = BitBuffer::from_iter( + std::iter::repeat_n(true, offset) + .chain((0..len).map(|index| Some(index) == expected)) + .chain(std::iter::repeat_n(true, 7)), + ); + let buffer = BitBuffer::new_with_offset(backing.inner().clone(), len, offset); + + assert_eq!(buffer.last_set_index(), expected); + } + #[rstest] #[case(5)] #[case(8)] @@ -971,12 +1050,21 @@ mod tests { #[case(65)] #[case(200)] #[case(1000)] - fn test_for_each_set_index_matches_set_indices(#[case] len: usize) { + fn test_set_index_visitors_match_set_indices(#[case] len: usize) { let buf = BitBuffer::collect_bool(len, |i| i % 5 == 0 || i % 7 == 0); let expected: Vec = buf.set_indices().collect(); + let mut got = Vec::new(); buf.for_each_set_index(|i| got.push(i)); assert_eq!(got, expected); + + let mut fallible_got = Vec::new(); + let result = buf.try_for_each_set_index(|i| { + fallible_got.push(i); + Ok::<(), ()>(()) + }); + assert_eq!(result, Ok(())); + assert_eq!(fallible_got, expected); } #[rstest] @@ -999,6 +1087,33 @@ mod tests { assert_eq!(got, (0..130).collect::>()); } + #[test] + fn test_try_for_each_set_index_stops_on_error() { + for (buffer, stop) in [ + (BitBuffer::new_set(130), 65), + (BitBuffer::collect_bool(130, |i| i % 3 == 0), 66), + ] { + let mut visited = Vec::new(); + let result = buffer.try_for_each_set_index(|index| { + visited.push(index); + if index == stop { + return Err(index); + } + + Ok(()) + }); + + assert_eq!(result, Err(stop)); + assert_eq!( + visited, + buffer + .set_indices() + .take_while(|&i| i <= stop) + .collect::>() + ); + } + } + #[test] fn test_map_cmp_conditional() { // map_cmp with conditional logic based on index and bit value diff --git a/vortex-buffer/src/bit/buf_mut.rs b/vortex-buffer/src/bit/buf_mut.rs index b3e7207a640..e496c34716e 100644 --- a/vortex-buffer/src/bit/buf_mut.rs +++ b/vortex-buffer/src/bit/buf_mut.rs @@ -6,6 +6,7 @@ use std::ops::Not; use bitvec::view::BitView; use crate::BitBuffer; +use crate::BufferAllocatorRef; use crate::BufferMut; use crate::ByteBufferMut; use crate::bit::collect_bool_words; @@ -17,6 +18,7 @@ use crate::bit::unset_bit_unchecked; use crate::buffer_mut; /// Sets all bits in the bit-range `[start_bit, end_bit)` of `slice` to `value`. +#[allow(clippy::inline_always)] #[inline(always)] pub(crate) fn fill_bits(slice: &mut [u8], start_bit: usize, end_bit: usize, value: bool) { if start_bit >= end_bit { @@ -102,6 +104,7 @@ pub struct BitBufferMut { impl BitBufferMut { /// Create new bit buffer from given byte buffer and logical bit length + #[inline] pub fn from_buffer(buffer: ByteBufferMut, offset: usize, len: usize) -> Self { assert!( len <= buffer.len() * 8, @@ -117,23 +120,36 @@ impl BitBufferMut { /// Creates a `BitBufferMut` from a [`BitBuffer`] by copying all of the data over. pub fn copy_from(bit_buffer: &BitBuffer) -> Self { + Self::copy_from_in(bit_buffer, bit_buffer.inner().allocator().clone()) + } + + /// Copies a bit buffer with the provided allocator. + pub fn copy_from_in(bit_buffer: &BitBuffer, allocator: BufferAllocatorRef) -> Self { Self { - buffer: ByteBufferMut::copy_from(bit_buffer.inner()), + buffer: ByteBufferMut::copy_from_in(bit_buffer.inner(), allocator), offset: bit_buffer.offset(), len: bit_buffer.len(), } } /// Create a new empty mutable bit buffer with requested capacity (in bits). + #[inline] pub fn with_capacity(capacity: usize) -> Self { + Self::with_capacity_in(capacity, BufferAllocatorRef::statically_allocated()) + } + + /// Create a mutable bit buffer with the provided allocator. + #[inline] + pub fn with_capacity_in(capacity: usize, allocator: BufferAllocatorRef) -> Self { Self { - buffer: BufferMut::with_capacity(capacity.div_ceil(8)), + buffer: BufferMut::with_capacity_in(capacity.div_ceil(8), allocator), offset: 0, len: 0, } } /// Create a new mutable buffer with requested `len` and all bits set to `true`. + #[inline] pub fn new_set(len: usize) -> Self { Self { buffer: buffer_mut![0xFF; len.div_ceil(8)], @@ -143,6 +159,7 @@ impl BitBufferMut { } /// Create a new mutable buffer with requested `len` and all bits set to `false`. + #[inline] pub fn new_unset(len: usize) -> Self { Self { buffer: BufferMut::zeroed(len.div_ceil(8)), @@ -152,12 +169,24 @@ impl BitBufferMut { } /// Create a new empty `BitBufferMut`. + #[allow(clippy::inline_always)] #[inline(always)] pub fn empty() -> Self { Self::with_capacity(0) } + /// Create an empty mutable bit buffer with the provided allocator. + pub fn empty_in(allocator: BufferAllocatorRef) -> Self { + Self::with_capacity_in(0, allocator) + } + + /// Returns the allocator that owns this buffer. + pub fn allocator(&self) -> &BufferAllocatorRef { + self.buffer.allocator() + } + /// Create a new mutable buffer with requested `len` and all bits set to `value`. + #[inline] pub fn full(value: bool, len: usize) -> Self { if value { Self::new_set(len) @@ -184,39 +213,17 @@ impl BitBufferMut { } } - /// Invokes `f` with indexes `0..len` collecting the boolean results into a new `BitBufferMut` - /// - /// `f` is invoked exactly once per index, in ascending order, and the results are packed - /// with the baseline SIMD byte→bit instruction of the target. - /// - /// # Performance + /// Mutable-buffer form of [`BitBuffer::collect_bool`]. /// - /// The packing is a few instructions per 64 bits, so evaluating `f` is usually the - /// bottleneck. In particular, a bounds-checked slice access in `f` (`|i| values[i] > x`) - /// blocks vectorization of the gather and can cost ~10x the packing itself. Since `f` only - /// ever sees indices `0..len`, callers reading from a slice with `len <= values.len()` may - /// soundly use `|i| unsafe { *values.get_unchecked(i) }`. - /// - /// Prefer this entry point for every predicate. Only switch to - /// [`Self::collect_bool_multiversioned`] after carefully checking that your specific `f` - /// meets its contract (a trivially cheap, bounds-check-free gather or comparison) — - /// ideally with a benchmark. + /// Calls `f` in the same order and uses the same packing path. #[inline] pub fn collect_bool bool>(len: usize, f: F) -> Self { Self::collect_words(len, |words| collect_bool_words(words, len, f)) } - /// Like [`Self::collect_bool`], but compiles the packing loop — with `f` inside it — once - /// per CPU feature level (AVX-512BW/AVX2/baseline) and selects a clone by runtime feature - /// detection. + /// Mutable-buffer form of [`BitBuffer::collect_bool_multiversioned`]. /// - /// Calling this asserts that `f` is small and simple enough (e.g. a bounds-check-free slice - /// gather or comparison) that duplicating it per feature level and paying a - /// `#[target_feature]` call boundary beats inlining it once into your function. For any - /// non-trivial `f` that assertion is false — the boundary deoptimizes the predicate — so - /// unless you have carefully checked (ideally benchmarked) that your specific `f` - /// qualifies, use [`Self::collect_bool`]. See - /// [`collect_bool_words_multiversioned`]. + /// Calls `f` in the same order and uses the same packing path. #[inline] pub fn collect_bool_multiversioned bool>(len: usize, f: F) -> Self { Self::collect_words(len, |words| { @@ -247,34 +254,40 @@ impl BitBufferMut { } /// Return the underlying byte buffer. + #[inline] pub fn inner(&self) -> &ByteBufferMut { &self.buffer } /// Consumes the buffer and return the underlying byte buffer. + #[inline] pub fn into_inner(self) -> ByteBufferMut { self.buffer } /// Get the current populated length of the buffer. + #[allow(clippy::inline_always)] #[inline(always)] pub fn len(&self) -> usize { self.len } /// True if the buffer has length 0. + #[allow(clippy::inline_always)] #[inline(always)] pub fn is_empty(&self) -> bool { self.len == 0 } /// Get the current bit offset of the buffer. + #[allow(clippy::inline_always)] #[inline(always)] pub fn offset(&self) -> usize { self.offset } /// Get the value at the requested index. + #[allow(clippy::inline_always)] #[inline(always)] pub fn value(&self, index: usize) -> bool { assert!(index < self.len); @@ -287,18 +300,21 @@ impl BitBufferMut { /// # Safety /// /// The caller must ensure that `index` is less than the length of the buffer. + #[allow(clippy::inline_always)] #[inline(always)] pub unsafe fn value_unchecked(&self, index: usize) -> bool { unsafe { get_bit_unchecked(self.buffer.as_ptr(), self.offset + index) } } /// Get the bit capacity of the buffer. + #[allow(clippy::inline_always)] #[inline(always)] pub fn capacity(&self) -> usize { (self.buffer.capacity() * 8) - self.offset } /// Reserve additional bit capacity for the buffer. + #[inline] pub fn reserve(&mut self, additional: usize) { let required_bits = self.offset + self.len + additional; let required_bytes = required_bits.div_ceil(8); // Rounds up. @@ -308,6 +324,7 @@ impl BitBufferMut { } /// Clears the bit buffer (but keeps any allocated memory). + #[inline] pub fn clear(&mut self) { // Also clear the byte buffer (not just `len`) so the "bits beyond len are zero" // invariant holds; `append_false` and `append_buffer` rely on it. @@ -398,6 +415,7 @@ impl BitBufferMut { /// /// - `new_len` must be less than or equal to [`capacity()`](Self::capacity) /// - The elements at `old_len..new_len` must be initialized + #[allow(clippy::inline_always)] #[inline(always)] pub unsafe fn set_len(&mut self, new_len: usize) { debug_assert!( @@ -415,6 +433,7 @@ impl BitBufferMut { /// Truncate the buffer to the given length. /// /// If the given length is greater than the current length, this is a no-op. + #[inline] pub fn truncate(&mut self, len: usize) { if len > self.len { return; @@ -514,6 +533,7 @@ impl BitBufferMut { /// /// This operates on an arbitrary range within the existing length of the buffer. /// Panics if `end > self.len` or `start > end`. + #[allow(clippy::inline_always)] #[inline(always)] pub fn fill_range(&mut self, start: usize, end: usize, value: bool) { assert!(end <= self.len, "end {end} exceeds len {}", self.len); @@ -529,6 +549,7 @@ impl BitBufferMut { /// # Safety /// /// The caller must ensure that `start <= end <= self.len`. + #[allow(clippy::inline_always)] #[inline(always)] pub unsafe fn fill_range_unchecked(&mut self, start: usize, end: usize, value: bool) { fill_bits( @@ -589,36 +610,20 @@ impl BitBufferMut { self.len += bit_len; } - /// Absorbs a mutable buffer that was previously split off. - /// - /// If the two buffers were previously contiguous and not mutated in a way that causes - /// re-allocation i.e., if other was created by calling split_off on this buffer, then this is - /// an O(1) operation that just decreases a reference count and sets a few indices. - /// - /// Otherwise, this method degenerates to self.append_buffer(&other). - pub fn unsplit(&mut self, other: Self) { - if (self.offset + self.len).is_multiple_of(8) && other.offset == 0 { - // We are aligned and can just append the buffers - self.buffer.unsplit(other.buffer); - self.len += other.len; - return; - } - - // Otherwise, we need to append the bits one by one - self.append_buffer(&other.freeze()) - } - /// Freeze the buffer in its current state into an immutable `BoolBuffer`. + #[inline] pub fn freeze(self) -> BitBuffer { BitBuffer::new_with_offset(self.buffer.freeze(), self.len, self.offset) } /// Get the underlying bytes as a slice + #[inline] pub fn as_slice(&self) -> &[u8] { self.buffer.as_slice() } /// Get the underlying bytes as a mutable slice + #[inline] pub fn as_mut_slice(&mut self) -> &mut [u8] { self.buffer.as_mut_slice() } diff --git a/vortex-buffer/src/bit/meta.rs b/vortex-buffer/src/bit/meta.rs index 866a308647d..d7197e912d8 100644 --- a/vortex-buffer/src/bit/meta.rs +++ b/vortex-buffer/src/bit/meta.rs @@ -1,6 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::Bound; +use std::ops::RangeBounds; + +use vortex_error::VortexExpect; + /// In-memory metadata describing a packed bitset: a normalized bit `offset` (always `< 8`) and a /// logical bit `len`. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] @@ -14,6 +19,7 @@ impl BitBufferMeta { /// /// Panics if `offset >= 8`. Use [`from_raw_offset`](Self::from_raw_offset) to normalize a /// larger offset. + #[inline] pub fn new(offset: usize, len: usize) -> Self { assert!(offset < 8, "BitBufferMeta offset must be < 8, got {offset}"); Self { offset, len } @@ -24,6 +30,7 @@ impl BitBufferMeta { /// /// Returns `(byte_offset, meta)` so the caller can slice its backing buffer by `byte_offset` /// and store the remaining sub-byte offset in `meta`. + #[inline] pub fn from_raw_offset(offset: usize, len: usize) -> (usize, Self) { ( offset / 8, @@ -34,19 +41,51 @@ impl BitBufferMeta { ) } + /// Return the leading byte offset and normalized metadata for a logical slice. + /// + /// # Panics + /// + /// Panics if the range is out of bounds or its end precedes its start. + #[inline] + pub fn slice(&self, range: impl RangeBounds) -> (usize, Self) { + let start = match range.start_bound() { + Bound::Included(&start) => start, + Bound::Excluded(&start) => start + .checked_add(1) + .vortex_expect("excluded slice start must not overflow"), + Bound::Unbounded => 0, + }; + let end = match range.end_bound() { + Bound::Included(&end) => end + .checked_add(1) + .vortex_expect("included slice end must not overflow"), + Bound::Excluded(&end) => end, + Bound::Unbounded => self.len, + }; + + assert!(start <= end); + assert!(start <= self.len); + assert!(end <= self.len); + + Self::from_raw_offset(self.offset + start, end - start) + } + /// The sub-byte bit offset. Always `< 8`. + #[allow(clippy::inline_always)] #[inline(always)] pub fn offset(&self) -> usize { self.offset } /// The logical length of the bitset in bits. + #[allow(clippy::inline_always)] #[inline(always)] pub fn len(&self) -> usize { self.len } /// Returns `true` if the bitset is empty. + #[allow(clippy::inline_always)] #[inline(always)] pub fn is_empty(&self) -> bool { self.len == 0 diff --git a/vortex-buffer/src/bit/mod.rs b/vortex-buffer/src/bit/mod.rs index 46d1bee89be..c0d82e10b3f 100644 --- a/vortex-buffer/src/bit/mod.rs +++ b/vortex-buffer/src/bit/mod.rs @@ -158,6 +158,7 @@ where /// # Panics /// /// Panics if `index` is not between 0 and length of `buf * 8`. +#[allow(clippy::inline_always)] #[inline(always)] pub fn get_bit(buf: &[u8], index: usize) -> bool { buf[index / 8] & (1 << (index % 8)) != 0 @@ -168,6 +169,7 @@ pub fn get_bit(buf: &[u8], index: usize) -> bool { /// # Safety /// /// `index` must be between 0 and length of `buf * 8`. +#[allow(clippy::inline_always)] #[inline(always)] pub unsafe fn get_bit_unchecked(buf: *const u8, index: usize) -> bool { (unsafe { *buf.add(index / 8) } & (1 << (index % 8))) != 0 @@ -178,6 +180,7 @@ pub unsafe fn get_bit_unchecked(buf: *const u8, index: usize) -> bool { /// # Safety /// /// `index` must be between 0 and length of `buf * 8`. +#[allow(clippy::inline_always)] #[inline(always)] pub unsafe fn set_bit_unchecked(buf: *mut u8, index: usize) { unsafe { *buf.add(index / 8) |= 1 << (index % 8) }; @@ -188,6 +191,7 @@ pub unsafe fn set_bit_unchecked(buf: *mut u8, index: usize) { /// # Safety /// /// `index` must be between 0 and length of `buf * 8`. +#[allow(clippy::inline_always)] #[inline(always)] pub unsafe fn unset_bit_unchecked(buf: *mut u8, index: usize) { unsafe { *buf.add(index / 8) &= !(1 << (index % 8)) }; diff --git a/vortex-buffer/src/bit/pack.rs b/vortex-buffer/src/bit/pack.rs index 565ce1e8345..128e27eb007 100644 --- a/vortex-buffer/src/bit/pack.rs +++ b/vortex-buffer/src/bit/pack.rs @@ -56,6 +56,7 @@ where /// the wider pack saves — and an indirect call per word is worse still (~4x on cheap /// predicates), since an opaque call target blocks fill/pack fusion regardless of how cheap /// the kernel *selection* is. For provably cheap predicates, use [`collect_bool_words_multiversioned`]. +#[allow(clippy::inline_always)] #[inline(always)] pub(crate) fn collect_bool_words_inline(words: &mut [u64], len: usize, f: F) where @@ -150,6 +151,7 @@ where /// /// Marked `#[inline(always)]` so each `#[target_feature]` wrapper gets its own fully-inlined /// copy compiled with that feature set. +#[allow(clippy::inline_always)] #[inline(always)] fn collect_bool_words_with(words: &mut [u64], len: usize, mut f: F, pack: P) where diff --git a/vortex-buffer/src/bit/view.rs b/vortex-buffer/src/bit/view.rs index 722119ee4ab..073df5943bc 100644 --- a/vortex-buffer/src/bit/view.rs +++ b/vortex-buffer/src/bit/view.rs @@ -63,6 +63,7 @@ impl<'a> BitBufferView<'a> { /// Create a new view over `buffer` with `len` bits, starting at bit zero. /// /// Panics if the buffer is not large enough to hold `len` bits. + #[inline] pub fn new(buffer: &'a [u8], len: usize) -> Self { Self::new_with_offset(buffer, len, 0) } @@ -70,6 +71,7 @@ impl<'a> BitBufferView<'a> { /// Create a new view over `buffer` with `len` bits, starting at the given bit `offset`. /// /// Panics if the buffer is not large enough to hold `len` bits after the offset. + #[inline] pub fn new_with_offset(buffer: &'a [u8], len: usize, offset: usize) -> Self { assert!( len.saturating_add(offset) <= buffer.len().saturating_mul(8), @@ -86,11 +88,13 @@ impl<'a> BitBufferView<'a> { } /// Create a new view over `buffer` described by `meta`. + #[inline] pub fn from_meta(buffer: &'a [u8], meta: BitBufferMeta) -> Self { Self::new_with_offset(buffer, meta.len(), meta.offset()) } /// Returns the [`BitBufferMeta`] (offset and length) describing this view. + #[inline] pub fn meta(&self) -> BitBufferMeta { BitBufferMeta::new(self.offset, self.len) } @@ -108,12 +112,14 @@ impl<'a> BitBufferView<'a> { } /// Offset of the start of the view in bits. Always `< 8`. + #[allow(clippy::inline_always)] #[inline(always)] pub fn offset(&self) -> usize { self.offset } /// Get a reference to the underlying byte slice. + #[allow(clippy::inline_always)] #[inline(always)] pub fn inner(&self) -> &'a [u8] { self.buffer @@ -142,6 +148,7 @@ impl<'a> BitBufferView<'a> { /// Create a new view over the range `[start, end)` of this view. /// /// Panics if the slice would extend beyond the end of the view. + #[inline] pub fn slice(&self, range: impl RangeBounds) -> BitBufferView<'a> { let (start, end) = resolve_range(range, self.len); BitBufferView::new_with_offset(self.buffer, end - start, self.offset + start) @@ -149,41 +156,49 @@ impl<'a> BitBufferView<'a> { /// Access chunks of the buffer aligned to an 8 byte boundary as /// `[prefix, , suffix]`. + #[inline] pub fn unaligned_chunks(&self) -> UnalignedBitChunk<'a> { UnalignedBitChunk::new(self.buffer, self.offset, self.len) } /// Access chunks of the underlying buffer as 8 byte chunks with a final trailer. + #[inline] pub fn chunks(&self) -> BitChunks<'a> { BitChunks::new(self.buffer, self.offset, self.len) } /// Get the number of set bits in the view. + #[inline] pub fn true_count(&self) -> usize { count_ones(self.buffer, self.offset, self.len) } /// Get the number of unset bits in the view. + #[inline] pub fn false_count(&self) -> usize { self.len - self.true_count() } /// Returns the position of the `nth` set bit (0-indexed), or `None` if out of range. + #[inline] pub fn select(&self, nth: usize) -> Option { bit_select(self.buffer, self.offset, self.len, nth) } /// Iterator over bits in the view. + #[inline] pub fn iter(&self) -> BitIterator<'a> { BitIterator::new(self.buffer, self.offset, self.len) } /// Iterator over set indices of the underlying buffer. + #[inline] pub fn set_indices(&self) -> BitIndexIterator<'a> { BitIndexIterator::new(self.buffer, self.offset, self.len) } /// Iterator over set slices of the underlying buffer. + #[inline] pub fn set_slices(&self) -> BitSliceIterator<'a> { BitSliceIterator::new(self.buffer, self.offset, self.len) } @@ -240,6 +255,7 @@ impl<'a> BitBufferMutView<'a> { /// Create a new mutable view over `buffer` with `len` bits, starting at bit zero. /// /// Panics if the buffer is not large enough to hold `len` bits. + #[inline] pub fn new(buffer: &'a mut [u8], len: usize) -> Self { Self::new_with_offset(buffer, len, 0) } @@ -247,6 +263,7 @@ impl<'a> BitBufferMutView<'a> { /// Create a new mutable view over `buffer` with `len` bits, starting at bit `offset`. /// /// Panics if the buffer is not large enough to hold `len` bits after the offset. + #[inline] pub fn new_with_offset(buffer: &'a mut [u8], len: usize, offset: usize) -> Self { assert!( len.saturating_add(offset) <= buffer.len().saturating_mul(8), @@ -286,6 +303,7 @@ impl<'a> BitBufferMutView<'a> { } /// Offset of the start of the view in bits. Always `< 8`. + #[allow(clippy::inline_always)] #[inline(always)] pub fn offset(&self) -> usize { self.offset @@ -324,16 +342,19 @@ impl<'a> BitBufferMutView<'a> { } /// Get the number of set bits in the view. + #[inline] pub fn true_count(&self) -> usize { self.as_view().true_count() } /// Get the number of unset bits in the view. + #[inline] pub fn false_count(&self) -> usize { self.as_view().false_count() } /// Iterator over bits in the view. + #[inline] pub fn iter(&self) -> BitIterator<'_> { self.as_view().iter() } @@ -341,6 +362,7 @@ impl<'a> BitBufferMutView<'a> { /// Set the bit at `index` to the given boolean value. /// /// Panics if `index` exceeds the view length. + #[inline] pub fn set_to(&mut self, index: usize, value: bool) { if value { self.set(index); @@ -354,6 +376,7 @@ impl<'a> BitBufferMutView<'a> { /// # Safety /// /// Caller must ensure that `index` is within the range of the view. + #[inline] pub unsafe fn set_to_unchecked(&mut self, index: usize, value: bool) { if value { // SAFETY: checked by caller @@ -367,6 +390,7 @@ impl<'a> BitBufferMutView<'a> { /// Set the bit at `index` to `true`. /// /// Panics if `index` exceeds the view length. + #[inline] pub fn set(&mut self, index: usize) { assert!(index < self.len, "index {index} exceeds len {}", self.len); // SAFETY: checked by assertion @@ -376,6 +400,7 @@ impl<'a> BitBufferMutView<'a> { /// Set the bit at `index` to `false`. /// /// Panics if `index` exceeds the view length. + #[inline] pub fn unset(&mut self, index: usize) { assert!(index < self.len, "index {index} exceeds len {}", self.len); // SAFETY: checked by assertion @@ -407,6 +432,7 @@ impl<'a> BitBufferMutView<'a> { /// Sets all bits in the range `[start, end)` to `value`. /// /// Panics if `end > self.len()` or `start > end`. + #[allow(clippy::inline_always)] #[inline(always)] pub fn fill_range(&mut self, start: usize, end: usize, value: bool) { assert!(end <= self.len, "end {end} exceeds len {}", self.len); @@ -420,6 +446,7 @@ impl<'a> BitBufferMutView<'a> { /// # Safety /// /// Caller must ensure that `start <= end <= self.len()`. + #[allow(clippy::inline_always)] #[inline(always)] pub unsafe fn fill_range_unchecked(&mut self, start: usize, end: usize, value: bool) { fill_bits(self.buffer, self.offset + start, self.offset + end, value); diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs index 4b587b94629..73d15991f0b 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -8,9 +8,10 @@ use std::fmt::Debug; use std::fmt::Formatter; use std::hash::Hash; use std::hash::Hasher; -use std::marker::PhantomData; use std::ops::Deref; use std::ops::RangeBounds; +use std::ptr::NonNull; +use std::sync::Arc; use bytes::Buf; use bytes::Bytes; @@ -18,72 +19,148 @@ use vortex_error::VortexExpect; use vortex_error::vortex_panic; use crate::Alignment; +use crate::Allocation; +use crate::BufferAllocatorRef; +use crate::BufferBacking; use crate::BufferMut; use crate::ByteBuffer; use crate::debug::TruncatedDebug; use crate::trusted_len::TrustedLen; /// An immutable buffer of items of `T`. +/// +/// Zero-sized element types are rejected at compile time when constructing a buffer. +/// +/// ```compile_fail +/// use vortex_buffer::Buffer; +/// let _ = Buffer::<()>::empty(); +/// ``` +/// +/// ```compile_fail +/// use vortex_buffer::Buffer; +/// let _ = Buffer::from(vec![(); 3]); +/// ``` +/// +/// ```compile_fail +/// use vortex_buffer::{Buffer, ByteBuffer}; +/// let _ = Buffer::<()>::from_byte_buffer(ByteBuffer::empty()); +/// ``` +/// +/// ```compile_fail +/// use bytes::Bytes; +/// use vortex_buffer::{Alignment, Buffer}; +/// let _ = Buffer::<()>::from_bytes_aligned(Bytes::new(), Alignment::none()); +/// ``` #[derive(Clone)] pub struct Buffer { - pub(crate) bytes: Bytes, + /// The first element in this view; may dangle for an empty buffer. + pub(crate) ptr: NonNull, + /// The number of initialized `T` values visible from `ptr`. pub(crate) length: usize, + /// The minimum alignment promised for `ptr` and preserved by aligned slices. pub(crate) alignment: Alignment, - pub(crate) _marker: PhantomData, + /// Shared ownership of the storage containing `ptr`, if any; `Buffer::empty` has no backing. + pub(crate) backing: Option>, } -/// Zero-length backing for empty buffers, "aligned" to [`Alignment::MAX`] so it satisfies any -/// valid alignment without allocating. A zero-length slice never reads memory, so it may use a -/// dangling pointer as long as it is non-null and aligned. -const EMPTY_BACKING: &[u8] = { - let addr = 1usize << (usize::BITS - 1); - assert!(Alignment::MAX.is_offset_aligned(addr)); - // SAFETY: the pointer is non-null and aligned, and the slice is zero-length. - unsafe { std::slice::from_raw_parts(std::ptr::without_provenance(addr), 0) } -}; +// SAFETY: Buffer is an immutable view over backing memory. Its pointer remains valid while the +// backing is live, and sharing elements follows the same bounds as sharing a slice. +unsafe impl Send for Buffer {} +// SAFETY: see the Send implementation above. +unsafe impl Sync for Buffer {} impl Default for Buffer { fn default() -> Self { Self { - bytes: Bytes::from_static(EMPTY_BACKING), + ptr: empty_ptr(), length: 0, alignment: Alignment::of::(), - _marker: PhantomData, + backing: None, } } } -impl PartialEq for Buffer { +impl PartialEq for Buffer { #[inline] fn eq(&self, other: &Self) -> bool { - self.bytes == other.bytes + self.as_slice() == other.as_slice() } } -impl Eq for Buffer {} +impl Eq for Buffer {} -impl Ord for Buffer { +impl Ord for Buffer { #[inline] fn cmp(&self, other: &Self) -> Ordering { - self.bytes.cmp(&other.bytes) + self.as_slice().cmp(other.as_slice()) } } -impl PartialOrd for Buffer { +impl PartialOrd for Buffer { #[inline] fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) + self.as_slice().partial_cmp(other.as_slice()) } } -impl Hash for Buffer { +impl Hash for Buffer { #[inline] fn hash(&self, state: &mut H) { - self.bytes.as_ref().hash(state) + self.as_slice().hash(state) } } impl Buffer { + pub(crate) fn from_allocation( + allocation: Allocation, + offset: usize, + length: usize, + alignment: Alignment, + ) -> Self { + // SAFETY: BufferMut keeps offset within allocation, including for empty buffers. + let ptr = unsafe { allocation.ptr().add(offset).cast() }; + Self { + ptr, + length, + alignment, + backing: Some(Arc::new(BufferBacking::Owned(allocation))), + } + } + + fn from_bytes(bytes: Bytes, alignment: Alignment) -> Self { + let length = bytes.len() / size_of::(); + if length == 0 { + return Self::empty_aligned(alignment); + } + let ptr = + NonNull::new(bytes.as_ptr().cast_mut().cast()).vortex_expect("Bytes pointer is null"); + Self { + ptr, + length, + alignment, + backing: Some(Arc::new(BufferBacking::Bytes(bytes))), + } + } + + #[cfg(feature = "arrow")] + pub(crate) fn from_arrow_owner( + arrow: arrow_buffer::Buffer, + length: usize, + alignment: Alignment, + ) -> Self { + if length == 0 { + return Self::empty_aligned(alignment); + } + let ptr = NonNull::new(arrow.as_ptr().cast_mut().cast()) + .vortex_expect("Arrow buffer pointer is null"); + Self { + ptr, + length, + alignment, + backing: Some(Arc::new(BufferBacking::Arrow(arrow))), + } + } + /// Returns a new `Buffer` copied from the provided `Vec`, `&[T]`, etc. /// /// Due to our underlying usage of `bytes::Bytes`, we are unable to take zero-copy ownership @@ -94,6 +171,11 @@ impl Buffer { BufferMut::copy_from(values).freeze() } + /// Returns a new `Buffer` copied with the provided allocator. + pub fn copy_from_in(values: impl AsRef<[T]>, allocator: BufferAllocatorRef) -> Self { + BufferMut::copy_from_in(values, allocator).freeze() + } + /// Returns a new `Buffer` copied from the provided slice and with the requested alignment. /// /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than @@ -121,6 +203,11 @@ impl Buffer { Self::zeroed_aligned(len, Alignment::of::()) } + /// Create a new zeroed `Buffer` with the provided allocator. + pub fn zeroed_in(len: usize, allocator: BufferAllocatorRef) -> Self { + BufferMut::zeroed_in(len, allocator).freeze() + } + /// Create a new zeroed `Buffer` with the requested alignment. /// /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than @@ -150,9 +237,9 @@ impl Buffer { /// Create a new empty `ByteBuffer` with the provided alignment. /// - /// This does not allocate: empty buffers are backed by a zero-length `Bytes` that is - /// aligned to [`Alignment::MAX`]. + /// This does not allocate. Empty buffers use an aligned dangling pointer. pub fn empty_aligned(alignment: Alignment) -> Self { + const { assert!(size_of::() != 0, "ZSTs are not supported") }; if !alignment.is_aligned_to(Alignment::of::()) { vortex_panic!( "Alignment {} must align to the scalar type's alignment {}", @@ -161,10 +248,10 @@ impl Buffer { ); } Self { - bytes: Bytes::from_static(EMPTY_BACKING), + ptr: empty_ptr(), length: 0, alignment, - _marker: PhantomData, + backing: None, } } @@ -176,6 +263,14 @@ impl Buffer { BufferMut::full(item, len).freeze() } + /// Create a full `Buffer` with the given value and allocator. + pub fn full_in(item: T, len: usize, allocator: BufferAllocatorRef) -> Self + where + T: Copy, + { + BufferMut::full_in(item, len, allocator).freeze() + } + /// Create a `Buffer` zero-copy from a `ByteBuffer`. /// /// ## Panics @@ -194,7 +289,30 @@ impl Buffer { /// Panics if the buffer is not aligned to the given alignment, if the length is not a multiple /// of the size of `T`, or if the given alignment is not aligned to that of `T`. pub fn from_byte_buffer_aligned(buffer: ByteBuffer, alignment: Alignment) -> Self { - Self::from_bytes_aligned(buffer.into_inner(), alignment) + const { assert!(size_of::() != 0, "ZSTs are not supported") }; + if !alignment.is_aligned_to(Alignment::of::()) { + vortex_panic!( + "Alignment {} must be compatible with the scalar type's alignment {}", + alignment, + Alignment::of::(), + ); + } + if !alignment.is_ptr_aligned(buffer.as_ptr()) { + vortex_panic!("Buffer must align to the requested alignment {}", alignment); + } + if !buffer.len().is_multiple_of(size_of::()) { + vortex_panic!( + "Buffer length {} must be a multiple of the scalar type's size {}", + buffer.len(), + size_of::() + ); + } + Self { + ptr: buffer.ptr.cast(), + length: buffer.length / size_of::(), + alignment, + backing: buffer.backing, + } } /// Create a `Buffer` zero-copy from a `Bytes`. @@ -204,6 +322,7 @@ impl Buffer { /// Panics if the buffer is not aligned to the size of `T`, or the length is not a multiple of /// the size of `T`. pub fn from_bytes_aligned(bytes: Bytes, alignment: Alignment) -> Self { + const { assert!(size_of::() != 0, "ZSTs are not supported") }; if !alignment.is_aligned_to(Alignment::of::()) { vortex_panic!( "Alignment {} must be compatible with the scalar type's alignment {}", @@ -224,24 +343,13 @@ impl Buffer { size_of::() ); } - let length = bytes.len() / size_of::(); - Self { - bytes, - length, - alignment, - _marker: Default::default(), - } + Self::from_bytes(bytes, alignment) } /// Create a buffer with values from the TrustedLen iterator. /// Should be preferred over `from_iter` when the iterator is known to be `TrustedLen`. pub fn from_trusted_len_iter>(iter: I) -> Self { - let (_, upper_bound) = iter.size_hint(); - let mut buffer = BufferMut::with_capacity( - upper_bound.vortex_expect("TrustedLen iterator has no upper bound"), - ); - buffer.extend_trusted(iter); - buffer.freeze() + BufferMut::from_trusted_len_iter(iter).freeze() } /// Map each element of the buffer with a closure. @@ -254,7 +362,8 @@ impl Buffer { Ok(mut_buf) => mut_buf.map_each_in_place(f), Err(buf) => { let len = buf.len(); - let mut out_buf = BufferMut::with_capacity(len); + let allocator = buf.allocator().clone(); + let mut out_buf = BufferMut::with_capacity_in(len, allocator); out_buf .spare_capacity_mut() .iter_mut() @@ -271,39 +380,63 @@ impl Buffer { /// Clear the buffer, preserving existing capacity. pub fn clear(&mut self) { - self.bytes.clear(); self.length = 0; } /// Returns the length of the buffer in elements of type T. + #[allow(clippy::inline_always)] #[inline(always)] pub fn len(&self) -> usize { self.length } /// Returns whether the buffer is empty. + #[allow(clippy::inline_always)] #[inline(always)] pub fn is_empty(&self) -> bool { self.length == 0 } /// Returns the alignment of the buffer. + #[allow(clippy::inline_always)] #[inline(always)] pub fn alignment(&self) -> Alignment { self.alignment } + /// Returns the allocator to use for derived buffers. + /// + /// External buffers use the static allocator. + pub fn allocator(&self) -> &BufferAllocatorRef { + match self.backing.as_deref() { + Some(backing) => backing.allocator(), + None => BufferAllocatorRef::static_ref(), + } + } + + /// Returns a raw pointer to the buffer's data. + #[allow(clippy::inline_always)] + #[inline(always)] + pub fn as_ptr(&self) -> *const T { + self.ptr.as_ptr() + } + /// Returns a slice over the buffer of elements of type T. + #[allow(clippy::inline_always)] #[inline(always)] pub fn as_slice(&self) -> &[T] { - // SAFETY: alignment of Buffer is checked on construction - unsafe { std::slice::from_raw_parts(self.bytes.as_ptr().cast(), self.length) } + // SAFETY: ptr points into the live backing and construction checks its alignment. + unsafe { std::slice::from_raw_parts(self.ptr.as_ptr(), self.length) } } /// Return a view over the buffer as an opaque byte slice. + #[allow(clippy::inline_always)] #[inline(always)] pub fn as_bytes(&self) -> &[u8] { - self.bytes.as_ref() + // SAFETY: the element range is initialized and remains live through backing. + unsafe { + std::slice::from_raw_parts(self.ptr.as_ptr().cast(), size_of_val(self.as_slice())) + } } /// Returns an iterator over the buffer of elements of type T. @@ -319,6 +452,7 @@ impl Buffer { /// /// Requires that `begin <= end` and `end <= self.len()`. /// Also requires that both `begin` and `end` are aligned to the buffer's required alignment. + #[allow(clippy::inline_always)] #[inline(always)] pub fn slice(&self, range: impl RangeBounds) -> Self { self.slice_with_alignment(range, self.alignment) @@ -330,6 +464,7 @@ impl Buffer { /// # Panics /// /// Requires that `begin <= end` and `end <= self.len()`. + #[allow(clippy::inline_always)] #[inline(always)] pub fn slice_unaligned(&self, range: impl RangeBounds) -> Self { self.slice_with_alignment(range, Alignment::of::()) @@ -377,8 +512,6 @@ impl Buffer { } let begin_byte = begin * size_of::(); - let end_byte = end * size_of::(); - if !alignment.is_offset_aligned(begin_byte) { vortex_panic!( "range start must be aligned to {alignment:?}, byte {}", @@ -390,10 +523,11 @@ impl Buffer { } Self { - bytes: self.bytes.slice(begin_byte..end_byte), + // SAFETY: begin is in bounds and the alignment check applies to the new pointer. + ptr: unsafe { self.ptr.add(begin) }, length: end - begin, alignment, - _marker: Default::default(), + backing: self.backing.clone(), } } @@ -405,6 +539,7 @@ impl Buffer { /// /// # Panics: /// Requires that the given sub slice is in fact contained within the Bytes buffer; otherwise this function will panic. + #[allow(clippy::inline_always)] #[inline(always)] pub fn slice_ref(&self, subset: &[T]) -> Self { self.slice_ref_with_alignment(subset, Alignment::of::()) @@ -432,74 +567,119 @@ impl Buffer { vortex_panic!("slice_ref subset must be aligned to {:?}", alignment); } - let subset_u8 = - unsafe { std::slice::from_raw_parts(subset.as_ptr().cast(), size_of_val(subset)) }; + let start = self.as_ptr().addr(); + let end = start + size_of_val(self.as_slice()); + let subset_start = subset.as_ptr().addr(); + let subset_end = subset_start + .checked_add(size_of_val(subset)) + .vortex_expect("slice_ref address overflow"); + if subset_start < start || subset_end > end { + vortex_panic!("slice_ref subset must be contained in the buffer"); + } Self { - bytes: self.bytes.slice_ref(subset_u8), + ptr: NonNull::new(subset.as_ptr().cast_mut()).vortex_expect("slice pointer is null"), length: subset.len(), alignment, - _marker: Default::default(), + backing: self.backing.clone(), } } - /// Returns the underlying aligned buffer. - pub fn inner(&self) -> &Bytes { - debug_assert_eq!( - self.length * size_of::(), - self.bytes.len(), - "Own length has to be the same as the underlying bytes length" - ); - &self.bytes - } - - /// Returns the underlying aligned buffer. - pub fn into_inner(self) -> Bytes { - debug_assert_eq!( - self.length * size_of::(), - self.bytes.len(), - "Own length has to be the same as the underlying bytes length" - ); - self.bytes + /// Returns the underlying bytes without copying. + pub fn into_bytes(self) -> Bytes { + if let Some(backing) = self.backing.as_ref() + && let BufferBacking::Bytes(bytes) = backing.as_ref() + { + let offset = self.ptr.cast::().addr().get() - bytes.as_ptr().addr(); + let length = self.length * size_of::(); + if offset == 0 && length == bytes.len() && Arc::strong_count(backing) == 1 { + return match self.backing { + Some(backing) => match Arc::try_unwrap(backing) { + Ok(BufferBacking::Bytes(bytes)) => bytes, + _ => unreachable!(), + }, + None => unreachable!(), + }; + } + return bytes.slice(offset..offset + length); + } + match self.backing { + Some(backing) => Bytes::from_owner(BufferBytesOwner { + ptr: self.ptr.cast(), + length: self.length * size_of::(), + backing, + }), + None => Bytes::new(), + } } /// Return the ByteBuffer for this `Buffer`. pub fn into_byte_buffer(self) -> ByteBuffer { ByteBuffer { - bytes: self.bytes, + ptr: self.ptr.cast(), length: self.length * size_of::(), alignment: self.alignment, - _marker: Default::default(), + backing: self.backing, } } /// Try to convert self into `BufferMut` if there is only a single strong reference. pub fn try_into_mut(self) -> Result, Self> { - self.bytes - .try_into_mut() - .map(|bytes| BufferMut { - bytes, - length: self.length, - alignment: self.alignment, - _marker: Default::default(), - }) - .map_err(|bytes| Self { - bytes, - length: self.length, - alignment: self.alignment, - _marker: Default::default(), - }) + let Self { + ptr, + length, + alignment, + backing, + } = self; + let Some(backing) = backing else { + return Ok(BufferMut::empty_aligned(alignment)); + }; + if !matches!(backing.as_ref(), BufferBacking::Owned(_)) { + return Err(Self { + ptr, + length, + alignment, + backing: Some(backing), + }); + } + match Arc::try_unwrap(backing) { + Ok(BufferBacking::Owned(allocation)) => { + let offset = ptr.addr().get() - allocation.ptr().addr().get(); + let capacity = if allocation.size() == 0 { + 0 + } else { + (allocation.size() - offset) / size_of::() + }; + Ok(BufferMut { + allocation, + ptr, + length, + capacity, + alignment, + _marker: Default::default(), + }) + } + Ok(_) => unreachable!(), + Err(backing) => Err(Self { + ptr, + length, + alignment, + backing: Some(backing), + }), + } } /// Convert self into `BufferMut`, cloning the data if there are multiple strong references. pub fn into_mut(self) -> BufferMut { - self.try_into_mut() - .unwrap_or_else(|buffer| BufferMut::::copy_from_aligned(&buffer, buffer.alignment)) + self.try_into_mut().unwrap_or_else(|buffer| { + let allocator = buffer.allocator().clone(); + BufferMut::::copy_from_aligned_in(&buffer, buffer.alignment, allocator) + }) } /// Returns whether a `Buffer` is aligned to the given alignment. pub fn is_aligned(&self, alignment: Alignment) -> bool { - alignment.is_ptr_aligned(self.bytes.as_ptr()) + alignment.is_ptr_aligned(self.as_ptr()) } /// Return a `Buffer` with the given alignment. Where possible, this will be zero-copy. @@ -515,7 +695,8 @@ impl Buffer { "Buffer is not aligned to requested alignment {alignment}, copying: {bt}" ) } - Self::copy_from_aligned(self, alignment) + let allocator = self.allocator().clone(); + BufferMut::copy_from_aligned_in(self, alignment, allocator).freeze() } } @@ -551,10 +732,10 @@ impl Buffer { ); Buffer { - bytes: self.bytes, + ptr: self.ptr.cast(), length: self.length, alignment: self.alignment, - _marker: PhantomData, + backing: self.backing, } } } @@ -635,48 +816,34 @@ impl FromIterator for Buffer { } } -// Helper struct to allow us to zero-copy any vec into a buffer -#[repr(transparent)] -struct Wrapper(Vec); - -impl AsRef<[u8]> for Wrapper { - fn as_ref(&self) -> &[u8] { - let data = self.0.as_ptr().cast::(); - let len = self.0.len() * size_of::(); - unsafe { std::slice::from_raw_parts(data, len) } - } -} - impl From> for Buffer where - T: Send + 'static, + T: Send + Sync + 'static, { fn from(value: Vec) -> Self { - let original_len = value.len(); - let wrapped_vec = Wrapper(value); - - let bytes = Bytes::from_owner(wrapped_vec); - - assert_eq!(bytes.as_ptr().align_offset(align_of::()), 0); - - Self { - bytes, - length: original_len, - alignment: Alignment::of::(), - _marker: PhantomData, + const { assert!(size_of::() != 0, "ZSTs are not supported") }; + let length = value.len(); + let alignment = Alignment::of::(); + if std::mem::needs_drop::() { + // Keep the typed owner so its elements are dropped. + Self { + ptr: NonNull::new(value.as_ptr().cast_mut()) + .vortex_expect("a Vec always has a non-null pointer"), + length, + alignment, + backing: Some(Arc::new(BufferBacking::External { + _owner: Box::new(value), + })), + } + } else { + Self::from_allocation(Allocation::from_vec(value), 0, length, alignment) } } } impl From for ByteBuffer { fn from(bytes: Bytes) -> Self { - let length = bytes.len(); - Self { - bytes, - length, - alignment: Alignment::of::(), - _marker: Default::default(), - } + Self::from_bytes(bytes, Alignment::of::()) } } @@ -700,11 +867,36 @@ impl Buf for ByteBuffer { self.alignment ); } - self.bytes.advance(cnt); + assert!(cnt <= self.length, "cannot advance past the buffer length"); + // SAFETY: cnt is within the initialized byte range. + self.ptr = unsafe { self.ptr.add(cnt) }; self.length -= cnt; } } +struct BufferBytesOwner { + ptr: NonNull, + length: usize, + backing: Arc, +} + +// SAFETY: the owner exposes immutable initialized bytes and keeps their backing live. +unsafe impl Send for BufferBytesOwner {} +unsafe impl Sync for BufferBytesOwner {} + +impl AsRef<[u8]> for BufferBytesOwner { + fn as_ref(&self) -> &[u8] { + let _ = &self.backing; + // SAFETY: ptr and length came from a live Buffer. + unsafe { std::slice::from_raw_parts(self.ptr.as_ptr(), self.length) } + } +} + +fn empty_ptr() -> NonNull { + let addr = 1usize << (usize::BITS - 1); + NonNull::new(std::ptr::without_provenance_mut(addr)).vortex_expect("empty pointer is non-null") +} + /// Owned iterator over a [`Buffer`]. pub struct BufferIterator { // Keep the buffer alive for the duration of the iteration. @@ -713,6 +905,11 @@ pub struct BufferIterator { end: *const T, } +// SAFETY: `BufferIterator` is a `Buffer` plus two cursors into it, so it can safely be +// `Send`/`Sync` exactly when `Buffer` is. Same bounds as `std::vec::IntoIter`. +unsafe impl Send for BufferIterator {} +unsafe impl Sync for BufferIterator {} + impl Iterator for BufferIterator { type Item = T; @@ -762,10 +959,17 @@ impl From> for Buffer { #[cfg(test)] mod test { + use std::mem::align_of; + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use bytes::Buf; + use bytes::Bytes; use crate::Alignment; use crate::Buffer; + use crate::BufferBacking; use crate::ByteBuffer; use crate::buffer; @@ -777,6 +981,17 @@ mod test { assert_eq!(aligned.as_slice(), &[0, 1, 2]); } + #[test] + fn buffer_iterator_send_sync() { + fn assert_send_sync(_: &T) {} + + let mut iter = buffer![0i32, 1, 2, 3].into_iter(); + assert_send_sync(&iter); + iter.next(); + let remaining: Vec = std::thread::spawn(move || iter.collect()).join().unwrap(); + assert_eq!(remaining, vec![1, 2, 3]); + } + #[test] fn slice() { let buf = buffer![0, 1, 2, 3, 4]; @@ -868,6 +1083,152 @@ mod test { assert_eq!(vec, buff.as_ref()); } + #[test] + fn from_vec_adopts_allocation() { + let mut vec = Vec::with_capacity(16); + vec.extend([1u32, 2, 3, 4, 5]); + let ptr = vec.as_ptr(); + let capacity = vec.capacity(); + + let buffer = Buffer::from(vec); + assert_eq!(buffer.as_ptr(), ptr); + + let Ok(mut buffer) = buffer.try_into_mut() else { + panic!("Vec-backed buffer should be uniquely owned") + }; + assert_eq!(buffer.capacity(), capacity); + assert_eq!(buffer.allocation.alignment(), align_of::()); + + buffer.extend(6..=32); + assert_eq!(buffer.as_slice(), (1..=32).collect::>()); + assert_eq!(buffer.allocation.alignment(), align_of::()); + } + + #[test] + fn byte_owner_preserves_slice_and_lifetime() { + struct Owner { + values: Vec, + drops: Arc, + } + + impl AsRef<[u8]> for Owner { + fn as_ref(&self) -> &[u8] { + &self.values[1..4] + } + } + + impl Drop for Owner { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::Relaxed); + } + } + + let drops = Arc::new(AtomicUsize::new(0)); + let owner = Owner { + values: vec![0, 1, 2, 3, 4], + drops: Arc::clone(&drops), + }; + let ptr = owner.as_ref().as_ptr(); + let buffer = ByteBuffer::from(Bytes::from_owner(owner)); + assert_eq!(buffer.as_ptr(), ptr); + assert_eq!(buffer.as_slice(), [1, 2, 3]); + let view = buffer.slice(1..); + drop(buffer); + assert_eq!(drops.load(Ordering::Relaxed), 0); + assert_eq!(view.as_slice(), [2, 3]); + drop(view); + assert_eq!(drops.load(Ordering::Relaxed), 1); + } + + #[test] + fn bytes_round_trip_reuses_owner() { + let bytes = Bytes::from_static(&[1, 2, 3, 4]); + let ptr = bytes.as_ptr(); + + let buffer = ByteBuffer::from(bytes); + assert!(matches!( + buffer.backing.as_deref(), + Some(BufferBacking::Bytes(_)) + )); + let bytes = buffer.into_bytes(); + + assert_eq!(bytes.as_ptr(), ptr); + assert_eq!(bytes.as_ref(), &[1, 2, 3, 4]); + } + + #[test] + fn external_try_into_mut_preserves_backing() { + let buffer = ByteBuffer::from(Bytes::from_static(&[1, 2, 3, 4])); + let Some(original_backing) = buffer.backing.as_ref() else { + panic!("external buffer has no backing") + }; + let backing = Arc::as_ptr(original_backing); + + let Err(buffer) = buffer.try_into_mut() else { + panic!("external buffer became mutable") + }; + + let Some(new_backing) = buffer.backing.as_ref() else { + panic!("external buffer has no backing") + }; + assert_eq!(Arc::as_ptr(new_backing), backing); + } + + #[test] + fn from_u8_vec_preserves_capacity() { + let mut vec = Vec::with_capacity(16); + vec.extend([1u8, 2, 3]); + + let buffer = Buffer::from(vec); + let Ok(buffer) = buffer.try_into_mut() else { + panic!("Vec-backed buffer should be uniquely owned") + }; + assert_eq!(buffer.capacity(), 16); + } + + #[test] + fn sliced_buffer_into_mut_has_safe_capacity() { + let mut original = crate::BufferMut::with_capacity(128); + original.extend(0u32..100); + let original = original.freeze(); + let sliced = original.slice(64..96); + drop(original); + + let Ok(mut sliced) = sliced.try_into_mut() else { + panic!("uniquely owned slice should become mutable") + }; + let ptr = sliced.as_ptr(); + let capacity = sliced.capacity(); + sliced.push_n(0, capacity - sliced.len()); + assert_eq!(sliced.len(), capacity); + assert_eq!(sliced.as_ptr(), ptr); + sliced.push(42); + assert_eq!(&sliced[..32], (64u32..96).collect::>()); + assert_eq!(&sliced[32..capacity], vec![0; capacity - 32]); + assert_eq!(sliced[capacity], 42); + } + + #[test] + fn from_vec_preserves_drop_glue() { + struct DropValue(Arc); + + impl Drop for DropValue { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + + let drops = Arc::new(AtomicUsize::new(0)); + let values = (0..3) + .map(|_| DropValue(Arc::clone(&drops))) + .collect::>(); + let buffer = Buffer::from(values); + + assert_eq!(drops.load(Ordering::Relaxed), 0); + drop(buffer); + assert_eq!(drops.load(Ordering::Relaxed), 3); + } + #[test] fn empty_aligned_max_alignment() { // Empty buffers are backed by a static and must satisfy any valid alignment. @@ -876,6 +1237,11 @@ mod test { assert!(buf.is_aligned(Alignment::MAX)); } + #[test] + fn empty_has_no_backing() { + assert!(Buffer::::empty().backing.is_none()); + } + #[test] fn empty_slice_preserves_alignment() { let buf = Buffer::::zeroed_aligned(8, Alignment::new(64)); diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index a7edfb00293..71e27c4f8ce 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -2,40 +2,68 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use core::mem::MaybeUninit; +use std::alloc::Layout; use std::any::type_name; use std::cmp::max; use std::fmt::Debug; use std::fmt::Formatter; -use std::io::Write; use std::ops::Deref; use std::ops::DerefMut; -use bytes::Buf; -use bytes::BufMut; -use bytes::BytesMut; -use bytes::buf::UninitSlice; +use itertools::Itertools; use vortex_error::VortexExpect; use vortex_error::vortex_panic; use crate::Alignment; +use crate::Allocation; use crate::Buffer; +use crate::BufferAllocatorRef; use crate::ByteBufferMut; use crate::debug::TruncatedDebug; use crate::trusted_len::TrustedLen; /// A mutable buffer that maintains a runtime-defined alignment through resizing operations. -#[derive(PartialEq, Eq)] +/// +/// Zero-sized element types are rejected at compile time when constructing a buffer. +/// +/// ```compile_fail +/// use vortex_buffer::BufferMut; +/// let _ = BufferMut::<()>::empty(); +/// ``` +/// +/// ```compile_fail +/// use vortex_buffer::BufferMut; +/// let _ = BufferMut::<()>::zeroed(3); +/// ``` pub struct BufferMut { - pub(crate) bytes: BytesMut, + /// The owned allocation, including any bytes before `ptr` used for alignment. + pub(crate) allocation: Allocation, + /// The first element, aligned to `alignment`; it may dangle for an empty buffer. + pub(crate) ptr: std::ptr::NonNull, + /// The number of initialized `T` values starting at `ptr`. pub(crate) length: usize, + /// The number of `T` values that fit from `ptr`. + pub(crate) capacity: usize, + /// The minimum alignment maintained for `ptr` across reallocations. pub(crate) alignment: Alignment, + /// Marks the buffer as logically owning values of `T` despite storing an erased allocation. pub(crate) _marker: std::marker::PhantomData, } +// SAFETY: BufferMut uniquely owns its allocation and only exposes T across threads. +unsafe impl Send for BufferMut {} +// SAFETY: shared access to BufferMut only exposes shared access to T. +unsafe impl Sync for BufferMut {} + impl BufferMut { /// Create a new `BufferMut` with the requested alignment and capacity. pub fn with_capacity(capacity: usize) -> Self { - Self::with_capacity_aligned(capacity, Alignment::of::()) + Self::with_capacity_in(capacity, BufferAllocatorRef::statically_allocated()) + } + + /// Create a new `BufferMut` with the requested capacity and allocator. + pub fn with_capacity_in(capacity: usize, allocator: BufferAllocatorRef) -> Self { + Self::with_capacity_aligned_in(capacity, Alignment::of::(), allocator) } /// Create a new `BufferMut` with the requested alignment and capacity. @@ -45,10 +73,24 @@ impl BufferMut { /// /// [`with_capacity_preferred_aligned`]: Self::with_capacity_preferred_aligned pub fn with_capacity_aligned(capacity: usize, alignment: Alignment) -> Self { - Self::with_capacity_preferred_aligned( + Self::with_capacity_aligned_in( + capacity, + alignment, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Create a new `BufferMut` with the requested alignment, capacity, and allocator. + pub fn with_capacity_aligned_in( + capacity: usize, + alignment: Alignment, + allocator: BufferAllocatorRef, + ) -> Self { + Self::with_capacity_preferred_aligned_in( capacity, alignment, Some(Alignment::DEFAULT_ALIGNMENT), + allocator, ) } @@ -61,6 +103,22 @@ impl BufferMut { alignment: Alignment, preferred_alignment: Option, ) -> Self { + Self::with_capacity_preferred_aligned_in( + capacity, + alignment, + preferred_alignment, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Create a new allocator-backed `BufferMut` with a requested and preferred alignment. + pub fn with_capacity_preferred_aligned_in( + capacity: usize, + alignment: Alignment, + preferred_alignment: Option, + allocator: BufferAllocatorRef, + ) -> Self { + const { assert!(size_of::() != 0, "ZSTs are not supported") }; let actual = max( alignment, preferred_alignment.unwrap_or(Alignment::of::()), @@ -74,12 +132,30 @@ impl BufferMut { ); } - let mut bytes = BytesMut::with_capacity((capacity * size_of::()) + *actual); - bytes.align_empty(actual); - + let size = capacity + .checked_mul(size_of::()) + .vortex_expect("buffer capacity overflow"); + let layout = if size == 0 { + Layout::from_size_align(0, actual.as_usize()) + .unwrap_or_else(|_| vortex_panic!("invalid empty buffer alignment")) + } else { + let allocation_size = size + .checked_add(actual.as_usize()) + .vortex_expect("buffer capacity overflow"); + Layout::from_size_align(allocation_size, 1).unwrap_or_else(|_| { + vortex_panic!("buffer capacity exceeds maximum allocation size") + }) + }; + let allocation = Allocation::allocate(layout, allocator); + let offset = allocation.ptr().as_ptr().align_offset(actual.as_usize()); + // SAFETY: the allocation includes enough padding to reach this aligned pointer. + let ptr = unsafe { allocation.ptr().add(offset).cast() }; + let capacity = (allocation.size() - offset) / size_of::(); Self { - bytes, + allocation, + ptr, length: 0, + capacity, alignment, _marker: Default::default(), } @@ -87,7 +163,12 @@ impl BufferMut { /// Create a new zeroed `BufferMut`. pub fn zeroed(len: usize) -> Self { - Self::zeroed_aligned(len, Alignment::of::()) + Self::zeroed_in(len, BufferAllocatorRef::statically_allocated()) + } + + /// Create a new zeroed `BufferMut` with the requested allocator. + pub fn zeroed_in(len: usize, allocator: BufferAllocatorRef) -> Self { + Self::zeroed_aligned_in(len, Alignment::of::(), allocator) } /// Create a new zeroed `BufferMut` with the requested alignment. @@ -97,7 +178,21 @@ impl BufferMut { /// /// [`zeroed_preferred_aligned`]: Self::zeroed_preferred_aligned pub fn zeroed_aligned(len: usize, alignment: Alignment) -> Self { - Self::zeroed_preferred_aligned(len, alignment, Some(Alignment::DEFAULT_ALIGNMENT)) + Self::zeroed_aligned_in(len, alignment, BufferAllocatorRef::statically_allocated()) + } + + /// Create a zeroed `BufferMut` with an alignment and allocator. + pub fn zeroed_aligned_in( + len: usize, + alignment: Alignment, + allocator: BufferAllocatorRef, + ) -> Self { + Self::zeroed_preferred_aligned_in( + len, + alignment, + Some(Alignment::DEFAULT_ALIGNMENT), + allocator, + ) } /// Create a new zeroed `BufferMut` with the requested alignment. @@ -109,15 +204,50 @@ impl BufferMut { alignment: Alignment, preferred_alignment: Option, ) -> Self { + Self::zeroed_preferred_aligned_in( + len, + alignment, + preferred_alignment, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Create a zeroed allocator-backed buffer with a requested and preferred alignment. + pub fn zeroed_preferred_aligned_in( + len: usize, + alignment: Alignment, + preferred_alignment: Option, + allocator: BufferAllocatorRef, + ) -> Self { + const { assert!(size_of::() != 0, "ZSTs are not supported") }; let preferred_alignment = preferred_alignment.unwrap_or(Alignment::of::()); let actual_alignment = max(preferred_alignment, alignment); - let mut bytes = BytesMut::zeroed((len * size_of::()) + *actual_alignment); - bytes.advance(bytes.as_ptr().align_offset(*actual_alignment)); - unsafe { bytes.set_len(len * size_of::()) }; - let actual_len = bytes.len().checked_div(size_of::()).unwrap_or(0); + let size = len + .checked_mul(size_of::()) + .vortex_expect("buffer length overflow"); + let layout = if size == 0 { + Layout::from_size_align(0, actual_alignment.as_usize()) + .unwrap_or_else(|_| vortex_panic!("invalid empty buffer alignment")) + } else { + let allocation_size = size + .checked_add(actual_alignment.as_usize()) + .vortex_expect("buffer length overflow"); + Layout::from_size_align(allocation_size, 1) + .unwrap_or_else(|_| vortex_panic!("buffer length exceeds maximum allocation size")) + }; + let allocation = Allocation::allocate_zeroed(layout, allocator); + let offset = allocation + .ptr() + .as_ptr() + .align_offset(actual_alignment.as_usize()); + // SAFETY: the allocation includes enough padding to reach this aligned pointer. + let ptr = unsafe { allocation.ptr().add(offset).cast() }; + let capacity = (allocation.size() - offset) / size_of::(); Self { - bytes, - length: actual_len, + allocation, + ptr, + length: len, + capacity, alignment, _marker: Default::default(), } @@ -135,7 +265,12 @@ impl BufferMut { /// /// [`empty_preferred_aligned`]: Self::empty_preferred_aligned pub fn empty_aligned(alignment: Alignment) -> Self { - Self::empty_preferred_aligned(alignment, Some(Alignment::DEFAULT_ALIGNMENT)) + Self::empty_aligned_in(alignment, BufferAllocatorRef::statically_allocated()) + } + + /// Create an empty `BufferMut` with an alignment and allocator. + pub fn empty_aligned_in(alignment: Alignment, allocator: BufferAllocatorRef) -> Self { + Self::with_capacity_aligned_in(0, alignment, allocator) } /// Create a new empty `BufferMut` with the provided alignment. @@ -146,7 +281,12 @@ impl BufferMut { alignment: Alignment, preferred_alignment: Option, ) -> Self { - BufferMut::with_capacity_preferred_aligned(0, alignment, preferred_alignment) + BufferMut::with_capacity_preferred_aligned_in( + 0, + alignment, + preferred_alignment, + BufferAllocatorRef::statically_allocated(), + ) } /// Create a new full `BufferMut` with the given value. @@ -154,14 +294,27 @@ impl BufferMut { where T: Copy, { - let mut buffer = BufferMut::::with_capacity(len); + Self::full_in(item, len, BufferAllocatorRef::statically_allocated()) + } + + /// Create a full `BufferMut` with the given value and allocator. + pub fn full_in(item: T, len: usize, allocator: BufferAllocatorRef) -> Self + where + T: Copy, + { + let mut buffer = BufferMut::::with_capacity_in(len, allocator); buffer.push_n(item, len); buffer } /// Create a mutable scalar buffer by copying the contents of the slice. pub fn copy_from(other: impl AsRef<[T]>) -> Self { - Self::copy_from_aligned(other, Alignment::of::()) + Self::copy_from_in(other, BufferAllocatorRef::statically_allocated()) + } + + /// Create a mutable scalar buffer by copying with the given allocator. + pub fn copy_from_in(other: impl AsRef<[T]>, allocator: BufferAllocatorRef) -> Self { + Self::copy_from_aligned_in(other, Alignment::of::(), allocator) } /// Create a mutable scalar buffer with the alignment by copying the contents of the slice. @@ -175,7 +328,21 @@ impl BufferMut { /// /// Panics when the requested alignment isn't itself aligned to type T. pub fn copy_from_aligned(other: impl AsRef<[T]>, alignment: Alignment) -> Self { - Self::copy_from_preferred_aligned(other, alignment, Some(Alignment::DEFAULT_ALIGNMENT)) + Self::copy_from_aligned_in(other, alignment, BufferAllocatorRef::statically_allocated()) + } + + /// Copy values into a mutable buffer with the given alignment and allocator. + pub fn copy_from_aligned_in( + other: impl AsRef<[T]>, + alignment: Alignment, + allocator: BufferAllocatorRef, + ) -> Self { + Self::copy_from_preferred_aligned_in( + other, + alignment, + Some(Alignment::DEFAULT_ALIGNMENT), + allocator, + ) } /// Create a mutable scalar buffer with the alignment by copying the contents of the slice. @@ -190,32 +357,58 @@ impl BufferMut { other: impl AsRef<[T]>, alignment: Alignment, preferred_alignment: Option, + ) -> Self { + Self::copy_from_preferred_aligned_in( + other, + alignment, + preferred_alignment, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Copy values with the given allocator, requested alignment, and preferred alignment. + pub fn copy_from_preferred_aligned_in( + other: impl AsRef<[T]>, + alignment: Alignment, + preferred_alignment: Option, + allocator: BufferAllocatorRef, ) -> Self { if !alignment.is_aligned_to(Alignment::of::()) { vortex_panic!("Given alignment is not aligned to type T") } let other = other.as_ref(); - let mut buffer = - Self::with_capacity_preferred_aligned(other.len(), alignment, preferred_alignment); + let mut buffer = Self::with_capacity_preferred_aligned_in( + other.len(), + alignment, + preferred_alignment, + allocator, + ); buffer.extend_from_slice(other); debug_assert_eq!(buffer.alignment(), alignment); buffer } /// Get the alignment of the buffer. + #[allow(clippy::inline_always)] #[inline(always)] pub fn alignment(&self) -> Alignment { self.alignment } + /// Returns the allocator that owns this buffer. + pub fn allocator(&self) -> &BufferAllocatorRef { + self.allocation.allocator() + } + /// Returns the length of the buffer. + #[allow(clippy::inline_always)] #[inline(always)] pub fn len(&self) -> usize { - debug_assert_eq!(self.length, self.bytes.len() / size_of::()); self.length } /// Returns whether the buffer is empty. + #[allow(clippy::inline_always)] #[inline(always)] pub fn is_empty(&self) -> bool { self.length == 0 @@ -224,29 +417,40 @@ impl BufferMut { /// Returns the capacity of the buffer. #[inline] pub fn capacity(&self) -> usize { - self.bytes.capacity() / size_of::() + self.capacity + } + + /// Returns a raw pointer to the buffer's data. + #[allow(clippy::inline_always)] + #[inline(always)] + pub fn as_ptr(&self) -> *const T { + self.ptr.as_ptr() + } + + /// Returns a mutable raw pointer to the buffer's data. + #[allow(clippy::inline_always)] + #[inline(always)] + pub fn as_mut_ptr(&mut self) -> *mut T { + self.ptr.as_ptr() } /// Returns a slice over the buffer of elements of type T. #[inline] pub fn as_slice(&self) -> &[T] { - let raw_slice = self.bytes.as_ref(); - // SAFETY: alignment of Buffer is checked on construction - unsafe { std::slice::from_raw_parts(raw_slice.as_ptr().cast(), self.length) } + // SAFETY: ptr is in the live allocation and construction checks its alignment. + unsafe { std::slice::from_raw_parts(self.as_ptr(), self.length) } } /// Returns a slice over the buffer of elements of type T. #[inline] pub fn as_mut_slice(&mut self) -> &mut [T] { - let raw_slice = self.bytes.as_mut(); - // SAFETY: alignment of Buffer is checked on construction - unsafe { std::slice::from_raw_parts_mut(raw_slice.as_mut_ptr().cast(), self.length) } + // SAFETY: BufferMut uniquely owns the allocation and the initialized range is in bounds. + unsafe { std::slice::from_raw_parts_mut(self.as_mut_ptr(), self.length) } } /// Clear the buffer, retaining any existing capacity. #[inline] pub fn clear(&mut self) { - unsafe { self.bytes.set_len(0) } self.length = 0; } @@ -268,8 +472,7 @@ impl BufferMut { /// Reserves capacity for at least `additional` more elements to be inserted in the buffer. #[inline] pub fn reserve(&mut self, additional: usize) { - let additional_bytes = additional * size_of::(); - if additional_bytes <= self.bytes.capacity() - self.bytes.len() { + if additional <= self.capacity() - self.length { // We can fit the additional bytes in the remaining capacity. Nothing to do. return; } @@ -278,17 +481,73 @@ impl BufferMut { self.reserve_allocate(additional); } - /// A separate function so we can inline the reserve call's fast path. According to `BytesMut` - /// this has significant performance implications. + /// A separate function so we can inline the reserve call's fast path. fn reserve_allocate(&mut self, additional: usize) { - let new_capacity: usize = ((self.length + additional) * size_of::()) + *self.alignment; - // Make sure we at least double in size each time we re-allocate to amortize the cost - let new_capacity = new_capacity.max(self.bytes.capacity() * 2); - - let mut bytes = BytesMut::with_capacity(new_capacity); - bytes.align_empty(self.alignment); - bytes.extend_from_slice(&self.bytes); - self.bytes = bytes; + let required = self + .length + .checked_add(additional) + .vortex_expect("buffer capacity overflow"); + let required_size = required + .checked_mul(size_of::()) + .vortex_expect("buffer capacity overflow"); + let alignment = self.alignment; + let current_size = self + .capacity + .checked_mul(size_of::()) + .vortex_expect("buffer capacity overflow"); + let logical_size = required_size + .max(current_size.saturating_mul(2)) + .max(Alignment::DEFAULT_ALIGNMENT.as_usize()); + let allocation_size = logical_size + .checked_add(alignment.as_usize()) + .vortex_expect("buffer capacity overflow"); + let allocation_alignment = if self.allocation.size() == 0 { + 1 + } else { + self.allocation.alignment() + }; + let layout = Layout::from_size_align(allocation_size, allocation_alignment) + .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); + + let old_offset = self.ptr.cast::().addr().get() - self.allocation.ptr().addr().get(); + let new_offset = if self.allocation.allocator().is_statically_allocated() { + let allocation = + Allocation::allocate(layout, BufferAllocatorRef::statically_allocated()); + let new_offset = allocation.ptr().as_ptr().align_offset(alignment.as_usize()); + // SAFETY: both allocations have room for the initialized elements and do not overlap. + unsafe { + std::ptr::copy_nonoverlapping( + self.ptr.cast::().as_ptr(), + allocation.ptr().as_ptr().add(new_offset), + self.length * size_of::(), + ); + } + self.allocation = allocation; + new_offset + } else { + self.allocation.grow(layout); + let new_offset = self + .allocation + .ptr() + .as_ptr() + .align_offset(alignment.as_usize()); + if new_offset != old_offset { + // SAFETY: grow preserved the initialized elements at old_offset. The new allocation + // has room for the requested elements plus alignment padding, and copy permits + // overlap. + unsafe { + std::ptr::copy( + self.allocation.ptr().as_ptr().add(old_offset), + self.allocation.ptr().as_ptr().add(new_offset), + self.length * size_of::(), + ); + } + } + new_offset + }; + // SAFETY: new_offset was computed within the allocation for alignment. + self.ptr = unsafe { self.allocation.ptr().add(new_offset).cast() }; + self.capacity = logical_size / size_of::(); } /// Returns the spare capacity of the buffer as a slice of `MaybeUninit`. @@ -328,13 +587,9 @@ impl BufferMut { /// ``` #[inline] pub fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit] { - let dst = self.bytes.spare_capacity_mut().as_mut_ptr(); - unsafe { - std::slice::from_raw_parts_mut( - dst as *mut MaybeUninit, - self.capacity() - self.length, - ) - } + // SAFETY: offset + length is within the allocation and points at spare capacity. + let dst = unsafe { self.as_mut_ptr().add(self.length) }.cast::>(); + unsafe { std::slice::from_raw_parts_mut(dst, self.capacity() - self.length) } } /// Sets the length of the buffer. @@ -348,7 +603,6 @@ impl BufferMut { #[inline] pub unsafe fn set_len(&mut self, len: usize) { debug_assert!(len <= self.capacity()); - unsafe { self.bytes.set_len(len * size_of::()) }; self.length = len; } @@ -368,9 +622,8 @@ impl BufferMut { pub unsafe fn push_unchecked(&mut self, item: T) { // SAFETY: the caller ensures we have sufficient capacity unsafe { - let dst: *mut T = self.bytes.spare_capacity_mut().as_mut_ptr().cast(); + let dst = self.as_mut_ptr().add(self.length); dst.write(item); - self.bytes.set_len(self.bytes.len() + size_of::()) } self.length += 1; } @@ -397,7 +650,8 @@ impl BufferMut { where T: Copy, { - let mut dst: *mut T = self.bytes.spare_capacity_mut().as_mut_ptr().cast(); + // SAFETY: the caller guarantees enough spare capacity. + let mut dst = unsafe { self.as_mut_ptr().add(self.length) }; // SAFETY: we checked the capacity in the reserve call unsafe { let end = dst.add(n); @@ -405,7 +659,6 @@ impl BufferMut { dst.write(item); dst = dst.add(1); } - self.bytes.set_len(self.bytes.len() + (n * size_of::())); } self.length += n; } @@ -425,73 +678,28 @@ impl BufferMut { #[inline] pub fn extend_from_slice(&mut self, slice: &[T]) { self.reserve(slice.len()); - let raw_slice = - unsafe { std::slice::from_raw_parts(slice.as_ptr().cast(), size_of_val(slice)) }; - self.bytes.extend_from_slice(raw_slice); - self.length += slice.len(); - } - - /// Splits the buffer into two at the given index. - /// - /// Afterward, self contains elements `[0, at)`, and the returned buffer contains elements - /// `[at, capacity)`. It’s guaranteed that the memory does not move, that is, the address of - /// self does not change, and the address of the returned slice is at bytes after that. - /// - /// This is an O(1) operation that just increases the reference count and sets a few indices. - /// - /// Panics if either half would have a length that is not a multiple of the alignment. - pub fn split_off(&mut self, at: usize) -> Self { - if at > self.capacity() { - vortex_panic!("Cannot split buffer of capacity {} at {}", self.len(), at); - } - - let bytes_at = at * size_of::(); - if !self.alignment.is_offset_aligned(bytes_at) { - vortex_panic!( - "Cannot split buffer at {}, resulting alignment is not {}", - at, - self.alignment - ); - } - - let new_bytes = self.bytes.split_off(bytes_at); - - // Adjust the lengths, given that length may be < at - let new_length = self.length.saturating_sub(at); - self.length = self.length.min(at); - - BufferMut { - bytes: new_bytes, - length: new_length, - alignment: self.alignment, - _marker: Default::default(), - } - } - - /// Absorbs a mutable buffer that was previously split off. - /// - /// If the two buffers were previously contiguous and not mutated in a way that causes - /// re-allocation i.e., if other was created by calling split_off on this buffer, then this is - /// an O(1) operation that just decreases a reference count and sets a few indices. - /// - /// Otherwise, this method degenerates to self.extend_from_slice(other.as_ref()). - pub fn unsplit(&mut self, other: Self) { - if self.alignment != other.alignment { - vortex_panic!( - "Cannot unsplit buffers with different alignments: {} and {}", - self.alignment, - other.alignment + // SAFETY: reserve made the destination valid and non-overlapping for slice.len() values. + unsafe { + std::ptr::copy_nonoverlapping( + slice.as_ptr(), + self.as_mut_ptr().add(self.length), + slice.len(), ); } - self.bytes.unsplit(other.bytes); - self.length += other.length; + self.length += slice.len(); } /// Return the [`ByteBufferMut`] for this [`BufferMut`]. pub fn into_byte_buffer(self) -> ByteBufferMut { + let capacity = self + .capacity + .checked_mul(size_of::()) + .vortex_expect("buffer capacity overflow"); ByteBufferMut { - bytes: self.bytes, + allocation: self.allocation, + ptr: self.ptr.cast(), length: self.length * size_of::(), + capacity, alignment: self.alignment, _marker: Default::default(), } @@ -499,12 +707,8 @@ impl BufferMut { /// Freeze the `BufferMut` into a `Buffer`. pub fn freeze(self) -> Buffer { - Buffer { - bytes: self.bytes.freeze(), - length: self.length, - alignment: self.alignment, - _marker: Default::default(), - } + let offset = self.ptr.cast::().addr().get() - self.allocation.ptr().addr().get(); + Buffer::from_allocation(self.allocation, offset, self.length, self.alignment) } /// Map each element of the buffer with a closure. @@ -531,15 +735,15 @@ impl BufferMut { /// /// If the data is not aligned, we copy it into a new allocation. pub fn aligned(self, alignment: Alignment) -> Self { - if self.as_ptr().align_offset(*alignment) == 0 { - Self { - bytes: self.bytes, - length: self.length, - alignment, - _marker: std::marker::PhantomData, - } + if self.as_ptr().align_offset(alignment.as_usize()) == 0 { + Self { alignment, ..self } } else { - Self::copy_from_aligned(self, alignment) + let capacity = self.capacity(); + let allocator = self.allocation.allocator().clone(); + let mut aligned = Self::with_capacity_aligned_in(capacity, alignment, allocator); + aligned.extend_from_slice(&self); + aligned.capacity = capacity; + aligned } } @@ -563,8 +767,10 @@ impl BufferMut { ); BufferMut { - bytes: self.bytes, + allocation: self.allocation, + ptr: self.ptr.cast(), length: self.length, + capacity: self.capacity, alignment: self.alignment, _marker: std::marker::PhantomData, } @@ -573,14 +779,24 @@ impl BufferMut { impl Clone for BufferMut { fn clone(&self) -> Self { - // NOTE(ngates): we cannot derive Clone since BytesMut copies on clone and the alignment - // might be messed up. - let mut buffer = BufferMut::::with_capacity_aligned(self.capacity(), self.alignment); + let mut buffer = BufferMut::::with_capacity_aligned_in( + self.capacity(), + self.alignment, + self.allocation.allocator().clone(), + ); buffer.extend_from_slice(self.as_slice()); buffer } } +impl PartialEq for BufferMut { + fn eq(&self, other: &Self) -> bool { + self.as_slice() == other.as_slice() + } +} + +impl Eq for BufferMut {} + impl Debug for BufferMut { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct(&format!("BufferMut<{}>", type_name::())) @@ -644,7 +860,7 @@ impl BufferMut { let unwritten = self.capacity() - self.len(); // We store `begin` in the case that the lower bound hint is incorrect. - let begin: *const T = self.bytes.spare_capacity_mut().as_mut_ptr().cast(); + let begin: *const T = self.spare_capacity_mut().as_mut_ptr().cast(); let mut dst: *mut T = begin.cast_mut(); // As a first step, we manually iterate the iterator up to the known capacity. @@ -683,16 +899,13 @@ impl BufferMut { /// The caller guarantees that the iterator will have a trusted upper bound, which allows the /// implementation to reserve all of the memory needed up front. pub fn extend_trusted>(&mut self, iter: I) { - // Since we know the exact upper bound (from `TrustedLen`), we can reserve all of the memory - // for this operation up front. let (_, upper_bound) = iter.size_hint(); self.reserve( upper_bound .vortex_expect("`TrustedLen` iterator somehow didn't have valid upper bound"), ); - // We store `begin` in the case that the upper bound hint is incorrect. - let begin: *const T = self.bytes.spare_capacity_mut().as_mut_ptr().cast(); + let begin: *const T = self.spare_capacity_mut().as_mut_ptr().cast(); let mut dst: *mut T = begin.cast_mut(); iter.for_each(|item| { @@ -702,13 +915,13 @@ impl BufferMut { // Note: We used to have `dst.add(iteration).write(item)`, here. However this was much // slower than just incrementing `dst`. - // SAFETY: The offsets fits in `isize`, and because we were able to reserve the memory + // SAFETY: The offset fits in `isize`, and because we were able to reserve the memory // we know that `add` will not overflow. unsafe { dst = dst.add(1) }; }); - // SAFETY: `dst` was derived from `begin`, which were both valid references to byte data, - // and since the only operation that `dst` has is `add`, we know that `dst >= begin`. + // SAFETY: `dst` starts at `begin` and advances by one for each item, so both pointers refer + // to the same allocation and `dst` is at or after `begin`. let items_written = unsafe { dst.offset_from_unsigned(begin) }; let length = self.len() + items_written; @@ -741,40 +954,7 @@ impl BufferMut { where I: TrustedLen>, { - let (_, upper_bound) = iter.size_hint(); - self.reserve( - upper_bound - .vortex_expect("`TrustedLen` iterator somehow didn't have valid upper bound"), - ); - - let begin: *const T = self.bytes.spare_capacity_mut().as_mut_ptr().cast(); - let mut dst: *mut T = begin.cast_mut(); - let mut result: Result<(), E> = Ok(()); - - for item in iter { - match item { - Ok(value) => { - // SAFETY: We reserved enough capacity to hold this item, and `dst` is a - // pointer derived from a valid reference to byte data. - unsafe { dst.write(value) }; - // SAFETY: The offset fits in `isize` because we reserved that much capacity. - unsafe { dst = dst.add(1) }; - } - Err(e) => { - result = Err(e); - break; - } - } - } - - // SAFETY: `dst` was derived from `begin`, both valid references to byte data, and - // `dst >= begin` since the only operation on `dst` is `add`. - let items_written = unsafe { dst.offset_from_unsigned(begin) }; - let length = self.len() + items_written; - // SAFETY: We have written valid items between the old length and the new length. - unsafe { self.set_len(length) }; - - result + iter.process_results(|values| self.extend_trusted(values)) } /// Like [`from_trusted_len_iter()`](Self::from_trusted_len_iter), but the iterator yields @@ -783,14 +963,7 @@ impl BufferMut { where I: TrustedLen>, { - let (_, upper_bound) = iter.size_hint(); - let mut buffer = Self::with_capacity( - upper_bound - .vortex_expect("`TrustedLen` iterator somehow didn't have valid upper bound"), - ); - - buffer.try_extend_trusted(iter)?; - Ok(buffer) + iter.process_results(|values| Self::from_trusted_len_iter(values)) } } @@ -813,147 +986,86 @@ where impl FromIterator for BufferMut { fn from_iter>(iter: I) -> Self { - // We don't infer the capacity here and just let the first call to `extend` do it for us. - let mut buffer = Self::with_capacity(0); + let iter = iter.into_iter(); + let mut buffer = Self::with_capacity(iter.size_hint().0); buffer.extend(iter); buffer } } -impl Buf for ByteBufferMut { - fn remaining(&self) -> usize { - self.len() - } +#[cfg(test)] +mod test { + use crate::Alignment; + use crate::BufferMut; + use crate::buffer_mut; - fn chunk(&self) -> &[u8] { - self.as_slice() - } + #[test] + fn capacity() { + let mut n = 57; + let mut buf = BufferMut::::with_capacity_aligned(n, Alignment::new(1024)); + assert!(buf.capacity() >= 57); - fn advance(&mut self, cnt: usize) { - if !self.alignment.is_offset_aligned(cnt) { - vortex_panic!( - "Cannot advance buffer by {} items, resulting alignment is not {}", - cnt, - self.alignment - ); + while n > 0 { + buf.push(0); + assert!(buf.capacity() >= n); + n -= 1 } - self.bytes.advance(cnt); - self.length -= cnt; - } -} -/// As per the BufMut implementation, we must support internal resizing when -/// asked to extend the buffer. -/// See: -unsafe impl BufMut for ByteBufferMut { - #[inline] - fn remaining_mut(&self) -> usize { - usize::MAX - self.len() - } - - #[inline] - unsafe fn advance_mut(&mut self, cnt: usize) { - if !self.alignment.is_offset_aligned(cnt) { - vortex_panic!( - "Cannot advance buffer by {} items, resulting alignment is not {}", - cnt, - self.alignment - ); - } - unsafe { self.bytes.advance_mut(cnt) }; - self.length -= cnt; + assert_eq!(buf.alignment(), Alignment::new(1024)); } - #[inline] - fn chunk_mut(&mut self) -> &mut UninitSlice { - self.bytes.chunk_mut() - } + #[test] + fn growth_preserves_alignment_and_values() { + let alignment = Alignment::new(4096); + let mut buffer = BufferMut::::with_capacity_aligned(1, alignment); - fn put(&mut self, mut src: T) - where - Self: Sized, - { - while src.has_remaining() { - let chunk = src.chunk(); - self.extend_from_slice(chunk); - src.advance(chunk.len()); + for value in 0..10_000 { + buffer.push(value); + assert!(alignment.is_offset_aligned(buffer.as_ptr().addr())); } - } - #[inline] - fn put_slice(&mut self, src: &[u8]) { - self.extend_from_slice(src); + assert_eq!(buffer.as_slice(), (0..10_000).collect::>()); } - #[inline] - fn put_bytes(&mut self, val: u8, cnt: usize) { - self.push_n(val, cnt) - } -} - -/// Extension trait for [`BytesMut`] that provides functions for aligning the buffer. -trait AlignedBytesMut { - /// Align an empty `BytesMut` to the specified alignment. - /// - /// ## Panics - /// - /// Panics if the buffer is not empty, or if there is not enough capacity to align the buffer. - fn align_empty(&mut self, alignment: Alignment); -} - -impl AlignedBytesMut for BytesMut { - fn align_empty(&mut self, alignment: Alignment) { - // TODO(joe): this is slow fixme - if !self.is_empty() { - vortex_panic!("ByteBufferMut must be empty"); - } - - let padding = self.as_ptr().align_offset(*alignment); - self.capacity() - .checked_sub(padding) - .vortex_expect("Not enough capacity to align buffer"); + #[test] + fn growth_seeds_and_doubles_logical_capacity() { + let alignment = Alignment::new(64); + let mut buffer = BufferMut::::empty_aligned(alignment); - // SAFETY: We know the buffer is empty, and we know we have enough capacity, so we can - // safely set the length to the padding and advance the buffer to the aligned offset. - unsafe { self.set_len(padding) }; - self.advance(padding); - } -} + buffer.push(0); + let capacity = buffer.capacity(); + assert_eq!(capacity, Alignment::DEFAULT_ALIGNMENT.as_usize()); -impl Write for ByteBufferMut { - fn write(&mut self, buf: &[u8]) -> std::io::Result { - self.extend_from_slice(buf); - Ok(buf.len()) + buffer.reserve(capacity); + assert_eq!(buffer.capacity(), capacity * 2); } - fn flush(&mut self) -> std::io::Result<()> { - Ok(()) - } -} + #[test] + fn static_growth_copies_live_data() { + let mut buffer = BufferMut::::with_capacity(1); + let capacity = buffer.capacity(); + buffer.extend(std::iter::repeat_n(7, capacity)); + let old_ptr = buffer.as_ptr(); -#[cfg(test)] -mod test { - use bytes::Buf; - use bytes::BufMut; + buffer.push(u32::MAX); - use crate::Alignment; - use crate::BufferMut; - use crate::ByteBufferMut; - use crate::buffer_mut; + assert_ne!(buffer.as_ptr(), old_ptr); + assert_eq!(&buffer[..capacity], vec![7; capacity]); + assert_eq!(buffer[capacity], u32::MAX); + } #[test] - fn capacity() { - let mut n = 57; - let mut buf = BufferMut::::with_capacity_aligned(n, Alignment::new(1024)); - assert!(buf.capacity() >= 57); + fn raising_logical_alignment_preserves_capacity() { + let buffer = + BufferMut::::with_capacity_preferred_aligned(1, Alignment::of::(), None); + let capacity = buffer.capacity(); - while n > 0 { - buf.push(0); - assert!(buf.capacity() >= n); - n -= 1 - } + let mut buffer = buffer.aligned(Alignment::new(2)); - assert_eq!(buf.alignment(), Alignment::new(1024)); + assert_eq!(buffer.capacity(), capacity); + buffer.extend(0..100); + assert!(Alignment::new(2).is_ptr_aligned(buffer.as_ptr())); + assert_eq!(buffer.as_slice(), (0..100).collect::>()); } #[test] @@ -981,6 +1093,15 @@ mod test { assert_eq!(result.err(), Some("bad")); } + #[test] + fn try_extend_trusted_retains_values_before_error() { + let mut buf = BufferMut::from_iter([0, 10]); + let result = buf.try_extend_trusted([Ok(20), Err("bad"), Ok(30)].into_iter()); + + assert_eq!(result, Err("bad")); + assert_eq!(buf.as_slice(), &[0, 10, 20]); + } + #[test] fn extend() { let mut buf = BufferMut::empty(); @@ -1023,34 +1144,16 @@ mod test { assert_eq!(buf.as_slice(), &[1u32, 2, 3]); } - #[test] - fn bytes_buf() { - let mut buf = ByteBufferMut::copy_from("helloworld".as_bytes()); - assert_eq!(buf.remaining(), 10); - assert_eq!(buf.chunk(), b"helloworld"); - - buf.advance(5); - assert_eq!(buf.remaining(), 5); - assert_eq!(buf.as_slice(), b"world"); - assert_eq!(buf.chunk(), b"world"); - } - - #[test] - fn bytes_buf_mut() { - let mut buf = ByteBufferMut::copy_from("hello".as_bytes()); - assert_eq!(BufMut::remaining_mut(&buf), usize::MAX - 5); - - buf.put_slice(b"world"); - assert_eq!(buf.as_slice(), b"helloworld"); - } - #[test] fn buffer_mut_zeroed() { const LEN: usize = 17; let mut buf = BufferMut::::zeroed(LEN); - assert_eq!(buf.as_ptr().align_offset(*Alignment::of::()), 0); + assert_eq!( + buf.as_ptr().align_offset(Alignment::of::().as_usize()), + 0 + ); assert_eq!(buf.as_slice(), &[0; LEN]); buf[3] = 7; @@ -1064,7 +1167,7 @@ mod test { let mut buf = BufferMut::::zeroed_aligned(LEN, alignment); - assert_eq!(buf.as_ptr().align_offset(*alignment), 0); + assert_eq!(buf.as_ptr().align_offset(alignment.as_usize()), 0); assert_eq!(buf.as_slice(), &[0; LEN]); buf[3] = 7; diff --git a/vortex-buffer/src/const.rs b/vortex-buffer/src/const.rs index bd696c371bb..d37631c3eea 100644 --- a/vortex-buffer/src/const.rs +++ b/vortex-buffer/src/const.rs @@ -30,6 +30,7 @@ impl ConstBuffer { } /// Returns a slice over the buffer of elements of type T. + #[allow(clippy::inline_always)] #[inline(always)] pub fn as_slice(&self) -> &[T] { self.0.as_slice() diff --git a/vortex-buffer/src/lib.rs b/vortex-buffer/src/lib.rs index ee113481353..99b1e7c3081 100644 --- a/vortex-buffer/src/lib.rs +++ b/vortex-buffer/src/lib.rs @@ -5,11 +5,10 @@ //! A library for working with custom aligned buffers of sized values. //! -//! The `vortex-buffer` crate is built around `bytes::Bytes` and therefore supports zero-copy -//! cloning and slicing, but differs in that it can define and maintain a custom alignment. +//! The `vortex-buffer` crate supports zero-copy cloning and slicing with a custom allocator and +//! runtime alignment. //! -//! * `Buffer` and `BufferMut` provide immutable and mutable wrappers around `bytes::Bytes` -//! and `bytes::BytesMut` respectively. +//! * `Buffer` and `BufferMut` provide immutable and mutable typed buffers. //! * `ByteBuffer` and `ByteBufferMut` are type aliases for `u8` buffers. //! * `BufferString` is a wrapper around a `ByteBuffer` that enforces utf-8 encoding. //! * `ConstBuffer` provides similar functionality to `Buffer` except with a @@ -47,6 +46,7 @@ //! `arrow_buffer::OffsetBuffer`. pub use alignment::*; +pub use allocation::*; pub use bit::*; pub use buffer::*; pub use buffer_mut::*; @@ -55,6 +55,7 @@ pub use r#const::*; pub use dispatch::*; pub use string::*; mod alignment; +mod allocation; #[cfg(feature = "arrow")] mod arrow; mod bit; diff --git a/vortex-buffer/src/serde.rs b/vortex-buffer/src/serde.rs index 9563236a50a..5d36f8da3d0 100644 --- a/vortex-buffer/src/serde.rs +++ b/vortex-buffer/src/serde.rs @@ -22,7 +22,7 @@ where where S: Serializer, { - serializer.serialize_bytes(self.inner().as_ref()) + serializer.serialize_bytes(self.as_bytes()) } } diff --git a/vortex-cloud/Cargo.toml b/vortex-cloud/Cargo.toml new file mode 100644 index 00000000000..7ac7de96104 --- /dev/null +++ b/vortex-cloud/Cargo.toml @@ -0,0 +1,83 @@ +[package] +name = "vortex-cloud" +description = "Cloud object store integration for Vortex: URL resolution and OpenDAL-backed services" +version = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } +authors = { workspace = true } +license = { workspace = true } +keywords = { workspace = true } +include = { workspace = true } +edition = { workspace = true } +readme = { workspace = true } +rust-version = { workspace = true } +categories = { workspace = true } + +[package.metadata.docs.rs] +all-features = true + +[dependencies] +http = { workspace = true, optional = true } +object_store = { workspace = true, features = ["fs"] } +object_store_opendal = { workspace = true, optional = true } +opendal = { workspace = true, optional = true } +parking_lot = { workspace = true, optional = true } +percent-encoding = { workspace = true, optional = true } +tracing = { workspace = true, optional = true } +url = { workspace = true } +vortex-utils = { workspace = true } + +[dev-dependencies] +rstest = { workspace = true } +tempfile = { workspace = true } +# So the registry's key resolution can be tested end to end against a real file open, which is +# where the key reaches the store. +tokio = { workspace = true, features = ["rt", "macros"] } +vortex-array = { workspace = true } +vortex-error = { workspace = true, features = ["object_store"] } +vortex-file = { workspace = true, features = ["object_store", "tokio"] } +vortex-io = { workspace = true, features = ["object_store", "tokio"] } +vortex-layout = { workspace = true } + +[features] +default = [] +# The URL -> ObjectStore registry, plus the cloud backends it resolves URLs to. Kept optional so +# that a consumer needing only an OpenDAL store does not pull in the native cloud backends. +registry = [ + "dep:parking_lot", + "object_store/aws", + "object_store/azure", + "object_store/gcp", + "object_store/http", +] +# The Hugging Face Hub, the `hf://` scheme. Served over `object_store`'s HTTP store, so it adds no +# cloud SDK of its own. +hf = ["dep:http", "dep:percent-encoding", "object_store/http"] +# Tencent Cloud COS, the `cos://` scheme. +cos = [ + "dep:opendal", + "dep:object_store_opendal", + "dep:tracing", + "object_store/cloud", + "opendal/services-cos", +] +# Alibaba Cloud OSS, the `oss://` scheme. +oss = [ + "dep:opendal", + "dep:object_store_opendal", + "dep:tracing", + "object_store/cloud", + "opendal/services-oss", +] +# Tencent Cloud GooseFS, the `goosefs://` scheme. +goosefs = [ + "dep:opendal", + "dep:object_store_opendal", + "dep:tracing", + "opendal/services-goosefs", +] +# Every OpenDAL-backed service. +opendal = ["cos", "oss", "goosefs"] + +[lints] +workspace = true diff --git a/vortex-cloud/src/hf/mod.rs b/vortex-cloud/src/hf/mod.rs new file mode 100644 index 00000000000..9a8caff8246 --- /dev/null +++ b/vortex-cloud/src/hf/mod.rs @@ -0,0 +1,420 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The Hugging Face Hub, served over [`object_store`]'s HTTP store. +//! +//! A Hub repository is a set of files behind an HTTP endpoint that honours range requests, so it +//! needs no cloud SDK of its own: a [`object_store::http::HttpStore`] rooted at the repository's +//! `resolve` prefix, carrying a bearer token when one is available, is the whole implementation. +//! Reads therefore keep every [`object_store::ClientOptions`] setting a caller passes — +//! connect/request timeouts, retries, proxy configuration, `allow_http` — unlike the OpenDAL-backed +//! schemes in the crate's `opendal` module, whose bridge owns its own HTTP client. +//! +//! # URL grammar +//! +//! ```text +//! hf://datasets//[@][/] +//! hf://spaces//[@][/] +//! hf:///[@][/] # a model repository +//! ``` +//! +//! matching the paths [`huggingface_hub`'s `HfFileSystem`][hffs] accepts. A revision containing `/` +//! (e.g. `refs/convert/parquet`) must be percent-encoded in the URL, as it must be there: +//! +//! ```text +//! hf://datasets/org/name@refs%2Fconvert%2Fparquet/data/train.vortex +//! ``` +//! +//! [hffs]: https://huggingface.co/docs/huggingface_hub/guides/hf_file_system +//! +//! # Configuration +//! +//! * `HF_TOKEN` — the API token. Falling back, as `huggingface_hub` does, to the token file at +//! `HF_TOKEN_PATH`, then `$HF_HOME/token`, then `$HOME/.cache/huggingface/token`. Without a token +//! the store reads anonymously, which is all a public repository needs. +//! * `HF_ENDPOINT` — the Hub endpoint, defaulting to `https://huggingface.co`. +//! +//! # Limitations +//! +//! The Hub does not implement WebDAV `PROPFIND`, which is how [`object_store`]'s HTTP store lists a +//! prefix, so [`object_store::ObjectStore::list`] fails against these URLs. Opening a file by name +//! works — that is `head` plus ranged `get`, both plain HTTP — so this serves reads of a known path. +//! Callers that need to expand a glob must list the repository through the Hub's own API first and +//! then open each returned path. + +use std::sync::Arc; + +use http::HeaderMap; +use http::HeaderValue; +use http::header::AUTHORIZATION; +use object_store::ClientOptions; +use object_store::ObjectStore; +use object_store::http::HttpBuilder; +use object_store::path::Path; +use percent_encoding::AsciiSet; +use percent_encoding::CONTROLS; +use percent_encoding::percent_decode_str; +use percent_encoding::utf8_percent_encode; +use url::Url; + +/// The URL scheme served by the Hugging Face Hub. +pub const HF_SCHEME: &str = "hf"; + +/// The Hub endpoint used when `HF_ENDPOINT` is unset. +const DEFAULT_ENDPOINT: &str = "https://huggingface.co"; + +/// The revision used when a URL does not name one. +const DEFAULT_REVISION: &str = "main"; + +const ENDPOINT_VAR: &str = "HF_ENDPOINT"; +const TOKEN_VAR: &str = "HF_TOKEN"; +const TOKEN_PATH_VAR: &str = "HF_TOKEN_PATH"; +const HF_HOME_VAR: &str = "HF_HOME"; +const HOME_VAR: &str = "HOME"; +const ALLOW_HTTP_VAR: &str = "ALLOW_HTTP"; + +/// The URL authority that marks a dataset repository. +const DATASETS_HOST: &str = "datasets"; +/// The URL authority that marks a Space repository. +const SPACES_HOST: &str = "spaces"; + +/// Characters escaped when a revision is spliced into the `resolve` path. +/// +/// `/` is the point of this: the Hub routes `refs/convert/parquet` only in its `%2F` form, since an +/// unescaped `/` would read as another path segment. +const REVISION_ESCAPES: &AsciiSet = &CONTROLS.add(b'/').add(b'%').add(b'?').add(b'#').add(b' '); + +/// Whether `scheme` is served by this module. +/// +/// The registry dispatches on this rather than comparing against [`HF_SCHEME`] itself, matching +/// how the OpenDAL module's `supports_scheme` is used. +pub(crate) fn supports_scheme(scheme: &str) -> bool { + scheme == HF_SCHEME +} + +/// Which kind of Hub repository a URL addresses. +/// +/// The Hub routes each kind under its own path prefix, which is the only thing that differs between +/// them here. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum HfRepoType { + /// A dataset repository, `hf://datasets//`. + Dataset, + /// A model repository, `hf:///`. + Model, + /// A Space repository, `hf://spaces//`. + Space, +} + +impl HfRepoType { + /// The path prefix, including its trailing `/`, that the Hub routes this kind under. + fn url_prefix(self) -> &'static str { + match self { + HfRepoType::Dataset => "datasets/", + HfRepoType::Model => "", + HfRepoType::Space => "spaces/", + } + } +} + +impl std::str::FromStr for HfRepoType { + type Err = HfStoreError; + + /// Parses either the singular name or the plural the Hub uses in its URLs, so a caller can pass + /// whichever it has. + fn from_str(s: &str) -> Result { + match s { + "dataset" | "datasets" => Ok(HfRepoType::Dataset), + "model" | "models" => Ok(HfRepoType::Model), + "space" | "spaces" => Ok(HfRepoType::Space), + other => Err(HfStoreError::UnknownRepoType(other.to_string())), + } + } +} + +/// Error type for building a Hugging Face Hub object store. +#[derive(Debug)] +pub enum HfStoreError { + /// The URL scheme is not one this module handles. + UnsupportedScheme(String), + /// The URL is not a well-formed `hf://` URL. + InvalidUrl(String), + /// The repository kind is not one the Hub serves. + UnknownRepoType(String), + /// The bearer token could not be used as an HTTP header value. + InvalidToken(http::header::InvalidHeaderValue), + /// The underlying HTTP store rejected the configuration. + Build(object_store::Error), +} + +impl std::fmt::Display for HfStoreError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + HfStoreError::UnsupportedScheme(s) => write!(f, "unsupported Hugging Face scheme: {s}"), + HfStoreError::InvalidUrl(url) => write!( + f, + "invalid Hugging Face URL {url}: expected hf://datasets//[@revision][/path], \ + hf://spaces//[@revision][/path] or hf:///[@revision][/path]" + ), + HfStoreError::UnknownRepoType(kind) => write!( + f, + "unknown Hugging Face repository type {kind}: expected one of dataset, model, space" + ), + HfStoreError::InvalidToken(e) => write!(f, "invalid Hugging Face token: {e}"), + HfStoreError::Build(e) => write!(f, "failed to build Hugging Face store: {e}"), + } + } +} + +impl std::error::Error for HfStoreError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + HfStoreError::InvalidToken(e) => Some(e), + HfStoreError::Build(e) => Some(e), + HfStoreError::UnsupportedScheme(_) + | HfStoreError::InvalidUrl(_) + | HfStoreError::UnknownRepoType(_) => None, + } + } +} + +impl From for object_store::Error { + fn from(e: HfStoreError) -> Self { + object_store::Error::Generic { + store: "HuggingFace", + source: Box::new(e), + } + } +} + +/// Configuration for a store serving one Hub repository at one revision. +/// +/// A store covers a single `(repository, revision)` pair because that pair is what the Hub's +/// `resolve` prefix names; the object key within the store is then the in-repository path. +#[derive(Debug, Clone)] +pub struct HfConfig { + /// Which kind of repository to address. + pub repo_type: HfRepoType, + /// The repository, as `/`. + pub repo_id: String, + /// The revision to read: a branch, tag or commit. Held decoded, so `refs/convert/parquet` is + /// spelled with `/`; it is percent-encoded when the `resolve` URL is built. + pub revision: String, + /// Bearer token for private and gated repositories. `None` reads anonymously. + pub token: Option, + /// The Hub endpoint, e.g. `https://huggingface.co`. + pub endpoint: String, + /// HTTP client configuration, applied to every read. + pub client_options: ClientOptions, +} + +impl Default for HfConfig { + fn default() -> Self { + Self { + repo_type: HfRepoType::Dataset, + repo_id: String::new(), + revision: DEFAULT_REVISION.to_string(), + token: None, + endpoint: DEFAULT_ENDPOINT.to_string(), + client_options: ClientOptions::default(), + } + } +} + +impl HfConfig { + /// Configuration for one repository, taking the token and endpoint from the process environment + /// exactly as resolving an `hf://` URL does. + /// + /// `revision` defaults to `main`. Callers that must read anonymously should clear + /// [`HfConfig::token`] afterwards, since this picks up whatever credentials the environment + /// offers. + pub fn from_env( + repo_type: HfRepoType, + repo_id: impl Into, + revision: Option, + ) -> Self { + Self::from_env_lookup(repo_type, repo_id, revision, |key| std::env::var(key).ok()) + } + + /// [`HfConfig::from_env`], reading environment configuration through `env_lookup`. + fn from_env_lookup( + repo_type: HfRepoType, + repo_id: impl Into, + revision: Option, + env_lookup: F, + ) -> Self + where + F: Fn(&str) -> Option, + { + Self { + repo_type, + repo_id: repo_id.into(), + revision: revision.unwrap_or_else(|| DEFAULT_REVISION.to_string()), + token: resolve_token(&env_lookup), + endpoint: env_lookup(ENDPOINT_VAR) + .filter(|endpoint| !endpoint.is_empty()) + .unwrap_or_else(|| DEFAULT_ENDPOINT.to_string()), + // Every other scheme picks `allow_http` out of the environment, because that is how + // `parse_url_opts` reads its configuration. Honour it here too, so a plain-HTTP + // `HF_ENDPOINT` (a test double, or a self-hosted Hub) behaves the same way. + client_options: ClientOptions::default().with_allow_http(allow_http(&env_lookup)), + } + } + + /// The URL this repository's files hang off, which is where the store is rooted. + fn resolve_prefix(&self) -> String { + let endpoint = self.endpoint.trim_end_matches('/'); + let prefix = self.repo_type.url_prefix(); + let revision = utf8_percent_encode(&self.revision, REVISION_ESCAPES); + format!("{endpoint}/{prefix}{}/resolve/{revision}", self.repo_id) + } +} + +/// Build an [`ObjectStore`] for a Hugging Face Hub repository from an [`HfConfig`]. +/// +/// This is the entry point for callers holding a strongly-typed configuration. The returned store is +/// rooted at the repository's `resolve` prefix, so its object keys are in-repository paths. +pub fn make_hf_store(config: HfConfig) -> Result, HfStoreError> { + if config.repo_id.is_empty() { + return Err(HfStoreError::InvalidUrl("".to_string())); + } + + let mut client_options = config.client_options.clone(); + if let Some(token) = config.token.as_deref() { + let mut headers = HeaderMap::new(); + let value = HeaderValue::from_str(&format!("Bearer {token}")) + .map_err(HfStoreError::InvalidToken)?; + // Sending the token as a default header, rather than per request, is what lets a plain + // `HttpStore` serve gated and private repositories. + headers.insert(AUTHORIZATION, value); + client_options = client_options.with_default_headers(headers); + } + + let store = HttpBuilder::new() + .with_url(config.resolve_prefix()) + .with_client_options(client_options) + .build() + .map_err(HfStoreError::Build)?; + + Ok(Arc::new(store)) +} + +/// Build an [`ObjectStore`] for an `hf://` URL, with the in-repository path of the addressed file. +/// +/// Configuration not carried by the URL is read from the process environment. The path is reported +/// the way [`object_store::parse_url_opts`] reports it — the key the returned store will see — so +/// that the caller can derive how deep into the URL the store is mounted. +pub fn make_hf_store_from_url(url: &Url) -> Result<(Arc, Path), HfStoreError> { + make_hf_store_from_url_with_env(url, |key| std::env::var(key).ok()) +} + +/// [`make_hf_store_from_url`], reading environment configuration through `env_lookup`. +/// +/// Tests pass a closure over a fixed map so they do not race against the process environment. +pub(crate) fn make_hf_store_from_url_with_env( + url: &Url, + env_lookup: F, +) -> Result<(Arc, Path), HfStoreError> +where + F: Fn(&str) -> Option, +{ + let (config, path) = url_to_config(url, env_lookup)?; + Ok((make_hf_store(config)?, path)) +} + +/// Translate an `hf://` URL into an [`HfConfig`] plus the in-repository path it addresses. +fn url_to_config(url: &Url, env_lookup: F) -> Result<(HfConfig, Path), HfStoreError> +where + F: Fn(&str) -> Option, +{ + if !supports_scheme(url.scheme()) { + return Err(HfStoreError::UnsupportedScheme(url.scheme().to_string())); + } + let invalid = || HfStoreError::InvalidUrl(url.to_string()); + + // `hf://datasets/org/name/...` parses with `datasets` as the authority and `/org/name/...` as + // the path, so the repository kind is the authority and a model repository's owner is too. + let host = url.host_str().ok_or_else(invalid)?; + let segments: Vec<&str> = url.path().split('/').filter(|s| !s.is_empty()).collect(); + + let (repo_type, owner, rest) = match host { + DATASETS_HOST | SPACES_HOST => { + let repo_type = if host == DATASETS_HOST { + HfRepoType::Dataset + } else { + HfRepoType::Space + }; + let (owner, rest) = segments.split_first().ok_or_else(invalid)?; + (repo_type, *owner, rest) + } + owner => (HfRepoType::Model, owner, segments.as_slice()), + }; + + let (name, rest) = rest.split_first().ok_or_else(invalid)?; + if owner.is_empty() || name.is_empty() { + return Err(invalid()); + } + + // A revision is appended to the repository name with `@`, and arrives percent-encoded when it + // contains `/`. Hold it decoded; `resolve_prefix` re-encodes it on the way back out. + let (name, revision) = match name.split_once('@') { + Some((name, revision)) => { + if name.is_empty() || revision.is_empty() { + return Err(invalid()); + } + let revision = percent_decode_str(revision) + .decode_utf8() + .map_err(|_| invalid())? + .into_owned(); + (name, revision) + } + None => (*name, DEFAULT_REVISION.to_string()), + }; + + let config = HfConfig::from_env_lookup( + repo_type, + format!("{owner}/{name}"), + Some(revision), + &env_lookup, + ); + + // The segments still carry their URL escapes, so decode them into the object key rather than + // joining them raw. + let path = Path::from_url_path(rest.join("/")).map_err(|_| invalid())?; + + Ok((config, path)) +} + +/// Whether plain-HTTP endpoints are permitted, from the `ALLOW_HTTP` variable the `object_store` +/// builders read. Anything other than a `true`-ish value leaves HTTPS as the requirement. +fn allow_http(env_lookup: &F) -> bool +where + F: Fn(&str) -> Option, +{ + env_lookup(ALLOW_HTTP_VAR) + .is_some_and(|value| matches!(value.trim().to_ascii_lowercase().as_str(), "true" | "1")) +} + +/// The bearer token to read with, following the same precedence as `huggingface_hub.get_token()`: +/// `HF_TOKEN`, then the token file at `HF_TOKEN_PATH`, `$HF_HOME/token` or +/// `$HOME/.cache/huggingface/token`. +fn resolve_token(env_lookup: &F) -> Option +where + F: Fn(&str) -> Option, +{ + if let Some(token) = env_lookup(TOKEN_VAR).filter(|token| !token.is_empty()) { + return Some(token); + } + + let token_path = env_lookup(TOKEN_PATH_VAR) + .or_else(|| env_lookup(HF_HOME_VAR).map(|home| format!("{home}/token"))) + .or_else(|| env_lookup(HOME_VAR).map(|home| format!("{home}/.cache/huggingface/token")))?; + + // A missing token file is the ordinary anonymous case, not an error. + let token = std::fs::read_to_string(token_path).ok()?; + let token = token.trim(); + (!token.is_empty()).then(|| token.to_string()) +} + +#[cfg(test)] +mod tests; diff --git a/vortex-cloud/src/hf/tests.rs b/vortex-cloud/src/hf/tests.rs new file mode 100644 index 00000000000..04b8fe40cb1 --- /dev/null +++ b/vortex-cloud/src/hf/tests.rs @@ -0,0 +1,364 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use url::Url; + +use super::DEFAULT_ENDPOINT; +use super::HfConfig; +use super::HfRepoType; +use super::HfStoreError; +use super::url_to_config; + +/// Configuration lookups over a fixed map, so these tests neither read nor mutate the process +/// environment. +fn env(vars: &[(&str, &str)]) -> impl Fn(&str) -> Option + use<> { + let vars: Vec<(String, String)> = vars + .iter() + .map(|(k, v)| ((*k).to_string(), (*v).to_string())) + .collect(); + move |key| { + vars.iter() + .find(|(k, _)| k == key) + .map(|(_, v)| v.to_string()) + } +} + +/// Parse a URL with no environment configuration at all, so no token is found and the default +/// endpoint applies. +fn parse(url: &str) -> Result<(HfConfig, String), Box> { + let url = Url::parse(url)?; + let (config, path) = url_to_config(&url, env(&[]))?; + Ok((config, path.as_ref().to_string())) +} + +#[rstest] +// A dataset repository, with and without an explicit revision. +#[case("hf://datasets/org/name", HfRepoType::Dataset, "org/name", "main", "")] +#[case( + "hf://datasets/org/name@main", + HfRepoType::Dataset, + "org/name", + "main", + "" +)] +#[case( + "hf://datasets/org/name/train.vortex", + HfRepoType::Dataset, + "org/name", + "main", + "train.vortex" +)] +#[case( + "hf://datasets/org/name/data/nested/train.vortex", + HfRepoType::Dataset, + "org/name", + "main", + "data/nested/train.vortex" +)] +#[case( + "hf://datasets/org/name@v1.0/train.vortex", + HfRepoType::Dataset, + "org/name", + "v1.0", + "train.vortex" +)] +// A revision containing `/` arrives percent-encoded and must be held decoded. +#[case( + "hf://datasets/org/name@refs%2Fconvert%2Fparquet/data/train.vortex", + HfRepoType::Dataset, + "org/name", + "refs/convert/parquet", + "data/train.vortex" +)] +// A bare owner/name is a model repository: the owner is the URL authority. +#[case( + "hf://org/name/model.vortex", + HfRepoType::Model, + "org/name", + "main", + "model.vortex" +)] +#[case( + "hf://org/name@abc123/model.vortex", + HfRepoType::Model, + "org/name", + "abc123", + "model.vortex" +)] +#[case( + "hf://spaces/org/name/app.vortex", + HfRepoType::Space, + "org/name", + "main", + "app.vortex" +)] +fn test_url_to_config( + #[case] url: &str, + #[case] repo_type: HfRepoType, + #[case] repo_id: &str, + #[case] revision: &str, + #[case] path: &str, +) -> Result<(), Box> { + let (config, parsed_path) = parse(url)?; + + assert_eq!(config.repo_type, repo_type); + assert_eq!(config.repo_id, repo_id); + assert_eq!(config.revision, revision); + assert_eq!(parsed_path, path); + Ok(()) +} + +#[rstest] +// A repository needs both an owner and a name. +#[case("hf://datasets/name-only")] +#[case("hf://org")] +// An `@` must have a name on the left and a revision on the right. +#[case("hf://datasets/org/@main")] +#[case("hf://datasets/org/name@")] +fn test_url_to_config_invalid(#[case] url: &str) -> Result<(), Box> { + let url = Url::parse(url)?; + assert!(matches!( + url_to_config(&url, env(&[])), + Err(HfStoreError::InvalidUrl(_)) + )); + Ok(()) +} + +#[test] +fn test_url_to_config_rejects_other_schemes() -> Result<(), Box> { + let url = Url::parse("s3://bucket/key.vortex")?; + assert!(matches!( + url_to_config(&url, env(&[])), + Err(HfStoreError::UnsupportedScheme(_)) + )); + Ok(()) +} + +#[rstest] +// Each repository kind hangs off its own prefix, and the store is rooted at the revision. +#[case( + HfRepoType::Dataset, + "main", + "https://huggingface.co/datasets/org/name/resolve/main" +)] +#[case( + HfRepoType::Model, + "main", + "https://huggingface.co/org/name/resolve/main" +)] +#[case( + HfRepoType::Space, + "main", + "https://huggingface.co/spaces/org/name/resolve/main" +)] +// The Hub only routes the escaped form of a revision containing `/`. +#[case( + HfRepoType::Dataset, + "refs/convert/parquet", + "https://huggingface.co/datasets/org/name/resolve/refs%2Fconvert%2Fparquet" +)] +fn test_resolve_prefix( + #[case] repo_type: HfRepoType, + #[case] revision: &str, + #[case] expected: &str, +) { + let config = HfConfig { + repo_type, + repo_id: "org/name".to_string(), + revision: revision.to_string(), + ..HfConfig::default() + }; + + assert_eq!(config.resolve_prefix(), expected); +} + +/// A percent-encoded revision must survive the URL -> config -> URL round trip, since that is how a +/// `refs/convert/parquet` read reaches the Hub. +#[test] +fn test_slash_revision_round_trips() -> Result<(), Box> { + let (config, path) = + parse("hf://datasets/org/name@refs%2Fconvert%2Fparquet/data/train.vortex")?; + + assert_eq!(config.revision, "refs/convert/parquet"); + assert_eq!( + config.resolve_prefix(), + "https://huggingface.co/datasets/org/name/resolve/refs%2Fconvert%2Fparquet" + ); + assert_eq!(path, "data/train.vortex"); + Ok(()) +} + +#[test] +fn test_endpoint_override() -> Result<(), Box> { + let url = Url::parse("hf://datasets/org/name/train.vortex")?; + let (config, _path) = url_to_config(&url, env(&[("HF_ENDPOINT", "https://hub.example.com/")]))?; + + assert_eq!(config.endpoint, "https://hub.example.com/"); + // The trailing slash on the endpoint must not double up in the resolve prefix. + assert_eq!( + config.resolve_prefix(), + "https://hub.example.com/datasets/org/name/resolve/main" + ); + Ok(()) +} + +#[test] +fn test_default_endpoint_when_unset_or_empty() -> Result<(), Box> { + let url = Url::parse("hf://datasets/org/name")?; + + let (unset, _) = url_to_config(&url, env(&[]))?; + assert_eq!(unset.endpoint, DEFAULT_ENDPOINT); + + let (empty, _) = url_to_config(&url, env(&[("HF_ENDPOINT", "")]))?; + assert_eq!(empty.endpoint, DEFAULT_ENDPOINT); + Ok(()) +} + +#[test] +fn test_token_from_env() -> Result<(), Box> { + let url = Url::parse("hf://datasets/org/name")?; + let (config, _path) = url_to_config(&url, env(&[("HF_TOKEN", "hf_from_env")]))?; + + assert_eq!(config.token.as_deref(), Some("hf_from_env")); + Ok(()) +} + +#[test] +fn test_token_from_token_file() -> Result<(), Box> { + let dir = tempfile::tempdir()?; + let token_file = dir.path().join("token"); + // The saved login file is newline-terminated, which must not reach the header value. + std::fs::write(&token_file, "hf_from_file\n")?; + + let url = Url::parse("hf://datasets/org/name")?; + let (config, _path) = url_to_config( + &url, + env(&[("HF_TOKEN_PATH", &token_file.to_string_lossy())]), + )?; + + assert_eq!(config.token.as_deref(), Some("hf_from_file")); + Ok(()) +} + +#[test] +fn test_token_env_wins_over_file() -> Result<(), Box> { + let dir = tempfile::tempdir()?; + let token_file = dir.path().join("token"); + std::fs::write(&token_file, "hf_from_file")?; + + let url = Url::parse("hf://datasets/org/name")?; + let (config, _path) = url_to_config( + &url, + env(&[ + ("HF_TOKEN", "hf_from_env"), + ("HF_TOKEN_PATH", &token_file.to_string_lossy()), + ]), + )?; + + assert_eq!(config.token.as_deref(), Some("hf_from_env")); + Ok(()) +} + +/// No token anywhere is the ordinary anonymous case, not an error. +#[test] +fn test_no_token_reads_anonymously() -> Result<(), Box> { + let url = Url::parse("hf://datasets/org/name")?; + let (config, _path) = url_to_config(&url, env(&[("HF_TOKEN_PATH", "/nonexistent/token")]))?; + + assert!(config.token.is_none()); + Ok(()) +} + +/// Building a store must succeed for both the anonymous and the credentialed path; a token that +/// cannot become a header value is reported rather than panicking. +#[test] +fn test_make_hf_store() -> Result<(), Box> { + let base = HfConfig { + repo_id: "org/name".to_string(), + ..HfConfig::default() + }; + + let _anonymous = super::make_hf_store(base.clone())?; + let _credentialed = super::make_hf_store(HfConfig { + token: Some("hf_token".to_string()), + ..base.clone() + })?; + + assert!(matches!( + super::make_hf_store(HfConfig { + token: Some("bad\nvalue".to_string()), + ..base + }), + Err(HfStoreError::InvalidToken(_)) + )); + assert!(matches!( + super::make_hf_store(HfConfig::default()), + Err(HfStoreError::InvalidUrl(_)) + )); + Ok(()) +} + +#[rstest] +// Both the singular name and the plural the Hub spells in its URLs are accepted, so a caller can +// pass whichever it has. +#[case("dataset", HfRepoType::Dataset)] +#[case("datasets", HfRepoType::Dataset)] +#[case("model", HfRepoType::Model)] +#[case("models", HfRepoType::Model)] +#[case("space", HfRepoType::Space)] +#[case("spaces", HfRepoType::Space)] +fn test_repo_type_from_str( + #[case] spelling: &str, + #[case] expected: HfRepoType, +) -> Result<(), Box> { + assert_eq!(spelling.parse::()?, expected); + Ok(()) +} + +#[test] +fn test_repo_type_from_str_rejects_unknown() { + assert!(matches!( + "notarepo".parse::(), + Err(HfStoreError::UnknownRepoType(_)) + )); +} + +/// `from_env` must resolve the token and endpoint the same way an `hf://` URL does, since the store +/// class and the URL registry are two doors onto one configuration. +#[test] +fn test_from_env_matches_url_resolution() -> Result<(), Box> { + let vars = &[ + ("HF_TOKEN", "hf_shared"), + ("HF_ENDPOINT", "https://hub.example.com"), + ]; + let url = Url::parse("hf://datasets/org/name/train.vortex")?; + + let (from_url, _path) = url_to_config(&url, env(vars))?; + let direct = HfConfig::from_env_lookup(HfRepoType::Dataset, "org/name", None, env(vars)); + + assert_eq!(from_url.token, direct.token); + assert_eq!(from_url.endpoint, direct.endpoint); + assert_eq!(from_url.resolve_prefix(), direct.resolve_prefix()); + Ok(()) +} + +/// `ALLOW_HTTP` must reach the client options, since that is how every other scheme is told plain +/// HTTP is acceptable and a self-hosted or stubbed Hub endpoint needs the same door. `ClientOptions` +/// exposes no getter, so the observable assertion is an actual plain-HTTP read; that lives in +/// `vortex-python/test/test_hf_datasets.py::test_hf_store_reads_a_repository_relative_path`. +#[test] +fn test_allow_http_is_configured_from_env() { + let allowed = HfConfig::from_env_lookup( + HfRepoType::Dataset, + "org/name", + None, + env(&[("ALLOW_HTTP", "true")]), + ); + let denied = HfConfig::from_env_lookup(HfRepoType::Dataset, "org/name", None, env(&[])); + + assert_ne!( + format!("{:?}", allowed.client_options), + format!("{:?}", denied.client_options) + ); +} diff --git a/vortex-cloud/src/lib.rs b/vortex-cloud/src/lib.rs new file mode 100644 index 00000000000..532ced1c81b --- /dev/null +++ b/vortex-cloud/src/lib.rs @@ -0,0 +1,45 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Cloud object store integration for Vortex. +//! +//! This crate owns everything Vortex needs to turn a URL into a credentialed +//! [`object_store::ObjectStore`]: +//! +//! * `Registry` resolves a URL to a store, caching one client per bucket. It reads configuration +//! out of environment variables case-insensitively, matching how the `object_store` `from_env` +//! builders behave. +//! * `opendal` supplies stores for cloud services the `object_store` crate does not implement +//! natively — Tencent Cloud COS, Alibaba Cloud OSS, and Tencent Cloud GooseFS — bridged +//! through `object_store_opendal`. +//! * `hf` serves Hugging Face Hub repositories over `object_store`'s HTTP store, adding no cloud +//! SDK of its own. +//! +//! Every Vortex language binding resolves URLs through this one crate, so a scheme added here is +//! reachable from Python, Java and DuckDB alike. +//! +//! Both are feature-gated, so this overview refers to them by name rather than by link: an +//! intra-doc link would dangle in a build that leaves the corresponding feature off. +//! +//! # Cargo features +//! +//! * `registry` — the `Registry`, plus the natively-supported cloud backends (S3, Azure, GCS, +//! HTTP) it resolves URLs to. +//! * `hf` — the Hugging Face Hub, the `hf://` scheme. +//! * `cos` — Tencent Cloud COS, the `cos://` scheme. +//! * `oss` — Alibaba Cloud OSS, the `oss://` scheme. +//! * `goosefs` — Tencent Cloud GooseFS, the `goosefs://` scheme. +//! * `opendal` — every OpenDAL-backed service above. +//! +//! The `registry` feature picks up whichever services are enabled, so a consumer that turns on +//! `oss` gets `oss://` resolution without touching its own scheme matching. + +#[cfg(feature = "hf")] +pub mod hf; +#[cfg(any(feature = "cos", feature = "goosefs", feature = "oss"))] +pub mod opendal; +#[cfg(feature = "registry")] +mod registry; + +#[cfg(feature = "registry")] +pub use registry::Registry; diff --git a/vortex-cloud/src/opendal/cos.rs b/vortex-cloud/src/opendal/cos.rs new file mode 100644 index 00000000000..b1a4f07f0a9 --- /dev/null +++ b/vortex-cloud/src/opendal/cos.rs @@ -0,0 +1,231 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Tencent Cloud COS, served over OpenDAL's `services::Cos`. + +use std::sync::Arc; + +use ::opendal::services; +use object_store::ObjectStore; +use object_store_opendal::OpendalStore; +use url::Url; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::opendal::OpenDALStoreError; +use crate::opendal::build_operator; +use crate::opendal::property_as_bool; +use crate::opendal::property_or_env; +use crate::opendal::warn_on_unknown_properties; + +/// The URL scheme served by Tencent Cloud COS. +pub const COS_SCHEME: &str = "cos"; + +/// Property keys recognized for `cos://` URLs. Anything else is warned about and dropped. +const KNOWN_PROPERTIES: &[&str] = &[ + "bucket", + "endpoint", + "secret_id", + "secret_key", + "root", + "disable_config_load", +]; + +/// Strongly-typed configuration for building an OpenDAL store against Tencent Cloud COS. +/// +/// The fields mirror the keyword arguments of the `CosStore` Python class. Building from a +/// `CosConfig` avoids the URL-round-trip that the [`crate::opendal::make_opendal_store`] entry point uses, +/// and is the preferred way to construct a COS store. +#[derive(Debug, Clone, Default)] +pub struct CosConfig { + /// COS bucket name. May be overridden by `properties["bucket"]` when adapting a URL. + pub bucket: String, + /// COS endpoint, e.g. `https://cos.ap-guangzhou.myqcloud.com`. + pub endpoint: String, + /// Tencent Cloud secret id (mapped to `TENCENTCLOUD_SECRET_ID`). + pub secret_id: Option, + /// Tencent Cloud secret key (mapped to `TENCENTCLOUD_SECRET_KEY`). + pub secret_key: Option, + /// Optional root prefix applied to all operations. + pub root: Option, + /// When `true`, disable OpenDAL's automatic config loading (so only the explicit + /// configuration is used). + pub disable_config_load: bool, +} + +/// Build an [`object_store::ObjectStore`] for Tencent Cloud COS directly from a [`CosConfig`]. +/// +/// This is the preferred entry point for callers that have a strongly-typed configuration object +/// (such as the `CosStore` pyclass in `vortex-python`). It does not synthesize a URL and so is not +/// fragile against reordering of `bucket` vs URL host precedence. +pub fn make_cos_store(config: CosConfig) -> Result, OpenDALStoreError> { + if config.bucket.is_empty() { + return Err(OpenDALStoreError::MissingConfig("bucket")); + } + if config.endpoint.is_empty() { + return Err(OpenDALStoreError::MissingConfig("endpoint")); + } + + let mut builder = services::Cos::default() + .bucket(&config.bucket) + .endpoint(&config.endpoint); + + if let Some(root) = config.root.as_deref() { + builder = builder.root(root); + } + if let Some(secret_id) = config.secret_id.as_deref() { + builder = builder.secret_id(secret_id); + } + if let Some(secret_key) = config.secret_key.as_deref() { + builder = builder.secret_key(secret_key); + } + if config.disable_config_load { + builder = builder.disable_config_load(); + } + + let operator = build_operator(builder)?; + Ok(Arc::new(OpendalStore::new(operator))) +} + +/// Translate a (`cos://` URL, properties) pair into a strongly-typed [`CosConfig`]. +/// +/// Bucket is taken from `properties["bucket"]` first and falls back to the URL host; endpoint is +/// taken from `properties["endpoint"]` first and falls back to `COS_ENDPOINT`; credentials fall +/// back to the variables OpenDAL's COS builder reads. +/// +/// `env_lookup` is the source of truth for environment-variable fallbacks. The production entry +/// point passes the real environment; tests pass a closure that returns from a fixed map, so they +/// do not race against the process environment. +pub(crate) fn url_and_properties_to_config( + url: &Url, + properties: &HashMap, + env_lookup: F, +) -> Result +where + F: Fn(&str) -> Option, +{ + warn_on_unknown_properties(properties, KNOWN_PROPERTIES); + + let bucket = properties + .get("bucket") + .cloned() + .or_else(|| url.host_str().map(str::to_string)) + .ok_or(OpenDALStoreError::MissingConfig("bucket"))?; + + let endpoint = property_or_env(properties, "endpoint", "COS_ENDPOINT", &env_lookup) + .ok_or(OpenDALStoreError::MissingConfig("endpoint"))?; + + Ok(CosConfig { + bucket, + endpoint, + secret_id: property_or_env( + properties, + "secret_id", + "TENCENTCLOUD_SECRET_ID", + &env_lookup, + ), + secret_key: property_or_env( + properties, + "secret_key", + "TENCENTCLOUD_SECRET_KEY", + &env_lookup, + ), + root: properties.get("root").cloned(), + disable_config_load: property_as_bool(properties, "disable_config_load"), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The endpoint is required. With a fixed env-lookup that returns `None`, the call must + /// fail with `MissingConfig("endpoint")` regardless of what the process environment + /// happens to contain. This makes the test deterministic under `cargo test`'s default + /// multi-threaded runner (and avoids the `unsafe std::env::set_var` that became unsound in + /// Rust 2024). + #[test] + fn cos_requires_endpoint() { + let url = Url::parse("cos://my-bucket/path").unwrap(); + let props = HashMap::new(); + let result = url_and_properties_to_config(&url, &props, |_| None).unwrap_err(); + assert!(matches!( + result, + OpenDALStoreError::MissingConfig("endpoint") + )); + } + + /// When `properties` does not contain `endpoint`, the env-lookup should be consulted. + #[test] + fn cos_falls_back_to_cos_endpoint_env() { + let url = Url::parse("cos://my-bucket/path").unwrap(); + let env = |key: &str| match key { + "COS_ENDPOINT" => Some("https://example.com".to_string()), + _ => None, + }; + let props = HashMap::new(); + let config = url_and_properties_to_config(&url, &props, env).expect("config"); + assert_eq!(config.endpoint, "https://example.com"); + } + + /// An explicit property must win over the environment fallback. + #[test] + fn cos_property_overrides_env() { + let url = Url::parse("cos://url-bucket/path").unwrap(); + let env = |key: &str| match key { + "COS_ENDPOINT" => Some("https://from-env.example.com".to_string()), + _ => None, + }; + let mut props = HashMap::new(); + props.insert("bucket".to_string(), "prop-bucket".to_string()); + props.insert( + "endpoint".to_string(), + "https://from-prop.example.com".to_string(), + ); + + let config = url_and_properties_to_config(&url, &props, env).expect("config"); + assert_eq!(config.bucket, "prop-bucket"); + assert_eq!(config.endpoint, "https://from-prop.example.com"); + } + + /// With a fixed env-lookup that returns explicit credentials, the strongly-typed + /// `make_cos_store` should build a store successfully. + #[test] + fn cos_builds_with_explicit_config() { + let url = Url::parse("cos://my-bucket/path/to/dataset.vortex").unwrap(); + let env = |key: &str| match key { + "COS_ENDPOINT" => Some("https://example.com".to_string()), + "TENCENTCLOUD_SECRET_ID" => Some("AKID".to_string()), + "TENCENTCLOUD_SECRET_KEY" => Some("secret".to_string()), + _ => None, + }; + let mut props = HashMap::new(); + props.insert("disable_config_load".to_string(), "true".to_string()); + + let config = url_and_properties_to_config(&url, &props, env).expect("config"); + assert_eq!(config.secret_id.as_deref(), Some("AKID")); + let store = make_cos_store(config).expect("store should build"); + // Sanity: the returned store is a non-null `Arc`. + assert!(Arc::strong_count(&store) >= 1); + } + + /// The strongly-typed `make_cos_store` entry point must reject an empty bucket or endpoint + /// with a `MissingConfig` error before consulting any environment or builder. + #[test] + fn cos_config_rejects_empty_fields() { + assert!(matches!( + make_cos_store(CosConfig { + bucket: String::new(), + ..CosConfig::default() + }), + Err(OpenDALStoreError::MissingConfig("bucket")) + )); + assert!(matches!( + make_cos_store(CosConfig { + bucket: "b".to_string(), + endpoint: String::new(), + ..CosConfig::default() + }), + Err(OpenDALStoreError::MissingConfig("endpoint")) + )); + } +} diff --git a/vortex-cloud/src/opendal/goosefs.rs b/vortex-cloud/src/opendal/goosefs.rs new file mode 100644 index 00000000000..2bec6cc0d37 --- /dev/null +++ b/vortex-cloud/src/opendal/goosefs.rs @@ -0,0 +1,266 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Tencent Cloud GooseFS, served over OpenDAL's `services::GooseFs`. +//! +//! GooseFS is a distributed caching file system accessed via native gRPC. Unlike COS/OSS it has +//! no bucket or HTTP endpoint — instead the client connects to a GooseFS master at +//! `host:port`. An HA cluster is expressed as a comma-separated list of master addresses. + +use std::sync::Arc; + +use ::opendal::services; +use object_store::ObjectStore; +use object_store_opendal::OpendalStore; +use url::Url; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::opendal::OpenDALStoreError; +use crate::opendal::build_operator; +use crate::opendal::property_or_env; +use crate::opendal::warn_on_unknown_properties; + +/// The URL scheme served by Tencent Cloud GooseFS. +pub const GOOSEFS_SCHEME: &str = "goosefs"; + +/// Property keys recognized for `goosefs://` URLs. Anything else is warned about and dropped. +const KNOWN_PROPERTIES: &[&str] = &[ + "master_addr", + "root", + "block_size", + "chunk_size", + "write_type", + "auth_type", + "auth_username", +]; + +/// Strongly-typed configuration for building an OpenDAL store against Tencent Cloud GooseFS. +/// +/// The fields mirror the keyword arguments of the `GoosefsStore` Python class. Building from a +/// [`GoosefsConfig`] avoids the URL-round-trip that the [`crate::opendal::make_opendal_store`] +/// entry point uses, and is the preferred way to construct a GooseFS store. +#[derive(Debug, Clone, Default)] +pub struct GoosefsConfig { + /// GooseFS master address(es). + /// + /// Single master: `"10.0.0.1:9200"` + /// HA (comma-separated): `"10.0.0.1:9200,10.0.0.2:9200,10.0.0.3:9200"` + /// + /// May be overridden by `properties["master_addr"]` when adapting a URL. Falls back to the + /// `GOOSEFS_MASTER_ADDR` environment variable. + pub master_addr: String, + /// Optional root prefix applied to all operations. + pub root: Option, + /// Block size in bytes for new files (default: 64 MiB). + pub block_size: Option, + /// Chunk size in bytes for streaming RPCs (default: 1 MiB). + pub chunk_size: Option, + /// Default write type: `"must_cache"`, `"cache_through"`, `"through"`, `"async_through"`. + pub write_type: Option, + /// Authentication type: `"nosasl"` or `"simple"` (default: `"simple"`). + pub auth_type: Option, + /// Authentication username (default: current OS user). + pub auth_username: Option, +} + +/// Build an [`object_store::ObjectStore`] for Tencent Cloud GooseFS directly from a +/// [`GoosefsConfig`]. +/// +/// This is the preferred entry point for callers that have a strongly-typed configuration object +/// (such as the `GoosefsStore` pyclass in `vortex-python`). It does not synthesize a URL and so is +/// not fragile against reordering of `master_addr` vs URL authority precedence. +pub fn make_goosefs_store( + config: GoosefsConfig, +) -> Result, OpenDALStoreError> { + if config.master_addr.is_empty() { + return Err(OpenDALStoreError::MissingConfig("master_addr")); + } + + let mut builder = services::GooseFs::default().master_addr(&config.master_addr); + + if let Some(root) = config.root.as_deref() { + builder = builder.root(root); + } + if let Some(block_size) = config.block_size { + builder = builder.block_size(block_size); + } + if let Some(chunk_size) = config.chunk_size { + builder = builder.chunk_size(chunk_size); + } + if let Some(write_type) = config.write_type.as_deref() { + builder = builder.write_type(write_type); + } + if let Some(auth_type) = config.auth_type.as_deref() { + builder = builder.auth_type(auth_type); + } + if let Some(auth_username) = config.auth_username.as_deref() { + builder = builder.auth_username(auth_username); + } + + let operator = build_operator(builder)?; + Ok(Arc::new(OpendalStore::new(operator))) +} + +/// Translate a (`goosefs://` URL, properties) pair into a strongly-typed [`GoosefsConfig`]. +/// +/// `master_addr` is taken from `properties["master_addr"]` first and falls back to the URL +/// authority (`host:port`); if neither is present it falls back to the `GOOSEFS_MASTER_ADDR` +/// environment variable. The remaining fields are taken from `properties` directly. +/// +/// `env_lookup` is the source of truth for environment-variable fallbacks. The production entry +/// point passes the real environment; tests pass a closure that returns from a fixed map, so they +/// do not race against the process environment. +pub(crate) fn url_and_properties_to_config( + url: &Url, + properties: &HashMap, + env_lookup: F, +) -> Result +where + F: Fn(&str) -> Option, +{ + warn_on_unknown_properties(properties, KNOWN_PROPERTIES); + + // master_addr: properties → URL authority → GOOSEFS_MASTER_ADDR env + let master_addr = properties + .get("master_addr") + .cloned() + .or_else(|| { + let host = url.host_str(); + let port = url.port(); + match (host, port) { + (Some(h), Some(p)) => Some(format!("{h}:{p}")), + (Some(h), None) => Some(h.to_string()), + _ => None, + } + }) + .or_else(|| { + property_or_env( + properties, + "master_addr", + "GOOSEFS_MASTER_ADDR", + &env_lookup, + ) + }) + .ok_or(OpenDALStoreError::MissingConfig("master_addr"))?; + + Ok(GoosefsConfig { + master_addr, + root: properties.get("root").cloned(), + block_size: properties.get("block_size").and_then(|s| s.parse().ok()), + chunk_size: properties.get("chunk_size").and_then(|s| s.parse().ok()), + write_type: properties.get("write_type").cloned(), + auth_type: properties.get("auth_type").cloned(), + auth_username: properties.get("auth_username").cloned(), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// `master_addr` is required. With a fixed env-lookup that returns `None` and a URL without + /// authority, the call must fail with `MissingConfig("master_addr")` regardless of what the + /// process environment happens to contain. + #[test] + fn goosefs_requires_master_addr() { + let url = Url::parse("goosefs:///path/to/file").unwrap(); + let props = HashMap::new(); + let result = url_and_properties_to_config(&url, &props, |_| None).unwrap_err(); + assert!(matches!( + result, + OpenDALStoreError::MissingConfig("master_addr") + )); + } + + /// The URL authority (`host:port`) should be used as `master_addr` when properties do not + /// contain it. + #[test] + fn goosefs_uses_url_authority_as_master_addr() { + let url = Url::parse("goosefs://10.0.0.1:9200/path/to/file").unwrap(); + let props = HashMap::new(); + let config = url_and_properties_to_config(&url, &props, |_| None).expect("config"); + assert_eq!(config.master_addr, "10.0.0.1:9200"); + } + + /// When `properties` does not contain `master_addr` and the URL has no authority, the + /// env-lookup should be consulted. + #[test] + fn goosefs_falls_back_to_env() { + let url = Url::parse("goosefs:///path/to/file").unwrap(); + let env = |key: &str| match key { + "GOOSEFS_MASTER_ADDR" => Some("10.0.0.1:9200".to_string()), + _ => None, + }; + let props = HashMap::new(); + let config = url_and_properties_to_config(&url, &props, env).expect("config"); + assert_eq!(config.master_addr, "10.0.0.1:9200"); + } + + /// An explicit property must win over the URL authority and the environment fallback. + #[test] + fn goosefs_property_overrides_url_and_env() { + let url = Url::parse("goosefs://url-host:9200/path").unwrap(); + let env = |key: &str| match key { + "GOOSEFS_MASTER_ADDR" => Some("from-env:9200".to_string()), + _ => None, + }; + let mut props = HashMap::new(); + props.insert("master_addr".to_string(), "from-prop:9200".to_string()); + + let config = url_and_properties_to_config(&url, &props, env).expect("config"); + assert_eq!(config.master_addr, "from-prop:9200"); + } + + /// The URL authority must win over the environment fallback when no explicit property is set. + #[test] + fn goosefs_url_authority_overrides_env() { + let url = Url::parse("goosefs://url-host:9200/path").unwrap(); + let env = |key: &str| match key { + "GOOSEFS_MASTER_ADDR" => Some("from-env:9200".to_string()), + _ => None, + }; + let props = HashMap::new(); + let config = url_and_properties_to_config(&url, &props, env).expect("config"); + assert_eq!(config.master_addr, "url-host:9200"); + } + + /// With a fixed env-lookup that returns an explicit master address, the strongly-typed + /// `make_goosefs_store` should build a store successfully. + #[test] + fn goosefs_builds_with_explicit_config() { + let config = GoosefsConfig { + master_addr: "127.0.0.1:9200".to_string(), + root: Some("/data".to_string()), + ..GoosefsConfig::default() + }; + let store = make_goosefs_store(config).expect("store should build"); + // Sanity: the returned store is a non-null `Arc`. + assert!(Arc::strong_count(&store) >= 1); + } + + /// The strongly-typed `make_goosefs_store` entry point must reject an empty `master_addr` + /// with a `MissingConfig` error before consulting any environment or builder. + #[test] + fn goosefs_config_rejects_empty_master_addr() { + assert!(matches!( + make_goosefs_store(GoosefsConfig { + master_addr: String::new(), + ..GoosefsConfig::default() + }), + Err(OpenDALStoreError::MissingConfig("master_addr")) + )); + } + + /// HA mode: a comma-separated list of master addresses should be accepted and build + /// successfully. + #[test] + fn goosefs_builds_ha_config() { + let config = GoosefsConfig { + master_addr: "10.0.0.1:9200,10.0.0.2:9200,10.0.0.3:9200".to_string(), + root: Some("/data".to_string()), + ..GoosefsConfig::default() + }; + let store = make_goosefs_store(config).expect("HA store should build"); + assert!(Arc::strong_count(&store) >= 1); + } +} diff --git a/vortex-cloud/src/opendal/mod.rs b/vortex-cloud/src/opendal/mod.rs new file mode 100644 index 00000000000..e4b01e46799 --- /dev/null +++ b/vortex-cloud/src/opendal/mod.rs @@ -0,0 +1,322 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! OpenDAL-backed [`object_store::ObjectStore`] implementations for cloud providers that are not +//! natively supported by the `object_store` crate: Tencent Cloud COS, Alibaba Cloud OSS, and +//! Tencent Cloud GooseFS. +//! +//! OpenDAL exposes each service as an `Operator`. We adapt an `Operator` into an +//! `object_store::ObjectStore` via the `object_store_opendal::OpendalStore` bridge. This lets +//! Vortex consume these services through its existing `ObjectStoreFileSystem` abstraction. +//! +//! Callers that dispatch on a URL scheme should ask [`supports_scheme`] rather than comparing +//! against [`COS_SCHEME`] / [`OSS_SCHEME`] / [`GOOSEFS_SCHEME`] themselves, so that enabling +//! another service does not require touching every call site. +//! +//! # Cargo features +//! +//! * `cos` — Tencent Cloud COS, the `cos://` scheme. +//! * `oss` — Alibaba Cloud OSS, the `oss://` scheme. +//! * `goosefs` — Tencent Cloud GooseFS, the `goosefs://` scheme. +//! +//! With a single service enabled the module still compiles: [`supports_scheme`] returns `false` +//! for every scheme it does not serve and [`make_opendal_store`] reports +//! [`OpenDALStoreError::UnsupportedScheme`]. +//! +//! # Limitations +//! +//! The `OpendalStore` bridge owns its own HTTP request client. Configuration that the JNI/Python +//! layers normally pass through [`object_store::ClientOptions`] — connect/request timeouts, +//! retries, proxy settings, `allow_http` — has no effect on URLs handled here. Properties that a +//! service does not recognize are logged at `warn` and dropped; callers that need strict +//! validation must pre-filter their property maps. + +#[cfg(feature = "cos")] +mod cos; +#[cfg(feature = "goosefs")] +mod goosefs; +#[cfg(feature = "oss")] +mod oss; + +use std::sync::Arc; + +#[cfg(any(feature = "cos", feature = "goosefs", feature = "oss"))] +use ::opendal::Operator; +use object_store::ObjectStore; +#[cfg(any(feature = "cos", feature = "goosefs", feature = "oss"))] +use tracing::warn; +use url::Url; +use vortex_utils::aliases::hash_map::HashMap; + +#[cfg(feature = "cos")] +pub use crate::opendal::cos::COS_SCHEME; +#[cfg(feature = "cos")] +pub use crate::opendal::cos::CosConfig; +#[cfg(feature = "cos")] +pub use crate::opendal::cos::make_cos_store; +#[cfg(feature = "goosefs")] +pub use crate::opendal::goosefs::GOOSEFS_SCHEME; +#[cfg(feature = "goosefs")] +pub use crate::opendal::goosefs::GoosefsConfig; +#[cfg(feature = "goosefs")] +pub use crate::opendal::goosefs::make_goosefs_store; +#[cfg(feature = "oss")] +pub use crate::opendal::oss::OSS_SCHEME; +#[cfg(feature = "oss")] +pub use crate::opendal::oss::OssConfig; +#[cfg(feature = "oss")] +pub use crate::opendal::oss::make_oss_store; + +/// Error type for building an OpenDAL-backed object store. +#[derive(Debug)] +pub enum OpenDALStoreError { + /// The URL scheme is not one this crate handles (e.g. `s3`, `gs`, ...), or its Cargo feature + /// is not enabled. + UnsupportedScheme(String), + /// A required configuration value (bucket and/or endpoint) was missing. + MissingConfig(&'static str), + /// The OpenDAL builder rejected the provided configuration. + Build(::opendal::Error), +} + +impl std::fmt::Display for OpenDALStoreError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + OpenDALStoreError::UnsupportedScheme(s) => { + write!(f, "unsupported OpenDAL scheme: {s}") + } + OpenDALStoreError::MissingConfig(k) => { + write!(f, "missing required OpenDAL store configuration: {k}") + } + OpenDALStoreError::Build(e) => write!(f, "failed to build OpenDAL store: {e}"), + } + } +} + +impl std::error::Error for OpenDALStoreError {} + +impl From for object_store::Error { + fn from(e: OpenDALStoreError) -> Self { + object_store::Error::Generic { + store: "OpenDAL", + source: Box::new(e), + } + } +} + +/// The URL schemes this crate can build stores for, given the enabled Cargo features. +pub const SUPPORTED_SCHEMES: &[&str] = &[ + #[cfg(feature = "cos")] + COS_SCHEME, + #[cfg(feature = "goosefs")] + GOOSEFS_SCHEME, + #[cfg(feature = "oss")] + OSS_SCHEME, +]; + +/// Returns `true` if `scheme` is served by an OpenDAL-backed store in this build. +/// +/// This is the dispatch predicate every caller should use: it tracks the enabled Cargo features, +/// so a consumer that adds a service feature picks it up without changing its own scheme matching. +/// +/// ``` +/// # use vortex_cloud::opendal::supports_scheme; +/// assert!(!supports_scheme("s3")); +/// ``` +pub fn supports_scheme(scheme: &str) -> bool { + SUPPORTED_SCHEMES.contains(&scheme) +} + +/// Build an [`object_store::ObjectStore`] for an OpenDAL-backed URL (`cos://`, `oss://`, +/// `goosefs://`). +/// +/// `properties` are per-request configuration overrides (matching the `HashMap` +/// passed through the JNI/Python layers). Missing values fall back to the environment variables +/// the corresponding service reads (e.g. `TENCENTCLOUD_SECRET_ID`, `ALIBABA_CLOUD_ACCESS_KEY_ID`, +/// `GOOSEFS_MASTER_ADDR`). +/// +/// Returns [`OpenDALStoreError::UnsupportedScheme`] if `url` uses a scheme this build does not +/// serve; test it up-front with [`supports_scheme`]. +pub fn make_opendal_store( + url: &Url, + properties: &HashMap, +) -> Result, OpenDALStoreError> { + make_opendal_store_with_env(url, properties, env_var_lookup) +} + +/// Build an OpenDAL-backed store, resolving environment fallbacks through `env_lookup` instead of +/// the process environment. +/// +/// Callers that already own a configuration source — such as a registry that resolves variables +/// case-insensitively — should use this so that store construction does not silently depend on +/// global state. [`make_opendal_store`] is the same call against the real environment. +pub fn make_opendal_store_with_env( + url: &Url, + properties: &HashMap, + env_lookup: F, +) -> Result, OpenDALStoreError> +where + F: Fn(&str) -> Option, +{ + match url.scheme() { + #[cfg(feature = "cos")] + COS_SCHEME => make_cos_store(cos::url_and_properties_to_config( + url, properties, env_lookup, + )?), + #[cfg(feature = "goosefs")] + GOOSEFS_SCHEME => make_goosefs_store(goosefs::url_and_properties_to_config( + url, properties, env_lookup, + )?), + #[cfg(feature = "oss")] + OSS_SCHEME => make_oss_store(oss::url_and_properties_to_config( + url, properties, env_lookup, + )?), + other => { + // Consumes the arguments so they count as used even in a build with no service + // features enabled, where every scheme lands here. + drop((properties, env_lookup)); + Err(OpenDALStoreError::UnsupportedScheme(other.to_string())) + } + } +} + +/// Default environment-variable lookup: reads `key` from the process environment. +/// +/// The lookup is factored out so tests can pass a fixed map instead of mutating the global +/// environment, which is unsound when `cargo test` runs tests on multiple threads within one +/// process (the `unsafe std::env::set_var` block became `unsafe` in Rust 2024 for exactly this +/// reason). +fn env_var_lookup(key: &str) -> Option { + std::env::var(key).ok() +} + +/// Take `key` from `properties`, falling back to `env_lookup(env_var)`. +#[cfg(any(feature = "cos", feature = "goosefs", feature = "oss"))] +pub(crate) fn property_or_env( + properties: &HashMap, + key: &str, + env_var: &str, + env_lookup: &F, +) -> Option +where + F: Fn(&str) -> Option, +{ + properties.get(key).cloned().or_else(|| env_lookup(env_var)) +} + +/// Take `key` from `properties` as a boolean, accepting the spellings `object_store` accepts in +/// its own configuration: `1`/`true`/`on`/`yes`/`y` and their negatives, case-insensitively. +/// +/// An absent key is `false`. A value that is not a boolean is warned about and read as `false`, +/// which is how [`warn_on_unknown_properties`] already treats a key the service cannot use. +#[cfg(any(feature = "cos", feature = "oss"))] +pub(crate) fn property_as_bool(properties: &HashMap, key: &str) -> bool { + let Some(value) = properties.get(key) else { + return false; + }; + match value.trim().to_ascii_lowercase().as_str() { + "1" | "true" | "on" | "yes" | "y" => true, + "0" | "false" | "off" | "no" | "n" => false, + _ => { + warn!("ignoring OpenDAL store property {key}: `{value}` is not a boolean"); + false + } + } +} + +/// Log a warning for every property key the service does not recognize. +#[cfg(any(feature = "cos", feature = "goosefs", feature = "oss"))] +pub(crate) fn warn_on_unknown_properties(properties: &HashMap, known: &[&str]) { + for key in properties.keys() { + if !known.contains(&key.as_str()) { + warn!("ignoring unknown OpenDAL store property: {key}"); + } + } +} + +/// Finish an OpenDAL builder into an [`Operator`], mapping builder errors into our error type. +#[cfg(any(feature = "cos", feature = "goosefs", feature = "oss"))] +pub(crate) fn build_operator(builder: B) -> Result +where + B: ::opendal::Builder, +{ + // OpenDAL 0.58+: Operator::new returns a finished Operator (no .finish()). + let operator: Operator = Operator::new(builder).map_err(OpenDALStoreError::Build)?; + Ok(operator) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rejects_unknown_scheme() { + let url = Url::parse("s3://bucket/path").unwrap(); + let props = HashMap::new(); + assert!(matches!( + make_opendal_store(&url, &props), + Err(OpenDALStoreError::UnsupportedScheme(_)) + )); + } + + #[test] + fn supports_scheme_tracks_enabled_features() { + assert!(!supports_scheme("s3")); + assert_eq!(supports_scheme("cos"), cfg!(feature = "cos")); + assert_eq!(supports_scheme("goosefs"), cfg!(feature = "goosefs")); + assert_eq!(supports_scheme("oss"), cfg!(feature = "oss")); + } + + /// Every scheme advertised by [`SUPPORTED_SCHEMES`] must actually reach a builder rather than + /// falling through to the `UnsupportedScheme` arm of [`make_opendal_store`]. Without a + /// configured endpoint the build fails, but it must fail with `MissingConfig`, not + /// `UnsupportedScheme` — that difference is what pins the dispatch table to the feature set. + #[test] + fn every_supported_scheme_dispatches() { + for scheme in SUPPORTED_SCHEMES { + let url = Url::parse(&format!("{scheme}://bucket/path")).unwrap(); + let props = HashMap::new(); + assert!( + !matches!( + make_opendal_store(&url, &props), + Err(OpenDALStoreError::UnsupportedScheme(_)) + ), + "{scheme} is advertised but not dispatched" + ); + } + } + + /// The spellings `object_store` accepts for its own boolean configuration, so that the same + /// property means the same thing whether a URL resolves to a native store or an OpenDAL one. + /// Values reach us verbatim from the caller's property map, hence the case and whitespace + /// cases; anything that is not a boolean stays `false` rather than becoming an error. + #[cfg(any(feature = "cos", feature = "oss"))] + #[rstest::rstest] + #[case("true", true)] + #[case("True", true)] + #[case("TRUE", true)] + #[case(" true ", true)] + #[case("1", true)] + #[case("on", true)] + #[case("yes", true)] + #[case("y", true)] + #[case("false", false)] + #[case("False", false)] + #[case("0", false)] + #[case("off", false)] + #[case("no", false)] + #[case("n", false)] + #[case("maybe", false)] + #[case("", false)] + fn property_as_bool_matches_object_store(#[case] value: &str, #[case] expected: bool) { + let mut props = HashMap::new(); + props.insert("skip_signature".to_string(), value.to_string()); + assert_eq!(property_as_bool(&props, "skip_signature"), expected); + } + + #[cfg(any(feature = "cos", feature = "oss"))] + #[test] + fn property_as_bool_is_false_when_absent() { + assert!(!property_as_bool(&HashMap::new(), "skip_signature")); + } +} diff --git a/vortex-cloud/src/opendal/oss.rs b/vortex-cloud/src/opendal/oss.rs new file mode 100644 index 00000000000..3d2ae4d41b3 --- /dev/null +++ b/vortex-cloud/src/opendal/oss.rs @@ -0,0 +1,258 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Alibaba Cloud OSS, served over OpenDAL's `services::Oss`. + +use std::sync::Arc; + +use ::opendal::services; +use object_store::ObjectStore; +use object_store_opendal::OpendalStore; +use url::Url; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::opendal::OpenDALStoreError; +use crate::opendal::build_operator; +use crate::opendal::property_as_bool; +use crate::opendal::property_or_env; +use crate::opendal::warn_on_unknown_properties; + +/// The URL scheme served by Alibaba Cloud OSS. +pub const OSS_SCHEME: &str = "oss"; + +/// Property keys recognized for `oss://` URLs. Anything else is warned about and dropped. +const KNOWN_PROPERTIES: &[&str] = &[ + "bucket", + "endpoint", + "access_key_id", + "access_key_secret", + "security_token", + "root", + "skip_signature", +]; + +/// Strongly-typed configuration for building an OpenDAL store against Alibaba Cloud OSS. +/// +/// This mirrors [`crate::opendal::CosConfig`], but OSS names its credentials `access_key_id` / +/// `access_key_secret` and has no equivalent of COS's `disable_config_load`; the closest knob is +/// [`OssConfig::skip_signature`], which sends unsigned requests for public buckets. +#[derive(Debug, Clone, Default)] +pub struct OssConfig { + /// OSS bucket name. May be overridden by `properties["bucket"]` when adapting a URL. + pub bucket: String, + /// OSS endpoint, e.g. `https://oss-cn-hangzhou.aliyuncs.com`. + pub endpoint: String, + /// Alibaba Cloud access key id (mapped to `ALIBABA_CLOUD_ACCESS_KEY_ID`). + pub access_key_id: Option, + /// Alibaba Cloud access key secret (mapped to `ALIBABA_CLOUD_ACCESS_KEY_SECRET`). + pub access_key_secret: Option, + /// STS security token for temporary credentials (mapped to `ALIBABA_CLOUD_SECURITY_TOKEN`). + pub security_token: Option, + /// Optional root prefix applied to all operations. + pub root: Option, + /// When `true`, send unsigned requests. Use this for public, read-only buckets. + pub skip_signature: bool, +} + +/// Build an [`object_store::ObjectStore`] for Alibaba Cloud OSS directly from an [`OssConfig`]. +/// +/// This is the preferred entry point for callers that have a strongly-typed configuration object. +/// It does not synthesize a URL and so is not fragile against reordering of `bucket` vs URL host +/// precedence. +pub fn make_oss_store(config: OssConfig) -> Result, OpenDALStoreError> { + if config.bucket.is_empty() { + return Err(OpenDALStoreError::MissingConfig("bucket")); + } + if config.endpoint.is_empty() { + return Err(OpenDALStoreError::MissingConfig("endpoint")); + } + + let mut builder = services::Oss::default() + .bucket(&config.bucket) + .endpoint(&config.endpoint); + + if let Some(root) = config.root.as_deref() { + builder = builder.root(root); + } + if let Some(access_key_id) = config.access_key_id.as_deref() { + builder = builder.access_key_id(access_key_id); + } + if let Some(access_key_secret) = config.access_key_secret.as_deref() { + builder = builder.access_key_secret(access_key_secret); + } + if let Some(security_token) = config.security_token.as_deref() { + builder = builder.security_token(security_token); + } + if config.skip_signature { + builder = builder.skip_signature(); + } + + let operator = build_operator(builder)?; + Ok(Arc::new(OpendalStore::new(operator))) +} + +/// Translate an (`oss://` URL, properties) pair into a strongly-typed [`OssConfig`]. +/// +/// Bucket is taken from `properties["bucket"]` first and falls back to the URL host; endpoint is +/// taken from `properties["endpoint"]` first and falls back to `OSS_ENDPOINT`; credentials fall +/// back to the `ALIBABA_CLOUD_*` variables OpenDAL's OSS builder reads. +/// +/// `env_lookup` is the source of truth for environment-variable fallbacks. The production entry +/// point passes the real environment; tests pass a closure that returns from a fixed map, so they +/// do not race against the process environment. +pub(crate) fn url_and_properties_to_config( + url: &Url, + properties: &HashMap, + env_lookup: F, +) -> Result +where + F: Fn(&str) -> Option, +{ + warn_on_unknown_properties(properties, KNOWN_PROPERTIES); + + let bucket = properties + .get("bucket") + .cloned() + .or_else(|| url.host_str().map(str::to_string)) + .ok_or(OpenDALStoreError::MissingConfig("bucket"))?; + + let endpoint = property_or_env(properties, "endpoint", "OSS_ENDPOINT", &env_lookup) + .ok_or(OpenDALStoreError::MissingConfig("endpoint"))?; + + Ok(OssConfig { + bucket, + endpoint, + access_key_id: property_or_env( + properties, + "access_key_id", + "ALIBABA_CLOUD_ACCESS_KEY_ID", + &env_lookup, + ), + access_key_secret: property_or_env( + properties, + "access_key_secret", + "ALIBABA_CLOUD_ACCESS_KEY_SECRET", + &env_lookup, + ), + security_token: property_or_env( + properties, + "security_token", + "ALIBABA_CLOUD_SECURITY_TOKEN", + &env_lookup, + ), + root: properties.get("root").cloned(), + skip_signature: property_as_bool(properties, "skip_signature"), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The endpoint is required, and the failure must be deterministic regardless of what the + /// process environment happens to contain — hence the fixed `|_| None` lookup. + #[test] + fn oss_requires_endpoint() { + let url = Url::parse("oss://my-bucket/path").unwrap(); + let props = HashMap::new(); + let result = url_and_properties_to_config(&url, &props, |_| None).unwrap_err(); + assert!(matches!( + result, + OpenDALStoreError::MissingConfig("endpoint") + )); + } + + /// OpenDAL's OSS service reads `ALIBABA_CLOUD_*` (with the underscore), not `ALIBABACLOUD_*`. + /// Pinning the exact names here keeps the fallbacks honest. + #[test] + fn oss_falls_back_to_env() { + let url = Url::parse("oss://my-bucket/path").unwrap(); + let env = |key: &str| match key { + "OSS_ENDPOINT" => Some("https://oss-cn-hangzhou.aliyuncs.com".to_string()), + "ALIBABA_CLOUD_ACCESS_KEY_ID" => Some("LTAI".to_string()), + "ALIBABA_CLOUD_ACCESS_KEY_SECRET" => Some("secret".to_string()), + "ALIBABA_CLOUD_SECURITY_TOKEN" => Some("token".to_string()), + _ => None, + }; + let props = HashMap::new(); + let config = url_and_properties_to_config(&url, &props, env).expect("config"); + + assert_eq!(config.bucket, "my-bucket"); + assert_eq!(config.endpoint, "https://oss-cn-hangzhou.aliyuncs.com"); + assert_eq!(config.access_key_id.as_deref(), Some("LTAI")); + assert_eq!(config.access_key_secret.as_deref(), Some("secret")); + assert_eq!(config.security_token.as_deref(), Some("token")); + } + + /// An explicit property must win over the environment fallback. + #[test] + fn oss_property_overrides_env() { + let url = Url::parse("oss://url-bucket/path").unwrap(); + let env = |key: &str| match key { + "OSS_ENDPOINT" => Some("https://from-env.example.com".to_string()), + "ALIBABA_CLOUD_ACCESS_KEY_ID" => Some("from-env".to_string()), + _ => None, + }; + let mut props = HashMap::new(); + props.insert("bucket".to_string(), "prop-bucket".to_string()); + props.insert( + "endpoint".to_string(), + "https://from-prop.example.com".to_string(), + ); + props.insert("access_key_id".to_string(), "from-prop".to_string()); + + let config = url_and_properties_to_config(&url, &props, env).expect("config"); + assert_eq!(config.bucket, "prop-bucket"); + assert_eq!(config.endpoint, "https://from-prop.example.com"); + assert_eq!(config.access_key_id.as_deref(), Some("from-prop")); + } + + /// A full config must reach the OpenDAL builder and produce a store. + #[test] + fn oss_builds_with_explicit_config() { + let url = Url::parse("oss://my-bucket/path/to/dataset.vortex").unwrap(); + let env = |key: &str| match key { + "OSS_ENDPOINT" => Some("https://oss-cn-hangzhou.aliyuncs.com".to_string()), + "ALIBABA_CLOUD_ACCESS_KEY_ID" => Some("LTAI".to_string()), + "ALIBABA_CLOUD_ACCESS_KEY_SECRET" => Some("secret".to_string()), + _ => None, + }; + let props = HashMap::new(); + let config = url_and_properties_to_config(&url, &props, env).expect("config"); + let store = make_oss_store(config).expect("store should build"); + assert!(Arc::strong_count(&store) >= 1); + } + + /// Public buckets need no credentials when `skip_signature` is set. + #[test] + fn oss_builds_anonymous() { + let store = make_oss_store(OssConfig { + bucket: "public-bucket".to_string(), + endpoint: "https://oss-cn-hangzhou.aliyuncs.com".to_string(), + skip_signature: true, + ..OssConfig::default() + }) + .expect("anonymous store should build"); + assert!(Arc::strong_count(&store) >= 1); + } + + /// `make_oss_store` must reject an empty bucket or endpoint before touching the builder. + #[test] + fn oss_config_rejects_empty_fields() { + assert!(matches!( + make_oss_store(OssConfig { + bucket: String::new(), + ..OssConfig::default() + }), + Err(OpenDALStoreError::MissingConfig("bucket")) + )); + assert!(matches!( + make_oss_store(OssConfig { + bucket: "b".to_string(), + endpoint: String::new(), + ..OssConfig::default() + }), + Err(OpenDALStoreError::MissingConfig("endpoint")) + )); + } +} diff --git a/vortex-cloud/src/registry/mod.rs b/vortex-cloud/src/registry/mod.rs new file mode 100644 index 00000000000..b854569959a --- /dev/null +++ b/vortex-cloud/src/registry/mod.rs @@ -0,0 +1,296 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Apache Software Foundation (ASF) + +//! A URL-to-[`ObjectStore`] registry. +//! +//! This is an adapted version of the `DefaultObjectStoreRegistry` from the `object_store` crate, +//! modified in two ways: +//! +//! 1. configuration is resolved out of environment variables case-insensitively, matching how the +//! various `Store::from_env` builders behave (see +//! ); +//! 2. schemes that `object_store` does not recognize natively — the OpenDAL-backed `cos://`, +//! `oss://`, and `goosefs://`, and the Hugging Face Hub's `hf://` — are served by the crate's +//! `opendal` and `hf` modules under the matching service feature. + +use std::sync::Arc; + +use object_store::ObjectStore; +use object_store::parse_url_opts; +use object_store::path::Path; +use object_store::registry::ObjectStoreRegistry; +use parking_lot::RwLock; +use url::Url; +use vortex_utils::aliases::hash_map::HashMap; + +#[derive(Debug, Default)] +struct PathEntry { + /// Store, if defined at this path + store: Option>, + /// Child [`PathEntry`], keyed by the next path segment in their path + children: HashMap, +} + +impl PathEntry { + /// Lookup a store based on URL path + /// + /// Returns the store and its path segment depth + fn lookup(&self, to_resolve: &Url) -> Option<(&Arc, usize)> { + let mut current = self; + let mut ret = self.store.as_ref().map(|store| (store, 0)); + let mut depth = 0; + // Traverse the PathEntry tree to find the longest match + for segment in path_segments(to_resolve.path()) { + match current.children.get(segment) { + Some(e) => { + current = e; + depth += 1; + if let Some(store) = ¤t.store { + ret = Some((store, depth)) + } + } + None => break, + } + } + ret + } +} + +/// An [`ObjectStoreRegistry`] that normalizes environment variables before doing lookups, and +/// resolves OpenDAL-backed schemes under the enabled service features. +/// +/// Stores are cached per URL authority, so repeated resolves against the same bucket share one +/// client. A store can also be registered explicitly at a path prefix with +/// [`ObjectStoreRegistry::register`], in which case it wins over the built-and-cached fallback. +/// +/// ``` +/// # use object_store::registry::ObjectStoreRegistry; +/// # use vortex_cloud::Registry; +/// # use url::Url; +/// # fn main() -> Result<(), Box> { +/// let registry = Registry::default(); +/// let store = object_store::memory::InMemory::new(); +/// registry.register(Url::parse("memory://bucket/")?, std::sync::Arc::new(store)); +/// +/// let (_store, path) = registry.resolve(&Url::parse("memory://bucket/a/b.vortex")?)?; +/// assert_eq!(path.as_ref(), "a/b.vortex"); +/// # Ok(()) +/// # } +/// ``` +#[derive(Debug, Default)] +pub struct Registry { + /// Mapping from [`url_key`] to [`PathEntry`] + map: RwLock>, + /// Where store configuration is read from when building a store. + env: EnvSource, +} + +/// Source of the configuration variables consulted when building a store. +#[derive(Debug, Default)] +enum EnvSource { + /// Read from the process environment. + #[default] + Process, + /// A fixed set of variables. + Fixed(Vec<(String, String)>), +} + +impl EnvSource { + /// Case-insensitive lookup of a single configuration variable. + #[cfg(any(feature = "cos", feature = "goosefs", feature = "hf", feature = "oss"))] + fn lookup(&self, key: &str) -> Option { + match self { + EnvSource::Process => std::env::var(key).ok(), + EnvSource::Fixed(vars) => vars + .iter() + .find(|(k, _)| k.eq_ignore_ascii_case(key)) + .map(|(_, v)| v.clone()), + } + } + + /// The configuration variables, with keys lowercased so that lookups are case-insensitive. + fn normalized_vars(&self) -> Vec<(String, String)> { + match self { + EnvSource::Process => std::env::vars() + .map(|(k, v)| (k.to_ascii_lowercase(), v)) + .collect(), + EnvSource::Fixed(vars) => vars + .iter() + .map(|(k, v)| (k.to_ascii_lowercase(), v.clone())) + .collect(), + } + } +} + +impl ObjectStoreRegistry for Registry { + fn register(&self, url: Url, store: Arc) -> Option> { + let mut map = self.map.write(); + let entry = entry_at( + &mut map, + url_key(&url), + url.path(), + num_segments(url.path()), + ); + entry.store.replace(store) + } + + fn resolve(&self, to_resolve: &Url) -> object_store::Result<(Arc, Path)> { + let key = url_key(to_resolve); + + // 1. Look up the user-registered map first. Every other scheme does this, so an explicit + // `Registry::register("cos://...", store)` should also win over the build-and-cache + // fallback below. This also means a previously-resolved store (built by us for the + // same URL) is served from the cache here, without rebuilding the client. + { + let map = self.map.read(); + if let Some((store, depth)) = map.get(key).and_then(|entry| entry.lookup(to_resolve)) { + let path = path_suffix(to_resolve, depth)?; + return Ok((Arc::clone(store), path)); + } + } + + // 2. Build the store and the path *within* that store. `build_store` reports the path + // the way `parse_url_opts` does — i.e. the path is what the store will see as the + // key — which is what keeps every scheme on the one caching rule in + // `cache_and_resolve`. The depth is then the difference in segment counts. + let (store, path) = self.build_store(to_resolve)?; + // The raw URL path counts a percent-encoded segment (e.g. `refs%2Fconvert%2Fparquet`) + // as one segment while `path` is percent-decoded and may count more, so saturate to + // mount the store at the URL root rather than underflowing. + let depth = num_segments(to_resolve.path()).saturating_sub(num_segments(path.as_ref())); + + // 3. Cache the store and return it alongside the path that lives inside the store. + self.cache_and_resolve(to_resolve, store, depth) + } +} + +impl Registry { + /// Create a registry that reads store configuration from the process environment. + pub fn new() -> Self { + Self::default() + } + + /// Create a registry over a fixed set of configuration variables, consulted instead of the + /// process environment when building stores. + /// + /// Lookups are case-insensitive. Pass each key at most once — the set's consumers disagree + /// on which duplicate wins. + pub fn with_vars(vars: I) -> Self + where + I: IntoIterator, + { + Self { + map: RwLock::default(), + env: EnvSource::Fixed(vars.into_iter().collect()), + } + } + + /// Caches `store` as the store serving the first `depth` path segments of `to_resolve`, and + /// returns it alongside the remaining path. + /// + /// If a racing `resolve` cached a store at the same position first, that store wins and the + /// one just built is dropped, so all callers of a given prefix share a single client. + fn cache_and_resolve( + &self, + to_resolve: &Url, + store: Arc, + depth: usize, + ) -> object_store::Result<(Arc, Path)> { + let path = path_suffix(to_resolve, depth)?; + + let mut map = self.map.write(); + let entry = entry_at(&mut map, url_key(to_resolve), to_resolve.path(), depth); + let stored = Arc::clone(match &entry.store { + None => entry.store.insert(store), + Some(existing) => existing, // Racing creation - use existing + }); + + Ok((stored, path)) + } + + /// Builds the [`ObjectStore`] that serves `to_resolve`, with the path within that store. + /// + /// This mirrors [`parse_url_opts`], extended with schemes the `object_store` crate does not + /// recognize. Reporting the path the way `parse_url_opts` does is what keeps every scheme on + /// the one caching rule in [`Registry::resolve`]: a scheme says where its store is rooted, and + /// the registry decides how to cache it. + fn build_store(&self, to_resolve: &Url) -> object_store::Result<(Arc, Path)> { + // OpenDAL-backed schemes (Tencent COS, Alibaba OSS, Tencent GooseFS) are not recognized + // by `object_store`, so build them from OpenDAL's own environment-variable configuration + // (e.g. `TENCENTCLOUD_SECRET_ID`, `GOOSEFS_MASTER_ADDR`). The operator is rooted at the + // bucket, which lives in the URL authority, so — exactly as for `s3://bucket/path` — + // the whole URL path is the object key. + #[cfg(any(feature = "cos", feature = "goosefs", feature = "oss"))] + if crate::opendal::supports_scheme(to_resolve.scheme()) { + let store = + crate::opendal::make_opendal_store_with_env(to_resolve, &HashMap::new(), |key| { + self.env.lookup(key) + })?; + return Ok((store, Path::from_url_path(to_resolve.path())?)); + } + + // The Hugging Face Hub is not recognized by `object_store` either. Unlike the OpenDAL + // schemes its store is not rooted at the URL authority — a repository and revision occupy + // path segments too — so it reports the in-repository path itself and the registry mounts + // the store as deep as that implies. + #[cfg(feature = "hf")] + if crate::hf::supports_scheme(to_resolve.scheme()) { + return Ok(crate::hf::make_hf_store_from_url_with_env( + to_resolve, + |key| self.env.lookup(key), + )?); + } + + let (store, path) = parse_url_opts(to_resolve, self.env.normalized_vars())?; + Ok((Arc::from(store), path)) + } +} + +/// Extracts the scheme and authority of a URL (components before the Path) +fn url_key(url: &Url) -> &str { + &url[..url::Position::AfterPort] +} + +/// Returns the non-empty segments of a path +/// +/// Note: We don't use [`Url::path_segments`] as we only want non-empty paths +fn path_segments(s: &str) -> impl Iterator { + s.split('/').filter(|x| !x.is_empty()) +} + +/// Returns the number of non-empty path segments in a path +fn num_segments(s: &str) -> usize { + path_segments(s).count() +} + +/// Returns the path of `url` skipping the first `depth` segments +fn path_suffix(url: &Url, depth: usize) -> Result { + // The segments come from a URL, so percent-decode them (one raw segment may decode to + // several path parts, e.g. `refs%2Fconvert%2Fparquet`). + let suffix = path_segments(url.path()) + .skip(depth) + .collect::>() + .join("/"); + Path::from_url_path(suffix).map_err(|e| object_store::Error::Generic { + store: "ObjectStoreRegistry", + source: Box::new(e), + }) +} + +/// Walks to the [`PathEntry`] for `key` sitting `depth` segments into `path`, creating the +/// intermediate entries as needed. +fn entry_at<'a>( + map: &'a mut HashMap, + key: &str, + path: &str, + depth: usize, +) -> &'a mut PathEntry { + let mut current = map.entry(key.to_string()).or_default(); + for segment in path_segments(path).take(depth) { + current = current.children.entry(segment.to_string()).or_default(); + } + current +} + +#[cfg(test)] +mod tests; diff --git a/vortex-cloud/src/registry/tests.rs b/vortex-cloud/src/registry/tests.rs new file mode 100644 index 00000000000..f65aeca7cd6 --- /dev/null +++ b/vortex-cloud/src/registry/tests.rs @@ -0,0 +1,304 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Write; +use std::sync::Arc; + +use object_store::ObjectStore; +use object_store::path::Path; +use object_store::registry::ObjectStoreRegistry; +use url::Url; + +use super::Registry; + +/// A registry whose S3 configuration comes from a fixed map rather than the process environment, +/// so these tests neither read nor mutate global state. +fn registry() -> Registry { + Registry::with_vars([("AWS_REGION".to_string(), "us-east-3".to_string())]) +} + +/// A percent-encoded segment (as HuggingFace dataset URLs use for `refs/convert/parquet`) decodes +/// to more path parts than the raw URL has segments. Both the build-and-cache branch and the +/// cached-store branch must decode it rather than underflowing the segment-count subtraction. +#[test] +fn test_resolve_percent_encoded_path() -> Result<(), Box> { + let registry = registry(); + let url = Url::parse( + "https://example.com/datasets/org/name/resolve/refs%2Fconvert%2Fparquet/dir/file.vortex", + )?; + let expected = Path::from("datasets/org/name/resolve/refs/convert/parquet/dir/file.vortex"); + + // First resolution parses and registers the store; it must decode the path, not panic. + let (_store, path) = registry.resolve(&url)?; + assert_eq!(path, expected); + + // Second resolution takes the cached-store branch and must agree. + let (_store, path) = registry.resolve(&url)?; + assert_eq!(path, expected); + Ok(()) +} + +#[test] +#[expect(clippy::use_debug)] +fn test_resolve_url() -> Result<(), Box> { + let (store, _) = registry().resolve(&Url::parse("s3://my-bucket/test")?)?; + + // NOTE(aduffy): object_store doesn't let us downcast stores, the only way to verify + // that a configuration was added was to validate that it ends up in the Debug + // output :/ + let mut debug_str = String::new(); + write!(&mut debug_str, "{store:?}")?; + + assert!(debug_str.contains("us-east-3")); + Ok(()) +} + +/// Two resolves of the same URL must return the same `Arc` (the cached client) and the same +/// path, and two different keys must get distinct stores. This pins the symmetry the registry +/// promises: cache hit returns the same client and the same key. +#[test] +fn test_resolve_url_caches_per_key() -> Result<(), Box> { + let registry = registry(); + let first = Url::parse("s3://my-bucket/first/second")?; + let second = Url::parse("s3://my-bucket/first/second")?; + + let (store_a, path_a) = registry.resolve(&first)?; + let (store_b, path_b) = registry.resolve(&second)?; + assert!(Arc::ptr_eq(&store_a, &store_b)); + assert_eq!(path_a, path_b); + + // A different bucket gets its own client. + let other = Url::parse("s3://other-bucket/first/second")?; + let (store_c, _) = registry.resolve(&other)?; + assert!(!Arc::ptr_eq(&store_a, &store_c)); + Ok(()) +} + +/// Two objects in the same bucket must share a cached client. This is the regression that +/// the bespoke "walk every segment then return the whole key" path caused: by walking all +/// segments it stored the store at full depth, so the next resolve saw the whole key +/// already consumed and returned an empty path. Pinning the path here guards against that. +#[test] +fn test_resolve_url_shared_client_same_bucket() -> Result<(), Box> { + let registry = registry(); + let a = Url::parse("s3://my-bucket/path/to/data.vortex")?; + let b = Url::parse("s3://my-bucket/other/data.vortex")?; + + let (store_a, path_a) = registry.resolve(&a)?; + let (store_b, path_b) = registry.resolve(&b)?; + assert!(Arc::ptr_eq(&store_a, &store_b)); + assert_eq!(path_a.as_ref(), "path/to/data.vortex"); + assert_eq!(path_b.as_ref(), "other/data.vortex"); + Ok(()) +} + +/// Pinning the `entry_at` helper at depth > 0 (the path-prefix case) protects the shared +/// `entry_at` / `cache_and_resolve` helpers from regressing prefix registration. +#[test] +fn test_register_at_prefix_shares_store() -> Result<(), Box> { + let registry = registry(); + let prefix = Url::parse("s3://my-bucket/prefix/")?; + let (inner, _) = registry.resolve(&prefix)?; + let cached: Arc = Arc::clone(&inner); + let replaced = registry.register(prefix.clone(), Arc::clone(&cached)); + // First registration at this prefix replaces nothing. + assert!(replaced.is_none()); + + // A resolve at a deeper key still walks through the registered prefix's store. + let deeper = Url::parse("s3://my-bucket/prefix/inner/key")?; + let (store_deeper, path_deeper) = registry.resolve(&deeper)?; + assert!(Arc::ptr_eq(&store_deeper, &cached)); + assert_eq!(path_deeper.as_ref(), "inner/key"); + Ok(()) +} + +/// A `file://` URL with its path cleared must resolve to a local store. `vortex-duckdb` resolves +/// exactly this shape — it clears the path off the glob's base URL before asking for a +/// filesystem — so this pins the most common local-file case. +#[test] +fn test_resolve_file_url_with_empty_path() -> Result<(), Box> { + let registry = registry(); + + let mut base = Url::parse("file:///data/part-0.vortex")?; + base.set_path(""); + let (store, path) = registry.resolve(&base)?; + assert!(format!("{store:?}").contains("LocalFileSystem")); + assert_eq!(path.as_ref(), ""); + + // The same store is reused for a second glob rooted at the same authority. + let (store_again, _) = registry.resolve(&base)?; + assert!(Arc::ptr_eq(&store, &store_again)); + Ok(()) +} + +/// An explicitly registered store must win over the build-and-cache fallback, including for +/// schemes `object_store` does not recognize natively. `memory://` stands in for any such scheme +/// so this holds regardless of which service features are enabled. +#[test] +fn test_registered_store_wins_over_build() -> Result<(), Box> { + let registry = registry(); + let registered: Arc = Arc::new(object_store::memory::InMemory::new()); + registry.register(Url::parse("memory://bucket/")?, Arc::clone(®istered)); + + let (store, path) = registry.resolve(&Url::parse("memory://bucket/a/b.vortex")?)?; + assert!(Arc::ptr_eq(&store, ®istered)); + assert_eq!(path.as_ref(), "a/b.vortex"); + Ok(()) +} + +/// An `hf://` URL must resolve to a store rooted at the repository revision, leaving the +/// in-repository path as the object key. The repository and revision occupy URL path segments, so +/// this is the one scheme whose store is not mounted at the URL authority — getting the depth wrong +/// would send the repository name to the Hub as part of the file path. +#[cfg(feature = "hf")] +#[rstest::rstest] +#[case("hf://datasets/org/name/data/train.vortex", "data/train.vortex")] +#[case( + "hf://datasets/org/name@refs%2Fconvert%2Fparquet/data/train.vortex", + "data/train.vortex" +)] +#[case("hf://org/name/model.vortex", "model.vortex")] +fn test_hf_scheme_mounts_at_the_repository( + #[case] url: &str, + #[case] expected: &str, +) -> Result<(), Box> { + let registry = registry(); + let url = Url::parse(url)?; + + // First resolution builds and caches the store; the second takes the cached-store branch. + // Both must report the same key, since the cached branch recomputes it from the mount depth. + let (_store, path) = registry.resolve(&url)?; + assert_eq!(path, Path::from(expected)); + let (_store, path) = registry.resolve(&url)?; + assert_eq!(path, Path::from(expected)); + Ok(()) +} + +/// Two revisions of one repository are different stores, since the revision is part of the prefix +/// the store is rooted at. +#[cfg(feature = "hf")] +#[test] +fn test_hf_revisions_do_not_share_a_store() -> Result<(), Box> { + let registry = registry(); + + let (main, _) = registry.resolve(&Url::parse("hf://datasets/org/name/train.vortex")?)?; + let (tagged, _) = registry.resolve(&Url::parse("hf://datasets/org/name@v2/train.vortex")?)?; + + assert!(!Arc::ptr_eq(&main, &tagged)); + Ok(()) +} + +/// The OpenDAL-backed schemes must resolve through the registry rather than falling through to +/// `parse_url_opts`, which does not recognize them. +/// +/// COS and OSS need an endpoint that is not derivable from the URL, so a bare +/// `scheme://bucket/key.vortex` URL fails the build; that failure must come from the OpenDAL +/// builder ("missing required OpenDAL store configuration"), not from `object_store`'s +/// unrecognized-scheme path. GooseFS is different: its `master_addr` is taken from the URL +/// authority, so the same URL builds a store successfully — which itself proves it reached the +/// OpenDAL builder, since `parse_url_opts` would have rejected `goosefs://`. +#[cfg(any(feature = "cos", feature = "oss", feature = "goosefs"))] +#[test] +fn test_opendal_schemes_reach_the_opendal_builder() -> Result<(), Box> { + for scheme in crate::opendal::SUPPORTED_SCHEMES { + let url = Url::parse(&format!("{scheme}://bucket/key.vortex"))?; + match registry().resolve(&url) { + // GooseFS derives `master_addr` from the URL authority, so the store builds + // successfully. `parse_url_opts` does not recognize `goosefs://`, so an `Ok` + // here proves the OpenDAL builder ran. + Ok(_) => {} + Err(err) => { + let message = err.to_string(); + assert!( + message.contains("OpenDAL"), + "{scheme} did not reach the OpenDAL builder: {message}" + ); + } + } + } + Ok(()) +} + +/// The key the registry reports must be the object's *literal* key. +/// +/// `~` is unreserved, so it travels through a URL verbatim; characters a URL has to escape decode +/// back to their literal form. +#[rstest::rstest] +#[case("repro~tilde~key/data.vortex", "repro~tilde~key/data.vortex")] +#[case("tablets/~~all~0/data.vortex", "tablets/~~all~0/data.vortex")] +#[case("dir/a[1].vortex", "dir/a[1].vortex")] +#[case("dir/a%23b.vortex", "dir/a#b.vortex")] +#[case("dir/a%20b.vortex", "dir/a b.vortex")] +#[case("dir/a%2520b.vortex", "dir/a%20b.vortex")] +fn test_resolve_reports_literal_key( + #[case] url_path: &str, + #[case] expected: &str, +) -> Result<(), Box> { + let registry = registry(); + registry.register( + Url::parse("memory://bucket/")?, + Arc::new(object_store::memory::InMemory::new()), + ); + + // Both the registered-store branch and the recomputed cached branch must agree. + let url = Url::parse(&format!("memory://bucket/{url_path}"))?; + let (_store, path) = registry.resolve(&url)?; + assert_eq!(path.as_ref(), expected); + let (_store, path) = registry.resolve(&url)?; + assert_eq!(path.as_ref(), expected); + Ok(()) +} + +#[rstest::rstest] +#[case::tilde("repro~tilde~key/data.vortex")] +#[case::double_tilde("tablets/~~all~0/data.vortex")] +#[case::brackets("dir/a[1].vortex")] +#[case::braces("dir/a{x}.vortex")] +#[case::caret("dir/a^b.vortex")] +#[case::pipe("dir/a|b.vortex")] +#[case::star("dir/a*b.vortex")] +#[case::space("dir/a b.vortex")] +#[case::plain("plain/data.vortex")] +#[tokio::test] +async fn test_resolved_key_reaches_the_object( + #[case] key: &str, +) -> Result<(), Box> { + use object_store::ObjectStoreExt; + use vortex_error::VortexError; + use vortex_file::OpenOptionsSessionExt; + + let session = vortex_array::array_session() + .with::() + .with::(); + + // Seed the store with the literal key, the way any other client would. + let store: Arc = Arc::new(object_store::memory::InMemory::new()); + let location = Path::parse(key)?; + assert_eq!( + location.as_ref(), + key, + "the store must hold the literal key" + ); + store.put(&location, vec![0u8; 1024].into()).await?; + + let registry = registry(); + registry.register(Url::parse("memory://bucket/")?, Arc::clone(&store)); + let (store, path) = registry.resolve(&Url::parse(&format!("memory://bucket/{key}"))?)?; + assert_eq!( + path.as_ref(), + key, + "the registry must report the literal key" + ); + + let Err(err) = session.open_options().open_object_store(&store, path).await else { + panic!("the payload is not a Vortex file, so opening it must fail") + }; + assert!( + !matches!( + err, + VortexError::ObjectStore(object_store::Error::NotFound { .. }, _) + ), + "opening {key:?} requested a key the store does not have: {err}" + ); + Ok(()) +} diff --git a/vortex-compressor/src/builtins/dict/binary.rs b/vortex-compressor/src/builtins/dict/binary.rs index 72b5f01141a..c407d0251c6 100644 --- a/vortex-compressor/src/builtins/dict/binary.rs +++ b/vortex-compressor/src/builtins/dict/binary.rs @@ -6,10 +6,13 @@ //! Vortex encoders must always produce unsigned integer codes; signed codes are only accepted //! for external compatibility. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::arrays::Dict; use vortex_array::arrays::DictArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::dict::DictArrayExt; @@ -45,6 +48,10 @@ impl Scheme for BinaryDictScheme { canonical.dtype().is_binary() } + fn produced_encodings(&self) -> Vec { + vec![Dict.id()] + } + fn stats_options(&self) -> GenerateStatsOptions { GenerateStatsOptions { count_distinct_values: true, diff --git a/vortex-compressor/src/builtins/dict/float.rs b/vortex-compressor/src/builtins/dict/float.rs index 074d3ce5e07..f962c3ff967 100644 --- a/vortex-compressor/src/builtins/dict/float.rs +++ b/vortex-compressor/src/builtins/dict/float.rs @@ -6,11 +6,14 @@ //! Vortex encoders must always produce unsigned integer codes; signed codes are only accepted for //! external compatibility. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::arrays::Dict; use vortex_array::arrays::DictArray; use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; @@ -51,6 +54,10 @@ impl Scheme for FloatDictScheme { canonical.dtype().is_float() } + fn produced_encodings(&self) -> Vec { + vec![Dict.id()] + } + fn stats_options(&self) -> GenerateStatsOptions { GenerateStatsOptions { count_distinct_values: true, diff --git a/vortex-compressor/src/builtins/dict/integer.rs b/vortex-compressor/src/builtins/dict/integer.rs index 4e91bace3ac..27a17ef94ad 100644 --- a/vortex-compressor/src/builtins/dict/integer.rs +++ b/vortex-compressor/src/builtins/dict/integer.rs @@ -6,11 +6,14 @@ //! Vortex encoders must always produce unsigned integer codes; signed codes are only accepted //! for external compatibility. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::arrays::Dict; use vortex_array::arrays::DictArray; use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; @@ -46,6 +49,10 @@ impl Scheme for IntDictScheme { canonical.dtype().is_int() } + fn produced_encodings(&self) -> Vec { + vec![Dict.id()] + } + fn stats_options(&self) -> GenerateStatsOptions { GenerateStatsOptions { count_distinct_values: true, diff --git a/vortex-compressor/src/builtins/dict/string.rs b/vortex-compressor/src/builtins/dict/string.rs index 64ed469dde9..f5cbcd54d89 100644 --- a/vortex-compressor/src/builtins/dict/string.rs +++ b/vortex-compressor/src/builtins/dict/string.rs @@ -6,10 +6,13 @@ //! Vortex encoders must always produce unsigned integer codes; signed codes are only accepted //! for external compatibility. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::arrays::Dict; use vortex_array::arrays::DictArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::dict::DictArrayExt; @@ -45,6 +48,10 @@ impl Scheme for StringDictScheme { canonical.dtype().is_utf8() } + fn produced_encodings(&self) -> Vec { + vec![Dict.id()] + } + fn stats_options(&self) -> GenerateStatsOptions { GenerateStatsOptions { count_distinct_values: true, diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 9dbbc611448..86d45d2c0d9 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -14,16 +14,19 @@ use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::Masked; use vortex_array::arrays::StructArray; +use vortex_array::arrays::UnionArray; use vortex_array::arrays::Variant; use vortex_array::arrays::VariantArray; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; -use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::arrays::listview::list_from_list_view; use vortex_array::arrays::masked::MaskedArraySlotsExt; -use vortex_array::arrays::scalar_fn::AnyScalarFn; use vortex_array::arrays::struct_::StructArrayExt; -use vortex_array::arrays::variant::VariantArrayExt; +use vortex_array::arrays::union::UnionArrayExt; +use vortex_array::arrays::union::UnionArraySlotsExt; +use vortex_array::arrays::variant::VariantArraySlotsExt; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; @@ -130,6 +133,18 @@ impl CascadingCompressor { )? .into_array()) } + Canonical::Union(union_array) => { + let type_ids = self.compress(union_array.type_ids(), exec_ctx)?; + let children = union_array + .iter_children() + .map(|child| self.compress(child, exec_ctx)) + .collect::, _>>()?; + + Ok( + UnionArray::try_new(type_ids, union_array.variants().clone(), children)? + .into_array(), + ) + } Canonical::List(list_view_array) => { if list_view_array.is_zero_copy_to_list() || list_view_array.elements().is_empty() { let list_array = list_from_list_view(list_view_array, exec_ctx)?; @@ -138,6 +153,7 @@ impl CascadingCompressor { self.compress_list_view_array(list_view_array, compress_ctx, exec_ctx) } } + Canonical::Map(map_array) => self.compress_map_array(map_array, compress_ctx, exec_ctx), Canonical::FixedSizeList(fsl_array) => { let compressed_elems = self.compress(fsl_array.elements(), exec_ctx)?; @@ -159,10 +175,6 @@ impl CascadingCompressor { compress_ctx, exec_ctx, )?; - // TODO(connor): HACK TO SUPPORT L2 DENORMALIZATION!!! - if scheme_compressed.is::() { - return Ok(scheme_compressed); - } // A constant extension array (that might be masked) is already in its terminal // representation, and compressing the storage separately cannot do better. @@ -305,8 +317,7 @@ impl CascadingCompressor { let after_nbytes = compressed.nbytes(); let actual_ratio = (after_nbytes != 0).then(|| before_nbytes as f64 / after_nbytes as f64); - // TODO(connor): HACK TO SUPPORT L2 DENORMALIZATION!!! - let accepted = after_nbytes < before_nbytes || compressed.is::(); + let accepted = after_nbytes < before_nbytes; trace::record_winner_compress_result( after_nbytes, diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index a661970950c..219b67e2519 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -65,6 +65,13 @@ impl CascadingCompressor { root_exclusions, } } + + /// Returns whether the compressor was configured with `scheme`. + pub fn has_scheme(&self, scheme: SchemeId) -> bool { + self.schemes + .iter() + .any(|candidate| candidate.id() == scheme) + } } // NB: Cascading compression logic is located in `vortex-compressor/src/compressor/cascade.rs`. diff --git a/vortex-compressor/src/compressor/structural.rs b/vortex-compressor/src/compressor/structural.rs index ec867c1c1dc..e0d9462ede1 100644 --- a/vortex-compressor/src/compressor/structural.rs +++ b/vortex-compressor/src/compressor/structural.rs @@ -9,12 +9,16 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ListArray; +use vortex_array::arrays::ListView; use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::MapArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::list::ListArrayExt; -use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::list::ListArraySlotsExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_error::VortexResult; +use vortex_error::vortex_err; use super::ROOT_SCHEME_ID; use crate::CascadingCompressor; @@ -105,6 +109,23 @@ impl CascadingCompressor { .into_array()) } + /// Compresses a [`MapArray`] by recursively compressing its entries [`ListViewArray`]. + pub(super) fn compress_map_array( + &self, + map_array: MapArray, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let data_parts = map_array.into_data_parts(); + let entries = self.compress_list_view_array(data_parts.entries, compress_ctx, exec_ctx)?; + let entries = entries + .as_opt::() + .ok_or_else(|| vortex_err!("Compressed map entries became {}", entries.encoding_id()))? + .into_owned(); + + Ok(MapArray::try_new(data_parts.map_dtype, entries)?.into_array()) + } + /// Compress very child slot of the array, then re-build it from them. pub(super) fn compress_physical_slots( &self, diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index 08eafa369de..3a2a6281047 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -4,6 +4,7 @@ use std::sync::LazyLock; use parking_lot::Mutex; +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -11,8 +12,16 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::Constant; +use vortex_array::arrays::Map; use vortex_array::arrays::NullArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::builders::MapBuilder; +use vortex_array::dtype::DType; +use vortex_array::dtype::MapDType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; @@ -63,6 +72,10 @@ impl Scheme for DirectRatioScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -95,6 +108,10 @@ impl Scheme for ImmediateAlwaysUseScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -127,6 +144,10 @@ impl Scheme for CallbackAlwaysUseScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -161,6 +182,10 @@ impl Scheme for CallbackSkipScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -195,6 +220,10 @@ impl Scheme for CallbackRatioScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -229,6 +258,10 @@ impl Scheme for HugeRatioScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -261,6 +294,10 @@ impl Scheme for ZeroBytesSamplingScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -470,6 +507,10 @@ impl Scheme for ThresholdObservingScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -507,6 +548,10 @@ impl Scheme for CallbackMatchingRatioScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -704,3 +749,99 @@ fn sampling_uses_scheme_stats_options() -> VortexResult<()> { assert!(matches!(score, EstimateScore::FiniteCompression(ratio) if ratio.is_finite())); Ok(()) } + +type MapEntryFixture<'a> = (i32, Option<&'a str>); +type MapRowFixture<'a> = Option>>; + +fn map_array_from_rows(rows: &[MapRowFixture<'_>], keys_sorted: bool) -> VortexResult { + let map_dtype = MapDType::try_new( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + keys_sorted, + )?; + let dtype = DType::Map(map_dtype.clone(), Nullability::Nullable); + let mut builder = MapBuilder::::with_capacity_in( + map_dtype, + Nullability::Nullable, + rows.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + + for row in rows { + let scalar = match row { + Some(entries) => { + let entries = entries + .iter() + .map(|(key, value)| { + let key = Scalar::primitive(*key, Nullability::NonNullable); + let value = value.map_or_else( + || Scalar::null(DType::Utf8(Nullability::Nullable)), + |value| Scalar::utf8(value, Nullability::Nullable), + ); + (key, value) + }) + .collect::>(); + Scalar::try_map(dtype.clone(), entries)? + } + None => Scalar::null(dtype.clone()), + }; + builder.append_value(scalar.as_map())?; + } + + Ok(builder.finish_into_map().into_array()) +} + +#[test] +fn map_compression_preserves_mixed_rows() -> VortexResult<()> { + let array = map_array_from_rows( + &[ + Some(vec![(1, Some("one")), (2, None)]), + None, + Some(vec![]), + Some(vec![(1, Some("dup-old")), (1, Some("dup-new"))]), + ], + false, + )?; + let mut exec_ctx = SESSION.create_execution_ctx(); + + let compressed = compressor().compress(&array, &mut exec_ctx)?; + + assert!(compressed.is::()); + assert_eq!(compressed.dtype(), array.dtype()); + assert_arrays_eq!(&compressed, &array, &mut exec_ctx); + Ok(()) +} + +#[test] +fn all_null_map_compression_preserves_values() -> VortexResult<()> { + let array = map_array_from_rows(&[None, None, None], false)?; + let mut exec_ctx = SESSION.create_execution_ctx(); + + let compressed = compressor().compress(&array, &mut exec_ctx)?; + + assert_eq!(compressed.dtype(), array.dtype()); + assert_arrays_eq!(&compressed, &array, &mut exec_ctx); + Ok(()) +} + +#[test] +fn map_compression_preserves_repeated_entry_children() -> VortexResult<()> { + let rows = (0..64) + .map(|idx| { + Some(vec![ + (idx % 4, Some("alpha")), + (idx % 4, Some("beta")), + (idx % 4, Some("alpha")), + ]) + }) + .collect::>(); + let array = map_array_from_rows(&rows, true)?; + let mut exec_ctx = SESSION.create_execution_ctx(); + + let compressed = compressor().compress(&array, &mut exec_ctx)?; + + assert!(compressed.is::()); + assert_eq!(compressed.dtype(), array.dtype()); + assert_arrays_eq!(&compressed, &array, &mut exec_ctx); + Ok(()) +} diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index 341acc8acbf..de9e67690d4 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -23,6 +23,7 @@ pub use estimate::EstimateVerdict; pub use exclusion::AncestorExclusion; pub use exclusion::ChildSelection; pub use exclusion::DescendantExclusion; +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -123,6 +124,13 @@ pub trait Scheme: Debug + Send + Sync { /// Whether this scheme can compress the given canonical array. fn matches(&self, canonical: &Canonical) -> bool; + /// The array encodings this scheme itself may introduce into its compressed output. + /// + /// Cascaded children are compressed by other schemes, which declare their own encodings, + /// so only encodings constructed directly by [`compress`](Scheme::compress) belong here. + /// Canonical arrays the scheme merely rearranges do not need to be declared. + fn produced_encodings(&self) -> Vec; + /// Returns the stats generation options this scheme requires. The compressor merges all /// eligible schemes' options before generating stats so that a single stats pass satisfies /// every scheme. diff --git a/vortex-compressor/src/stats/integer.rs b/vortex-compressor/src/stats/integer.rs index a399c1060f9..cded1a76aa4 100644 --- a/vortex-compressor/src/stats/integer.rs +++ b/vortex-compressor/src/stats/integer.rs @@ -381,17 +381,12 @@ where }; let sliced = buffer.slice(head_idx..array.len()); - let mut chunks = sliced.as_slice().chunks_exact(64); + let (chunks, remainder) = sliced.as_slice().as_chunks::<64>(); match validity.bit_buffer() { AllOr::All => { - for chunk in &mut chunks { - inner_loop_nonnull( - chunk.try_into().ok().vortex_expect("chunk size must be 64"), - count_distinct_values, - &mut loop_state, - ) + for chunk in chunks { + inner_loop_nonnull(chunk, count_distinct_values, &mut loop_state) } - let remainder = chunks.remainder(); inner_loop_naive( remainder, count_distinct_values, @@ -403,7 +398,7 @@ where AllOr::Some(v) => { let mask = v.slice(head_idx..array.len()); let mut offset = 0; - for chunk in &mut chunks { + for chunk in chunks { let validity = mask.slice(offset..(offset + 64)); offset += 64; @@ -411,14 +406,10 @@ where // All nulls -> no stats to update. 0 => continue, // Inner loop for when validity check can be elided. - 64 => inner_loop_nonnull( - chunk.try_into().ok().vortex_expect("chunk size must be 64"), - count_distinct_values, - &mut loop_state, - ), + 64 => inner_loop_nonnull(chunk, count_distinct_values, &mut loop_state), // Inner loop for when we need to check validity. _ => inner_loop_nullable( - chunk.try_into().ok().vortex_expect("chunk size must be 64"), + chunk, count_distinct_values, &validity, &mut loop_state, @@ -426,7 +417,6 @@ where } } // Final iteration, run naive loop. - let remainder = chunks.remainder(); inner_loop_naive( remainder, count_distinct_values, @@ -489,6 +479,7 @@ struct LoopState { } /// Inner loop for non-null chunks of 64 values. +#[allow(clippy::inline_always)] #[inline(always)] fn inner_loop_nonnull( values: &[T; 64], @@ -510,6 +501,7 @@ fn inner_loop_nonnull( } /// Inner loop for nullable chunks of 64 values. +#[allow(clippy::inline_always)] #[inline(always)] fn inner_loop_nullable( values: &[T; 64], @@ -534,6 +526,7 @@ fn inner_loop_nullable( } /// Fallback inner loop for remainder values. +#[allow(clippy::inline_always)] #[inline(always)] fn inner_loop_naive( values: &[T], diff --git a/vortex-compute/Cargo.toml b/vortex-compute/Cargo.toml index f223d455395..f8b710aacc5 100644 --- a/vortex-compute/Cargo.toml +++ b/vortex-compute/Cargo.toml @@ -28,6 +28,7 @@ arrow-schema = { workspace = true } divan = { workspace = true } num-traits = { workspace = true } rand = { workspace = true } +vortex-bench-support = { workspace = true } [lints] workspace = true diff --git a/vortex-compute/benches/lane_kernels.rs b/vortex-compute/benches/lane_kernels.rs index af8e27d6199..2e9e145add7 100644 --- a/vortex-compute/benches/lane_kernels.rs +++ b/vortex-compute/benches/lane_kernels.rs @@ -12,6 +12,11 @@ //! Each Vortex kernel bench has a sibling `arrow_*` baseline bench using the //! equivalent arrow-rs kernel over the same data shape, so the divan report //! lines up side-by-side. +//! +//! The checked-add pair carries `#[cpu_features]`, so both are measured on every +//! walltime CPU-feature leg rather than in simulation: they are one lane loop +//! compiled differently per leg, and comparing them against arrow-rs is only +//! meaningful under the same build flags. The cast benches stay in simulation. #![expect(clippy::unwrap_used)] #![expect(clippy::clone_on_ref_ptr)] @@ -49,7 +54,7 @@ fn main() { divan::main(); } -const SIZES: &[usize] = &[65_536]; +const SIZES: &[usize] = &[16_384]; // ----------------------------------------------------------------------------- // Cast fixture (u64/u16/i32 lanes + a single validity mask). @@ -324,6 +329,7 @@ fn add_fixture(n: usize) -> AddFixture { } } +#[vortex_bench_support::cpu_features] #[divan::bench(args = SIZES)] fn lanezip_checked_add_u32(bencher: Bencher, n: usize) { let f = add_fixture(n); @@ -346,6 +352,7 @@ fn lanezip_checked_add_u32(bencher: Bencher, n: usize) { }); } +#[vortex_bench_support::cpu_features] #[divan::bench(args = SIZES)] fn arrow_checked_add_u32(bencher: Bencher, n: usize) { let f = add_fixture(n); diff --git a/vortex-compute/src/lane_kernels/map_in_place.rs b/vortex-compute/src/lane_kernels/map_in_place.rs index 335e1c456df..7ca58ac1609 100644 --- a/vortex-compute/src/lane_kernels/map_in_place.rs +++ b/vortex-compute/src/lane_kernels/map_in_place.rs @@ -37,6 +37,7 @@ pub trait IndexedSinkExt: IndexedSink + Sized { where F: FnMut(Self::Item) -> Self::Write, { + #[allow(clippy::inline_always)] #[inline(always)] fn chunk(values: &mut S, f: &mut F, base: usize, count: usize) where @@ -88,6 +89,7 @@ pub trait IndexedSinkExt: IndexedSink + Sized { Self::Write: Default, F: FnMut(Self::Item) -> Option, { + #[allow(clippy::inline_always)] #[inline(always)] fn chunk(values: &mut S, base: usize, count: usize, f: &mut F) -> Option where @@ -162,6 +164,7 @@ pub trait IndexedSinkExt: IndexedSink + Sized { /// Bit-pack `is_none()` flags per lane, then AND with `src_chunk` post-loop to /// drop null-lane failures. The per-lane attribution work is `OR + shift` /// (no `min`/`csel`), giving LLVM more freedom to vectorize the value pipeline. + #[allow(clippy::inline_always)] #[inline(always)] fn chunk( values: &mut S, diff --git a/vortex-compute/src/lane_kernels/map_into.rs b/vortex-compute/src/lane_kernels/map_into.rs index 258913e9fc7..ee8c10b1477 100644 --- a/vortex-compute/src/lane_kernels/map_into.rs +++ b/vortex-compute/src/lane_kernels/map_into.rs @@ -5,6 +5,7 @@ //! caller-provided `&mut [MaybeUninit]`. use std::mem::MaybeUninit; +use std::ops::BitOrAssign; use vortex_buffer::BitBuffer; @@ -17,6 +18,9 @@ use crate::lane_kernels::source::IndexedSource; /// `impl IndexedSourceExt for S` below. Bring the trait into /// scope (`use vortex_compute::lane_kernels::IndexedSourceExt;`) to call /// them with method syntax: `values.try_map_masked_into(&mask, &mut out, f)`. +/// +/// Callbacks implement [`Fn`] because each lane must be independent. A callback that mutates +/// captured state introduces a loop-carried dependency that can prevent vectorization. pub trait IndexedSourceExt: IndexedSource + Sized { /// Fallible map with mask-aware error attribution. `f` returns `Option`; /// `None` indicates a per-lane failure (e.g. range overflow on a narrowing cast). @@ -27,7 +31,7 @@ pub trait IndexedSourceExt: IndexedSource + Sized { /// `is_none()` flags are bit-packed into a `u64` at the lane's position, then /// AND-combined with the chunk's validity bitmap — null-lane bits vanish. /// - /// The closure shape is the same as [`try_map_into`] (`FnMut(Item) -> Option`); + /// The closure shape is the same as [`try_map_into`] (`Fn(Item) -> Option`); /// the mask parameter is what makes this kernel mask-aware. Callers that need to /// distinguish null lanes inside the closure (e.g. to short-circuit an expensive /// computation) should construct their own per-lane validity check externally; for @@ -47,17 +51,18 @@ pub trait IndexedSourceExt: IndexedSource + Sized { self, mask: &BitBuffer, out: &mut [MaybeUninit], - mut f: F, + f: F, ) -> Result<(), usize> where R: Copy + Default, - F: FnMut(Self::Item) -> Option, + F: Fn(Self::Item) -> Option, { + #[allow(clippy::inline_always)] #[inline(always)] fn chunk( values: &S, out: &mut [MaybeUninit], - f: &mut F, + f: &F, src_chunk: u64, base: usize, count: usize, @@ -65,7 +70,7 @@ pub trait IndexedSourceExt: IndexedSource + Sized { where S: IndexedSource, R: Copy + Default, - F: FnMut(S::Item) -> Option, + F: Fn(S::Item) -> Option, { let mut fail_bits: u64 = 0; for bit_idx in 0..count { @@ -91,7 +96,7 @@ pub trait IndexedSourceExt: IndexedSource + Sized { let remainder = len % 64; for (chunk_idx, src_chunk) in chunks.iter().enumerate() { - if let Some(idx) = chunk(&values, out, &mut f, src_chunk, chunk_idx * 64, 64) { + if let Some(idx) = chunk(&values, out, &f, src_chunk, chunk_idx * 64, 64) { return Err(idx); } } @@ -99,7 +104,7 @@ pub trait IndexedSourceExt: IndexedSource + Sized { && let Some(idx) = chunk( &values, out, - &mut f, + &f, chunks.remainder_bits(), chunks_count * 64, remainder, @@ -118,20 +123,16 @@ pub trait IndexedSourceExt: IndexedSource + Sized { /// /// Panics if `out.len() != self.len()`. #[inline] - fn map_into(self, out: &mut [MaybeUninit], mut f: F) + fn map_into(self, out: &mut [MaybeUninit], f: F) where - F: FnMut(Self::Item) -> R, + F: Fn(Self::Item) -> R, { + #[allow(clippy::inline_always)] #[inline(always)] - fn chunk( - values: &S, - out: &mut [MaybeUninit], - f: &mut F, - base: usize, - count: usize, - ) where + fn chunk(values: &S, out: &mut [MaybeUninit], f: &F, base: usize, count: usize) + where S: IndexedSource, - F: FnMut(S::Item) -> R, + F: Fn(S::Item) -> R, { for bit_idx in 0..count { let idx = base + bit_idx; @@ -149,10 +150,10 @@ pub trait IndexedSourceExt: IndexedSource + Sized { let remainder = len % CHUNK_LEN; for chunk_idx in 0..chunks_count { - chunk(&values, out, &mut f, chunk_idx * CHUNK_LEN, CHUNK_LEN); + chunk(&values, out, &f, chunk_idx * CHUNK_LEN, CHUNK_LEN); } if remainder != 0 { - chunk(&values, out, &mut f, chunks_count * CHUNK_LEN, remainder); + chunk(&values, out, &f, chunks_count * CHUNK_LEN, remainder); } } @@ -178,15 +179,16 @@ pub trait IndexedSourceExt: IndexedSource + Sized { /// /// Panics if `words.len() < self.len().div_ceil(64)`. #[inline] - fn map_bits_into(self, words: &mut [u64], mut f: F) + fn map_bits_into(self, words: &mut [u64], f: F) where - F: FnMut(Self::Item) -> bool, + F: Fn(Self::Item) -> bool, { + #[allow(clippy::inline_always)] #[inline(always)] - fn chunk(values: &S, f: &mut F, base: usize, count: usize) -> u64 + fn chunk(values: &S, f: &F, base: usize, count: usize) -> u64 where S: IndexedSource, - F: FnMut(S::Item) -> bool, + F: Fn(S::Item) -> bool, { let mut packed: u64 = 0; for bit_idx in 0..count { @@ -210,13 +212,65 @@ pub trait IndexedSourceExt: IndexedSource + Sized { let remainder = len % 64; for word_idx in 0..full { - words[word_idx] = chunk(&values, &mut f, word_idx * 64, 64); + words[word_idx] = chunk(&values, &f, word_idx * 64, 64); } if remainder != 0 { - words[full] = chunk(&values, &mut f, full * 64, remainder); + words[full] = chunk(&values, &f, full * 64, remainder); } } + /// Split value/failure map with **no validity awareness at all**: write every lane's value + /// unconditionally and OR-reduce its failure evidence into the return. + /// + /// The fastest checked shape, running at the speed of the unchecked [`map_into`] in exchange + /// for reporting only _that_ some lane failed and never exiting early. Re-run the now known + /// cold input through [`try_map_into`] or [`try_map_masked_into`] to attribute the failure or + /// to drop the null-lane ones. The evidence reduces inside the kernel because a captured `&mut` + /// becomes a loop-carried memory dependence that blocks vectorization. + /// + /// Anything other than [`Default`] means failure, and `bool` is the ordinary `Fail`. Wider + /// words exist for operations where deriving a `bool` costs the vectorization it guards. + /// **`Fail` must be no wider than `R`**, asserted below, or the reduction rather than the + /// operation decides how many lanes fit in a vector. + /// + /// [`map_into`]: IndexedSourceExt::map_into + /// [`try_map_into`]: IndexedSourceExt::try_map_into + /// [`try_map_masked_into`]: IndexedSourceExt::try_map_masked_into + /// + /// # Panics + /// + /// Panics if `out.len() != self.len()`. + #[inline] + fn map_checked_into(self, out: &mut [MaybeUninit], apply: Apply) -> Fail + where + Fail: Copy + Default + BitOrAssign, + Apply: Fn(Self::Item) -> (R, Fail), + { + const { + assert!( + size_of::() <= size_of::(), + "failure evidence must be no wider than the value, or it bounds the vector width" + ) + }; + + let values = self; + let len = values.len(); + assert_eq!(out.len(), len, "out must have the same length as values"); + + let mut failed = Fail::default(); + for idx in 0..len { + // SAFETY: idx < len by the loop bound, and out.len() == len. + let val = unsafe { values.get_unchecked(idx) }; + + let (result, failure) = apply(val); + failed |= failure; + + // SAFETY: idx < len == out.len(). + unsafe { out.get_unchecked_mut(idx).write(result) }; + } + failed + } + /// Fallible map with **no validity awareness at all** — every `None` returned /// by the closure is treated as a failure, even at null lanes. /// @@ -240,26 +294,27 @@ pub trait IndexedSourceExt: IndexedSource + Sized { /// /// Panics if `out.len() != self.len()`. #[inline] - fn try_map_into(self, out: &mut [MaybeUninit], mut f: F) -> Result<(), usize> + fn try_map_into(self, out: &mut [MaybeUninit], f: F) -> Result<(), usize> where R: Copy + Default, - F: FnMut(Self::Item) -> Option, + F: Fn(Self::Item) -> Option, { /// Returns `true` if any lane in `[base, base+count)` failed (OR-reduced); /// the cold attribution path is called at the kernel level so it can be /// inlined separately for full vs remainder. + #[allow(clippy::inline_always)] #[inline(always)] fn chunk( values: &S, out: &mut [MaybeUninit], - f: &mut F, + f: &F, base: usize, count: usize, ) -> bool where S: IndexedSource, R: Copy + Default, - F: FnMut(S::Item) -> Option, + F: Fn(S::Item) -> Option, { let mut fail_acc: u64 = 0; for bit_idx in 0..count { @@ -283,14 +338,14 @@ pub trait IndexedSourceExt: IndexedSource + Sized { for chunk_idx in 0..chunks_count { let base = chunk_idx * CHUNK_LEN; - if chunk(&values, out, &mut f, base, CHUNK_LEN) { - return Err(attribute_failure_no_mask(&values, base, CHUNK_LEN, &mut f)); + if chunk(&values, out, &f, base, CHUNK_LEN) { + return Err(attribute_failure_no_mask(&values, base, CHUNK_LEN, &f)); } } if remainder != 0 { let base = chunks_count * CHUNK_LEN; - if chunk(&values, out, &mut f, base, remainder) { - return Err(attribute_failure_no_mask(&values, base, remainder, &mut f)); + if chunk(&values, out, &f, base, remainder) { + return Err(attribute_failure_no_mask(&values, base, remainder, &f)); } } Ok(()) @@ -310,7 +365,7 @@ fn cold_scan( values: &S, base: usize, chunk_len: usize, - mut lane_fails: impl FnMut(usize /* bit_idx */, S::Item) -> bool, + lane_fails: impl Fn(usize /* bit_idx */, S::Item) -> bool, ) -> usize where S: IndexedSource, @@ -329,10 +384,10 @@ where /// Cold attribution for the no-mask variant. Replays `f` over the chunk to find /// the first lane that returns `None`. #[inline] -fn attribute_failure_no_mask(values: &S, base: usize, chunk_len: usize, f: &mut F) -> usize +fn attribute_failure_no_mask(values: &S, base: usize, chunk_len: usize, f: &F) -> usize where S: IndexedSource, - F: FnMut(S::Item) -> Option, + F: Fn(S::Item) -> Option, { cold_scan(values, base, chunk_len, |_bit_idx, val| f(val).is_none()) } @@ -546,6 +601,26 @@ mod tests { assert!(res.is_ok(), "null lane should bypass the range check"); } + #[test] + fn map_checked_into_writes_all_lanes_and_reduces_flag() { + let mut values: Vec = (0..130).collect(); + let mut out = vec![MaybeUninit::::uninit(); 130]; + let failed = values + .as_slice() + .map_checked_into(&mut out, |v| (v as u32, v > u32::MAX as u64)); + assert!(!failed); + assert_eq!(write_t(out), (0..130u32).collect::>()); + + values[77] = (u32::MAX as u64) + 1; + let mut out = vec![MaybeUninit::::uninit(); 130]; + let failed = values + .as_slice() + .map_checked_into(&mut out, |v| (v as u32, v > u32::MAX as u64)); + assert!(failed); + // Failing lanes still write their (wrapped) value. + assert_eq!(write_t(out)[76], 76); + } + #[test] fn map_bits_into_packs_full_and_remainder_words() { let values: Vec = (0..130).collect(); diff --git a/vortex-compute/src/lane_kernels/source.rs b/vortex-compute/src/lane_kernels/source.rs index 44c9cb27096..f578d7a2952 100644 --- a/vortex-compute/src/lane_kernels/source.rs +++ b/vortex-compute/src/lane_kernels/source.rs @@ -10,9 +10,8 @@ /// reads carry no inter-iteration data dependency — the autovectorizer treats each /// lane independently. pub trait IndexedSource { - /// The per-lane item type. Must be `Copy` so the kernels can pass it through - /// the closure by value without extra moves. - type Item: Copy; + /// The per-lane item type passed through the kernel by value. + type Item; /// Logical lane count. fn len(&self) -> usize; /// Returns true when there are no lanes. @@ -55,9 +54,10 @@ impl IndexedSource for &mut [T] { /// Pair of two [`IndexedSource`]s of equal length. Yields `(A::Item, B::Item)` per lane. /// -/// Use this to drive a binary kernel from two columns. Length equality is enforced -/// at construction. -pub struct LaneZip(pub A, pub B); +/// Use this to drive a binary kernel from two columns. Length equality is enforced at +/// construction, and the private fields prevent callers from bypassing that check. +#[derive(Clone, Copy)] +pub struct LaneZip(A, B); impl LaneZip { /// Build a `LaneZip` from two equal-length sources. @@ -79,7 +79,6 @@ impl IndexedSource for LaneZip { type Item = (A::Item, B::Item); #[inline] fn len(&self) -> usize { - debug_assert_eq!(self.0.len(), self.1.len()); self.0.len() } #[inline] @@ -88,3 +87,14 @@ impl IndexedSource for LaneZip { unsafe { (self.0.get_unchecked(i), self.1.get_unchecked(i)) } } } + +#[cfg(test)] +mod tests { + use super::LaneZip; + + #[test] + #[should_panic(expected = "LaneZip operands must have the same length")] + fn rejects_mismatched_lengths() { + _ = LaneZip::new(&[1_u8][..], &[2_u8, 3][..]); + } +} diff --git a/vortex-cuda/Cargo.toml b/vortex-cuda/Cargo.toml index be6faa77e7b..d5237fd1f49 100644 --- a/vortex-cuda/Cargo.toml +++ b/vortex-cuda/Cargo.toml @@ -20,7 +20,6 @@ workspace = true [features] default = [] _test-harness = [] -unstable_encodings = ["vortex/unstable_encodings"] [dependencies] arc-swap = { workspace = true } @@ -43,6 +42,10 @@ vortex-cub = { path = "cub" } vortex-cuda-macros = { workspace = true } vortex-error = { workspace = true, features = ["object_store"] } vortex-nvcomp = { path = "nvcomp" } +vortex-onpair = { workspace = true } + +[target.'cfg(target_os = "linux")'.dependencies] +rustix = { workspace = true } [dev-dependencies] criterion = { workspace = true } @@ -61,6 +64,10 @@ fastlanes = { workspace = true } name = "for_cuda" harness = false +[[bench]] +name = "delta_cuda" +harness = false + [[bench]] name = "dict_cuda" harness = false @@ -105,6 +112,10 @@ harness = false name = "fsst_cuda" harness = false +[[bench]] +name = "onpair_cuda" +harness = false + [[bench]] name = "list_view_cuda" harness = false diff --git a/vortex-cuda/benches/arrow_binary_cuda.rs b/vortex-cuda/benches/arrow_binary_cuda.rs index 6668abd3866..91d625abf6f 100644 --- a/vortex-cuda/benches/arrow_binary_cuda.rs +++ b/vortex-cuda/benches/arrow_binary_cuda.rs @@ -5,7 +5,6 @@ #![expect(clippy::cast_possible_truncation)] -#[allow(dead_code)] mod bench_config; mod timed_launch_strategy; diff --git a/vortex-cuda/benches/bench_config/mod.rs b/vortex-cuda/benches/bench_config/mod.rs index 75905fee19d..909018a434f 100644 --- a/vortex-cuda/benches/bench_config/mod.rs +++ b/vortex-cuda/benches/bench_config/mod.rs @@ -9,6 +9,10 @@ use criterion::Criterion; /// /// 100M elements keeps every kernel above ~500 µs, well above the /// ~15 µs CUDA driver noise floor that caused 15-45% swings at 10M. +// Each bench binary includes this module textually, and the string-heavy +// benches (fsst, onpair, arrow_binary, list_view) define their own smaller +// sizes instead of this const — allow it to go unused in those binaries. +#[allow(dead_code)] pub const BENCH_SIZES: &[(usize, &str)] = &[(100_000_000, "100M")]; /// Returns a [`Criterion`] configuration tuned for CUDA benchmarks. diff --git a/vortex-cuda/benches/delta_cuda.rs b/vortex-cuda/benches/delta_cuda.rs new file mode 100644 index 00000000000..1fa72b90d30 --- /dev/null +++ b/vortex-cuda/benches/delta_cuda.rs @@ -0,0 +1,114 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! CUDA benchmarks for FastLanes delta decompression. +//! +//! Each element width is its own case: the lane count is `1024 / bit-width`, so the width sets +//! both how a chunk is split (128 lanes of 8 rows for `u8` through 16 lanes of 64 rows for +//! `u64`) and how much of a block is busy during the scan. + +#![expect(clippy::unwrap_used)] + +mod bench_config; +mod timed_launch_strategy; + +use std::mem::size_of; +use std::sync::Arc; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use criterion::BenchmarkId; +use criterion::Criterion; +use criterion::Throughput; +use cudarc::driver::DeviceRepr; +use futures::executor::block_on; +use vortex::array::IntoArray; +use vortex::array::VortexSessionExecute; +use vortex::array::array_session; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::validity::Validity; +use vortex::buffer::Buffer; +use vortex::dtype::NativePType; +use vortex::encodings::fastlanes::Delta; +use vortex::encodings::fastlanes::DeltaArray; +use vortex::error::VortexExpect; +use vortex_cuda::CudaDispatchMode; +use vortex_cuda::CudaSession; +use vortex_cuda::executor::CudaArrayExt; +use vortex_cuda_macros::cuda_available; +use vortex_cuda_macros::cuda_not_available; + +use crate::bench_config::BENCH_SIZES; +use crate::timed_launch_strategy::TimedLaunchStrategy; + +/// Builds a delta-encoded array of `len` values that stay inside `T`. +fn make_delta_array(len: usize) -> DeltaArray +where + T: NativePType + From, +{ + // A small repeating step keeps every delta narrow, which is the shape delta is chosen for. + let data: Vec = (0..len) + .map(|i| >::from(u8::try_from(i % 251).vortex_expect("modulo fits u8"))) + .collect(); + let primitive = PrimitiveArray::new(Buffer::from(data), Validity::NonNullable); + + let mut ctx = array_session().create_execution_ctx(); + Delta::try_from_primitive_array(&primitive, &mut ctx).vortex_expect("failed to delta encode") +} + +fn benchmark_delta_typed(c: &mut Criterion, type_name: &str) +where + T: NativePType + DeviceRepr + From, +{ + let mut group = c.benchmark_group("cuda"); + + for &(len, len_str) in BENCH_SIZES { + group.throughput(Throughput::Bytes((len * size_of::()) as u64)); + + let delta = make_delta_array::(len); + + group.bench_with_input( + BenchmarkId::new(format!("cuda/delta/{type_name}"), len_str), + &delta, + |b, delta| { + b.iter_custom(|iters| { + let timed = TimedLaunchStrategy::default(); + let timer = timed.timer(); + + let mut cuda_ctx = + CudaSession::create_execution_ctx(&vortex_cuda::cuda_session()) + .vortex_expect("failed to create execution context") + .with_dispatch_mode(CudaDispatchMode::StandaloneOnly) + .with_launch_strategy(Arc::new(timed)); + + for _ in 0..iters { + block_on(delta.clone().into_array().execute_cuda(&mut cuda_ctx)).unwrap(); + } + + Duration::from_nanos(timer.load(Ordering::Relaxed)) + }); + }, + ); + } + + group.finish(); +} + +fn benchmark_delta(c: &mut Criterion) { + benchmark_delta_typed::(c, "u8"); + benchmark_delta_typed::(c, "u16"); + benchmark_delta_typed::(c, "u32"); + benchmark_delta_typed::(c, "u64"); +} + +criterion::criterion_group! { + name = benches; + config = bench_config::cuda_bench_config(); + targets = benchmark_delta +} + +#[cuda_available] +criterion::criterion_main!(benches); + +#[cuda_not_available] +fn main() {} diff --git a/vortex-cuda/benches/dynamic_dispatch_cuda.rs b/vortex-cuda/benches/dynamic_dispatch_cuda.rs index e4726d7c46d..2d9db5ddbbd 100644 --- a/vortex-cuda/benches/dynamic_dispatch_cuda.rs +++ b/vortex-cuda/benches/dynamic_dispatch_cuda.rs @@ -8,7 +8,6 @@ mod bench_config; use std::f64::consts::PI; -use std::marker::PhantomData; use std::mem::size_of; use std::os::raw::c_void; use std::ptr; @@ -50,13 +49,13 @@ use vortex::encodings::fastlanes::BitPackedArrayExt; use vortex::encodings::fastlanes::BitPackedData; use vortex::encodings::fastlanes::FoR; use vortex::encodings::fastlanes::FoRArrayExt; +use vortex::encodings::fastlanes::FoRArraySlotsExt; use vortex::encodings::fastlanes::FoRData; use vortex::encodings::runend::RunEnd; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_err; use vortex_cuda::CudaBufferExt; -use vortex_cuda::CudaDeviceBuffer; use vortex_cuda::CudaDispatchMode; use vortex_cuda::CudaExecutionCtx; use vortex_cuda::CudaSession; @@ -78,14 +77,12 @@ use crate::bench_config::BENCH_SIZES; fn run_timed( cuda_ctx: &mut CudaExecutionCtx, array_len: usize, - output_buf: &CudaDeviceBuffer, + output_buf: &mut CudaSlice, device_plan: &Arc>, shared_mem_bytes: u32, ) -> VortexResult { let cuda_function = cuda_ctx.load_function("dynamic_dispatch", &[T::PTYPE])?; let array_len_u64 = array_len as u64; - let output_view = output_buf.as_view::(); - let (output_ptr, record_output) = output_view.device_ptr(cuda_ctx.stream()); let (plan_ptr, record_plan) = device_plan.device_ptr(cuda_ctx.stream()); let stream = cuda_ctx.stream(); @@ -98,7 +95,7 @@ fn run_timed( .map_err(|e| vortex_err!("{e:?}"))?; let mut launch_builder = cuda_ctx.stream().launch_builder(&cuda_function); - launch_builder.arg(&output_ptr); + launch_builder.arg(output_buf); launch_builder.arg(&array_len_u64); launch_builder.arg(&plan_ptr); @@ -114,7 +111,7 @@ fn run_timed( .launch(config) .map_err(|e| vortex_err!("kernel launch failed: {e}"))?; } - drop((record_output, record_plan)); + drop(record_plan); let stream = cuda_ctx.stream(); let ctx = stream.context(); @@ -139,9 +136,8 @@ struct BenchRunner { smem_bytes: u32, len: usize, device_plan: Arc>, - output_buf: CudaDeviceBuffer, + output_buf: CudaSlice, _plan_buffers: Vec, - _phantom: PhantomData, } impl BenchRunner { @@ -171,22 +167,19 @@ impl BenchRunner { smem_bytes: shared_mem_bytes, len, device_plan, - output_buf: CudaDeviceBuffer::new( - cuda_ctx - .device_alloc::(len.next_multiple_of(1024)) - .expect("alloc output"), - ), + output_buf: cuda_ctx + .device_alloc::(len.next_multiple_of(1024)) + .expect("alloc output"), _plan_buffers: device_buffers, - _phantom: PhantomData, } } - fn run(&self, cuda_ctx: &mut CudaExecutionCtx) -> Duration { + fn run(&mut self, cuda_ctx: &mut CudaExecutionCtx) -> Duration { cuda_ctx.stream().synchronize().unwrap(); run_timed::( cuda_ctx, self.len, - &self.output_buf, + &mut self.output_buf, &self.device_plan, self.smem_bytes, ) @@ -226,7 +219,7 @@ fn bench_for_bitpacked(c: &mut Criterion) { let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session()).vortex_expect("ctx"); - let bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); + let mut bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); b.iter_custom(|iters| { let mut total_time = Duration::ZERO; @@ -271,7 +264,7 @@ fn bench_dict_bp_codes(c: &mut Criterion) { let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session()).vortex_expect("ctx"); - let bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); + let mut bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); b.iter_custom(|iters| { let mut total_time = Duration::ZERO; @@ -315,7 +308,7 @@ fn bench_runend(c: &mut Criterion) { let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session()).vortex_expect("ctx"); - let bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); + let mut bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); b.iter_custom(|iters| { let mut total_time = Duration::ZERO; @@ -393,7 +386,7 @@ fn bench_dict_bp_codes_alp_for_bp_values_dynanmic_dispatch(c: &mut Criterion) { let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session()).vortex_expect("ctx"); - let bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); + let mut bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); b.iter_custom(|iters| { let mut total_time = Duration::ZERO; @@ -423,6 +416,7 @@ mod standalone { offset_within_chunk: u32, num_patches: u32, n_chunks: u32, + indices_base: u32, } unsafe impl DeviceRepr for NullGpuPatches {} @@ -437,16 +431,17 @@ mod standalone { offset_within_chunk: 0, num_patches: 0, n_chunks: 0, + indices_base: 0, }; } pub(super) struct DictAlpForBitpackedRunner { values_packed: BufferHandle, codes_packed: BufferHandle, - values_for_buf: CudaDeviceBuffer, - values_alp_buf: CudaDeviceBuffer, - codes_buf: CudaDeviceBuffer, - output_buf: CudaDeviceBuffer, + values_for_buf: CudaSlice, + values_alp_buf: CudaSlice, + codes_buf: CudaSlice, + output_buf: CudaSlice, bit_unpack_fn: CudaFunction, alp_fn: CudaFunction, dict_fn: CudaFunction, @@ -484,26 +479,18 @@ mod standalone { .vortex_expect("codes packed"); let values_len = values_bp.len(); - let values_for_buf = CudaDeviceBuffer::new( - cuda_ctx - .device_alloc::(values_len.next_multiple_of(1024)) - .vortex_expect("alloc values for"), - ); - let values_alp_buf = CudaDeviceBuffer::new( - cuda_ctx - .device_alloc::(values_len.next_multiple_of(1024)) - .vortex_expect("alloc values alp"), - ); - let codes_buf = CudaDeviceBuffer::new( - cuda_ctx - .device_alloc::(len.next_multiple_of(1024)) - .vortex_expect("alloc codes"), - ); - let output_buf = CudaDeviceBuffer::new( - cuda_ctx - .device_alloc::(len) - .vortex_expect("alloc output"), - ); + let values_for_buf = cuda_ctx + .device_alloc::(values_len.next_multiple_of(1024)) + .vortex_expect("alloc values for"); + let values_alp_buf = cuda_ctx + .device_alloc::(values_len.next_multiple_of(1024)) + .vortex_expect("alloc values alp"); + let codes_buf = cuda_ctx + .device_alloc::(len.next_multiple_of(1024)) + .vortex_expect("alloc codes"); + let output_buf = cuda_ctx + .device_alloc::(len) + .vortex_expect("alloc output"); cuda_ctx.stream().synchronize().expect("setup sync"); @@ -533,15 +520,11 @@ mod standalone { // Setup owns all allocations and H2D copies; this times only the standalone // kernel sequence needed to mirror the fused dynamic-dispatch plan. - pub(super) fn run(&self, cuda_ctx: &mut CudaExecutionCtx) -> Duration { + pub(super) fn run(&mut self, cuda_ctx: &mut CudaExecutionCtx) -> Duration { cuda_ctx.stream().synchronize().unwrap(); let values_packed_view = self.values_packed.cuda_view::().unwrap(); let codes_packed_view = self.codes_packed.cuda_view::().unwrap(); - let values_for_view = self.values_for_buf.as_view::(); - let values_alp_view = self.values_alp_buf.as_view::(); - let codes_view = self.codes_buf.as_view::(); - let output_view = self.output_buf.as_view::(); let patches = NullGpuPatches::NULL; cuda_ctx.stream().synchronize().unwrap(); @@ -560,7 +543,7 @@ mod standalone { { let mut launch = cuda_ctx.stream().launch_builder(&self.bit_unpack_fn); launch.arg(&values_packed_view); - launch.arg(&values_for_view); + launch.arg(&mut self.values_for_buf); launch.arg(&self.values_reference); launch.arg(&patches); unsafe { @@ -577,8 +560,8 @@ mod standalone { { let mut launch = cuda_ctx.stream().launch_builder(&self.alp_fn); let values_len = self.values_len as u64; - launch.arg(&values_for_view); - launch.arg(&values_alp_view); + launch.arg(&self.values_for_buf); + launch.arg(&mut self.values_alp_buf); launch.arg(&self.alp_f); launch.arg(&self.alp_e); launch.arg(&values_len); @@ -598,7 +581,7 @@ mod standalone { let mut launch = cuda_ctx.stream().launch_builder(&self.bit_unpack_fn); let reference = 0u32; launch.arg(&codes_packed_view); - launch.arg(&codes_view); + launch.arg(&mut self.codes_buf); launch.arg(&reference); launch.arg(&patches); unsafe { @@ -615,10 +598,10 @@ mod standalone { { let mut launch = cuda_ctx.stream().launch_builder(&self.dict_fn); let codes_len = self.codes_len as u64; - launch.arg(&codes_view); + launch.arg(&self.codes_buf); launch.arg(&codes_len); - launch.arg(&values_alp_view); - launch.arg(&output_view); + launch.arg(&self.values_alp_buf); + launch.arg(&mut self.output_buf); unsafe { launch .launch(LaunchConfig { @@ -710,7 +693,7 @@ fn bench_dict_bp_codes_alp_for_bp_values_composed_standalone(c: &mut Criterion) let mut cuda_ctx = CudaSession::create_execution_ctx(&session) .vortex_expect("ctx") .with_dispatch_mode(CudaDispatchMode::StandaloneOnly); - let bench_runner = standalone::DictAlpForBitpackedRunner::new( + let mut bench_runner = standalone::DictAlpForBitpackedRunner::new( values_bp, *values_reference, exponents, @@ -794,7 +777,7 @@ fn bench_alp_for_bitpacked_f64(c: &mut Criterion) { let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session()).vortex_expect("ctx"); - let bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); + let mut bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); b.iter_custom(|iters| { let mut total_time = Duration::ZERO; @@ -851,7 +834,7 @@ fn bench_dict_bp_codes_bp_for_values(c: &mut Criterion) { let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session()).vortex_expect("ctx"); - let bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); + let mut bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); b.iter_custom(|iters| { let mut total_time = Duration::ZERO; @@ -917,7 +900,7 @@ fn bench_alp_for_bitpacked(c: &mut Criterion) { let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session()).vortex_expect("ctx"); - let bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); + let mut bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); b.iter_custom(|iters| { let mut total_time = Duration::ZERO; @@ -964,7 +947,7 @@ fn bench_dict_bp_u8_codes_u32_values(c: &mut Criterion) { let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session()).vortex_expect("ctx"); - let bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); + let mut bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); b.iter_custom(|iters| { let mut total_time = Duration::ZERO; @@ -1007,7 +990,7 @@ fn bench_dict_bp_u16_codes_u32_values(c: &mut Criterion) { let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session()).vortex_expect("ctx"); - let bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); + let mut bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); b.iter_custom(|iters| { let mut total_time = Duration::ZERO; @@ -1050,7 +1033,7 @@ fn bench_dict_bp_u32_codes_u32_values(c: &mut Criterion) { let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session()).vortex_expect("ctx"); - let bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); + let mut bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); b.iter_custom(|iters| { let mut total_time = Duration::ZERO; diff --git a/vortex-cuda/benches/fsst_cuda.rs b/vortex-cuda/benches/fsst_cuda.rs index 49930c9a30a..5e2256fbcd9 100644 --- a/vortex-cuda/benches/fsst_cuda.rs +++ b/vortex-cuda/benches/fsst_cuda.rs @@ -5,7 +5,6 @@ #![expect(clippy::unwrap_used)] -#[allow(dead_code)] mod bench_config; mod timed_launch_strategy; @@ -25,6 +24,7 @@ use vortex::dtype::DType; use vortex::dtype::Nullability; use vortex::encodings::fsst::FSST; use vortex::encodings::fsst::FSSTArrayExt; +use vortex::encodings::fsst::FSSTArraySlotsExt; use vortex::error::VortexExpect; use vortex_cuda::CudaDispatchMode; use vortex_cuda::CudaSession; @@ -34,6 +34,7 @@ use vortex_cuda::arrow::release_device_array; use vortex_cuda::executor::CudaArrayExt; use vortex_cuda_macros::cuda_available; use vortex_cuda_macros::cuda_not_available; +use vortex_fsst::FSSTSymbolTable; use vortex_fsst::test_utils::make_fsst_clickbench_urls; use crate::timed_launch_strategy::TimedLaunchStrategy; @@ -72,10 +73,16 @@ fn make_fixture(n: usize) -> FSSTBenchFixture { lens.as_slice::

().iter().map(|x| *x as u64).sum() }); - let binary = FSST::try_new( + let binary = FSST::try_new_with_symbol_table( DType::Binary(Nullability::NonNullable), - fsst.symbols().clone(), - fsst.symbol_lengths().clone(), + Arc::new( + FSSTSymbolTable::new_padded( + fsst.padded_symbols().clone(), + fsst.padded_symbol_lengths().clone(), + fsst.n_symbols(), + ) + .vortex_expect("construction"), + ), fsst.codes(), fsst.uncompressed_lengths().clone(), setup_ctx.execution_ctx(), diff --git a/vortex-cuda/benches/list_view_cuda.rs b/vortex-cuda/benches/list_view_cuda.rs index 4ce365210d6..9fb84f7ff49 100644 --- a/vortex-cuda/benches/list_view_cuda.rs +++ b/vortex-cuda/benches/list_view_cuda.rs @@ -5,7 +5,6 @@ #![expect(clippy::cast_possible_truncation)] -#[allow(dead_code)] mod bench_config; mod timed_launch_strategy; diff --git a/vortex-cuda/benches/onpair_cuda.rs b/vortex-cuda/benches/onpair_cuda.rs new file mode 100644 index 00000000000..0aa716f0c0a --- /dev/null +++ b/vortex-cuda/benches/onpair_cuda.rs @@ -0,0 +1,110 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! CUDA benchmarks for OnPair decompression. + +#![expect(clippy::unwrap_used)] + +mod bench_config; +mod timed_launch_strategy; + +use std::sync::Arc; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use criterion::BenchmarkId; +use criterion::Criterion; +use criterion::Throughput; +use futures::executor::block_on; +use vortex::array::ArrayRef; +use vortex::array::IntoArray; +use vortex::array::arrays::VarBinArray; +use vortex::dtype::DType; +use vortex::dtype::Nullability; +use vortex::error::VortexExpect; +use vortex_cuda::CudaDispatchMode; +use vortex_cuda::CudaSession; +use vortex_cuda::executor::CudaArrayExt; +use vortex_cuda_macros::cuda_available; +use vortex_cuda_macros::cuda_not_available; +use vortex_onpair::DEFAULT_CONFIG; +use vortex_onpair::onpair_compress; + +use crate::timed_launch_strategy::TimedLaunchStrategy; + +// Bench-local size instead of the workspace 100M default: each input is a +// URL-shaped string, much heavier per-element than the fixed-width primitives +// other kernels benchmark. +const BENCH_SIZES: &[(usize, &str)] = &[(10_000_000, "10M")]; + +struct OnPairBenchFixture { + array: ArrayRef, + uncompressed_size: u64, +} + +fn make_fixture(n: usize) -> OnPairBenchFixture { + let mut setup_ctx = CudaSession::create_execution_ctx(&vortex_cuda::cuda_session()) + .vortex_expect("failed to create execution context"); + + let strings: Vec = (0..n) + .map(|i| format!("https://www.example.com/path/{i}/segment?q={}", i % 97)) + .collect(); + let uncompressed_size = strings.iter().map(|s| s.len() as u64).sum(); + let varbin = VarBinArray::from_iter( + strings.iter().map(|s| Some(s.as_str())), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + let array = onpair_compress(&varbin, DEFAULT_CONFIG, setup_ctx.execution_ctx()) + .vortex_expect("OnPair compression failed"); + + OnPairBenchFixture { + array, + uncompressed_size, + } +} + +fn benchmark_onpair_cuda_decompress(c: &mut Criterion) { + let mut group = c.benchmark_group("cuda"); + + for &(n, len_str) in BENCH_SIZES { + let fixture = make_fixture(n); + + group.throughput(Throughput::Bytes(fixture.uncompressed_size)); + group.bench_with_input( + BenchmarkId::new("cuda/onpair/decompress_to_varbinview", len_str), + &fixture.array, + |b, onpair_array| { + b.iter_custom(|iters| { + let timed = TimedLaunchStrategy::default(); + let timer = timed.timer(); + + let mut cuda_ctx = + CudaSession::create_execution_ctx(&vortex_cuda::cuda_session()) + .vortex_expect("failed to create execution context") + .with_dispatch_mode(CudaDispatchMode::StandaloneOnly) + .with_launch_strategy(Arc::new(timed)); + + for _ in 0..iters { + block_on(onpair_array.clone().execute_cuda(&mut cuda_ctx)).unwrap(); + } + Duration::from_nanos(timer.load(Ordering::Relaxed)) + }); + }, + ); + } + + group.finish(); +} + +criterion::criterion_group! { + name = benches; + config = bench_config::cuda_bench_config(); + targets = benchmark_onpair_cuda_decompress +} + +#[cuda_available] +criterion::criterion_main!(benches); + +#[cuda_not_available] +fn main() {} diff --git a/vortex-cuda/benches/zstd_cuda.rs b/vortex-cuda/benches/zstd_cuda.rs index c0f9afaadd1..9b58808c2ad 100644 --- a/vortex-cuda/benches/zstd_cuda.rs +++ b/vortex-cuda/benches/zstd_cuda.rs @@ -82,6 +82,7 @@ async fn execute_zstd_kernel( .record(stream) .map_err(|e| vortex_err!("Failed to record start event: {:?}", e))?; + let (_device_output_ptr, record_device_output) = exec.device_output.device_ptr_mut(stream); let (device_actual_sizes_ptr, record_actual_sizes) = exec.device_actual_sizes.device_ptr_mut(stream); let (nvcomp_temp_buffer_ptr, record_temp) = exec.nvcomp_temp_buffer.device_ptr_mut(stream); @@ -103,7 +104,12 @@ async fn execute_zstd_kernel( ) .map_err(|e| vortex_err!("nvcomp decompress_async failed: {}", e))?; } - drop((record_actual_sizes, record_temp, record_statuses)); + drop(( + record_device_output, + record_actual_sizes, + record_temp, + record_statuses, + )); let end_event = ctx .new_event(Some(CUevent_flags::CU_EVENT_BLOCKING_SYNC)) diff --git a/vortex-cuda/build.rs b/vortex-cuda/build.rs index 0621a2884b3..8880392ed55 100644 --- a/vortex-cuda/build.rs +++ b/vortex-cuda/build.rs @@ -28,19 +28,24 @@ fn main() { // Source directory for kernels (hand-written and generated .cu/.cuh files) let kernels_src = Path::new(&manifest_dir).join("kernels/src"); - // Output directory for compiled .ptx files - separate by profile. - let kernels_gen = Path::new(&manifest_dir).join("kernels/gen").join(&profile); - - std::fs::create_dir_all(&kernels_gen).expect("Failed to create kernels/gen directory"); - - println!("cargo:rerun-if-env-changed=PROFILE"); - - // nvcc availability depends on PATH (e.g. entering/leaving a nix shell that - // provides cuda_nvcc). Without this, cargo caches a stale build.rs result - // from an environment without nvcc, and switching to one with nvcc does not - // trigger PTX recompilation — producing a binary with an empty embedded_ptx - // table that silently falls back to CPU at runtime. - println!("cargo:rerun-if-env-changed=PATH"); + let out_dir = PathBuf::from(env::var("OUT_DIR").expect("OUT_DIR not set")); + // Keep different architecture selections isolated across build trees. + let kernels_gen = out_dir.join("kernels"); + + // NVCC's output depends on its flags and, with native, the visible GPUs. + for name in [ + "PROFILE", + "PATH", + "CUDA_PATH", + "VORTEX_CUDA_ARCH_FLAGS", + "CUDA_VISIBLE_DEVICES", + "CUDA_DEVICE_ORDER", + "NVCC_PREPEND_FLAGS", + "NVCC_APPEND_FLAGS", + "NVCC_CCBIN", + ] { + println!("cargo:rerun-if-env-changed={name}"); + } // Regenerate bit_unpack kernels only when the generator changes println!( @@ -54,18 +59,19 @@ fn main() { generate_unpack::(&kernels_src, 32).expect("Failed to generate unpack for u32"); generate_unpack::(&kernels_src, 16).expect("Failed to generate unpack for u64"); - let out_dir = PathBuf::from(env::var("OUT_DIR").expect("OUT_DIR not set")); generate_arrow_device_array_bindings(Path::new(&manifest_dir), &out_dir); generate_dynamic_dispatch_bindings(&kernels_src, &out_dir); generate_patches_bindings(&kernels_src, &out_dir); - generate_embedded_ptx(&out_dir, &kernels_gen).expect("Failed to generate embedded PTX source"); - if !is_cuda_available() { + // The kernel loader unconditionally includes embedded_kernels.rs, so emit a Rust stub + // even without nvcc, replacing any stale table from a previous CUDA-enabled build. + generate_embedded_kernels(&out_dir, &[]).expect("Failed to generate empty kernel table"); return; } - // Watch and compile .cu and .cuh files from kernels/src to PTX in kernels/gen + std::fs::create_dir_all(&kernels_gen).expect("Failed to create kernel output directory"); + let mut kernel_files = Vec::new(); if let Ok(entries) = std::fs::read_dir(&kernels_src) { for path in entries.flatten().map(|entry| entry.path()) { let is_generated = path @@ -74,12 +80,10 @@ fn main() { .is_some_and(|n| n.starts_with("bit_unpack_")); match path.extension().and_then(|e| e.to_str()) { - Some("cuh") | Some("h") => { - // Only watch hand-written .cuh/.h files, not generated ones - // (generated files are rebuilt when cuda_kernel_generator changes) - if !is_generated { - println!("cargo:rerun-if-changed={}", path.display()); - } + // Only watch hand-written .cuh/.h files, not generated ones + // (generated files are rebuilt when cuda_kernel_generator changes) + Some("cuh") | Some("h") if !is_generated => { + println!("cargo:rerun-if-changed={}", path.display()); } Some("cu") => { // Only watch hand-written .cu files, not generated ones @@ -87,61 +91,45 @@ fn main() { if !is_generated { println!("cargo:rerun-if-changed={}", path.display()); } - // Compile all .cu files to PTX in gen directory - nvcc_compile_ptx(&kernels_src, &kernels_gen, &path, &profile) + let kernel = nvcc_compile_kernel(&kernels_src, &kernels_gen, &path, &profile) .map_err(|e| { format!("Failed to compile CUDA kernel {}: {}", path.display(), e) }) .unwrap(); + kernel_files.push(kernel); } _ => {} } } } - // Refresh embedded_ptx.rs after nvcc has compiled PTX so the binary embeds the latest kernels. - generate_embedded_ptx(&out_dir, &kernels_gen).expect("Failed to generate embedded PTX source"); + kernel_files.sort(); + generate_embedded_kernels(&out_dir, &kernel_files) + .expect("Failed to generate embedded kernels"); } -/// Generates `embedded_ptx.rs`, a module-name lookup table of `include_str!`-embedded PTX. -/// -/// Called before the `nvcc` check to create an empty table for no-toolkit builds, then again after -/// PTX compilation so runtime loading never depends on build-machine filesystem paths. -fn generate_embedded_ptx(out_dir: &Path, kernels_gen: &Path) -> io::Result<()> { - let mut ptx_files = Vec::new(); - if let Ok(entries) = std::fs::read_dir(kernels_gen) { - for entry in entries.flatten() { - let path = entry.path(); - if path.extension().and_then(|e| e.to_str()) == Some("ptx") { - ptx_files.push(path); - } - } - } - ptx_files.sort(); - - let mut file = File::create(out_dir.join("embedded_ptx.rs"))?; - if ptx_files.is_empty() { +/// Embeds only this build's fatbins, excluding stale outputs from previous configurations. +fn generate_embedded_kernels(out_dir: &Path, kernel_files: &[PathBuf]) -> io::Result<()> { + let mut file = File::create(out_dir.join("embedded_kernels.rs"))?; + if kernel_files.is_empty() { writeln!( file, - "pub(crate) fn embedded_ptx(_module_name: &str) -> Option<&'static str> {{" + "pub(crate) fn embedded_kernel(_module_name: &str) -> Option<&'static [u8]> {{ None }}" )?; - writeln!(file, " None")?; - writeln!(file, "}}")?; return Ok(()); } - writeln!( file, - "pub(crate) fn embedded_ptx(module_name: &str) -> Option<&'static str> {{" + "pub(crate) fn embedded_kernel(module_name: &str) -> Option<&'static [u8]> {{" )?; writeln!(file, " match module_name {{")?; - for path in ptx_files { + for path in kernel_files { let Some(module_name) = path.file_stem().and_then(|stem| stem.to_str()) else { continue; }; writeln!( file, - " {module_name:?} => Some(include_str!({:?})),", + " {module_name:?} => Some(include_bytes!({:?})),", path.to_string_lossy() )?; } @@ -167,13 +155,16 @@ fn generate_unpack(output_dir: &Path, thread_count: usize) -> io:: Ok(cu_path) } -fn nvcc_compile_ptx( +fn nvcc_compile_kernel( include_dir: &Path, output_dir: &Path, cu_path: &Path, profile: &str, -) -> io::Result<()> { +) -> io::Result { + let architecture_flags = + env::var("VORTEX_CUDA_ARCH_FLAGS").unwrap_or_else(|_| "-arch=native".to_owned()); let mut cmd = Command::new("nvcc"); + cmd.args(architecture_flags.split_whitespace()); if profile == "debug" { cmd.arg("-O0"); @@ -200,23 +191,21 @@ fn nvcc_compile_ptx( cmd.arg("-O3"); } - // Output PTX file goes to output_dir with same base name - let ptx_path = output_dir + let fatbin_path = output_dir .join(cu_path.file_name().unwrap()) - .with_extension("ptx"); + .with_extension("fatbin"); cmd.arg("-std=c++20") - .arg("-arch=native") // Flags forwarded to Clang. .arg("--compiler-options=-Wall -Wextra -Wpedantic -Werror") .arg("--restrict") - .arg("--ptx") + // Fatbins preserve real/virtual architecture lists; PTX alone cannot. + .arg("--fatbin") .arg("--include-path") .arg(include_dir) - .arg("-c") .arg(cu_path) .arg("-o") - .arg(&ptx_path); + .arg(&fatbin_path); let res = cmd.output()?; @@ -246,7 +235,7 @@ fn nvcc_compile_ptx( cu_path.display() ))); } - Ok(()) + Ok(fatbin_path) } /// Generate bindings for the vendored Arrow C Device ABI header. diff --git a/vortex-cuda/cub/build.rs b/vortex-cuda/cub/build.rs index 5888d9234dd..fc74de26e8f 100644 --- a/vortex-cuda/cub/build.rs +++ b/vortex-cuda/cub/build.rs @@ -17,7 +17,19 @@ use std::path::PathBuf; use std::process::Command; fn main() { - println!("cargo:rerun-if-env-changed=CUDA_PATH"); + // Keep CUB's architecture selection in sync with the embedded kernels. + for name in [ + "CUDA_PATH", + "PATH", + "VORTEX_CUDA_ARCH_FLAGS", + "CUDA_VISIBLE_DEVICES", + "CUDA_DEVICE_ORDER", + "NVCC_PREPEND_FLAGS", + "NVCC_APPEND_FLAGS", + "NVCC_CCBIN", + ] { + println!("cargo:rerun-if-env-changed={name}"); + } let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap()); let kernels_dir = manifest_dir.join("kernels"); @@ -56,8 +68,11 @@ fn is_cuda_available() -> bool { fn compile_shared_library(kernel_dir: &Path, sources: &[PathBuf], out_dir: &Path) { let lib_path = out_dir.join("libvortex_cub.so"); + let architecture_flags = + env::var("VORTEX_CUDA_ARCH_FLAGS").unwrap_or_else(|_| "-arch=native".to_owned()); let mut cmd = Command::new("nvcc"); - cmd.args(["-std=c++20", "-arch=native"]); + cmd.arg("-std=c++20") + .args(architecture_flags.split_whitespace()); if env::var("PROFILE").unwrap() == "debug" { cmd.args(["-O0", "-g", "-G", "-lineinfo"]); @@ -102,6 +117,7 @@ fn generate_rust_bindings(kernels_dir: &Path, out_dir: &Path) { .allowlist_function("filter_bytemask_.*") .allowlist_function("filter_bitmask_.*") .allowlist_function("scan_exclusive_sum_.*") + .allowlist_function("onpair_.*") // Allow CUDA types .allowlist_type("cudaError_t") // Blocklist cudaStream_t and define it manually as an opaque pointer diff --git a/vortex-cuda/cub/kernels/filter.h b/vortex-cuda/cub/kernels/filter.h index c49dc62faed..0fae1b639a8 100644 --- a/vortex-cuda/cub/kernels/filter.h +++ b/vortex-cuda/cub/kernels/filter.h @@ -99,6 +99,26 @@ cudaError_t scan_exclusive_sum_i64(void *d_temp, int64_t num_items, cudaStream_t stream); +// Fused OnPair per-batch offsets regeneration: a single sweep kernel reduces +// each 128-token batch's decoded size and exclusive-scans the sizes in-kernel +// via decoupled look-back, writing `num_batches + 1` offsets whose last +// element is the total decoded byte count. `code_width` selects the code +// stream's element size in bytes (1 or 2). A code outside the dictionary +// raises `*status` to 1 and contributes zero bytes. +cudaError_t onpair_batch_offsets_temp_size(size_t *temp_bytes, int64_t num_batches); + +cudaError_t onpair_batch_offsets(void *d_temp, + size_t temp_bytes, + const void *codes, + uint32_t code_width, + const uint8_t *lens, + uint32_t dict_size, + uint64_t total_tokens, + uint64_t *chunk_offsets, + uint32_t *status, + int64_t num_batches, + cudaStream_t stream); + #ifdef __cplusplus } #endif diff --git a/vortex-cuda/cub/kernels/onpair.cu b/vortex-cuda/cub/kernels/onpair.cu new file mode 100644 index 00000000000..65336ec0886 --- /dev/null +++ b/vortex-cuda/cub/kernels/onpair.cu @@ -0,0 +1,218 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +// OnPair per-batch output-offsets regeneration, fused into a single sweep +// kernel: every warp reduces one 128-token batch's decoded size from the codes +// and the length LUT, the first warp scans the block's batch sums, and CUB's +// decoupled look-back (`ScanTileState` + `TilePrefixCallbackOp` — the same +// machinery `cub::DeviceScan` is built on) resolves the running prefix of all +// preceding tiles in-kernel. The per-batch sizes live only in registers and +// shared memory; only the offsets are written to global memory. +// +// `onpair_batch_offsets` is the sole entry point. The tile-state init launch +// it performs first is CUB's own requirement — `DeviceScan` runs the identical +// init internally before its sweep. + +#include +#include +#include +#include + +namespace { + +// One warp reduces one 128-token batch (4 tokens per lane); a 512-thread block +// covers a tile of 16 batches. The decode kernel shares this geometry. +constexpr uint32_t ONPAIR_TOKENS_PER_BATCH = 128; +constexpr uint32_t ONPAIR_BLOCK_THREADS = 512; +constexpr uint32_t ONPAIR_WARPS_PER_BLOCK = ONPAIR_BLOCK_THREADS / 32; + +using OnPairTileState = cub::ScanTileState; + +// Version-stable u64 addition functor (avoids deprecated cub thread operators). +struct SumU64 { + __host__ __device__ inline uint64_t operator()(uint64_t a, uint64_t b) const { + return a + b; + } +}; + +using OnPairPrefixOp = cub::TilePrefixCallbackOp; + +__global__ void onpair_batch_offsets_init(OnPairTileState tile_state, int num_tiles) { + tile_state.InitializeStatus(num_tiles); +} + +// The fused sweep, instantiated for the two code widths OnPair stores (u16 +// natively, u8 when the compressor narrowed the codes). A code outside the +// dictionary raises `status` to 1 and contributes zero bytes: the host must +// check the flag before trusting the offsets and before launching the decode +// kernel, whose dictionary gathers are unchecked. +template +__global__ +__launch_bounds__(ONPAIR_BLOCK_THREADS) void onpair_batch_offsets_sweep(const CodeT *__restrict codes, + const uint8_t *__restrict lens, + uint32_t dict_size, + uint64_t total_tokens, + uint64_t *__restrict chunk_offsets, + uint32_t *__restrict status, + OnPairTileState tile_state, + int64_t num_batches) { + const uint32_t lane = threadIdx.x & 31; + const uint32_t warp = threadIdx.x >> 5; + const int tile_idx = static_cast(blockIdx.x); + const int64_t batch = (int64_t)tile_idx * ONPAIR_WARPS_PER_BLOCK + (int64_t)warp; + + // Warp-parallel reduction of this batch's decoded size. Code reads are + // lane-consecutive (coalesced); the length LUT is small and + // cache-resident. Batches past the end (last tile) contribute zero. + uint32_t batch_bytes = 0; + if (batch < num_batches) { + const uint64_t base = (uint64_t)batch * (uint64_t)ONPAIR_TOKENS_PER_BATCH; +#pragma unroll + for (uint8_t token = 0; token < 4; ++token) { + const uint64_t i = base + (uint64_t)lane + (uint64_t)(token * 32); + if (i < total_tokens) { + const uint32_t code = (uint32_t)codes[i]; + if (code < dict_size) { + batch_bytes += (uint32_t)lens[code]; + } else { + atomicMax(status, 1u); + } + } + } + } +#pragma unroll + for (uint8_t offset = 16; offset > 0; offset >>= 1) { + batch_bytes += __shfl_down_sync(0xffffffffu, batch_bytes, offset); + } + + __shared__ uint64_t warp_sums[ONPAIR_WARPS_PER_BLOCK]; + __shared__ uint64_t warp_excl[ONPAIR_WARPS_PER_BLOCK]; + __shared__ typename OnPairPrefixOp::TempStorage prefix_storage; + if (lane == 0) { + warp_sums[warp] = (uint64_t)batch_bytes; + } + __syncthreads(); + + // The first warp scans the tile's batch sums and resolves the running + // prefix of all preceding tiles via decoupled look-back. + if (warp == 0) { + const uint64_t lane_sum = (lane < ONPAIR_WARPS_PER_BLOCK) ? warp_sums[lane] : 0; + uint64_t inclusive = lane_sum; +#pragma unroll + for (uint8_t offset = 1; offset < 32; offset <<= 1) { + const uint64_t shifted = __shfl_up_sync(0xffffffffu, inclusive, offset); + if (lane >= offset) { + inclusive += shifted; + } + } + const uint64_t aggregate = __shfl_sync(0xffffffffu, inclusive, 31); + + uint64_t prefix = 0; + if (tile_idx == 0) { + if (lane == 0) { + tile_state.SetInclusive(0, aggregate); + } + } else { + // Collective over the first warp, as BlockScan would invoke it. + // The callback's return value is only defined in lane 0 (its + // look-back window reduction is a WarpReduce); broadcast it, the + // same way BlockScan shares the prefix before applying it. + OnPairPrefixOp prefix_op(tile_state, prefix_storage, SumU64(), tile_idx); + const uint64_t lane0_prefix = prefix_op(aggregate); + prefix = __shfl_sync(0xffffffffu, lane0_prefix, 0); + } + if (lane < ONPAIR_WARPS_PER_BLOCK) { + warp_excl[lane] = prefix + inclusive - lane_sum; + } + } + __syncthreads(); + + if (lane == 0 && batch < num_batches) { + chunk_offsets[batch] = warp_excl[warp]; + // The trailing slot holds the total decoded byte count. + if (batch == num_batches - 1) { + chunk_offsets[num_batches] = warp_excl[warp] + warp_sums[warp]; + } + } +} + +int onpair_num_tiles(int64_t num_batches) { + return static_cast((num_batches + ONPAIR_WARPS_PER_BLOCK - 1) / ONPAIR_WARPS_PER_BLOCK); +} + +} // namespace + +// Query the look-back tile-state storage for `num_batches` batches. +extern "C" cudaError_t onpair_batch_offsets_temp_size(size_t *temp_bytes, int64_t num_batches) { + if (num_batches < 0 || num_batches / ONPAIR_WARPS_PER_BLOCK >= INT_MAX) { + return cudaErrorInvalidValue; + } + return OnPairTileState::AllocationSize(onpair_num_tiles(num_batches), *temp_bytes); +} + +// Regenerate the OnPair decode kernel's per-batch output offsets on `stream` +// in one fused sweep, writing `chunk_offsets[0..num_batches]` where +// `chunk_offsets[b]` is the decoded byte count preceding batch `b` and +// `chunk_offsets[num_batches]` is the total. `code_width` selects the code +// stream's element size in bytes (1 or 2), so narrowed codes never need a +// widening pass. A code outside the dictionary raises `*status` to 1 and +// contributes zero bytes; the caller must check the flag before trusting the +// offsets. +extern "C" cudaError_t onpair_batch_offsets(void *d_temp, + size_t temp_bytes, + const void *codes, + uint32_t code_width, + const uint8_t *lens, + uint32_t dict_size, + uint64_t total_tokens, + uint64_t *chunk_offsets, + uint32_t *status, + int64_t num_batches, + cudaStream_t stream) { + if (num_batches <= 0 || num_batches / ONPAIR_WARPS_PER_BLOCK >= INT_MAX) { + return cudaErrorInvalidValue; + } + const int num_tiles = onpair_num_tiles(num_batches); + + OnPairTileState tile_state; + cudaError_t err = tile_state.Init(num_tiles, d_temp, temp_bytes); + if (err != cudaSuccess) { + return err; + } + + constexpr int INIT_THREADS = 128; + const int init_blocks = (num_tiles + INIT_THREADS - 1) / INIT_THREADS; + onpair_batch_offsets_init<<>>(tile_state, num_tiles); + err = cudaGetLastError(); + if (err != cudaSuccess) { + return err; + } + + switch (code_width) { + case 1: + onpair_batch_offsets_sweep + <<>>(static_cast(codes), + lens, + dict_size, + total_tokens, + chunk_offsets, + status, + tile_state, + num_batches); + break; + case 2: + onpair_batch_offsets_sweep + <<>>(static_cast(codes), + lens, + dict_size, + total_tokens, + chunk_offsets, + status, + tile_state, + num_batches); + break; + default: + return cudaErrorInvalidValue; + } + return cudaGetLastError(); +} diff --git a/vortex-cuda/cub/src/lib.rs b/vortex-cuda/cub/src/lib.rs index ae8e3b52b0f..7221f902e11 100644 --- a/vortex-cuda/cub/src/lib.rs +++ b/vortex-cuda/cub/src/lib.rs @@ -7,13 +7,15 @@ //! which is used for GPU-accelerated filtering in Vortex. //! //! The library is compiled at build time and loaded at runtime via libloading. -//! This avoids link-time dependencies on CUDA. +//! This avoids link-time dependencies on CUDA. Relocated executables can bundle the compiled +//! shared library in the same directory. //! //! # Platform Support //! //! CUB is part of the CUDA Toolkit. This crate requires CUDA (nvcc) to be //! available at build time. +use std::env; use std::path::PathBuf; use std::sync::OnceLock; @@ -23,6 +25,7 @@ pub mod sys; mod error; pub mod filter; +pub mod onpair; pub mod scan; pub use error::CubError; @@ -32,8 +35,13 @@ static CUB_LIB: OnceLock> = OnceLock::new(); fn load_cub() -> Result { let lib_name = "libvortex_cub.so"; - let build_lib_dir = env!("OUT_DIR"); - let lib_path = PathBuf::from(build_lib_dir).join(lib_name); + let bundled_lib_path = env::current_exe() + .ok() + .and_then(|path| path.parent().map(|parent| parent.join(lib_name))); + let build_lib_path = PathBuf::from(env!("OUT_DIR")).join(lib_name); + let lib_path = bundled_lib_path + .filter(|path| path.is_file()) + .unwrap_or(build_lib_path); // SAFETY: The library at the path is a valid vortex_cub shared library // compiled during the build process. diff --git a/vortex-cuda/cub/src/onpair.rs b/vortex-cuda/cub/src/onpair.rs new file mode 100644 index 00000000000..614b69c42b2 --- /dev/null +++ b/vortex-cuda/cub/src/onpair.rs @@ -0,0 +1,71 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Rust wrappers around the fused OnPair per-batch offsets regeneration. + +use std::ffi::c_void; + +use crate::cub_library; +use crate::error::CubError; +use crate::error::check_cuda_error; +pub use crate::sys::cudaStream_t; + +/// Get temporary storage size (the look-back tile state) for +/// [`batch_offsets`]. +pub fn batch_offsets_temp_size(num_batches: i64) -> Result { + let lib = cub_library()?; + let mut temp_bytes: usize = 0; + let err = unsafe { (lib.onpair_batch_offsets_temp_size)(&raw mut temp_bytes, num_batches) }; + check_cuda_error(err, "onpair_batch_offsets_temp_size")?; + Ok(temp_bytes) +} + +/// Regenerate the OnPair decode kernel's per-batch output offsets in one fused +/// sweep: each warp reduces one 128-token batch's decoded size and the +/// exclusive scan over the sizes runs in-kernel via decoupled look-back. +/// Writes `num_batches + 1` offsets; the last is the total decoded byte count. +/// `code_width` selects the code stream's element size in bytes (1 or 2). A +/// code outside the dictionary raises `*status` to 1 and contributes zero +/// bytes. +/// +/// # Safety +/// +/// All device pointers must be valid and properly sized: +/// - `d_temp` must have at least `temp_bytes` bytes allocated. +/// - `codes` must have at least `total_tokens` elements of `code_width` bytes +/// each, with `total_tokens <= num_batches * 128`. +/// - `lens` must have at least `dict_size` bytes. +/// - `chunk_offsets` must have at least `num_batches + 1` `u64` values. +/// - `status` must point to a valid device `u32`. +#[allow(clippy::too_many_arguments)] +pub unsafe fn batch_offsets( + d_temp: *mut c_void, + temp_bytes: usize, + codes: *const c_void, + code_width: u32, + lens: *const u8, + dict_size: u32, + total_tokens: u64, + chunk_offsets: *mut u64, + status: *mut u32, + num_batches: i64, + stream: cudaStream_t, +) -> Result<(), CubError> { + let lib = cub_library()?; + let err = unsafe { + (lib.onpair_batch_offsets)( + d_temp, + temp_bytes, + codes, + code_width, + lens, + dict_size, + total_tokens, + chunk_offsets, + status, + num_batches, + stream, + ) + }; + check_cuda_error(err, "onpair_batch_offsets") +} diff --git a/vortex-cuda/ffi/README.md b/vortex-cuda/ffi/README.md index 1b07e12f0df..d83974921ec 100644 --- a/vortex-cuda/ffi/README.md +++ b/vortex-cuda/ffi/README.md @@ -23,3 +23,7 @@ Push host-resident arrays and close the sink using the standard `vx_array_sink_* Use `vx_cuda_scan_path_arrow_device_stream` to read such a local file through pinned host buffers and receive an Arrow C Device stream. Reuse the same CUDA session across scans so the pinned buffer pool and CUDA state are reused as well. + +On Linux, use `vx_cuda_scan_path_arrow_device_stream_with_options` with +`vx_cuda_scan_options.flags = VX_CUDA_SCAN_FLAG_DIRECT_IO` to bypass the operating system page +cache for pooled data-plane reads. Footer and zone-map reads remain buffered on the host. diff --git a/vortex-cuda/ffi/cinclude/vortex_cuda.h b/vortex-cuda/ffi/cinclude/vortex_cuda.h index 971ad9872f4..3c83b9b9ffd 100644 --- a/vortex-cuda/ffi/cinclude/vortex_cuda.h +++ b/vortex-cuda/ffi/cinclude/vortex_cuda.h @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors #pragma once +#include #include #include "vortex.h" @@ -75,6 +76,40 @@ vx_array_sink *vx_cuda_array_sink_open_file(const vx_session *session, const vx_dtype *dtype, vx_error **error_out); +/** + * Open a CUDA-readable Vortex file sink with a fixed row block size. + * + * `block_rows` controls the row granularity of CUDA-flat data blocks. Passing zero uses the default + * writer strategy: 8,192-row blocks may be coalesced into data blocks targeting 1 MiB. Passing any + * nonzero value disables this byte-size coalescing, so passing 8,192 is not equivalent to passing + * zero. + * + * Write and scan sizing are independent. To align on-disk row blocks with scan batches, pass the + * same nonzero value to this function and `vx_cuda_scan_path_arrow_device_stream_batch_rows`; the + * API does not enforce a match. + */ +vx_array_sink *vx_cuda_array_sink_open_file_block_rows(const vx_session *session, + vx_view path, + const vx_dtype *dtype, + size_t block_rows, + vx_error **error_out); + +/** + * Options for scanning a CUDA-compatible Vortex file. + * + * Zero-initialize this struct to use buffered file I/O and layout-derived batch splitting. + */ +/** Bypass the operating system page cache for pooled data-plane reads. + * Footer and zone-map reads remain buffered. Supported only on Linux. */ +#define VX_CUDA_SCAN_FLAG_DIRECT_IO (UINT32_C(1) << 0) + +typedef struct vx_cuda_scan_options { + /** Bitwise combination of `VX_CUDA_SCAN_FLAG_*` values. */ + uint32_t flags; + /** Number of rows in each output ArrowDeviceArray. Zero uses layout-derived splitting. */ + size_t batch_rows; +} vx_cuda_scan_options; + /** * Scan a local CUDA-compatible Vortex file as an Arrow C Device stream. * @@ -90,6 +125,35 @@ int vx_cuda_scan_path_arrow_device_stream(const vx_session *session, struct ArrowDeviceArrayStream *out_stream, vx_error **error_out); +/** + * Scan a local CUDA-compatible Vortex file with fixed-size row batches. + * + * `batch_rows` controls the number of rows in each output `ArrowDeviceArray`. Pass zero to use the + * layout-derived splitting of `vx_cuda_scan_path_arrow_device_stream`. + * + * Scan and write sizing are independent. To align scan batches with on-disk row blocks, pass the + * same nonzero value to this function and `vx_cuda_array_sink_open_file_block_rows`; the API does + * not enforce a match. + */ +int vx_cuda_scan_path_arrow_device_stream_batch_rows(const vx_session *session, + vx_view path, + size_t batch_rows, + struct ArrowDeviceArrayStream *out_stream, + vx_error **error_out); + +/** + * Scan a local CUDA-compatible Vortex file with explicit options. + * + * This has the same ownership and file compatibility requirements as + * `vx_cuda_scan_path_arrow_device_stream`. Pass NULL or a zero-initialized options struct to use + * buffered file I/O and layout-derived batch splitting. + */ +int vx_cuda_scan_path_arrow_device_stream_with_options(const vx_session *session, + vx_view path, + const vx_cuda_scan_options *options, + struct ArrowDeviceArrayStream *out_stream, + vx_error **error_out); + /** * Export a borrowed Vortex array for cuDF's Arrow Device import path. * diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 12bc584dad6..44fc7d87938 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -15,17 +15,19 @@ use std::sync::Arc; use arrow_schema::ffi::FFI_ArrowSchema; use vortex::array::stream::ArrayStreamExt; use vortex::compressor::BtrBlocksCompressorBuilder; +use vortex::editions::ComponentKind; +use vortex::editions::EditionSessionExt; use vortex::error::VortexResult; use vortex::error::vortex_ensure; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteStrategyBuilder; -use vortex::io::VortexReadAt; use vortex::io::runtime::BlockingRuntime; -use vortex::io::session::RuntimeSessionExt; +use vortex::layout::scan::split_by::SplitBy; use vortex::session::SessionExt; use vortex::session::VortexSession; +use vortex_cuda::CudaOpenOptionsExt; use vortex_cuda::CudaSession; -use vortex_cuda::PooledFileReadAt; +use vortex_cuda::PooledFileReadAtOptions; use vortex_cuda::arrow::ArrowDeviceArray; use vortex_cuda::arrow::ArrowDeviceArrayStream; use vortex_cuda::arrow::DeviceArrayExt; @@ -50,6 +52,22 @@ use vortex_ffi::vx_view; const VX_CUDA_OK: c_int = 0; const VX_CUDA_ERR: c_int = 1; +/// Enable direct I/O for pooled CUDA file reads. +pub const VX_CUDA_SCAN_FLAG_DIRECT_IO: u32 = 1 << 0; +const VX_CUDA_SCAN_KNOWN_FLAGS: u32 = VX_CUDA_SCAN_FLAG_DIRECT_IO; + +/// Options for scanning a CUDA-compatible Vortex file. +/// +/// Zero-initialize this struct to use buffered file I/O and layout-derived batch splitting. +#[repr(C)] +#[derive(Default)] +pub struct vx_cuda_scan_options { + /// A bitwise combination of `VX_CUDA_SCAN_FLAG_*` values. + pub flags: u32, + /// Number of rows in each output batch. Zero uses layout-derived splitting. + pub batch_rows: usize, +} + /// Return a Vortex session with a [`CudaSession`] session variable. /// /// If `session` already has CUDA support, this returns a clone of it. Otherwise it @@ -99,18 +117,63 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file( path: vx_view, dtype: *const vx_dtype, error_out: *mut *mut vx_error, +) -> *mut vx_array_sink { + unsafe { vx_cuda_array_sink_open_file_block_rows(session, path, dtype, 0, error_out) } +} + +/// Open a CUDA-readable Vortex file sink with a fixed row block size. +/// +/// `block_rows` controls the row granularity of CUDA-flat data blocks. Passing zero preserves the +/// default writer strategy used by [`vx_cuda_array_sink_open_file`]. Any nonzero value disables +/// byte-size coalescing so data blocks retain the requested row granularity. +/// +/// Write and scan sizing are independent. To align on-disk row blocks with scan batches, pass the +/// same nonzero value to this function and [`vx_cuda_scan_path_arrow_device_stream_batch_rows`]. +/// +/// # Safety +/// +/// `session`, `path`, and `dtype` must satisfy the same requirements as +/// `vx_array_sink_open_file`. If `error_out` is non-null, it must be valid for writing one error +/// pointer. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file_block_rows( + session: *const vx_session, + path: vx_view, + dtype: *const vx_dtype, + block_rows: usize, + error_out: *mut *mut vx_error, ) -> *mut vx_array_sink { try_or(error_out, ptr::null_mut(), || { - session_with_cuda(unsafe { vx_session_ref(session) }?)?; - let strategy = WriteStrategyBuilder::default() - .with_btrblocks_builder(BtrBlocksCompressorBuilder::default().only_cuda_compatible()) - .with_flat_strategy(Arc::new(CudaFlatLayoutStrategy::default())) - .build(); - unsafe { vx_array_sink_open_file_with_strategy(session, path, dtype, strategy) } + let vortex_session = unsafe { vx_session_ref(session) }?; + session_with_cuda(vortex_session)?; + let allowed_encodings = vortex_session + .enabled_component_ids(ComponentKind::Array) + .into_iter() + .collect(); + let mut strategy = WriteStrategyBuilder::default() + .with_btrblocks_builder( + BtrBlocksCompressorBuilder::default() + .only_cuda_compatible() + .retain_allowed_encodings(&allowed_encodings), + ) + .with_flat_strategy(Arc::new(CudaFlatLayoutStrategy::default())); + if block_rows > 0 { + // The default byte-size target can coalesce several row blocks into one data block. + // A scan using the same row count would then split inside that data block, defeating + // the requested alignment. The explicit block row count already defines the desired + // granularity for this opt-in path, so a separate byte-size target is unnecessary. + strategy = strategy + .with_row_block_size(block_rows) + .with_data_block_target_bytes(None); + } + unsafe { vx_array_sink_open_file_with_strategy(session, path, dtype, strategy.build()) } }) } -/// Scan a local Vortex file through pinned host buffers and export an Arrow C Device stream. +/// Scan a local Vortex file with buffered I/O and export an Arrow C Device stream. +/// +/// Footer and zone-map reads remain on the host. Data segments are staged through pinned host +/// buffers and transferred directly to the GPU. /// /// The file must use encodings and layouts supported by the CUDA execution path, such as files /// written by [`vx_cuda_array_sink_open_file`]. Pinned staging buffers are reused across scans made @@ -134,26 +197,94 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream( path: vx_view, out_stream: *mut ArrowDeviceArrayStream, error_out: *mut *mut vx_error, +) -> c_int { + unsafe { + vx_cuda_scan_path_arrow_device_stream_with_options( + session, + path, + ptr::null(), + out_stream, + error_out, + ) + } +} + +/// Scan a local Vortex file and export an Arrow C Device stream with fixed-size row batches. +/// +/// `batch_rows` controls the number of rows in each output batch. Passing zero preserves the +/// layout-derived splitting used by [`vx_cuda_scan_path_arrow_device_stream`]. +/// +/// Scan and write sizing are independent. To align scan batches with on-disk row blocks, pass the +/// same nonzero value to this function and [`vx_cuda_array_sink_open_file_block_rows`]. +/// +/// # Safety +/// +/// `session` must be a valid borrowed handle created by `vortex-ffi`. `path` must be valid for the +/// duration of this call and contain UTF-8. `out_stream` must be a valid writable pointer. If +/// `error_out` is non-null, it must be valid for writing one error pointer. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_batch_rows( + session: *const vx_session, + path: vx_view, + batch_rows: usize, + out_stream: *mut ArrowDeviceArrayStream, + error_out: *mut *mut vx_error, +) -> c_int { + let options = vx_cuda_scan_options { + batch_rows, + ..Default::default() + }; + unsafe { + vx_cuda_scan_path_arrow_device_stream_with_options( + session, + path, + &raw const options, + out_stream, + error_out, + ) + } +} + +/// Scan a local Vortex file with explicit options and export an Arrow C Device stream. +/// +/// This has the same ownership and file compatibility requirements as +/// [`vx_cuda_scan_path_arrow_device_stream`]. Pass a null `options` pointer or a zero-initialized +/// [`vx_cuda_scan_options`] to use buffered file I/O and layout-derived batch splitting. +/// +/// # Safety +/// +/// `session` must be a valid borrowed handle created by `vortex-ffi`. `path` must be valid for the +/// duration of this call and contain UTF-8. `options`, when non-null, must point to a valid +/// [`vx_cuda_scan_options`]. `out_stream` must be a valid writable pointer. If `error_out` is +/// non-null, it must be valid for writing one error pointer. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_with_options( + session: *const vx_session, + path: vx_view, + options: *const vx_cuda_scan_options, + out_stream: *mut ArrowDeviceArrayStream, + error_out: *mut *mut vx_error, ) -> c_int { try_or(error_out, VX_CUDA_ERR, || { vortex_ensure!(!out_stream.is_null(), "null ArrowDeviceArrayStream output"); let path = unsafe { path.as_str() }?.to_owned(); let session = session_with_cuda(unsafe { vx_session_ref(session) }?)?; - let cuda_session = session.get::(); - let stream = cuda_session.stream()?; - let pool = Arc::clone(cuda_session.pinned_buffer_pool()); - drop(cuda_session); - - let reader: Arc = Arc::new(PooledFileReadAt::open( - path, - session.handle(), - pool, - stream, - )?); + let options = unsafe { scan_options(options) }?; let array_stream = ffi_runtime().block_on(async { - let file = session.open_options().open(reader).await?; - Ok::<_, vortex::error::VortexError>(file.scan()?.into_array_stream()?.boxed()) + let file = session + .open_options() + .with_cuda() + .with_read_at_options(options.read_at_options) + .open_path(path) + .await?; + let scan = file.scan()?; + let scan = if options.batch_rows == 0 { + scan + } else { + scan.with_split_by(SplitBy::RowCount(options.batch_rows)) + }; + Ok::<_, vortex::error::VortexError>(scan.into_array_stream()?.boxed()) })?; let device_stream = array_stream.export_device_array_stream(&session, ffi_runtime())?; @@ -162,6 +293,46 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream( }) } +struct CudaScanOptions { + read_at_options: PooledFileReadAtOptions, + batch_rows: usize, +} + +unsafe fn scan_options(options: *const vx_cuda_scan_options) -> VortexResult { + let (flags, batch_rows) = if options.is_null() { + (0, 0) + } else { + let options = unsafe { &*options }; + (options.flags, options.batch_rows) + }; + vortex_ensure!( + flags & !VX_CUDA_SCAN_KNOWN_FLAGS == 0, + "unsupported CUDA scan option flags: {:#x}", + flags & !VX_CUDA_SCAN_KNOWN_FLAGS + ); + + let read_at_options = PooledFileReadAtOptions::default(); + let read_at_options = if flags & VX_CUDA_SCAN_FLAG_DIRECT_IO == 0 { + read_at_options + } else { + #[cfg(target_os = "linux")] + { + read_at_options.with_direct_io() + } + #[cfg(not(target_os = "linux"))] + { + return Err(vortex::error::vortex_err!( + "direct CUDA file I/O is only supported on Linux" + )); + } + }; + + Ok(CudaScanOptions { + read_at_options, + batch_rows, + }) +} + /// Export a borrowed Vortex array for cuDF's Arrow Device import path. /// /// On success returns `0` and writes independently releasable `out_schema` and `out_array`; the @@ -265,6 +436,49 @@ mod tests { use super::*; + #[test] + fn scan_options_default_to_buffered_io() { + let options = vx_cuda_scan_options::default(); + assert_eq!(options.flags, 0); + assert_eq!(options.batch_rows, 0); + } + + #[test] + fn rejects_unknown_scan_option_flags() { + let options = vx_cuda_scan_options { + flags: 1 << 31, + ..Default::default() + }; + assert!(unsafe { scan_options(&raw const options) }.is_err()); + } + + #[cfg(target_os = "linux")] + #[test] + fn maps_direct_io_scan_option_to_pooled_reader() -> VortexResult<()> { + let options = vx_cuda_scan_options { + flags: VX_CUDA_SCAN_FLAG_DIRECT_IO, + ..Default::default() + }; + assert_eq!( + unsafe { scan_options(&raw const options) }?.read_at_options, + PooledFileReadAtOptions::default().with_direct_io() + ); + Ok(()) + } + + #[test] + fn maps_batch_rows_scan_option() -> VortexResult<()> { + let options = vx_cuda_scan_options { + batch_rows: 8192, + ..Default::default() + }; + assert_eq!( + unsafe { scan_options(&raw const options) }?.batch_rows, + 8192 + ); + Ok(()) + } + fn test_session(session: VortexSession) -> *mut vx_session { Box::into_raw(Box::new(session)).cast::() } diff --git a/vortex-cuda/gpu-scan-cli/Cargo.toml b/vortex-cuda/gpu-scan-cli/Cargo.toml index 9234d303952..468158a35ba 100644 --- a/vortex-cuda/gpu-scan-cli/Cargo.toml +++ b/vortex-cuda/gpu-scan-cli/Cargo.toml @@ -12,6 +12,10 @@ repository = { workspace = true } rust-version = { workspace = true } version = { workspace = true } +[[bin]] +name = "gpu-scan-cli" +test = false + [lints] workspace = true diff --git a/vortex-cuda/gpu-scan-cli/src/main.rs b/vortex-cuda/gpu-scan-cli/src/main.rs index 00f6cd1ac1c..2aa32729ef3 100644 --- a/vortex-cuda/gpu-scan-cli/src/main.rs +++ b/vortex-cuda/gpu-scan-cli/src/main.rs @@ -24,6 +24,8 @@ use vortex::array::arrays::StructArray; use vortex::array::arrays::struct_::StructArrayExt; use vortex::buffer::ByteBufferMut; use vortex::compressor::BtrBlocksCompressorBuilder; +use vortex::editions::ComponentKind; +use vortex::editions::EditionSessionExt; use vortex::error::VortexResult; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; @@ -31,9 +33,9 @@ use vortex::file::WriteStrategyBuilder; use vortex::io::session::RuntimeSessionExt; use vortex::session::SessionExt; use vortex::session::VortexSession; +use vortex_cuda::CudaOpenOptionsExt; use vortex_cuda::CudaSession; use vortex_cuda::PooledByteBufferReadAt; -use vortex_cuda::PooledFileReadAt; use vortex_cuda::TracingLaunchStrategy; use vortex_cuda::executor::CudaArrayExt; use vortex_cuda::layout::CudaFlatLayoutStrategy; @@ -92,9 +94,17 @@ async fn main() -> VortexResult<()> { /// Build the write strategy used for CUDA-compatible file output. #[cuda_available] -fn cuda_write_strategy() -> Arc { +fn cuda_write_strategy(session: &VortexSession) -> Arc { + let allowed_encodings = session + .enabled_component_ids(ComponentKind::Array) + .into_iter() + .collect(); WriteStrategyBuilder::default() - .with_btrblocks_builder(BtrBlocksCompressorBuilder::default().only_cuda_compatible()) + .with_btrblocks_builder( + BtrBlocksCompressorBuilder::default() + .only_cuda_compatible() + .retain_allowed_encodings(&allowed_encodings), + ) .with_flat_strategy(Arc::new(CudaFlatLayoutStrategy::default())) .build() } @@ -111,7 +121,7 @@ async fn cmd_convert(input: PathBuf, output: PathBuf) -> VortexResult<()> { let mut out = tokio::fs::File::create(&output).await?; session .write_options() - .with_strategy(cuda_write_strategy()) + .with_strategy(cuda_write_strategy(&session)) .write(&mut out, scan) .await?; @@ -158,11 +168,9 @@ async fn cmd_scan(path: PathBuf, gpu_file: bool, json_output: bool) -> VortexRes let pool = Arc::clone(cuda_session.pinned_buffer_pool()); let cuda_stream = cuda_session.stream()?; drop(cuda_session); - let handle = session.handle(); let gpu_file_handle = if gpu_file { - let reader = PooledFileReadAt::open(&path, handle, Arc::clone(&pool), cuda_stream)?; - session.open_options().open(Arc::new(reader)).await? + session.open_options().with_cuda().open_path(&path).await? } else { let (recompressed, footer) = recompress_for_gpu(&path, &session).await?; let reader = PooledByteBufferReadAt::new(recompressed, Arc::clone(&pool), cuda_stream); @@ -225,7 +233,7 @@ async fn recompress_for_gpu( let mut out = ByteBufferMut::empty(); let result = session .write_options() - .with_strategy(cuda_write_strategy()) + .with_strategy(cuda_write_strategy(session)) .write(&mut out, scan) .await?; diff --git a/vortex-cuda/kernels/src/arrow_offsets.cu b/vortex-cuda/kernels/src/arrow_offsets.cu new file mode 100644 index 00000000000..79bbb2435c7 --- /dev/null +++ b/vortex-cuda/kernels/src/arrow_offsets.cu @@ -0,0 +1,63 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "config.cuh" + +#include +#include + +// Convert integer lengths to CUB scan input. The final zero makes an exclusive scan produce an +// Arrow offsets array of length len + 1. +#define GENERATE_FROM_LENGTHS_KERNEL(suffix, LengthT, is_signed) \ + extern "C" __global__ void arrow_offsets_from_lengths_##suffix(const LengthT *__restrict lengths, \ + int32_t *__restrict scan, \ + uint32_t *status, \ + uint64_t len) { \ + const uint64_t scan_len = len + 1; \ + const uint64_t elements_per_block = (uint64_t)blockDim.x * ELEMENTS_PER_THREAD; \ + const uint64_t block_start = (uint64_t)blockIdx.x * elements_per_block; \ + const uint64_t block_stop = min(block_start + elements_per_block, scan_len); \ + for (uint64_t idx = block_start + threadIdx.x; idx < block_stop; idx += blockDim.x) { \ + if (idx == len) { \ + scan[idx] = 0; \ + continue; \ + } \ + const LengthT raw_length = lengths[idx]; \ + if (is_signed && raw_length < 0) { \ + scan[idx] = 0; \ + atomicMax(status, 1u); \ + continue; \ + } \ + const uint64_t length = (uint64_t)raw_length; \ + if (length > (uint64_t)INT32_MAX) { \ + scan[idx] = 0; \ + atomicMax(status, 2u); \ + continue; \ + } \ + scan[idx] = (int32_t)length; \ + } \ + } + +// The first i32 prefix-sum overflow must be negative because every scan input is nonnegative and at +// most INT32_MAX. Retaining this validation as a separate pass keeps malformed inputs from reaching +// consumers with wrapped Arrow offsets. +extern "C" __global__ void +arrow_offsets_validate(const int32_t *__restrict offsets, uint32_t *status, uint64_t scan_len) { + const uint64_t elements_per_block = (uint64_t)blockDim.x * ELEMENTS_PER_THREAD; + const uint64_t block_start = (uint64_t)blockIdx.x * elements_per_block; + const uint64_t block_stop = min(block_start + elements_per_block, scan_len); + for (uint64_t idx = block_start + threadIdx.x; idx < block_stop; idx += blockDim.x) { + if (offsets[idx] < 0) { + atomicMax(status, 2u); + } + } +} + +GENERATE_FROM_LENGTHS_KERNEL(i8, int8_t, true) +GENERATE_FROM_LENGTHS_KERNEL(i16, int16_t, true) +GENERATE_FROM_LENGTHS_KERNEL(i32, int32_t, true) +GENERATE_FROM_LENGTHS_KERNEL(i64, int64_t, true) +GENERATE_FROM_LENGTHS_KERNEL(u8, uint8_t, false) +GENERATE_FROM_LENGTHS_KERNEL(u16, uint16_t, false) +GENERATE_FROM_LENGTHS_KERNEL(u32, uint32_t, false) +GENERATE_FROM_LENGTHS_KERNEL(u64, uint64_t, false) diff --git a/vortex-cuda/kernels/src/bit_unpack_16.cu b/vortex-cuda/kernels/src/bit_unpack_16.cu index a784df201d3..43acff3d08d 100644 --- a/vortex-cuda/kernels/src/bit_unpack_16.cu +++ b/vortex-cuda/kernels/src/bit_unpack_16.cu @@ -13,11 +13,16 @@ __device__ void _bit_unpack_16_device(const uint16_t *__restrict in, uint16_t *_ } __syncwarp(); - // Step 2: Apply patches to shared memory in parallel + // Step 2: Apply patches to shared memory in parallel. + // + // Patch values are stored in the same frame-of-reference domain as the packed values, so + // they take the same `+ reference` the lane decoder applies to every unpacked value. For a + // plain bit-packed array the reference is zero and this is a no-op; for FoR-over-bit-packed + // it is what keeps patched positions from coming out short by the reference. PatchesCursor cursor(patches, blockIdx.x, thread_idx, 32); auto patch = cursor.next(); while (patch.index != FL_CHUNK) { - shared_out[patch.index] = patch.value; + shared_out[patch.index] = patch.value + reference; patch = cursor.next(); } __syncwarp(); diff --git a/vortex-cuda/kernels/src/bit_unpack_32.cu b/vortex-cuda/kernels/src/bit_unpack_32.cu index 3f8fcb5c227..b9f29c72c3a 100644 --- a/vortex-cuda/kernels/src/bit_unpack_32.cu +++ b/vortex-cuda/kernels/src/bit_unpack_32.cu @@ -13,11 +13,16 @@ __device__ void _bit_unpack_32_device(const uint32_t *__restrict in, uint32_t *_ } __syncwarp(); - // Step 2: Apply patches to shared memory in parallel + // Step 2: Apply patches to shared memory in parallel. + // + // Patch values are stored in the same frame-of-reference domain as the packed values, so + // they take the same `+ reference` the lane decoder applies to every unpacked value. For a + // plain bit-packed array the reference is zero and this is a no-op; for FoR-over-bit-packed + // it is what keeps patched positions from coming out short by the reference. PatchesCursor cursor(patches, blockIdx.x, thread_idx, 32); auto patch = cursor.next(); while (patch.index != FL_CHUNK) { - shared_out[patch.index] = patch.value; + shared_out[patch.index] = patch.value + reference; patch = cursor.next(); } __syncwarp(); diff --git a/vortex-cuda/kernels/src/bit_unpack_64.cu b/vortex-cuda/kernels/src/bit_unpack_64.cu index ebe0b125369..aebe500c653 100644 --- a/vortex-cuda/kernels/src/bit_unpack_64.cu +++ b/vortex-cuda/kernels/src/bit_unpack_64.cu @@ -13,11 +13,16 @@ __device__ void _bit_unpack_64_device(const uint64_t *__restrict in, uint64_t *_ } __syncwarp(); - // Step 2: Apply patches to shared memory in parallel + // Step 2: Apply patches to shared memory in parallel. + // + // Patch values are stored in the same frame-of-reference domain as the packed values, so + // they take the same `+ reference` the lane decoder applies to every unpacked value. For a + // plain bit-packed array the reference is zero and this is a no-op; for FoR-over-bit-packed + // it is what keeps patched positions from coming out short by the reference. PatchesCursor cursor(patches, blockIdx.x, thread_idx, 16); auto patch = cursor.next(); while (patch.index != FL_CHUNK) { - shared_out[patch.index] = patch.value; + shared_out[patch.index] = patch.value + reference; patch = cursor.next(); } __syncwarp(); diff --git a/vortex-cuda/kernels/src/bit_unpack_8.cu b/vortex-cuda/kernels/src/bit_unpack_8.cu index b2fcfd26f04..cabea862e59 100644 --- a/vortex-cuda/kernels/src/bit_unpack_8.cu +++ b/vortex-cuda/kernels/src/bit_unpack_8.cu @@ -13,11 +13,16 @@ __device__ void _bit_unpack_8_device(const uint8_t *__restrict in, uint8_t *__re } __syncwarp(); - // Step 2: Apply patches to shared memory in parallel + // Step 2: Apply patches to shared memory in parallel. + // + // Patch values are stored in the same frame-of-reference domain as the packed values, so + // they take the same `+ reference` the lane decoder applies to every unpacked value. For a + // plain bit-packed array the reference is zero and this is a no-op; for FoR-over-bit-packed + // it is what keeps patched positions from coming out short by the reference. PatchesCursor cursor(patches, blockIdx.x, thread_idx, 32); auto patch = cursor.next(); while (patch.index != FL_CHUNK) { - shared_out[patch.index] = patch.value; + shared_out[patch.index] = patch.value + reference; patch = cursor.next(); } __syncwarp(); diff --git a/vortex-cuda/kernels/src/date_time_parts.cu b/vortex-cuda/kernels/src/date_time_parts.cu index ccb3e614991..a556ee5c57e 100644 --- a/vortex-cuda/kernels/src/date_time_parts.cu +++ b/vortex-cuda/kernels/src/date_time_parts.cu @@ -43,22 +43,34 @@ __device__ void date_time_parts(const DaysT *__restrict days, } #define EXPAND_DAYS(X) \ + X(u8, uint8_t) \ + X(u16, uint16_t) \ + X(u32, uint32_t) \ + X(u64, uint64_t) \ X(i8, int8_t) \ X(i16, int16_t) \ X(i32, int32_t) \ X(i64, int64_t) #define EXPAND_SUBSECONDS(d, DT, s, ST) \ + GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, u8, uint8_t) \ + GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, u16, uint16_t) \ + GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, u32, uint32_t) \ + GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, u64, uint64_t) \ GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, i8, int8_t) \ GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, i16, int16_t) \ GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, i32, int32_t) \ GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, i64, int64_t) #define EXPAND_SECONDS(d, DT) \ + EXPAND_SUBSECONDS(d, DT, u8, uint8_t) \ + EXPAND_SUBSECONDS(d, DT, u16, uint16_t) \ + EXPAND_SUBSECONDS(d, DT, u32, uint32_t) \ + EXPAND_SUBSECONDS(d, DT, u64, uint64_t) \ EXPAND_SUBSECONDS(d, DT, i8, int8_t) \ EXPAND_SUBSECONDS(d, DT, i16, int16_t) \ EXPAND_SUBSECONDS(d, DT, i32, int32_t) \ EXPAND_SUBSECONDS(d, DT, i64, int64_t) -// Generate all 64 kernels (4³) +// Generate all 512 kernels (8³: every signed and unsigned integer width per component) EXPAND_DAYS(EXPAND_SECONDS) diff --git a/vortex-cuda/kernels/src/delta.cu b/vortex-cuda/kernels/src/delta.cu new file mode 100644 index 00000000000..9e7cd4e74aa --- /dev/null +++ b/vortex-cuda/kernels/src/delta.cu @@ -0,0 +1,78 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "fastlanes_common.cuh" +#include "types.cuh" + +// FastLanes delta decode. +// +// Delta is stored in the FastLanes transposed layout, where a 1024-element chunk is +// FL_LANES independent columns of (1024 / FL_LANES) rows. Each column carries its own +// running total seeded from that lane's base, so the chunk decodes as FL_LANES independent +// sequential scans rather than one scan over 1024 elements. One thread owns one lane. +// +// This mirrors `fastlanes::Delta::undelta` followed by `Transpose::untranspose`, which is what +// the CPU decoder (`delta_decompress`) runs. Both steps must stay in step with that crate. + +/// Maps a position in the transposed layout to its position in natural order. +/// +/// Mirrors `fastlanes::transpose`; `untranspose` is `output[transpose(i)] = input[i]`. +__device__ inline uint32_t fl_transpose_index(uint32_t idx) { + const uint32_t lane = idx % 16; + const uint32_t order = (idx / 16) % 8; + const uint32_t row = idx / 128; + return (lane * 64) + (FL_ORDER[order] * 8) + row; +} + +/// Decodes `num_chunks` full 1024-element delta chunks. +/// +/// `deltas` and `output` hold `num_chunks * FL_CHUNK` elements; `bases` holds +/// `num_chunks * FL_LANES`. Only unsigned types are instantiated: the CPU decoder +/// reinterprets signed input through its unsigned counterpart so that the wrapping add here +/// inverts the wrapping subtract done at compress time. +template +__device__ void delta_decode_kernel(const T *const __restrict deltas, + const T *const __restrict bases, + T *const __restrict output, + uint64_t num_chunks) { + constexpr uint32_t LANES = FL_LANES; + constexpr uint32_t ROWS = FL_CHUNK / LANES; + + // The undelta pass writes the chunk in transposed order; the untranspose pass then reads it + // back in a different order, so the whole chunk is staged in shared memory between them. + __shared__ T transposed[FL_CHUNK]; + + for (uint64_t chunk = blockIdx.x; chunk < num_chunks; chunk += gridDim.x) { + const T *const in = deltas + chunk * FL_CHUNK; + const T *const base = bases + chunk * LANES; + + // Each lane accumulates down its own column, seeded by that lane's base. + for (uint32_t lane = threadIdx.x; lane < LANES; lane += blockDim.x) { + T running = base[lane]; + for (uint32_t row = 0; row < ROWS; ++row) { + const uint32_t idx = INDEX(row, lane); + running = static_cast(running + in[idx]); + transposed[idx] = running; + } + } + __syncthreads(); + + T *const out = output + chunk * FL_CHUNK; + for (uint32_t i = threadIdx.x; i < FL_CHUNK; i += blockDim.x) { + out[fl_transpose_index(i)] = transposed[i]; + } + // Guard the staging buffer before the next chunk overwrites it. + __syncthreads(); + } +} + +#define GENERATE_DELTA_KERNEL(suffix, Type) \ + extern "C" __global__ void delta_##suffix(const Type *const __restrict deltas, \ + const Type *const __restrict bases, \ + Type *const __restrict output, \ + uint64_t num_chunks) { \ + delta_decode_kernel(deltas, bases, output, num_chunks); \ + } + +// Signed input is reinterpreted to its unsigned counterpart before the launch. +FOR_EACH_UNSIGNED_INT(GENERATE_DELTA_KERNEL) diff --git a/vortex-cuda/kernels/src/list.cu b/vortex-cuda/kernels/src/list.cu new file mode 100644 index 00000000000..3e87e267299 --- /dev/null +++ b/vortex-cuda/kernels/src/list.cu @@ -0,0 +1,36 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "config.cuh" +#include "types.cuh" + +// Converts Arrow-style `List` offsets into `ListView` offset/size pairs. +// +// `List` stores `list_len + 1` monotonically increasing offsets; a `ListView` stores one offset +// and one size per list. Both outputs are written by the same thread so the two views of a list +// are always produced together. +template +__device__ void list_views(const OffsetT *const __restrict offsets, + OffsetT *const __restrict out_offsets, + OffsetT *const __restrict out_sizes, + uint64_t list_len) { + const uint32_t elements_per_block = blockDim.x * ELEMENTS_PER_THREAD; + const uint64_t block_start = static_cast(blockIdx.x) * elements_per_block; + const uint64_t block_end = min(block_start + elements_per_block, list_len); + + for (uint64_t idx = block_start + threadIdx.x; idx < block_end; idx += blockDim.x) { + const OffsetT start = offsets[idx]; + out_offsets[idx] = start; + out_sizes[idx] = static_cast(offsets[idx + 1] - start); + } +} + +#define GENERATE_LIST_VIEWS_KERNEL(offset_suffix, OffsetT) \ + extern "C" __global__ void list_views_##offset_suffix(const OffsetT *const __restrict offsets, \ + OffsetT *const __restrict out_offsets, \ + OffsetT *const __restrict out_sizes, \ + uint64_t list_len) { \ + list_views(offsets, out_offsets, out_sizes, list_len); \ + } + +FOR_EACH_INTEGER(GENERATE_LIST_VIEWS_KERNEL) diff --git a/vortex-cuda/kernels/src/onpair.cu b/vortex-cuda/kernels/src/onpair.cu new file mode 100644 index 00000000000..905583b9bb6 --- /dev/null +++ b/vortex-cuda/kernels/src/onpair.cu @@ -0,0 +1,165 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "config.cuh" + +#include + +// Support kernels for OnPair GPU decompression. The per-batch output-offsets +// regeneration lives in the CUB shim (`cub/kernels/onpair.cu`), where the +// reduction and exclusive scan run as one fused sweep; this module holds the +// window-bounds and view-construction kernels launched around the decode. + +// Tokens per decode batch. Must match `ONPAIR_TOKENS_PER_BATCH` in +// `cub/kernels/onpair.cu` and `TOKENS_PER_BATCH` in the Rust launch code. +constexpr uint32_t ONPAIR_TOKENS_PER_BATCH = 128; + +// Token and byte positions of the visible code window's bounds, resolved +// entirely on device. A sliced array keeps its whole `codes` child, so the +// decode runs over the full stream; the window's token bounds are read from +// the (possibly slice-narrowed, possibly device-resident) `codes_offsets` +// child — the offsets are nondecreasing, so the window's min and max are its +// first and last elements — and each boundary's byte position is the +// whole-batch prefix from `chunk_offsets` plus a warp reduction over the +// boundary batch's head `[batch_start, boundary)`. Launched after the offsets +// sweep with one 32-thread block per boundary: block 0 resolves +// `codes_offsets[0]`, block 1 resolves `codes_offsets[last]`. The raw token +// boundary is written to `scratch[blockIdx.x]` (a signed negative offset +// sign-extends huge and fails the host's post-readback range validation; the +// reduction clamps it to `total_tokens` so a corrupt offset cannot read out +// of bounds) and the byte position to `scratch[2 + blockIdx.x]`. Block 0 +// also packs the sweep's outputs into the scratch — the full decoded heap +// size (`chunk_offsets`'s last entry) into `scratch[4]` and the corruption +// status flag into `scratch[5]` — so the host gates the decode kernel on a +// single readback. A code outside the dictionary contributes zero bytes +// (the sweep already raised the status flag the host checks before trusting +// the scratch). `CodeT` is the code stream's element type (u16 natively, u8 +// when the compressor narrowed the codes); `OffsetT` is the `codes_offsets` +// element type. +template +__device__ inline void onpair_window_offsets_body(const CodeT *__restrict codes, + const uint8_t *__restrict lens, + uint32_t dict_size, + const uint64_t *__restrict chunk_offsets, + uint64_t total_tokens, + const OffsetT *__restrict codes_offsets, + uint64_t last, + const uint32_t *__restrict status, + uint64_t *__restrict scratch) { + const uint64_t requested = (uint64_t)codes_offsets[blockIdx.x == 0 ? 0 : last]; + const uint64_t boundary = requested < total_tokens ? requested : total_tokens; + const uint64_t batch = boundary / ONPAIR_TOKENS_PER_BATCH; + const uint64_t batch_base = batch * ONPAIR_TOKENS_PER_BATCH; + const uint32_t lane = threadIdx.x & 31u; + + uint32_t partial = 0; +#pragma unroll + for (uint8_t token = 0; token < 4; ++token) { + const uint64_t i = batch_base + lane + (uint64_t)(token * 32u); + if (i < boundary) { + const uint32_t code = (uint32_t)codes[i]; + if (code < dict_size) { + partial += (uint32_t)lens[code]; + } + } + } +#pragma unroll + for (uint8_t offset = 16; offset > 0; offset >>= 1) { + partial += __shfl_down_sync(0xffffffffu, partial, offset); + } + if (lane == 0) { + scratch[blockIdx.x] = requested; + scratch[2 + blockIdx.x] = chunk_offsets[batch] + (uint64_t)partial; + if (blockIdx.x == 0) { + const uint64_t num_batches = + (total_tokens + ONPAIR_TOKENS_PER_BATCH - 1) / ONPAIR_TOKENS_PER_BATCH; + scratch[4] = chunk_offsets[num_batches]; + scratch[5] = (uint64_t)*status; + } + } +} + +// One entry point per (code width, codes_offsets ptype) pair: +// `onpair_window_offsets_{code}_{offset}`. +#define GENERATE_WINDOW_OFFSETS_KERNEL(code_suffix, CodeT, offset_suffix, OffsetT) \ + extern "C" __global__ void onpair_window_offsets_##code_suffix##_##offset_suffix( \ + const CodeT *__restrict codes, \ + const uint8_t *__restrict lens, \ + uint32_t dict_size, \ + const uint64_t *__restrict chunk_offsets, \ + uint64_t total_tokens, \ + const OffsetT *__restrict codes_offsets, \ + uint64_t last, \ + const uint32_t *__restrict status, \ + uint64_t *__restrict scratch) { \ + onpair_window_offsets_body(codes, \ + lens, \ + dict_size, \ + chunk_offsets, \ + total_tokens, \ + codes_offsets, \ + last, \ + status, \ + scratch); \ + } + +#define GENERATE_WINDOW_OFFSETS_KERNELS(code_suffix, CodeT) \ + GENERATE_WINDOW_OFFSETS_KERNEL(code_suffix, CodeT, i8, int8_t) \ + GENERATE_WINDOW_OFFSETS_KERNEL(code_suffix, CodeT, i16, int16_t) \ + GENERATE_WINDOW_OFFSETS_KERNEL(code_suffix, CodeT, i32, int32_t) \ + GENERATE_WINDOW_OFFSETS_KERNEL(code_suffix, CodeT, i64, int64_t) \ + GENERATE_WINDOW_OFFSETS_KERNEL(code_suffix, CodeT, u8, uint8_t) \ + GENERATE_WINDOW_OFFSETS_KERNEL(code_suffix, CodeT, u16, uint16_t) \ + GENERATE_WINDOW_OFFSETS_KERNEL(code_suffix, CodeT, u32, uint32_t) \ + GENERATE_WINDOW_OFFSETS_KERNEL(code_suffix, CodeT, u64, uint64_t) + +GENERATE_WINDOW_OFFSETS_KERNELS(u8, uint8_t) +GENERATE_WINDOW_OFFSETS_KERNELS(u16, uint16_t) + +// Arrow/Vortex variable-length view records are 16 bytes. Values up to 12 bytes +// are stored inline after the u32 length. Longer values store their first four +// bytes, backing-buffer index, and byte offset. +constexpr uint8_t MAX_INLINED_SIZE = 12; + +// Build one BinaryView over the flat decoded byte stream. Row `rid`'s bytes +// are `output_bytes[row_offsets[rid]..row_offsets[rid + 1])`; the offsets are +// the Arrow i32 offsets built on device from the lengths child. The Rust +// caller only launches this when the window fits a single backing buffer +// (`MAX_BUFFER_LEN`, i32::MAX), so every offset is non-negative and fits the +// view's u32 fields. +__device__ inline void onpair_write_view(const int32_t *__restrict row_offsets, + const uint8_t *__restrict output_bytes, + uint4 *__restrict views, + uint64_t rid) { + const uint32_t start = (uint32_t)row_offsets[rid]; + const uint32_t len = (uint32_t)row_offsets[rid + 1] - start; + if (len <= MAX_INLINED_SIZE) { + uint32_t words[3] = {0, 0, 0}; +#pragma unroll + for (uint8_t i = 0; i < MAX_INLINED_SIZE; i++) { + if (i < len) { + words[i >> 2] |= (uint32_t)output_bytes[start + i] << (8u * (i & 3u)); + } + } + views[rid] = make_uint4(len, words[0], words[1], words[2]); + return; + } + + const uint32_t prefix = (uint32_t)output_bytes[start] | ((uint32_t)output_bytes[start + 1] << 8u) | + ((uint32_t)output_bytes[start + 2] << 16u) | + ((uint32_t)output_bytes[start + 3] << 24u); + views[rid] = make_uint4(len, prefix, 0, start); +} + +extern "C" __global__ void onpair_build_views(const int32_t *__restrict row_offsets, + const uint8_t *__restrict output_bytes, + uint4 *__restrict views, + uint64_t num_rows) { + const uint64_t elements_per_block = (uint64_t)blockDim.x * ELEMENTS_PER_THREAD; + const uint64_t block_start = (uint64_t)blockIdx.x * elements_per_block; + const uint64_t block_end = + (block_start + elements_per_block < num_rows) ? (block_start + elements_per_block) : num_rows; + for (uint64_t rid = block_start + threadIdx.x; rid < block_end; rid += blockDim.x) { + onpair_write_view(row_offsets, output_bytes, views, rid); + } +} diff --git a/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu b/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu new file mode 100644 index 00000000000..2d22023c8d1 --- /dev/null +++ b/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu @@ -0,0 +1,263 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include +#include +#include +#include + +// OnPair decompress — 4 tokens/thread, split-read dictionary. +// +// One warp decodes one 128-token batch (4 tokens per lane) in four phases: +// +// 1. Load — each lane fetches its 4 codes, their first-8 dictionary bytes +// (`dict_s8`), and their lengths. +// 2. Scan — a warp prefix-scan of the lengths positions every token within +// the batch and yields the batch's total decoded size. +// 3. Stage — token bytes are gathered into a per-warp shared staging buffer; +// only the rare `len > 8` tokens touch the full 16-byte-row +// `dict_padded`. +// 4. Drain — the staged bytes stream to global output as an aligned `uint4` +// body between a byte head and tail. +// +// Baseline `onpair_shmem_4tpt` is L1/TEX-cache-request bound on the per-token +// 16-byte `uint4` gather into the 64 KB padded dict, where the dict L1 hit rate +// is only ~31% (the 64 KB dict thrashes against the streaming codes/output). +// +// Most tokens are short (mean dict len ~6). This variant reads the common case +// from the **32 KB** `dict_s8` array (first 8 bytes/entry, `uint2`) and only +// touches the 64 KB `dict_padded` for the rare `len > 8` tokens. Halving the +// hot dict working set aims to raise the dict L1 hit rate, cutting L2 sectors +// and L1/TEX-request pressure. As a bonus, holding `uint2 lo[4]` (32 B) instead +// of `uint4 t[4]` (64 B) lowers register pressure. + +#ifndef WARPS_PER_BLOCK_MAX +#define WARPS_PER_BLOCK_MAX 16u +#endif +// The launch bounds are an exception, not a pattern: this kernel was tuned on +// GH200 where NCU showed it register-capped at 64 (50% occupancy); bounding to +// 2 blocks/SM was benchmarked as the best trade against spills. Re-evaluate on +// new architectures rather than copying this to other kernels. +#ifndef ONPAIR_LAUNCH_BOUNDS +#define ONPAIR_LAUNCH_BOUNDS __launch_bounds__(512, 2) +#endif +#define WARP_BUF_BYTES 2080u + +__device__ inline uint32_t onpair_warp_inclusive_scan_u32(uint32_t x, uint32_t lane) { + constexpr unsigned mask = 0xffffffffu; +#pragma unroll + for (uint8_t offset = 1; offset < 32; offset <<= 1) { + uint32_t y = __shfl_up_sync(mask, x, offset); + if (lane >= offset) { + x += y; + } + } + return x; +} + +// One lane's slice of a 128-token batch: 4 tokens, strided a warp apart. +struct OnPairTokens { + // First 8 dictionary bytes of each token (the common-case read). + uint2 lo[4]; + // Dictionary codes, kept for the rare `len > 8` high-byte gather. + uint32_t code[4]; + // Decoded token lengths. + uint32_t len[4]; +}; + +// Phase 1 — load this lane's 4 (code, dict_s8 bytes, length) triples. Tokens +// outside the visible window load as empty. `CodeT` is the code stream's +// element type (u16 natively, u8 when the compressor narrowed the codes). +template +__device__ inline OnPairTokens onpair_load_tokens(const CodeT *__restrict codes, + const uint8_t *__restrict dict_s8, + const uint8_t *__restrict lens, + uint64_t base_i, + uint64_t token_start, + uint64_t token_end) { + OnPairTokens t; +#pragma unroll + for (uint8_t k = 0; k < 4; ++k) { + const uint64_t i = base_i + (uint64_t)(k * 32); + if (i >= token_start && i < token_end) { + const uint32_t code = (uint32_t)codes[i]; + t.code[k] = code; + t.lo[k] = *reinterpret_cast(dict_s8 + (size_t)code * 8u); + t.len[k] = (uint32_t)lens[code]; + } else { + t.code[k] = 0u; + t.lo[k] = make_uint2(0u, 0u); + t.len[k] = 0u; + } + } + return t; +} + +// Phase 2 — position every token within the batch: `excl[k]` is the exclusive +// prefix (the token's staging offset) via 4 chained warp scans of the lengths. +// Returns the batch's total decoded byte count. +__device__ inline uint32_t onpair_scan_offsets(const uint32_t (&len)[4], uint32_t lane, uint32_t (&excl)[4]) { + constexpr unsigned mask = 0xffffffffu; + uint32_t acc_base = 0u; +#pragma unroll + for (uint8_t k = 0; k < 4; ++k) { + const uint32_t incl = onpair_warp_inclusive_scan_u32(len[k], lane); + excl[k] = acc_base + (incl - len[k]); + acc_base += __shfl_sync(mask, incl, 31); + } + return acc_base; +} + +// Phase 3 — gather each token's bytes into the warp's shared staging buffer at +// its scanned offset. The common case writes only the 8 `dict_s8` bytes +// already in registers; tokens longer than 8 bytes take the rare path through +// the full padded dictionary. +__device__ inline void onpair_stage_tokens(const OnPairTokens &t, + const uint32_t (&excl)[4], + const uint8_t *__restrict dict_padded, + uint8_t *__restrict s_buf) { +#pragma unroll + for (uint8_t k = 0; k < 4; ++k) { + const uint32_t len = t.len[k]; + if (len == 0u) { + continue; + } + const uint32_t base = excl[k]; + const uint8_t *lob = reinterpret_cast(&t.lo[k]); + const uint32_t nlo = len < 8u ? len : 8u; +#pragma unroll + for (uint8_t j = 0; j < 8; ++j) { + if (j < nlo) { + s_buf[base + j] = lob[j]; + } + } + if (len > 8u) { + // Rare path: high bytes from the full padded dict. + const uint2 hi = *reinterpret_cast(dict_padded + (size_t)t.code[k] * 16u + 8u); + const uint8_t *hib = reinterpret_cast(&hi); +#pragma unroll + for (uint8_t j = 0; j < 8; ++j) { + if (8u + j < len) { + s_buf[base + 8 + j] = hib[j]; + } + } + } + } +} + +// Phase 4 — copy the staged batch to global output: a byte head up to the +// first 16-aligned output address, an aligned `uint4` body with streaming +// stores, and a byte tail. The caller offset `s_buf` so shared and global +// 16-byte alignment phases match. +__device__ inline void onpair_drain(const uint8_t *__restrict s_buf, + uint8_t *__restrict output_bytes, + uint64_t out_start, + uint32_t head_pre, + uint32_t warp_total, + uint32_t lane) { + const uint32_t head = head_pre < warp_total ? head_pre : warp_total; + if (lane < head) { + output_bytes[out_start + (uint64_t)lane] = s_buf[lane]; + } + if (head >= warp_total) { + return; + } + + const uint32_t body_chunks = (warp_total - head) >> 4; + for (uint8_t k = (uint8_t)lane; k < body_chunks; k += 32u) { + const uint32_t off = head + k * 16u; + const uint4 v = *reinterpret_cast(s_buf + off); + __stcs(reinterpret_cast(output_bytes + out_start + off), v); + } + + const uint32_t tail_start = head + (body_chunks << 4); + if (lane < warp_total - tail_start) { + output_bytes[out_start + (uint64_t)tail_start + (uint64_t)lane] = s_buf[tail_start + lane]; + } +} + +template +__device__ inline void onpair_decode_body(const CodeT *__restrict codes, + const uint64_t *__restrict chunk_offsets, + const uint8_t *__restrict dict_s8, + const uint8_t *__restrict dict_padded, + const uint8_t *__restrict lens, + uint8_t *__restrict output_bytes, + uint64_t token_start, + uint64_t token_end, + uint64_t first_batch, + uint64_t byte_start) { + const uint32_t lane = threadIdx.x & 31; + const uint32_t warp_id = threadIdx.x >> 5; + const uint64_t window_chunk = (uint64_t)blockIdx.x * (uint64_t)(blockDim.x >> 5) + (uint64_t)warp_id; + const uint64_t chunk = first_batch + window_chunk; + if (chunk * 128u >= token_end) { + return; + } + + __shared__ __align__(16) uint8_t s_buf_all[WARPS_PER_BLOCK_MAX * WARP_BUF_BYTES]; + uint8_t *s_buf_base = &s_buf_all[warp_id * WARP_BUF_BYTES]; + + const OnPairTokens t = + onpair_load_tokens(codes, dict_s8, lens, chunk * 128u + (uint64_t)lane, token_start, token_end); + + uint32_t excl[4]; + const uint32_t warp_total = onpair_scan_offsets(t.len, lane, excl); + + // Offset the staging buffer by (out_start % 16) so the drain's global + // 16-byte stores land aligned when copied from 16-aligned shared reads. + const uint64_t out_start = chunk == first_batch ? 0 : chunk_offsets[chunk] - byte_start; + const uint32_t head_pre = (16u - (uint32_t)(out_start & 15u)) & 15u; + uint8_t *s_buf = s_buf_base + ((16u - head_pre) & 15u); + + onpair_stage_tokens(t, excl, dict_padded, s_buf); + __syncwarp(); + + onpair_drain(s_buf, output_bytes, out_start, head_pre, warp_total, lane); +} + +extern "C" __global__ ONPAIR_LAUNCH_BOUNDS void +onpair_shmem_4tpt_split8read_u8(const uint8_t *__restrict codes, + const uint64_t *__restrict chunk_offsets, + const uint8_t *__restrict dict_s8, + const uint8_t *__restrict dict_padded, + const uint8_t *__restrict lens, + uint8_t *__restrict output_bytes, + uint64_t token_start, + uint64_t token_end, + uint64_t first_batch, + uint64_t byte_start) { + onpair_decode_body(codes, + chunk_offsets, + dict_s8, + dict_padded, + lens, + output_bytes, + token_start, + token_end, + first_batch, + byte_start); +} + +extern "C" __global__ ONPAIR_LAUNCH_BOUNDS void +onpair_shmem_4tpt_split8read_u16(const uint16_t *__restrict codes, + const uint64_t *__restrict chunk_offsets, + const uint8_t *__restrict dict_s8, + const uint8_t *__restrict dict_padded, + const uint8_t *__restrict lens, + uint8_t *__restrict output_bytes, + uint64_t token_start, + uint64_t token_end, + uint64_t first_batch, + uint64_t byte_start) { + onpair_decode_body(codes, + chunk_offsets, + dict_s8, + dict_padded, + lens, + output_bytes, + token_start, + token_end, + first_batch, + byte_start); +} diff --git a/vortex-cuda/kernels/src/runend.cu b/vortex-cuda/kernels/src/runend.cu index a3f1d245dbe..7472e7dafba 100644 --- a/vortex-cuda/kernels/src/runend.cu +++ b/vortex-cuda/kernels/src/runend.cu @@ -127,6 +127,45 @@ __device__ void runend_decode_kernel(const EndsT *const __restrict ends, } } +// Expands run-end encoded validity bits into a packed output bitmap. +// +// Mirrors `runend_decode_kernel`, but each thread owns one complete output byte so that +// threads never race on bits within the same byte. Runs are located with a global binary +// search per element rather than the shared-memory cache: validity expansion runs once per +// array and is not the decode hot path. +template +__device__ void runend_bool_kernel(const EndsT *const __restrict ends, + uint64_t num_runs, + const uint8_t *const __restrict values, + uint64_t values_bit_offset, + uint64_t offset, + uint64_t output_len, + uint8_t *const __restrict output) { + const uint64_t output_bytes = (output_len + 7) / 8; + const uint32_t elements_per_block = blockDim.x * ELEMENTS_PER_THREAD; + const uint64_t block_start = static_cast(blockIdx.x) * elements_per_block; + const uint64_t block_end = min(block_start + elements_per_block, output_bytes); + + for (uint64_t byte_idx = block_start + threadIdx.x; byte_idx < block_end; byte_idx += blockDim.x) { + const uint64_t row_start = byte_idx * 8; + uint8_t packed = 0; +#pragma unroll + for (uint32_t bit = 0; bit < 8; ++bit) { + const uint64_t row = row_start + bit; + if (row < output_len) { + uint64_t run_idx = upper_bound(ends, num_runs, row + offset); + if (run_idx >= num_runs) { + run_idx = num_runs - 1; + } + const uint64_t value_idx = values_bit_offset + run_idx; + const uint8_t value = (values[value_idx / 8] >> (value_idx % 8)) & 1; + packed |= static_cast(value << bit); + } + } + output[byte_idx] = packed; + } +} + #define GENERATE_RUNEND_KERNEL(value_suffix, ValueType, ends_suffix, EndsType) \ extern "C" __global__ void runend_##value_suffix##_##ends_suffix( \ const EndsType *const __restrict ends, \ @@ -155,3 +194,21 @@ GENERATE_RUNEND_KERNELS_FOR_VALUE(i64, int64_t) GENERATE_RUNEND_KERNELS_FOR_VALUE(f16, __half) GENERATE_RUNEND_KERNELS_FOR_VALUE(f32, float) GENERATE_RUNEND_KERNELS_FOR_VALUE(f64, double) + +#define GENERATE_RUNEND_BOOL_KERNEL(ends_suffix, EndsType) \ + extern "C" __global__ void runend_bool_##ends_suffix(const EndsType *const __restrict ends, \ + uint64_t num_runs, \ + const uint8_t *const __restrict values, \ + uint64_t values_bit_offset, \ + uint64_t offset, \ + uint64_t output_len, \ + uint8_t *const __restrict output) { \ + runend_bool_kernel(ends, num_runs, values, values_bit_offset, offset, output_len, output); \ + } + +// Validity bitmaps use a different physical layout and launch unit, but dispatch over the +// same run-end index types. +GENERATE_RUNEND_BOOL_KERNEL(u8, uint8_t) +GENERATE_RUNEND_BOOL_KERNEL(u16, uint16_t) +GENERATE_RUNEND_BOOL_KERNEL(u32, uint32_t) +GENERATE_RUNEND_BOOL_KERNEL(u64, uint64_t) diff --git a/vortex-cuda/kernels/src/zigzag.cu b/vortex-cuda/kernels/src/zigzag.cu index 4059452da56..f1ba98e00da 100644 --- a/vortex-cuda/kernels/src/zigzag.cu +++ b/vortex-cuda/kernels/src/zigzag.cu @@ -8,26 +8,20 @@ // Converts unsigned integers back to signed using the ZigZag encoding scheme. // Formula: decoded = (encoded >> 1) ^ -(encoded & 1) // This interleaves positive and negative numbers: 0, -1, 1, -2, 2, -3, ... -template +template struct ZigZagOp { - __device__ inline SignedT operator()(UnsignedT value) const { - // ZigZag decode: (n >> 1) ^ -(n & 1) - // The -(n & 1) is equivalent to: if (n & 1) then -1 else 0 - return static_cast((value >> 1) ^ (~(value & 1) + 1)); + __device__ inline UnsignedT operator()(UnsignedT value) const { + return (value >> 1) ^ (UnsignedT(0) - (value & 1)); } }; -// Macro to generate ZigZag kernel for each type. -// In-place operation: unsigned input, signed output (same size, reinterpret). -#define GENERATE_ZIGZAG_KERNEL(suffix, UnsignedType, SignedType) \ +// Macro to generate the in-place ZigZag kernel for each type. +#define GENERATE_ZIGZAG_KERNEL(suffix, UnsignedType) \ extern "C" __global__ void zigzag_##suffix(UnsignedType *__restrict values, uint64_t array_len) { \ - scalar_kernel(values, \ - reinterpret_cast(values), \ - array_len, \ - ZigZagOp {}); \ + scalar_kernel_inplace(values, array_len, ZigZagOp {}); \ } -GENERATE_ZIGZAG_KERNEL(u8, uint8_t, int8_t) -GENERATE_ZIGZAG_KERNEL(u16, uint16_t, int16_t) -GENERATE_ZIGZAG_KERNEL(u32, uint32_t, int32_t) -GENERATE_ZIGZAG_KERNEL(u64, uint64_t, int64_t) +GENERATE_ZIGZAG_KERNEL(u8, uint8_t) +GENERATE_ZIGZAG_KERNEL(u16, uint16_t) +GENERATE_ZIGZAG_KERNEL(u32, uint32_t) +GENERATE_ZIGZAG_KERNEL(u64, uint64_t) diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index 66c50a0413c..4a96187a9db 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -38,7 +38,7 @@ use vortex::array::arrays::extension::ExtensionArrayExt; use vortex::array::arrays::fixed_size_list::FixedSizeListArrayExt; use vortex::array::arrays::fixed_size_list::FixedSizeListDataParts; use vortex::array::arrays::list::ListDataParts; -use vortex::array::arrays::listview::ListViewArrayExt; +use vortex::array::arrays::listview::ListViewArraySlotsExt; use vortex::array::arrays::listview::list_from_list_view; use vortex::array::arrays::primitive::PrimitiveDataParts; use vortex::array::arrays::struct_::StructDataParts; @@ -55,12 +55,12 @@ use vortex::dtype::Nullability; use vortex::dtype::PType; use vortex::dtype::i256; use vortex::encodings::fsst::FSST; -use vortex::encodings::fsst::FSSTArray; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::error::vortex_ensure; use vortex::error::vortex_err; use vortex::extension::datetime::AnyTemporal; +use vortex_onpair::OnPair; use crate::CudaBufferExt; use crate::CudaDeviceBuffer; @@ -81,8 +81,9 @@ use crate::cub::exclusive_sum_i32; use crate::device_buffer::CUDF_VALIDITY_BUFFER_PADDING; use crate::executor::CudaArrayExt; use crate::executor::execute_validity_cuda; -use crate::kernel::FSSTVarBin; +use crate::kernel::DecodedVarBin; use crate::kernel::decode_fsst_varbin; +use crate::kernel::decode_onpair_varbin; /// An implementation of `ExportDeviceArray` that exports Vortex arrays to `ArrowDeviceArray` by /// first decoding the array on the GPU and then converting the canonical type to the nearest @@ -229,11 +230,23 @@ fn export_array( // `CudaDispatchMode` only governs `execute_cuda`'s fused-vs-standalone planning. let array = match array.try_downcast::() { Ok(fsst) if ctx.cuda_session().varbin_export_layout() == VarBinExportLayout::VarBin => { - return export_fsst_varbin(fsst, ctx).await; + let decoded = decode_fsst_varbin(fsst, ctx).await?; + return export_decoded_varbin(decoded, ctx).await; } Ok(fsst) => fsst.into_array(), Err(array) => array, }; + // OnPair takes the same offset-based export shortcut as FSST. + let array = match array.try_downcast::() { + Ok(onpair) + if ctx.cuda_session().varbin_export_layout() == VarBinExportLayout::VarBin => + { + let decoded = decode_onpair_varbin(onpair, ctx).await?; + return export_decoded_varbin(decoded, ctx).await; + } + Ok(onpair) => onpair.into_array(), + Err(array) => array, + }; let cuda_array = array.execute_cuda(ctx).await?; export_canonical(cuda_array, ctx).await @@ -487,11 +500,9 @@ where .await; } - let output_buffer = ctx.device_alloc::(len)?; - let output_device = CudaDeviceBuffer::new(output_buffer); + let mut output_buffer = ctx.device_alloc::(len)?; let values_view = values.cuda_view::()?; - let output_view = output_device.as_view::(); let len_u64 = len as u64; let cuda_function = ctx.load_function_with_suffixes( "decimal_cast", @@ -499,10 +510,12 @@ where )?; ctx.launch_kernel(&cuda_function, len, |args| { - args.arg(&values_view).arg(&output_view).arg(&len_u64); + args.arg(&values_view).arg(&mut output_buffer).arg(&len_u64); })?; - Ok(BufferHandle::new_device(Arc::new(output_device))) + Ok(BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new( + output_buffer, + )))) } /// Export Vortex binary views as an Arrow Device array with `Utf8View`/`BinaryView` layout. @@ -581,20 +594,22 @@ async fn export_varbin( export_varbin_buffers(len, validity_buffer, null_count, offsets, values, ctx) } -async fn export_fsst_varbin( - fsst: FSSTArray, +/// Export an offset-based decompression result (FSST or OnPair) with the +/// standard Arrow `Utf8`/`Binary` layout. +async fn export_decoded_varbin( + decoded: DecodedVarBin, ctx: &mut CudaExecutionCtx, ) -> VortexResult<(ArrowArray, SyncEvent)> { - let FSSTVarBin { + let DecodedVarBin { dtype, len, offsets, values, validity, - } = decode_fsst_varbin(fsst, ctx).await?; + } = decoded; vortex_ensure!( matches!(dtype, DType::Utf8(_) | DType::Binary(_)), - "FSST produced invalid variable-length dtype {dtype}" + "offset-based decode produced invalid variable-length dtype {dtype}" ); let (validity_buffer, null_count) = export_arrow_validity_buffer(validity, len, 0, ctx).await?; export_varbin_buffers(len, validity_buffer, null_count, offsets, values, ctx) @@ -655,21 +670,25 @@ async fn export_binary_buffers( } let data_buffer_ptrs = device_buffer_from(ptr_values, ctx).await?; let data_buffer_lens = device_buffer_from(len_values, ctx).await?; - let status = device_buffer_from(vec![0u32], ctx).await?; + let mut status = ctx.device_alloc::(1)?; + ctx.stream() + .memset_zeros(&mut status) + .map_err(|err| vortex_err!("Failed to zero Arrow Binary status buffer: {err}"))?; let scan_input = init_binary_scan( views, validity, &data_buffer_lens, device_data_buffers.len(), - &status, + &mut status, len, ctx, )?; let output_offsets = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new( exclusive_sum_i32(&scan_input, len + 1, ctx)?, ))); - validate_binary_offsets(&output_offsets, len, &status, ctx)?; + validate_binary_offsets(&output_offsets, len, &mut status, ctx)?; + let status = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(status))); // One status read covers init_scan and offset validation. Both must pass before gather may // dereference view payloads through the scanned offsets. Enqueue both copies up front so the @@ -726,7 +745,7 @@ fn init_binary_scan( validity: Option<&BufferHandle>, data_buffer_lens: &BufferHandle, data_buffer_count: usize, - status: &BufferHandle, + status: &mut CudaSlice, len: usize, ctx: &mut CudaExecutionCtx, ) -> VortexResult> { @@ -738,18 +757,17 @@ fn init_binary_scan( .transpose()? .unwrap_or(0); let lens_view = data_buffer_lens.cuda_view::()?; - let status_view = status.cuda_view::()?; let data_buffer_count_u64 = data_buffer_count as u64; let len_u64 = len as u64; - let scan_input = ctx.device_alloc::(scan_len)?; + let mut scan_input = ctx.device_alloc::(scan_len)?; let kernel = ctx.load_function_with_suffixes("arrow_binary", &["init_scan"])?; ctx.launch_kernel(&kernel, scan_len, |args| { args.arg(&views_view) .arg(&validity_ptr) .arg(&lens_view) - .arg(&scan_input) - .arg(&status_view) + .arg(&mut scan_input) + .arg(&mut *status) .arg(&data_buffer_count_u64) .arg(&len_u64); })?; @@ -763,17 +781,16 @@ fn init_binary_scan( fn validate_binary_offsets( offsets: &BufferHandle, len: usize, - status: &BufferHandle, + status: &mut CudaSlice, ctx: &mut CudaExecutionCtx, ) -> VortexResult<()> { let scan_len = len + 1; let offsets_view = offsets.cuda_view::()?; - let status_view = status.cuda_view::()?; let scan_len_u64 = scan_len as u64; let kernel = ctx.load_function_with_suffixes("arrow_binary", &["validate_offsets"])?; ctx.launch_kernel(&kernel, scan_len, |args| { - args.arg(&offsets_view).arg(&status_view).arg(&scan_len_u64); + args.arg(&offsets_view).arg(&mut *status).arg(&scan_len_u64); }) } @@ -785,7 +802,7 @@ fn gather_binary_values( len: usize, ctx: &mut CudaExecutionCtx, ) -> VortexResult { - let output_values = ctx.device_alloc::(total_bytes.max(1))?; + let mut output_values = ctx.device_alloc::(total_bytes.max(1))?; if total_bytes != 0 { let views_view = views.cuda_view::()?; @@ -799,7 +816,7 @@ fn gather_binary_values( args.arg(&views_view) .arg(&ptrs_view) .arg(&offsets_view) - .arg(&output_values) + .arg(&mut output_values) .arg(&len_u64) .arg(&total_bytes_u64); })?; @@ -979,10 +996,8 @@ pub fn count_arrow_validity_nulls( ctx.stream() .memset_zeros(&mut count) .map_err(|err| vortex_err!("Failed to zero Arrow validity count buffer: {err}"))?; - let count = CudaDeviceBuffer::new(count); let input_view = bitmap.cuda_view::()?; - let output_view = count.as_view::(); let len = u64::try_from(len)?; let arrow_offset = u64::try_from(arrow_offset)?; @@ -998,14 +1013,14 @@ pub fn count_arrow_validity_nulls( }; ctx.launch_kernel_config(&kernel, config, expected_bytes, |args| { args.arg(&input_view) - .arg(&output_view) + .arg(&mut count) .arg(&len) .arg(&arrow_offset); })?; let valid_count = ctx .stream() - .clone_dtoh(&output_view) + .clone_dtoh(&count) .map_err(|err| vortex_err!("Failed to copy Arrow validity count to host: {err}"))? .into_iter() .next() @@ -1055,12 +1070,9 @@ pub fn repack_arrow_validity_buffer( ctx.stream() .memset_zeros(&mut output) .map_err(|err| vortex_err!("Failed to zero Arrow validity buffer padding: {err}"))?; - // The memset above zeroed all allocation bytes after the logical output. - let output_device = CudaDeviceBuffer::new_with_zeroed_tail(output, output_bytes)?; if output_words > 0 { let input_view = input_buffer.cuda_view::()?; - let output_view = output_device.as_view::(); let len = u64::try_from(len)?; let input_offset = u64::try_from(input_offset)?; let arrow_offset = u64::try_from(arrow_offset)?; @@ -1076,7 +1088,7 @@ pub fn repack_arrow_validity_buffer( }; ctx.launch_kernel_config(&kernel, config, output_words, |args| { args.arg(&input_view) - .arg(&output_view) + .arg(&mut output) .arg(&len) .arg(&input_offset) .arg(&arrow_offset) @@ -1084,6 +1096,8 @@ pub fn repack_arrow_validity_buffer( })?; } + // The memset above zeroed all allocation bytes after the logical output. + let output_device = CudaDeviceBuffer::new_with_zeroed_tail(output, output_bytes)?; Ok(BufferHandle::new_device(Arc::new(output_device)).slice(0..output_bytes)) } @@ -1255,13 +1269,13 @@ fn fixed_size_list_offsets( let output_len = len .checked_add(1) .ok_or_else(|| vortex_err!("FixedSizeList Arrow List offsets length overflows usize"))?; - let offsets = ctx.device_alloc::(output_len)?; + let mut offsets = ctx.device_alloc::(output_len)?; let base = 0i32; let output_len_u64 = output_len as u64; let kernel = ctx.load_function_with_suffixes("sequence", &["i32"])?; ctx.launch_kernel(&kernel, output_len, |args| { - args.arg(&offsets) + args.arg(&mut offsets) .arg(&base) .arg(&list_size) .arg(&output_len_u64); @@ -1440,6 +1454,8 @@ mod tests { use rstest::rstest; use vortex::array::ArrayRef; use vortex::array::IntoArray; + use vortex::array::VortexSessionExecute; + use vortex::array::array_session; use vortex::array::arrays::BoolArray; use vortex::array::arrays::ChunkedArray; use vortex::array::arrays::DecimalArray; @@ -1502,7 +1518,7 @@ mod tests { } fn cuda_ctx_with_varbin_layout(layout: VarBinExportLayout) -> VortexResult { - let session = vortex::array::array_session() + let session = array_session() .with_some(CudaSession::try_default()?.with_varbin_export_layout(layout)); CudaSession::create_execution_ctx(&session) } @@ -1670,6 +1686,7 @@ mod tests { Arc::from([first, second]), dtype, Validity::NonNullable, + &mut array_session().create_execution_ctx(), ) .vortex_expect("valid multi-buffer VarBinViewArray") .into_array(); @@ -2520,10 +2537,9 @@ mod tests { let fsst = fsst_compress(&varbin, &compressor, ctx.execution_ctx())?; // Same codes, but uncompressed lengths whose sum exceeds i32::MAX. - let oversized = FSST::try_new( + let oversized = FSST::try_new_with_symbol_table( DType::Utf8(Nullability::NonNullable), - fsst.symbols().clone(), - fsst.symbol_lengths().clone(), + fsst.symbol_table(), fsst.codes(), PrimitiveArray::from_iter([i32::MAX, i32::MAX]).into_array(), ctx.execution_ctx(), @@ -2543,7 +2559,7 @@ mod tests { assert!( error .to_string() - .contains("FSST decoded size exceeds Arrow i32 offset range") + .contains("length sum exceeds Arrow i32 offset range") ); Ok(()) } @@ -2805,9 +2821,8 @@ mod tests { Ok(()) } - // Regression test: with an average list size >= 128 the host list-view rebuild picks its - // list-by-list strategy, which may canonicalize Dict elements. The schema must describe the - // rebuilt child layout. + // Regression test: the host list-view rebuild uses a take for large lists, which preserves + // dictionary elements. The schema must describe the rebuilt child layout. #[crate::test] async fn test_export_host_large_lists_dictionary_list_view_schema_matches_rebuilt_child() -> VortexResult<()> { @@ -2836,7 +2851,17 @@ mod tests { field, Field::new_list( "", - Field::new(Field::LIST_FIELD_DEFAULT_NAME, DataType::Int32, true), + Field::new( + Field::LIST_FIELD_DEFAULT_NAME, + DataType::Dictionary( + Box::new(DataType::Int64), + Box::new(DataType::Dictionary( + Box::new(DataType::Int16), + Box::new(DataType::Int32), + )), + ), + true, + ), false, ) ); @@ -2846,9 +2871,11 @@ mod tests { ); let list_children = unsafe { std::slice::from_raw_parts(exported.array.array.children, 1) }; let child = unsafe { &*list_children[0] }; - assert!(child.dictionary.is_null()); + assert!(!child.dictionary.is_null()); assert_eq!(child.length, 256); assert_eq!(child.n_buffers, 2); + let nested_dict = unsafe { &*child.dictionary }; + assert!(!nested_dict.dictionary.is_null()); unsafe { release_exported_array(&raw mut exported.array.array) }; Ok(()) diff --git a/vortex-cuda/src/arrow/list_view.rs b/vortex-cuda/src/arrow/list_view.rs index 6e3f25a939f..2ce3b3b5669 100644 --- a/vortex-cuda/src/arrow/list_view.rs +++ b/vortex-cuda/src/arrow/list_view.rs @@ -162,16 +162,16 @@ where O: NativePType + DeviceRepr + Send + Sync + 'static, S: NativePType + DeviceRepr + Send + Sync + 'static, { - let status = new_list_view_status(ctx).await?; + let mut status = new_list_view_status(ctx)?; let scan_len = list_len + 1; - let scan_input = init_list_view_rebuild_scan::(&sizes, &status, list_len, ctx)?; + let scan_input = init_list_view_rebuild_scan::(&sizes, &mut status, list_len, ctx)?; let output_offsets = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new( exclusive_sum_i32(&scan_input, scan_len, ctx)?, ))); - validate_list_view_rebuild_offsets::(&sizes, &output_offsets, &status, list_len, ctx)?; - check_list_view_rebuild_status(&status).await?; + validate_list_view_rebuild_offsets::(&sizes, &output_offsets, &mut status, list_len, ctx)?; + check_list_view_rebuild_status(&status, ctx)?; let total_values = total_values_from_offsets(&output_offsets, list_len).await?; let value_width = values_ptype.byte_width(); @@ -188,10 +188,10 @@ where elements_len, list_len, value_width, - &status, + &mut status, ctx, )?; - check_list_view_rebuild_status(&status).await?; + check_list_view_rebuild_status(&status, ctx)?; let values = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(output_values))) .slice(0..output_values_bytes); @@ -199,16 +199,24 @@ where } /// Allocate the device status word used by list-view rebuild kernels. -async fn new_list_view_status(ctx: &mut CudaExecutionCtx) -> VortexResult { - ctx.ensure_on_device(BufferHandle::new_host( - Buffer::from(vec![0u32]).into_byte_buffer(), - )) - .await +fn new_list_view_status(ctx: &mut CudaExecutionCtx) -> VortexResult> { + let mut status = ctx.device_alloc::(1)?; + ctx.stream() + .memset_zeros(&mut status) + .map_err(|err| vortex_err!("Failed to zero list-view rebuild status buffer: {err}"))?; + Ok(status) } /// Convert the list-view rebuild status word into a Vortex error. -async fn check_list_view_rebuild_status(status: &BufferHandle) -> VortexResult<()> { - match Buffer::::from_byte_buffer(status.try_to_host()?.await?)[0] { +fn check_list_view_rebuild_status( + status: &CudaSlice, + ctx: &CudaExecutionCtx, +) -> VortexResult<()> { + let status = ctx + .stream() + .clone_dtoh(status) + .map_err(|err| vortex_err!("Failed to copy list-view rebuild status to host: {err}"))?[0]; + match status { 0 => Ok(()), 1 => vortex_bail!( "cannot export device-resident ListViewArray as Arrow List: offsets/sizes are invalid for the child elements" @@ -223,7 +231,7 @@ async fn check_list_view_rebuild_status(status: &BufferHandle) -> VortexResult<( /// Initialize the exclusive-scan input for rebuilt Arrow List offsets. fn init_list_view_rebuild_scan( sizes: &BufferHandle, - status: &BufferHandle, + status: &mut CudaSlice, list_len: usize, ctx: &mut CudaExecutionCtx, ) -> VortexResult> @@ -232,17 +240,16 @@ where { let scan_len = list_len + 1; let sizes_view = sizes.cuda_view::()?; - let status_view = status.cuda_view::()?; let list_len_u64 = list_len as u64; let scan_len_u64 = scan_len as u64; - let scan_input = ctx.device_alloc::(scan_len)?; + let mut scan_input = ctx.device_alloc::(scan_len)?; let init_kernel = ctx .load_function_with_suffixes("list_view", &["rebuild_init_scan", &S::PTYPE.to_string()])?; ctx.launch_kernel(&init_kernel, scan_len, |args| { args.arg(&sizes_view) - .arg(&scan_input) - .arg(&status_view) + .arg(&mut scan_input) + .arg(&mut *status) .arg(&list_len_u64) .arg(&scan_len_u64); })?; @@ -254,7 +261,7 @@ where fn validate_list_view_rebuild_offsets( sizes: &BufferHandle, output_offsets: &BufferHandle, - status: &BufferHandle, + status: &mut CudaSlice, list_len: usize, ctx: &mut CudaExecutionCtx, ) -> VortexResult<()> @@ -263,7 +270,6 @@ where { let sizes_view = sizes.cuda_view::()?; let output_offsets_view = output_offsets.cuda_view::()?; - let status_view = status.cuda_view::()?; let list_len_u64 = list_len as u64; let validate_kernel = ctx.load_function_with_suffixes( "list_view", @@ -273,7 +279,7 @@ where ctx.launch_kernel(&validate_kernel, list_len, |args| { args.arg(&sizes_view) .arg(&output_offsets_view) - .arg(&status_view) + .arg(&mut *status) .arg(&list_len_u64); }) } @@ -304,7 +310,7 @@ fn gather_rebuilt_primitive_values( elements_len: usize, list_len: usize, value_width: usize, - status: &BufferHandle, + status: &mut CudaSlice, ctx: &mut CudaExecutionCtx, ) -> VortexResult> where @@ -315,8 +321,7 @@ where let sizes_view = sizes.cuda_view::()?; let values_view = values.cuda_view::()?; let output_offsets_view = output_offsets.cuda_view::()?; - let output_values = ctx.device_alloc::(output_values_bytes.max(1))?; - let status_view = status.cuda_view::()?; + let mut output_values = ctx.device_alloc::(output_values_bytes.max(1))?; let list_len_u64 = list_len as u64; let elements_len_u64 = elements_len as u64; let value_width_u64 = value_width as u64; @@ -334,8 +339,8 @@ where .arg(&sizes_view) .arg(&output_offsets_view) .arg(&values_view) - .arg(&output_values) - .arg(&status_view) + .arg(&mut output_values) + .arg(&mut *status) .arg(&list_len_u64) .arg(&elements_len_u64) .arg(&value_width_u64); @@ -512,17 +517,11 @@ where S: NativePType + DeviceRepr + Send + Sync + 'static, { let output_len = len + 1; - let output = ctx.device_alloc::(output_len)?; - - let status = ctx - .ensure_on_device(BufferHandle::new_host( - Buffer::from(vec![0u32]).into_byte_buffer(), - )) - .await?; + let mut output = ctx.device_alloc::(output_len)?; + let mut status = new_list_view_status(ctx)?; let offsets_view = offsets.cuda_view::()?; let sizes_view = sizes.cuda_view::()?; - let status_view = status.cuda_view::()?; let list_len_u64 = len as u64; let kernel = ctx.load_function_with_suffixes( @@ -532,12 +531,16 @@ where ctx.launch_kernel(&kernel, len, |args| { args.arg(&offsets_view) .arg(&sizes_view) - .arg(&output) - .arg(&status_view) + .arg(&mut output) + .arg(&mut status) .arg(&list_len_u64); })?; - match Buffer::::from_byte_buffer(status.try_to_host()?.await?)[0] { + let status = ctx + .stream() + .clone_dtoh(&status) + .map_err(|err| vortex_err!("Failed to copy list-view offsets status to host: {err}"))?[0]; + match status { 0 => Ok(DeviceListViewOffsets::Contiguous(BufferHandle::new_device( Arc::new(CudaDeviceBuffer::new(output)), ))), diff --git a/vortex-cuda/src/arrow/mod.rs b/vortex-cuda/src/arrow/mod.rs index 071e22085c4..f606e623d4b 100644 --- a/vortex-cuda/src/arrow/mod.rs +++ b/vortex-cuda/src/arrow/mod.rs @@ -10,6 +10,7 @@ mod canonical; mod list_view; +mod offsets; use std::ffi::CString; use std::ffi::c_char; @@ -29,7 +30,10 @@ use async_trait::async_trait; pub(crate) use canonical::CanonicalDeviceArrayExport; use cudarc::driver::CudaEvent; use cudarc::driver::CudaStream; +use cudarc::driver::DevicePtr; use cudarc::runtime::sys::cudaEvent_t; +pub(crate) use offsets::I32Offsets; +pub(crate) use offsets::i32_offsets_from_lengths; use vortex::array::ArrayRef; use vortex::array::arrays::Dict; use vortex::array::arrays::FixedSizeList; @@ -37,9 +41,9 @@ use vortex::array::arrays::List; use vortex::array::arrays::ListView; use vortex::array::arrays::Struct; use vortex::array::arrays::dict::DictArraySlotsExt; -use vortex::array::arrays::fixed_size_list::FixedSizeListArrayExt; -use vortex::array::arrays::list::ListArrayExt; -use vortex::array::arrays::listview::ListViewArrayExt; +use vortex::array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use vortex::array::arrays::list::ListArraySlotsExt; +use vortex::array::arrays::listview::ListViewArraySlotsExt; use vortex::array::arrays::struct_::StructArrayExt; use vortex::array::buffer::BufferHandle; use vortex::array::stream::SendableArrayStream; @@ -166,9 +170,18 @@ impl PrivateData { // null pointer Ok(ptr::null()) } - Some(handle) => usize::try_from(handle.cuda_device_ptr()?) - .map(|ptr| ptr as *const c_void) - .map_err(|_| vortex_err!("CUDA device pointer does not fit in usize")), + Some(handle) => { + // The buffer may have been populated on a different stream (for example, + // pooled file reads use a dedicated H2D stream). Access it through cudarc's + // stream-aware API so this export stream waits for pending writes before + // its Arrow sync event is recorded. + let view = handle.cuda_view::()?; + let (device_ptr, record_read) = view.device_ptr(ctx.stream()); + drop(record_read); + usize::try_from(device_ptr) + .map(|ptr| ptr as *const c_void) + .map_err(|_| vortex_err!("CUDA device pointer does not fit in usize")) + } } }) .collect::>>()? diff --git a/vortex-cuda/src/arrow/offsets.rs b/vortex-cuda/src/arrow/offsets.rs new file mode 100644 index 00000000000..4eda482cb41 --- /dev/null +++ b/vortex-cuda/src/arrow/offsets.rs @@ -0,0 +1,101 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Device construction of standard Arrow `i32` offsets. + +use std::sync::Arc; + +use cudarc::driver::DeviceRepr; +use cudarc::driver::PushKernelArg; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::primitive::PrimitiveDataParts; +use vortex::array::buffer::BufferHandle; +use vortex::array::match_each_integer_ptype; +use vortex::buffer::Buffer; +use vortex::dtype::NativePType; +use vortex::error::VortexResult; +use vortex::error::vortex_bail; +use vortex::error::vortex_err; + +use crate::CudaBufferExt; +use crate::CudaDeviceBuffer; +use crate::cub::exclusive_sum_i32; +use crate::executor::CudaExecutionCtx; + +pub(crate) struct I32Offsets { + pub(crate) buffer: BufferHandle, + pub(crate) total: usize, +} + +/// Build Arrow-compatible offsets from a canonical integer length array. +/// +/// Length conversion, prefix sum, and overflow validation stay on the active CUDA stream. Only the +/// final offset and status word are copied to the host because callers need the total to size their +/// output allocation. +pub(crate) async fn i32_offsets_from_lengths( + lengths: PrimitiveArray, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + let len = lengths.len(); + let ptype = lengths.ptype(); + let PrimitiveDataParts { buffer, .. } = lengths.into_data_parts(); + let lengths = ctx.ensure_on_device(buffer).await?; + + match_each_integer_ptype!(ptype, |L| { + i32_offsets_from_lengths_typed::(&lengths, len, ctx).await + }) +} + +async fn i32_offsets_from_lengths_typed( + lengths: &BufferHandle, + len: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult +where + L: NativePType + DeviceRepr + Send + Sync + 'static, +{ + let scan_len = len + .checked_add(1) + .ok_or_else(|| vortex_err!("Arrow offset count overflow"))?; + let mut status = ctx.device_alloc::(1)?; + ctx.stream() + .memset_zeros(&mut status) + .map_err(|err| vortex_err!("Failed to zero Arrow offset status buffer: {err}"))?; + let lengths_view = lengths.cuda_view::()?; + let mut scan_input = ctx.device_alloc::(scan_len)?; + let ptype = L::PTYPE.to_string(); + let scan_kernel = + ctx.load_function_with_suffixes("arrow_offsets", &["from", "lengths", &ptype])?; + let len_u64 = u64::try_from(len)?; + + ctx.launch_kernel(&scan_kernel, scan_len, |args| { + args.arg(&lengths_view) + .arg(&mut scan_input) + .arg(&mut status) + .arg(&len_u64); + })?; + + let offsets = exclusive_sum_i32(&scan_input, scan_len, ctx)?; + let scan_len_u64 = u64::try_from(scan_len)?; + let validate_kernel = ctx.load_function_with_suffixes("arrow_offsets", &["validate"])?; + ctx.launch_kernel(&validate_kernel, scan_len, |args| { + args.arg(&offsets).arg(&mut status).arg(&scan_len_u64); + })?; + + let offsets = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(offsets))); + let status = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(status))); + let status_copy = status.try_to_host()?; + let total_copy = offsets.slice_typed::(len..scan_len).try_to_host()?; + let (status_bytes, total_bytes) = futures::try_join!(status_copy, total_copy)?; + match Buffer::::from_byte_buffer(status_bytes)[0] { + 0 => {} + 1 => vortex_bail!("cannot build Arrow offsets from a negative length"), + 2 => vortex_bail!("length sum exceeds Arrow i32 offset range"), + status => vortex_bail!("unexpected Arrow offset status {status}"), + } + + Ok(I32Offsets { + buffer: offsets, + total: usize::try_from(Buffer::::from_byte_buffer(total_bytes)[0])?, + }) +} diff --git a/vortex-cuda/src/bit_unpack_gen.rs b/vortex-cuda/src/bit_unpack_gen.rs index 2482c0996b8..7b3ce148753 100644 --- a/vortex-cuda/src/bit_unpack_gen.rs +++ b/vortex-cuda/src/bit_unpack_gen.rs @@ -152,11 +152,16 @@ __device__ void _bit_unpack_{bits}_device(const uint{bits}_t *__restrict in, uin }} __syncwarp(); - // Step 2: Apply patches to shared memory in parallel + // Step 2: Apply patches to shared memory in parallel. + // + // Patch values are stored in the same frame-of-reference domain as the packed values, so + // they take the same `+ reference` the lane decoder applies to every unpacked value. For a + // plain bit-packed array the reference is zero and this is a no-op; for FoR-over-bit-packed + // it is what keeps patched positions from coming out short by the reference. PatchesCursor cursor(patches, blockIdx.x, thread_idx, {thread_count}); auto patch = cursor.next(); while (patch.index != FL_CHUNK) {{ - shared_out[patch.index] = patch.value; + shared_out[patch.index] = patch.value + reference; patch = cursor.next(); }} __syncwarp(); diff --git a/vortex-cuda/src/canonical.rs b/vortex-cuda/src/canonical.rs index 9f1fce7e68e..5a3d1e224f2 100644 --- a/vortex-cuda/src/canonical.rs +++ b/vortex-cuda/src/canonical.rs @@ -5,29 +5,46 @@ use std::sync::Arc; use async_trait::async_trait; use futures::future::try_join_all; +use vortex::array::ArrayRef; use vortex::array::Canonical; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; use vortex::array::arrays::BoolArray; use vortex::array::arrays::DecimalArray; use vortex::array::arrays::ExtensionArray; +use vortex::array::arrays::ListViewArray; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::StructArray; use vortex::array::arrays::VarBinViewArray; use vortex::array::arrays::bool::BoolDataParts; use vortex::array::arrays::decimal::DecimalDataParts; use vortex::array::arrays::extension::ExtensionArrayExt; +use vortex::array::arrays::listview::ListViewDataParts; use vortex::array::arrays::primitive::PrimitiveDataParts; use vortex::array::arrays::struct_::StructDataParts; use vortex::array::arrays::varbinview::BinaryView; use vortex::array::arrays::varbinview::VarBinViewDataParts; use vortex::array::buffer::BufferHandle; use vortex::array::legacy_session; +use vortex::array::validity::Validity; use vortex::buffer::BitBuffer; use vortex::buffer::Buffer; use vortex::buffer::ByteBuffer; use vortex::error::VortexResult; +/// Copy a canonical child array to the host. +async fn child_into_host(child: ArrayRef) -> VortexResult { + #[expect( + clippy::disallowed_methods, + reason = "CanonicalCudaExt threads no session through" + )] + Ok(child + .execute::(&mut legacy_session().create_execution_ctx())? + .into_host() + .await? + .into_array()) +} + /// Move all canonical data from to_host from device. #[async_trait] pub trait CanonicalCudaExt { @@ -36,9 +53,32 @@ pub trait CanonicalCudaExt { Self: Sized; } +/// Copies an array-backed validity mask back to the host. +/// +/// Only [`Validity::Array`] owns a buffer; the other variants are metadata and pass through. +/// Migrating the values of a nullable array without its validity leaves the mask on the +/// device, and the first host read of it — canonicalising to Arrow, say — panics in +/// `BufferHandle::unwrap_host`. +async fn validity_into_host(validity: Validity) -> VortexResult { + let Validity::Array(array) = validity else { + return Ok(validity); + }; + #[expect( + clippy::disallowed_methods, + reason = "CanonicalCudaExt threads no session through" + )] + let mut ctx = legacy_session().create_execution_ctx(); + Ok(Validity::Array( + array + .execute::(&mut ctx)? + .into_host() + .await? + .into_array(), + )) +} + #[async_trait] impl CanonicalCudaExt for Canonical { - #[allow(clippy::disallowed_methods)] async fn into_host(self) -> VortexResult { match self { Canonical::Struct(struct_array) => { @@ -53,10 +93,15 @@ impl CanonicalCudaExt for Canonical { let mut host_fields = vec![]; for field in fields.iter() { + #[expect( + clippy::disallowed_methods, + reason = "CanonicalCudaExt threads no session through" + )] + let mut ctx = legacy_session().create_execution_ctx(); host_fields.push( field .clone() - .execute::(&mut legacy_session().create_execution_ctx())? + .execute::(&mut ctx)? .into_host() .await? .into_array(), @@ -67,15 +112,13 @@ impl CanonicalCudaExt for Canonical { struct_fields.names().clone(), host_fields, len, - validity, + validity_into_host(validity).await?, ))) } n @ Canonical::Null(_) => Ok(n), Canonical::Bool(bool) => { - // NOTE: update to copy to host when adding buffer handle. - // Also update other method to copy validity to host. let len = bool.len(); - let validity = bool.validity()?; + let validity = validity_into_host(bool.validity()?).await?; let BoolDataParts { bits, meta } = bool.into_data().into_parts(len); let bits = BitBuffer::new_with_offset( @@ -95,7 +138,7 @@ impl CanonicalCudaExt for Canonical { Ok(Canonical::Primitive(PrimitiveArray::from_byte_buffer( buffer.try_into_host()?.await?, ptype, - validity, + validity_into_host(validity).await?, ))) } Canonical::Decimal(decimal) => { @@ -106,6 +149,7 @@ impl CanonicalCudaExt for Canonical { validity, .. } = decimal.into_data_parts(); + let validity = validity_into_host(validity).await?; Ok(Canonical::Decimal(unsafe { DecimalArray::new_unchecked_handle( BufferHandle::new_host(values.try_into_host()?.await?), @@ -122,6 +166,7 @@ impl CanonicalCudaExt for Canonical { validity, dtype, } = varbinview.into_data_parts(); + let validity = validity_into_host(validity).await?; // Copy all device views to host let host_views = views.try_into_host()?.await?; @@ -140,12 +185,34 @@ impl CanonicalCudaExt for Canonical { VarBinViewArray::new_unchecked(host_views, host_buffers, dtype, validity) })) } + Canonical::List(list) => { + let ListViewDataParts { + elements, + offsets, + sizes, + validity, + .. + } = list.into_data_parts(); + let validity = validity_into_host(validity).await?; + + Ok(Canonical::List(ListViewArray::try_new( + child_into_host(elements).await?, + child_into_host(offsets).await?, + child_into_host(sizes).await?, + validity, + )?)) + } Canonical::Extension(ext) => { // Copy the storage array to host and rewrap in ExtensionArray. + #[expect( + clippy::disallowed_methods, + reason = "CanonicalCudaExt threads no session through" + )] + let mut ctx = legacy_session().create_execution_ctx(); let host_storage = ext .storage_array() .clone() - .execute::(&mut legacy_session().create_execution_ctx())? + .execute::(&mut ctx)? .into_host() .await? .into_array(); diff --git a/vortex-cuda/src/cub.rs b/vortex-cuda/src/cub.rs index 4b6009e7cea..0612a13fc07 100644 --- a/vortex-cuda/src/cub.rs +++ b/vortex-cuda/src/cub.rs @@ -8,13 +8,70 @@ use std::ffi::c_void; use cudarc::driver::CudaSlice; use cudarc::driver::DevicePtr; use cudarc::driver::DevicePtrMut; +use vortex::array::buffer::BufferHandle; use vortex::error::VortexResult; use vortex::error::vortex_err; +use vortex_cub::onpair; use vortex_cub::scan; use vortex_cub::scan::cudaStream_t; +use crate::CudaBufferExt; use crate::CudaExecutionCtx; +/// Regenerate the OnPair decode kernel's per-batch output offsets in one +/// fused sweep (see `cub/kernels/onpair.cu`): the per-batch decoded-size +/// reduction and the exclusive scan over the sizes run in a single kernel via +/// decoupled look-back. `code_width` selects the code stream's element size in +/// bytes (1 or 2). Returns `num_batches + 1` offsets; the last is the total +/// decoded byte count. A code outside the dictionary raises `status` to 1; +/// the caller must check the flag before trusting the offsets. +#[allow(clippy::too_many_arguments)] +pub(crate) fn onpair_batch_offsets( + codes: &BufferHandle, + code_width: u32, + lens: &BufferHandle, + dict_size: u32, + num_tokens: usize, + num_batches: usize, + status: &mut CudaSlice, + ctx: &mut CudaExecutionCtx, +) -> VortexResult> { + let num_batches_i64 = i64::try_from(num_batches)?; + let temp_bytes = onpair::batch_offsets_temp_size(num_batches_i64) + .map_err(|err| vortex_err!("CUB onpair_batch_offsets_temp_size failed: {err}"))?; + + let mut temp = ctx.device_alloc::(temp_bytes.max(1))?; + let mut chunk_offsets = ctx.device_alloc::(num_batches + 1)?; + let codes_ptr = codes.cuda_device_ptr()?; + let lens_ptr = lens.cuda_device_ptr()?; + let total_tokens = u64::try_from(num_tokens)?; + let stream = ctx.stream(); + let stream_ptr = stream.cu_stream() as cudaStream_t; + let (status_ptr, record_status) = status.device_ptr_mut(stream); + let (offsets_ptr, record_offsets) = chunk_offsets.device_ptr_mut(stream); + let (temp_ptr, record_temp) = temp.device_ptr_mut(stream); + + ctx.launch_external(num_tokens, || unsafe { + onpair::batch_offsets( + temp_ptr as *mut c_void, + temp_bytes, + codes_ptr as *const c_void, + code_width, + lens_ptr as *const u8, + dict_size, + total_tokens, + offsets_ptr as *mut u64, + status_ptr as *mut u32, + num_batches_i64, + stream_ptr, + ) + .map_err(|err| vortex_err!("CUB onpair_batch_offsets failed: {err}")) + })?; + drop((record_status, record_offsets, record_temp)); + + Ok(chunk_offsets) +} + pub(crate) fn exclusive_sum_i32( input: &CudaSlice, len: usize, @@ -47,3 +104,138 @@ pub(crate) fn exclusive_sum_i32( Ok(output) } + +#[cfg(test)] +mod tests { + use vortex::error::VortexExpect; + + use super::*; + use crate::session::CudaSession; + + /// Upload synthetic codes and lengths, regenerate the chunk offsets, and + /// read them back together with the status flag. The code width follows + /// the element type. + async fn batch_offsets_roundtrip( + codes: Vec, + lens: Vec, + ) -> VortexResult<(Vec, u32)> + where + C: cudarc::driver::DeviceRepr + + cudarc::driver::ValidAsZeroBits + + std::fmt::Debug + + Send + + Sync + + 'static, + { + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let num_tokens = codes.len(); + let num_batches = num_tokens.div_ceil(128); + let dict_size = u32::try_from(lens.len())?; + let code_width = u32::try_from(size_of::())?; + + let codes_dev = ctx.copy_to_device(codes)?.await?; + let lens_dev = ctx.copy_to_device(lens)?.await?; + let mut status = ctx.device_alloc::(1)?; + ctx.stream() + .memset_zeros(&mut status) + .map_err(|e| vortex_err!("Failed to zero status flag: {e}"))?; + + let offsets = onpair_batch_offsets( + &codes_dev, + code_width, + &lens_dev, + dict_size, + num_tokens, + num_batches, + &mut status, + &mut ctx, + )?; + + let offsets = ctx + .stream() + .clone_dtoh(&offsets) + .map_err(|e| vortex_err!("Failed to copy offsets to host: {e}"))?; + let status = ctx + .stream() + .clone_dtoh(&status) + .map_err(|e| vortex_err!("Failed to copy status to host: {e}"))?; + Ok((offsets, status[0])) + } + + /// The exclusive prefix at 128-token boundaries, plus the trailing total. + fn host_reference(codes: &[u16], lens: &[u8]) -> Vec { + let mut expected = Vec::with_capacity(codes.len().div_ceil(128) + 1); + expected.push(0u64); + let mut acc = 0u64; + for (i, &code) in codes.iter().enumerate() { + acc += u64::from(lens[code as usize]); + if (i + 1) % 128 == 0 { + expected.push(acc); + } + } + if !codes.len().is_multiple_of(128) { + expected.push(acc); + } + expected + } + + /// A single partial batch: one tile, no look-back. + #[crate::test] + async fn test_onpair_batch_offsets_single_batch() -> VortexResult<()> { + let lens: Vec = (1..=16).collect(); + let codes: Vec = (0..100u16).map(|i| i % 16).collect(); + let expected = host_reference(&codes, &lens); + + let (offsets, status) = batch_offsets_roundtrip(codes, lens).await?; + assert_eq!(status, 0); + assert_eq!(offsets, expected); + Ok(()) + } + + /// Many look-back tiles with a ragged tail batch; the offsets must match + /// a host prefix sum sampled at 128-token boundaries. Regression test for + /// the look-back prefix being defined only in lane 0. + #[crate::test] + async fn test_onpair_batch_offsets_multi_tile() -> VortexResult<()> { + let lens: Vec = (1..=16u8).cycle().take(300).collect(); + let codes: Vec = (0..2000u32 * 128 - 57) + .map(|i| u16::try_from(i * 31 % 300).vortex_expect("bounded by dictionary size")) + .collect(); + let expected = host_reference(&codes, &lens); + + let (offsets, status) = batch_offsets_roundtrip(codes, lens).await?; + assert_eq!(status, 0); + assert_eq!(offsets, expected); + Ok(()) + } + + /// u8 codes dispatch the narrow sweep instantiation. + #[crate::test] + async fn test_onpair_batch_offsets_u8_codes() -> VortexResult<()> { + let lens: Vec = (1..=16).collect(); + let codes: Vec = (0..300u32) + .map(|i| u8::try_from(i * 7 % 16).vortex_expect("bounded by dictionary size")) + .collect(); + let widened: Vec = codes.iter().map(|&c| u16::from(c)).collect(); + let expected = host_reference(&widened, &lens); + + let (offsets, status) = batch_offsets_roundtrip(codes, lens).await?; + assert_eq!(status, 0); + assert_eq!(offsets, expected); + Ok(()) + } + + /// A code outside the dictionary raises the status flag and contributes + /// zero bytes. + #[crate::test] + async fn test_onpair_batch_offsets_flags_out_of_range_code() -> VortexResult<()> { + let lens = vec![2u8; 4]; + let mut codes = vec![1u16; 200]; + codes[130] = 9; + + let (offsets, status) = batch_offsets_roundtrip(codes, lens).await?; + assert_eq!(status, 1); + assert_eq!(offsets, vec![0, 256, 256 + 71 * 2]); + Ok(()) + } +} diff --git a/vortex-cuda/src/device_buffer.rs b/vortex-cuda/src/device_buffer.rs index 954731fab44..602fa16ba94 100644 --- a/vortex-cuda/src/device_buffer.rs +++ b/vortex-cuda/src/device_buffer.rs @@ -12,6 +12,7 @@ use std::sync::Arc; use cudarc::driver::CudaSlice; use cudarc::driver::CudaView; +use cudarc::driver::CudaViewMut; use cudarc::driver::DevicePtr; use cudarc::driver::DeviceRepr; use cudarc::driver::sys; @@ -55,7 +56,6 @@ pub struct CudaDeviceBuffer { /// reads without copying or repacking. zeroed_tail_start: Option, } - mod private { use std::fmt::Debug; use std::sync::Arc; @@ -63,6 +63,7 @@ mod private { use cudarc::driver::CudaSlice; use cudarc::driver::CudaStream; use cudarc::driver::CudaView; + use cudarc::driver::CudaViewMut; use cudarc::driver::DeviceRepr; use vortex::buffer::Alignment; use vortex::error::VortexExpect; @@ -74,8 +75,11 @@ mod private { /// Get a reference to the underlying cuStream. fn stream(&self) -> &Arc; - /// Access the values as a bytes view + /// Access the values as a bytes view. fn as_bytes_view(&self) -> CudaView<'_, u8>; + + /// Access the values as a mutable bytes view. + fn as_bytes_view_mut(&mut self) -> CudaViewMut<'_, u8>; } // CudaSlice needs to be held by the CudaDeviceBuffer @@ -90,10 +94,17 @@ mod private { fn as_bytes_view(&self) -> CudaView<'_, u8> { let bytes_len = self.len() * size_of::(); - // SAFETY: all types can be reinterpreted as a byte slice + // SAFETY: all types can be reinterpreted as a byte slice. let result = unsafe { self.as_view().transmute::(bytes_len) }; result.vortex_expect("Downcasting CudaSlice => CudaSlice must succeed") } + + fn as_bytes_view_mut(&mut self) -> CudaViewMut<'_, u8> { + let bytes_len = self.len() * size_of::(); + // SAFETY: all types can be reinterpreted as a mutable byte slice. + let result = unsafe { self.transmute_mut::(bytes_len) }; + result.vortex_expect("Downcasting CudaSlice => CudaSlice must succeed") + } } } @@ -146,8 +157,8 @@ impl CudaDeviceBuffer { // Return a new &[T] let new_len = self.len / size_of::(); - // SAFETY: All DeviecRepr types are aligned to < 256 bytes, which is what CUDA allocator - // gives us back. So we should not suffer any alignment issues at runtime. + // SAFETY: All DeviceRepr types are aligned to < 256 bytes, which is what CUDA allocator + // gives us back. So we should not suffer any alignment issues at runtime. unsafe { self.allocation .as_bytes_view() @@ -156,6 +167,21 @@ impl CudaDeviceBuffer { .vortex_expect("Failed to transmute from CudaView to CudaView") } } + + pub(crate) fn with_view_mut( + &mut self, + function: impl FnOnce(&mut CudaViewMut<'_, T>) -> R, + ) -> Option { + let new_len = self.len / size_of::(); + let allocation = Arc::get_mut(&mut self.allocation)?; + let mut bytes = allocation.as_bytes_view_mut(); + let mut bytes = bytes.slice_mut(self.offset..self.offset + self.len); + // SAFETY: The same alignment and size requirements as `as_view` apply, and unique + // ownership of both allocation Arcs guarantees no other device view can alias this one. + let mut values = unsafe { bytes.transmute_mut::(new_len) } + .vortex_expect("Failed to transmute from CudaViewMut to CudaViewMut"); + Some(function(&mut values)) + } } #[cfg(test)] @@ -180,10 +206,30 @@ pub(crate) fn cuda_backing_allocation(handle: &BufferHandle) -> VortexResult( + handle: BufferHandle, + function: impl FnOnce(&mut CudaViewMut<'_, T>) -> R, +) -> VortexResult<(BufferHandle, Option)> { + if !handle.is_on_device() { + return Err(vortex_err!("Buffer is not on device")); + } + + let device_buffer = handle.unwrap_device(); + if !device_buffer.as_any().is::() { + return Err(vortex_err!( + "expected CudaDeviceBuffer, was {device_buffer:?}" + )); + } + + let device_buffer: Arc = device_buffer; + let mut cuda_buf = Arc::downcast::(device_buffer) + .map_err(|_| vortex_err!("CudaDeviceBuffer downcast failed after type check"))?; + let result = Arc::get_mut(&mut cuda_buf).and_then(|buffer| buffer.with_view_mut(function)); + + Ok((BufferHandle::new_device(cuda_buf), result)) +} + +/// Extension trait for getting CUDA views from a [`BufferHandle`]. pub trait CudaBufferExt { /// Returns a readonly [`CudaView`] for the buffer handle. /// @@ -340,9 +386,10 @@ impl DeviceBuffer for CudaDeviceBuffer { alignment: Alignment, ) -> VortexResult>> { let stream = self.allocation.stream(); - - // Add offset to device pointer to account for any previous slicing operations. - let src_ptr = self.device_ptr + self.offset as u64; + let source = self.as_view::(); + // Use cudarc's read guard so this allocation stream waits for writes submitted on any + // other managed stream before scheduling the device-to-host copy. + let (src_ptr, record_read) = source.device_ptr(stream); let mut host_buffer: ByteBufferMut = ByteBufferMut::with_capacity_aligned(self.len, alignment); @@ -365,6 +412,7 @@ impl DeviceBuffer for CudaDeviceBuffer { .result() .map_err(|e| vortex_err!("Failed to schedule async copy to host: {}", e))?; } + drop(record_read); let cuda_slice = Arc::clone(&self.allocation); @@ -425,7 +473,7 @@ impl DeviceBuffer for CudaDeviceBuffer { fn aligned(self: Arc, alignment: Alignment) -> VortexResult> { let effective_ptr = self.device_ptr + self.offset as u64; - if effective_ptr.is_multiple_of(*alignment as u64) { + if effective_ptr.is_multiple_of(alignment.as_usize() as u64) { Ok(Arc::new(CudaDeviceBuffer { allocation: Arc::clone(&self.allocation), offset: self.offset, @@ -441,3 +489,69 @@ impl DeviceBuffer for CudaDeviceBuffer { } } } + +#[cfg(test)] +mod tests { + use cudarc::driver::CudaContext; + use cudarc::driver::PushKernelArg; + use vortex::buffer::Buffer; + use vortex::dtype::PType; + use vortex::error::VortexResult; + use vortex::error::vortex_bail; + + use super::*; + use crate::CudaSession; + + #[crate::test] + async fn mutable_view_requires_unique_ownership() -> VortexResult<()> { + let ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let allocation = ctx.device_alloc::(4)?; + let mut handle = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(allocation))); + + let (next, result) = with_cuda_view_mut::(handle, |view| view.len())?; + handle = next; + assert_eq!(result, Some(4)); + + let alias = handle.clone(); + let (next, result) = with_cuda_view_mut::(handle, |view| view.len())?; + handle = next; + assert!(result.is_none()); + drop(alias); + + let (_, result) = with_cuda_view_mut::(handle, |view| view.len())?; + assert_eq!(result, Some(4)); + Ok(()) + } + + #[crate::test] + async fn copy_to_host_waits_for_write_on_another_stream() -> VortexResult<()> { + let context = CudaContext::new(0) + .map_err(|err| vortex_err!("failed to create CUDA context: {err}"))?; + let cuda_session = CudaSession::with_stream_pool_capacity(context, 2); + let session = vortex::array::array_session().with_some(cuda_session); + let allocation_ctx = CudaSession::create_execution_ctx(&session)?; + let mut launch_ctx = CudaSession::create_execution_ctx(&session)?; + assert_ne!( + allocation_ctx.stream().cu_stream(), + launch_ctx.stream().cu_stream() + ); + + let handle = allocation_ctx.copy_to_device(vec![0u32; 4])?.await?; + let kernel = launch_ctx.load_function("constant_numeric", &[PType::U32])?; + let value = 42u32; + let len = 4u64; + let (handle, launch) = with_cuda_view_mut::(handle, |view| { + launch_ctx.launch_kernel(&kernel, 4, |args| { + args.arg(view).arg(&value).arg(&len); + }) + })?; + let Some(launch) = launch else { + vortex_bail!("fresh CUDA allocation was unexpectedly shared"); + }; + launch?; + + let values = Buffer::::from_byte_buffer(handle.try_to_host()?.await?); + assert_eq!(values.as_slice(), &[42; 4]); + Ok(()) + } +} diff --git a/vortex-cuda/src/dynamic_dispatch/mod.rs b/vortex-cuda/src/dynamic_dispatch/mod.rs index 329ea19c061..ad9c02ba6a1 100644 --- a/vortex-cuda/src/dynamic_dispatch/mod.rs +++ b/vortex-cuda/src/dynamic_dispatch/mod.rs @@ -471,7 +471,7 @@ impl MaterializedPlan { ))); } - let output_buf = CudaDeviceBuffer::new(ctx.device_alloc::(len.next_multiple_of(1024))?); + let mut output = ctx.device_alloc::(len.next_multiple_of(1024))?; // Upload the packed plan bytes to the device. let device_plan = Arc::new( @@ -503,20 +503,20 @@ impl MaterializedPlan { .map(|view| view.device_ptr(&stream).1) .collect::>(); - let output_ptr = output_buf.offset_ptr(); let (plan_ptr, plan_read_record) = device_plan.device_ptr(&stream); let array_len_u64 = len as u64; ctx.launch_kernel_config(&cuda_function, config, len, |args| { - args.arg(&output_ptr); + args.arg(&mut output); args.arg(&array_len_u64); args.arg(&plan_ptr); })?; drop((device_buffer_read_records, plan_read_record)); + let output = CudaDeviceBuffer::new(output); Ok(Canonical::Primitive(PrimitiveArray::from_buffer_handle( - BufferHandle::new_device(output_buf.slice_typed::(0..len)), + BufferHandle::new_device(output.slice_typed::(0..len)), output_ptype, self.validity, ))) @@ -560,6 +560,7 @@ mod tests { use vortex::encodings::fastlanes::BitPackedArrayExt; use vortex::encodings::fastlanes::FoR; use vortex::encodings::fastlanes::FoRArrayExt; + use vortex::encodings::fastlanes::FoRArraySlotsExt; use vortex::encodings::runend::RunEnd; use vortex::encodings::sequence::Sequence; use vortex::encodings::zigzag::ZigZag; @@ -576,7 +577,6 @@ mod tests { use super::*; use crate::CanonicalCudaExt; use crate::CudaBufferExt; - use crate::CudaDeviceBuffer; use crate::CudaExecutionCtx; use crate::cuda_session; use crate::executor::CudaArrayExt; @@ -787,12 +787,9 @@ mod tests { plan: &CudaDispatchPlan, shared_mem_bytes: u32, ) -> VortexResult> { - let output_slice = cuda_ctx + let mut output = cuda_ctx .device_alloc::(output_len) .vortex_expect("alloc output"); - let output_buf = CudaDeviceBuffer::new(output_slice); - let output_view = output_buf.as_view::(); - let (output_ptr, record_output) = output_view.device_ptr(cuda_ctx.stream()); let device_plan = Arc::new( cuda_ctx @@ -812,7 +809,7 @@ mod tests { .load_function("dynamic_dispatch", &[PType::U32]) .vortex_expect("load kernel"); let mut launch_builder = cuda_ctx.launch_builder(&cuda_function); - launch_builder.arg(&output_ptr); + launch_builder.arg(&mut output); launch_builder.arg(&array_len_u64); launch_builder.arg(&plan_ptr); @@ -827,11 +824,11 @@ mod tests { .launch(config) .map_err(|e| vortex_err!("kernel launch: {e}"))?; } - drop((record_output, record_plan)); + drop(record_plan); cuda_ctx .stream() - .clone_dtoh(&output_buf.as_view::()) + .clone_dtoh(&output) .map_err(|e| vortex_err!("copy back: {e}")) } diff --git a/vortex-cuda/src/dynamic_dispatch/plan_builder.rs b/vortex-cuda/src/dynamic_dispatch/plan_builder.rs index d7c5a58712d..43f6a5a3712 100644 --- a/vortex-cuda/src/dynamic_dispatch/plan_builder.rs +++ b/vortex-cuda/src/dynamic_dispatch/plan_builder.rs @@ -18,7 +18,7 @@ use vortex::array::arrays::ScalarFn; use vortex::array::arrays::Slice; use vortex::array::arrays::dict::DictArraySlotsExt; use vortex::array::arrays::scalar_fn::ScalarFnArrayExt; -use vortex::array::arrays::slice::SliceArrayExt; +use vortex::array::arrays::slice::SliceArraySlotsExt; use vortex::array::buffer::BufferHandle; use vortex::array::patches::Patches; use vortex::array::validity::Validity; @@ -32,11 +32,13 @@ use vortex::encodings::fastlanes::BitPacked; use vortex::encodings::fastlanes::BitPackedArrayExt; use vortex::encodings::fastlanes::FoR; use vortex::encodings::fastlanes::FoRArrayExt; +use vortex::encodings::fastlanes::FoRArraySlotsExt; use vortex::encodings::runend::RunEnd; use vortex::encodings::runend::RunEndArrayExt; +use vortex::encodings::runend::RunEndArraySlotsExt; use vortex::encodings::sequence::Sequence; use vortex::encodings::zigzag::ZigZag; -use vortex::encodings::zigzag::ZigZagArrayExt; +use vortex::encodings::zigzag::ZigZagArraySlotsExt; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::error::vortex_err; diff --git a/vortex-cuda/src/executor.rs b/vortex-cuda/src/executor.rs index 6a03bfc2dca..1c56bdab090 100644 --- a/vortex-cuda/src/executor.rs +++ b/vortex-cuda/src/executor.rs @@ -197,7 +197,7 @@ impl CudaExecutionCtx { /// /// # Arguments /// - /// * `module_name` - Name of the module (`kernels/{module_name}.ptx`) + /// * `module_name` - Kernel source name without the `.cu` extension /// * `ptypes` - List of ptype strings for the kernel name /// /// # Errors @@ -222,7 +222,7 @@ impl CudaExecutionCtx { /// /// # Arguments /// - /// * `module_name` - Name of the module (`kernels/{module_name}.ptx`) + /// * `module_name` - Kernel source name without the `.cu` extension /// * `type_suffixes` - List of type suffix strings for the kernel name /// /// # Errors @@ -239,7 +239,8 @@ impl CudaExecutionCtx { /// Returns a launch builder for a CUDA kernel function. /// - /// Arguments can be added to the kernel launch with `.arg(buffer)`. + /// Arguments can be added to the kernel launch with `.arg(&buffer)`. Buffers written by the + /// kernel must be passed with `.arg(&mut buffer)` so their writes are tracked. /// /// # Arguments /// @@ -312,7 +313,6 @@ impl CudaExecutionCtx { } /// Returns the Vortex session backing this CUDA execution context. - #[cfg(feature = "unstable_encodings")] pub(crate) fn session(&self) -> &vortex::session::VortexSession { self.ctx.session() } diff --git a/vortex-cuda/src/file.rs b/vortex-cuda/src/file.rs new file mode 100644 index 00000000000..2a9806729f7 --- /dev/null +++ b/vortex-cuda/src/file.rs @@ -0,0 +1,254 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! CUDA-aware Vortex file opening. + +use std::path::Path; +use std::sync::Arc; + +use vortex::error::VortexResult; +use vortex::error::vortex_bail; +use vortex::file::VortexFile; +use vortex::file::VortexOpenOptions; +use vortex::io::session::RuntimeSessionExt; +use vortex::layout::DynLayout; +use vortex::layout::layouts::zoned::LegacyStats; +use vortex::layout::layouts::zoned::Zoned; +use vortex::layout::segments::SegmentFuture; +use vortex::layout::segments::SegmentId; +use vortex::layout::segments::SegmentSource; + +use crate::CudaSessionExt; +use crate::PooledFileReadAt; +use crate::PooledFileReadAtOptions; +use crate::layout::register_cuda_layout; + +/// Extension trait for opening CUDA-readable files from [`VortexOpenOptions`]. +pub trait CudaOpenOptionsExt { + /// Configure the opener to keep control-plane segments on the host and read data-plane + /// segments into CUDA device memory. + fn with_cuda(self) -> CudaOpenOptions; +} + +impl CudaOpenOptionsExt for VortexOpenOptions { + fn with_cuda(self) -> CudaOpenOptions { + CudaOpenOptions { + inner: self, + read_at_options: PooledFileReadAtOptions::default(), + } + } +} + +/// CUDA-aware Vortex file open options. +/// +/// Construct this adapter with [`CudaOpenOptionsExt::with_cuda`]. Standard file options should be +/// configured before calling `with_cuda`. +pub struct CudaOpenOptions { + inner: VortexOpenOptions, + read_at_options: PooledFileReadAtOptions, +} + +impl CudaOpenOptions { + /// Configure how the pooled data-plane reader opens and reads the local file. + /// + /// Footer and zone-map reads continue to use the standard buffered host reader. + pub fn with_read_at_options(mut self, options: PooledFileReadAtOptions) -> Self { + self.read_at_options = options; + self + } + + /// Open a local Vortex file for CUDA execution. + /// + /// The footer and zone-map segments are read through the ordinary host path. All other file + /// segments are staged through the session-scoped pinned buffer pool and transferred to the + /// device. + pub async fn open_path(self, path: impl AsRef) -> VortexResult { + let path = path.as_ref().to_owned(); + register_cuda_layout(self.inner.session()); + + // A host cache cannot front the data source: cached host buffers would violate the CUDA + // source's device-placement contract. It remains enabled on the control-plane source. + let data_options = self.inner.clone().without_segment_cache(); + let control_file = self.inner.open_path(&path).await?; + let footer = control_file.footer().clone(); + + let session = control_file.session().clone(); + let cuda_session = session.cuda_session(); + let stream = cuda_session.stream()?; + let pool = Arc::clone(cuda_session.pinned_buffer_pool()); + drop(cuda_session); + + let reader = Arc::new(PooledFileReadAt::open_with_options( + &path, + session.handle(), + pool, + stream, + self.read_at_options, + )?); + let data_file = data_options + .with_footer(footer.clone()) + .open(reader) + .await?; + + let control_segments = + control_plane_segments(footer.layout().as_ref(), footer.segment_map().len())?; + let routed = Arc::new(RoutingSegmentSource { + control_segments: control_segments.into(), + control: control_file.segment_source(), + data: data_file.segment_source(), + }); + + Ok(data_file.with_segment_source(routed)) + } +} + +struct RoutingSegmentSource { + control_segments: Arc<[bool]>, + control: Arc, + data: Arc, +} + +impl SegmentSource for RoutingSegmentSource { + fn request(&self, id: SegmentId) -> SegmentFuture { + let Ok(idx) = usize::try_from(*id) else { + return self.data.request(id); + }; + if self.control_segments.get(idx).copied().unwrap_or(false) { + self.control.request(id) + } else { + self.data.request(id) + } + } +} + +fn control_plane_segments(layout: &dyn DynLayout, segment_count: usize) -> VortexResult> { + let mut control_segments = vec![false; segment_count]; + mark_zone_map_segments(layout, &mut control_segments)?; + Ok(control_segments) +} + +fn mark_zone_map_segments( + layout: &dyn DynLayout, + control_segments: &mut [bool], +) -> VortexResult<()> { + if layout.is::() || layout.is::() { + let (Some(data), Some(zones)) = (layout.slot(0)?, layout.slot(1)?) else { + vortex_bail!("zone map layout is missing its data or zones child"); + }; + mark_layout_segments(zones.as_ref(), control_segments)?; + mark_zone_map_segments(data.as_ref(), control_segments)?; + return Ok(()); + } + + for child in layout.children()? { + mark_zone_map_segments(child.as_ref(), control_segments)?; + } + Ok(()) +} + +fn mark_layout_segments(layout: &dyn DynLayout, control_segments: &mut [bool]) -> VortexResult<()> { + for id in layout.segment_ids() { + let idx = usize::try_from(*id)?; + let Some(is_control) = control_segments.get_mut(idx) else { + vortex_bail!("layout references out-of-range segment {id}"); + }; + *is_control = true; + } + + for child in layout.children()? { + mark_layout_segments(child.as_ref(), control_segments)?; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::num::NonZeroUsize; + + use futures::FutureExt; + use parking_lot::Mutex; + use vortex::array::buffer::BufferHandle; + use vortex::array::dtype::DType; + use vortex::array::dtype::Nullability::NonNullable; + use vortex::array::dtype::PType; + use vortex::array::dtype::StructFields; + use vortex::buffer::ByteBuffer; + use vortex::layout::layouts::flat::FlatLayout; + use vortex::layout::layouts::zoned::ZonedLayout; + use vortex::session::registry::ReadContext; + + use super::*; + + #[derive(Clone)] + struct RecordingSource { + requests: Arc>>, + value: u8, + } + + impl RecordingSource { + fn new(value: u8) -> Self { + Self { + requests: Arc::new(Mutex::new(Vec::new())), + value, + } + } + } + + impl SegmentSource for RecordingSource { + fn request(&self, id: SegmentId) -> SegmentFuture { + self.requests.lock().push(id); + let value = self.value; + async move { Ok(BufferHandle::new_host(ByteBuffer::from(vec![value]))) }.boxed() + } + } + + #[tokio::test] + async fn routes_control_and_data_segments() { + let control = RecordingSource::new(1); + let data = RecordingSource::new(2); + let source = RoutingSegmentSource { + control_segments: Arc::from([false, true, false]), + control: Arc::new(control.clone()), + data: Arc::new(data.clone()), + }; + + let control_buffer = source + .request(SegmentId::from(1)) + .await + .unwrap() + .unwrap_host(); + let data_buffer = source + .request(SegmentId::from(2)) + .await + .unwrap() + .unwrap_host(); + + assert_eq!(control_buffer.as_ref(), &[1]); + assert_eq!(data_buffer.as_ref(), &[2]); + assert_eq!(*control.requests.lock(), [SegmentId::from(1)]); + assert_eq!(*data.requests.lock(), [SegmentId::from(2)]); + } + + #[test] + fn marks_only_zone_map_subtree_segments_as_control_plane() -> VortexResult<()> { + let read_ctx = ReadContext::new([]); + let data_dtype = DType::Primitive(PType::I32, NonNullable); + let data = + FlatLayout::new(16, data_dtype, SegmentId::from(0), read_ctx.clone()).into_layout(); + let zones = FlatLayout::new( + 1, + DType::Struct(StructFields::empty(), NonNullable), + SegmentId::from(1), + read_ctx, + ) + .into_layout(); + let layout = + ZonedLayout::try_new(data, zones, NonZeroUsize::new(16).unwrap(), Arc::new([]))? + .into_layout(); + + let control_segments = control_plane_segments(layout.as_ref(), 2)?; + + assert_eq!(control_segments, [false, true]); + Ok(()) + } +} diff --git a/vortex-cuda/src/hybrid_dispatch/mod.rs b/vortex-cuda/src/hybrid_dispatch/mod.rs index b3a0151a727..d607f9da7d9 100644 --- a/vortex-cuda/src/hybrid_dispatch/mod.rs +++ b/vortex-cuda/src/hybrid_dispatch/mod.rs @@ -297,7 +297,6 @@ mod tests { /// Dict(values=ZstdBuffers(FoR(BP)), codes=FoR(BP)) — ZstdBuffers is /// executed separately, then Dict+FoR+BP fuses with its output as a LOAD. /// 3 launches: nvcomp + fused FoR+BP + fused LOAD+FoR+BP+DICT. - #[cfg(feature = "unstable_encodings")] #[crate::test] async fn test_partial_fusion() -> VortexResult<()> { let mut cpu_ctx = array_session().create_execution_ctx(); diff --git a/vortex-cuda/src/kernel/arrays/constant.rs b/vortex-cuda/src/kernel/arrays/constant.rs index d385333a024..b52f1a0cb13 100644 --- a/vortex-cuda/src/kernel/arrays/constant.rs +++ b/vortex-cuda/src/kernel/arrays/constant.rs @@ -106,8 +106,7 @@ where .ok_or_else(|| vortex_err!("Expected non-null primitive scalar value"))?; // Allocate output buffer on device - let output_buffer = ctx.device_alloc::

(array_len)?; - let output_view = output_buffer.as_view(); + let mut output_buffer = ctx.device_alloc::

()?; + let mut device_buffer = ctx.ensure_on_device(buffer).await?; let array_len_u64 = array_len as u64; - // Load kernel function - let kernel_ptypes = [P::PTYPE]; - let cuda_function = ctx.load_function("for", &kernel_ptypes)?; + let cuda_function = ctx.load_function("for", &[P::PTYPE])?; + let (next, launch) = with_cuda_view_mut::(device_buffer, |view| { + ctx.launch_kernel(&cuda_function, array_len, |args| { + args.arg(view).arg(&reference).arg(&array_len_u64); + }) + })?; + device_buffer = next; + if let Some(launch) = launch { + launch?; + return Ok(Canonical::Primitive(PrimitiveArray::from_buffer_handle( + device_buffer, + P::PTYPE, + validity, + ))); + } + // Preserve aliased inputs by using the out-of-place kernel only when unique mutable access is + // unavailable. + let input_view = device_buffer.cuda_view::

()?; + let mut output = ctx.device_alloc::

(array_len)?; + let ptype_suffix = P::PTYPE.to_string(); + let cuda_function = + ctx.load_function_with_suffixes("for", &["in", "out", ptype_suffix.as_str()])?; ctx.launch_kernel(&cuda_function, array_len, |args| { - args.arg(&cuda_view).arg(&reference).arg(&array_len_u64); + args.arg(&input_view) + .arg(&mut output) + .arg(&reference) + .arg(&array_len_u64); })?; - // Build result - in-place reuses the same buffer + let output = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(output))); Ok(Canonical::Primitive(PrimitiveArray::from_buffer_handle( - device_buffer, + output, P::PTYPE, validity, ))) @@ -134,6 +158,7 @@ mod tests { use vortex::buffer::Buffer; use vortex::dtype::NativePType; use vortex::encodings::fastlanes::BitPacked; + use vortex::encodings::fastlanes::BitPackedArrayExt; use vortex::encodings::fastlanes::FoR; use vortex::encodings::fastlanes::FoRArray; use vortex::error::VortexExpect; @@ -204,4 +229,50 @@ mod tests { assert_arrays_eq!(for_array, gpu_result, &mut ctx); } + + /// Patched positions must pick up the frame of reference, exactly like unpacked ones. + /// + /// The bit-packed exceptions are stored reference-relative, so a decoder that writes them + /// straight into the output leaves every patched value short by the reference. A plain + /// bit-packed array cannot catch that: its reference is zero. + #[rstest] + #[case::u32(100_000u32)] + #[case::u64(1_000_000u64)] + #[crate::test] + async fn test_ffor_patched_values_include_reference(#[case] reference: T) -> VortexResult<()> + where + T: NativePType + Into + From, + { + let mut ctx = array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + // Values that fit in 8 bits, with a handful that do not and so become patches. + let mut values = (0..2048u32) + .map(|i| >::from(i % 200)) + .collect::>(); + for index in [7, 1023, 1024, 2047] { + values[index] = >::from(1u32 << 17); + } + + let values = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); + let packed = BitPacked::encode(&values, 8, &mut array_session().create_execution_ctx())?; + assert!( + packed.patches().is_some(), + "test setup expects the exceptions to be stored as patches" + ); + let for_array = FoR::try_new(packed.into_array(), reference.into())?; + + let gpu_result = FoRExecutor + .execute(for_array.clone().into_array(), &mut cuda_ctx) + .await + .vortex_expect("GPU decompression failed") + .into_host() + .await? + .into_array(); + + assert_arrays_eq!(for_array, gpu_result, &mut ctx); + + Ok(()) + } } diff --git a/vortex-cuda/src/kernel/encodings/fsst.rs b/vortex-cuda/src/kernel/encodings/fsst.rs index de46d0b8fa1..74cfcfadf7d 100644 --- a/vortex-cuda/src/kernel/encodings/fsst.rs +++ b/vortex-cuda/src/kernel/encodings/fsst.rs @@ -13,10 +13,12 @@ use cudarc::driver::PushKernelArg; use tracing::instrument; use vortex::array::ArrayRef; use vortex::array::Canonical; +use vortex::array::arrays::Bool; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::VarBinViewArray; +use vortex::array::arrays::bool::BoolDataParts; use vortex::array::arrays::primitive::PrimitiveDataParts; -use vortex::array::arrays::varbin::VarBinArrayExt; +use vortex::array::arrays::varbin::VarBinArraySlotsExt; use vortex::array::arrays::varbinview::build_views::MAX_BUFFER_LEN; use vortex::array::arrays::varbinview::build_views::build_views; use vortex::array::buffer::BufferHandle; @@ -25,23 +27,31 @@ use vortex::array::match_each_integer_ptype; use vortex::array::match_each_unsigned_integer_ptype; use vortex::array::validity::Validity; use vortex::buffer::Alignment; -use vortex::buffer::ByteBuffer; +use vortex::buffer::Buffer; use vortex::dtype::DType; use vortex::dtype::NativePType; use vortex::encodings::fsst::FSST; use vortex::encodings::fsst::FSSTArray; use vortex::encodings::fsst::FSSTArrayExt; +use vortex::encodings::fsst::FSSTArraySlotsExt; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_err; +use crate::CanonicalCudaExt; use crate::CudaBufferExt; use crate::CudaDeviceBuffer; +use crate::arrow::I32Offsets; +use crate::arrow::i32_offsets_from_lengths; +use crate::executor::CudaArrayExt; use crate::executor::CudaExecute; use crate::executor::CudaExecutionCtx; +use crate::executor::execute_validity_cuda; -/// Device-resident offset-based result of FSST decompression. -pub(crate) struct FSSTVarBin { +/// Device-resident offset-based (Arrow `Utf8`/`Binary`) decompression result: +/// i32 offsets plus a contiguous values heap. Produced by the FSST and OnPair +/// varbin decoders for the offset-based Arrow export path. +pub(crate) struct DecodedVarBin { pub(crate) dtype: DType, pub(crate) len: usize, pub(crate) offsets: BufferHandle, @@ -53,19 +63,30 @@ pub(crate) struct FSSTVarBin { /// /// `BitBuffer` slices normalize whole-byte offsets but can retain a sub-byte offset. Passing that /// offset to CUDA lets the kernels address sliced validity directly without repacking it on host. -fn cuda_validity( +async fn cuda_validity( validity: &Validity, len: usize, ctx: &mut CudaExecutionCtx, -) -> VortexResult<(u64, ByteBuffer)> { - let bits = validity - .clone() - .execute_mask(len, ctx.execution_ctx())? - .into_bit_buffer(); - let (bit_offset, bit_len, bytes) = bits.into_inner(); - debug_assert_eq!(bit_len, len); - let byte_len = (bit_offset + bit_len).div_ceil(8); - Ok((bit_offset as u64, bytes.slice(0..byte_len))) +) -> VortexResult<(u64, BufferHandle)> { + match execute_validity_cuda(validity.clone(), len, ctx).await? { + Validity::NonNullable | Validity::AllValid => Ok(( + 0, + BufferHandle::new_host(Buffer::from(vec![u8::MAX; len.div_ceil(8)]).into_byte_buffer()), + )), + Validity::AllInvalid => Ok(( + 0, + BufferHandle::new_host(Buffer::from(vec![0u8; len.div_ceil(8)]).into_byte_buffer()), + )), + Validity::Array(array) => { + let bool_array = array.try_downcast::().map_err(|array| { + vortex_err!("CUDA validity execution produced {}", array.dtype()) + })?; + let BoolDataParts { bits, meta } = bool_array.into_data().into_parts(len); + let bit_offset = meta.offset(); + let byte_len = (bit_offset + len).div_ceil(8); + Ok((u64::try_from(bit_offset)?, bits.slice(0..byte_len))) + } + } } /// CUDA decoder for FSST. @@ -105,12 +126,18 @@ impl CudaExecute for FSSTExecutor { let lens = fsst .uncompressed_lengths() .clone() - .execute::(ctx.execution_ctx())?; + .execute_cuda(ctx) + .await? + .into_host() + .await? + .into_primitive(); let codes_offsets = fsst .codes() .offsets() .clone() - .execute::(ctx.execution_ctx())?; + .execute_cuda(ctx) + .await? + .into_primitive(); // Prefix-sum lens to per-string u64 output offsets so the kernel // knows where to write each decoded string. @@ -138,52 +165,35 @@ impl CudaExecute for FSSTExecutor { pub(crate) async fn decode_fsst_varbin( fsst: FSSTArray, ctx: &mut CudaExecutionCtx, -) -> VortexResult { +) -> VortexResult { let dtype = fsst.dtype().clone(); let validity = fsst.codes().validity()?; let len = fsst.len(); let lens = fsst .uncompressed_lengths() .clone() - .execute::(ctx.execution_ctx())?; + .execute_cuda(ctx) + .await? + .into_primitive(); let codes_offsets = fsst .codes() .offsets() .clone() - .execute::(ctx.execution_ctx())?; - - let output_offsets = match_each_integer_ptype!(lens.ptype(), |P| { - let mut offsets = Vec::with_capacity(lens.len() + 1); - let mut acc = 0u64; - offsets.push(0i32); - #[allow(clippy::unnecessary_cast)] - for &length in lens.as_slice::

() { - let length = u64::try_from(length as i128) - .map_err(|_| vortex_err!("FSST uncompressed length cannot be negative"))?; - acc = acc - .checked_add(length) - .ok_or_else(|| vortex_err!("FSST decoded size overflow"))?; - offsets - .push(i32::try_from(acc).map_err(|_| { - vortex_err!("FSST decoded size exceeds Arrow i32 offset range") - })?); - } - VortexResult::Ok(offsets) - })?; - let total_size = usize::try_from( - *output_offsets - .last() - .vortex_expect("output_offsets has at least one entry"), - )?; + .execute_cuda(ctx) + .await? + .into_primitive(); + let I32Offsets { + buffer: output_offsets, + total: total_size, + } = i32_offsets_from_lengths(lens, ctx).await?; if total_size == 0 { - let offsets = ctx.copy_to_device(output_offsets)?.await?; let allocation = CudaDeviceBuffer::new(ctx.device_alloc::(1)?); let values = BufferHandle::new_device(allocation.slice(0..0)); - return Ok(FSSTVarBin { + return Ok(DecodedVarBin { dtype, len, - offsets, + offsets: output_offsets, values, validity, }); @@ -197,10 +207,10 @@ pub(crate) async fn decode_fsst_varbin( async fn decode_fsst_varbin_typed( fsst: FSSTArray, codes_offsets: PrimitiveArray, - output_offsets: Vec, + output_offsets: BufferHandle, total_size: usize, ctx: &mut CudaExecutionCtx, -) -> VortexResult +) -> VortexResult where U: NativePType + DeviceRepr + Send + Sync + 'static, { @@ -208,27 +218,30 @@ where let validity = fsst.codes().validity()?; let len = fsst.len(); let len_u64 = len as u64; - let symbols_u64: Vec = fsst.symbols().iter().map(|s| s.to_u64()).collect(); - let symbol_lengths = fsst.symbol_lengths().clone(); + let symbols_u64 = fsst + .symbols() + .iter() + .map(|s| s.to_u64()) + .collect::>(); + let symbol_lengths = fsst.padded_symbol_lengths().slice(0..fsst.n_symbols()); let codes_bytes_handle = fsst.codes_bytes_handle().clone(); let PrimitiveDataParts { buffer: codes_offsets_buffer, .. } = codes_offsets.into_data_parts(); - let (validity_bit_offset, validity_bits) = cuda_validity(&validity, len, ctx)?; + let (validity_bit_offset, validity_bits) = cuda_validity(&validity, len, ctx).await?; - let (symbols, symbol_lengths, output_offsets, validity_device, codes_bytes, codes_offsets) = futures::try_join!( + let (symbols, symbol_lengths, validity_device, codes_bytes, codes_offsets) = futures::try_join!( ctx.copy_to_device(symbols_u64)?, ctx.copy_to_device(symbol_lengths)?, - ctx.copy_to_device(output_offsets)?, - ctx.copy_to_device(validity_bits)?, + ctx.ensure_on_device(validity_bits), ctx.ensure_on_device(codes_bytes_handle), ctx.ensure_on_device(codes_offsets_buffer), )?; // The kernel gates store widths on `out_pos % N` relative to the base, so the base must // satisfy the widest store (u128 → 16). - let output = ctx.device_alloc::(total_size)?; + let mut output = ctx.device_alloc::(total_size)?; let (output_base_ptr, _) = output.device_ptr(ctx.stream()); assert_eq!( output_base_ptr % 16, @@ -253,11 +266,11 @@ where .arg(&output_offsets_view) .arg(&validity_view) .arg(&validity_bit_offset) - .arg(&output) + .arg(&mut output) .arg(&len_u64); })?; - Ok(FSSTVarBin { + Ok(DecodedVarBin { dtype, len, offsets: output_offsets, @@ -294,29 +307,33 @@ where })); } - let symbols_u64: Vec = fsst.symbols().iter().map(|s| s.to_u64()).collect(); - let symbol_lengths = fsst.symbol_lengths().clone(); + let symbols_u64 = fsst + .symbols() + .iter() + .map(|s| s.to_u64()) + .collect::>(); + let symbol_lengths = fsst.padded_symbol_lengths().slice(0..fsst.n_symbols()); let codes_bytes_handle = fsst.codes_bytes_handle().clone(); let PrimitiveDataParts { buffer: codes_offsets_buffer, .. } = codes_offsets.into_data_parts(); - let (validity_bit_offset, validity_bits) = cuda_validity(&validity, num_strings, ctx)?; + let (validity_bit_offset, validity_bits) = cuda_validity(&validity, num_strings, ctx).await?; let (symbols, symbol_lengths, output_offsets, validity_device, codes_bytes, codes_offsets) = futures::try_join!( ctx.copy_to_device(symbols_u64)?, ctx.copy_to_device(symbol_lengths)?, ctx.copy_to_device(output_offsets)?, - ctx.copy_to_device(validity_bits)?, + ctx.ensure_on_device(validity_bits), ctx.ensure_on_device(codes_bytes_handle), ctx.ensure_on_device(codes_offsets_buffer), )?; // The kernel checks store alignment relative to the base via // `out_pos % N`, so the base must satisfy the widest store (u128 → 16). - let device_output = ctx.device_alloc::(total_size)?; - let device_views = (total_size <= MAX_BUFFER_LEN) + let mut device_output = ctx.device_alloc::(total_size)?; + let mut device_views = (total_size <= MAX_BUFFER_LEN) .then(|| ctx.device_alloc::(num_strings)) .transpose()?; let (output_base_ptr, _) = device_output.device_ptr(ctx.stream()); @@ -343,8 +360,8 @@ where .arg(&output_offsets_view) .arg(&validity_view) .arg(&validity_bit_offset) - .arg(&device_output); - if let Some(device_views) = device_views.as_ref() { + .arg(&mut device_output); + if let Some(device_views) = device_views.as_mut() { args.arg(device_views); } else { args.arg(&null_views); @@ -371,12 +388,7 @@ where let host_bytes = host_bytes.slice(0..total_size); let (buffers, views) = match_each_integer_ptype!(lens.ptype(), |P| { - build_views( - 0, - MAX_BUFFER_LEN, - host_bytes.into_mut(), - lens.as_slice::

(), - ) + build_views(0, MAX_BUFFER_LEN, host_bytes, lens.as_slice::

()) }); Ok(Canonical::VarBinView(unsafe { diff --git a/vortex-cuda/src/kernel/encodings/mod.rs b/vortex-cuda/src/kernel/encodings/mod.rs index 6ad37e7ead1..fb292bc445c 100644 --- a/vortex-cuda/src/kernel/encodings/mod.rs +++ b/vortex-cuda/src/kernel/encodings/mod.rs @@ -5,13 +5,14 @@ mod alp; mod bitpacked; mod date_time_parts; mod decimal_byte_parts; +mod delta; mod for_; mod fsst; +mod onpair; mod runend; mod sequence; mod zigzag; mod zstd; -#[cfg(feature = "unstable_encodings")] mod zstd_buffers; pub(crate) use alp::ALPExecutor; @@ -19,15 +20,17 @@ pub(crate) use bitpacked::BitPackedExecutor; pub(crate) use bitpacked::bitpacked_slice_view; pub(crate) use date_time_parts::DateTimePartsExecutor; pub(crate) use decimal_byte_parts::DecimalBytePartsExecutor; +pub(crate) use delta::DeltaExecutor; pub(crate) use for_::FoRExecutor; +pub(crate) use fsst::DecodedVarBin; pub(crate) use fsst::FSSTExecutor; -pub(crate) use fsst::FSSTVarBin; pub(crate) use fsst::decode_fsst_varbin; +pub(crate) use onpair::OnPairExecutor; +pub(crate) use onpair::decode_onpair_varbin; pub(crate) use runend::RunEndExecutor; pub(crate) use sequence::SequenceExecutor; pub(crate) use zigzag::ZigZagExecutor; pub(crate) use zstd::ZstdExecutor; pub use zstd::ZstdKernelPrep; pub use zstd::zstd_kernel_prepare; -#[cfg(feature = "unstable_encodings")] pub(crate) use zstd_buffers::ZstdBuffersExecutor; diff --git a/vortex-cuda/src/kernel/encodings/onpair.rs b/vortex-cuda/src/kernel/encodings/onpair.rs new file mode 100644 index 00000000000..8298048b9d9 --- /dev/null +++ b/vortex-cuda/src/kernel/encodings/onpair.rs @@ -0,0 +1,1074 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! CUDA executor for OnPair decompression. +//! +//! Decoding runs entirely on the GPU. A sliced array keeps its whole `codes` +//! child, so the lightweight prefix sweep covers the retained stream, while +//! the dictionary gather/scatter and output allocation cover only the visible +//! token window: +//! +//! 1. `onpair_batch_offsets` (in the CUB shim) regenerates the per-batch +//! output offsets (`chunk_offsets`) the decode kernel positions its writes +//! with: one fused sweep reduces every 128-token batch's decoded size and +//! exclusive-scans the sizes in-kernel via decoupled look-back. +//! 2. `onpair_window_offsets` reads this array's token window from the +//! device-resident `codes_offsets` child — the offsets are nondecreasing, +//! so the window's min and max are its first and last elements — and +//! resolves the window's byte positions inside the decoded stream in the +//! same launch: the whole-batch prefix from `chunk_offsets` plus a +//! partial-batch reduction over each boundary batch's head. No boundary is +//! read on host before the single gating readback. +//! 3. `onpair_shmem_4tpt_split8read` gathers each token's bytes from the +//! split dictionary layout and scatters only the visible window into an +//! exact-sized output byte stream. +//! +//! Every kernel that reads the codes is instantiated for the two widths +//! OnPair stores (u16 natively, u8 when the compressor narrowed the codes), +//! so the code stream is decompressed on device and never widened. +//! +//! The heap size and window bounds come from the GPU; each output path builds +//! its row offsets from the lengths child with [`i32_offsets_from_lengths`] +//! on device and validates them against the window size before they index the +//! decoded bytes. The result is exposed either as a canonical `VarBinView` +//! (views built on-device by `onpair_build_views`, or on host for windows +//! that exceed a single backing buffer — the only path that materialises the +//! lengths) or as Arrow-compatible i32 offsets plus values via +//! [`decode_onpair_varbin`] — mirroring the FSST varbin path. + +use std::fmt::Debug; +use std::sync::Arc; + +use async_trait::async_trait; +use cudarc::driver::CudaSlice; +use cudarc::driver::DevicePtr; +use cudarc::driver::DeviceRepr; +use cudarc::driver::LaunchConfig; +use cudarc::driver::PushKernelArg; +use num_traits::AsPrimitive; +use tracing::instrument; +use vortex::array::ArrayRef; +use vortex::array::Canonical; +use vortex::array::IntoArray; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::VarBinViewArray; +use vortex::array::arrays::primitive::PrimitiveDataParts; +use vortex::array::arrays::varbinview::build_views::MAX_BUFFER_LEN; +use vortex::array::arrays::varbinview::build_views::build_views; +use vortex::array::buffer::BufferHandle; +use vortex::array::buffer::DeviceBuffer; +use vortex::array::builtins::ArrayBuiltins; +use vortex::array::match_each_integer_ptype; +use vortex::array::validity::Validity; +use vortex::buffer::Buffer; +use vortex::dtype::DType; +use vortex::dtype::NativePType; +use vortex::dtype::PType; +use vortex::error::VortexExpect; +use vortex::error::VortexResult; +use vortex::error::vortex_bail; +use vortex::error::vortex_ensure; +use vortex::error::vortex_err; +use vortex_array::ArrayView; +use vortex_onpair::DictionaryView; +use vortex_onpair::MAX_TOKEN_SIZE; +use vortex_onpair::OnPair; +use vortex_onpair::OnPairArray; +use vortex_onpair::OnPairArrayExt; +use vortex_onpair::OnPairArraySlotsExt; +use vortex_onpair::OnPairDictionary; + +use crate::CanonicalCudaExt; +use crate::CudaBufferExt; +use crate::CudaDeviceBuffer; +use crate::arrow::I32Offsets; +use crate::arrow::i32_offsets_from_lengths; +use crate::cub::onpair_batch_offsets; +use crate::executor::CudaArrayExt; +use crate::executor::CudaExecute; +use crate::executor::CudaExecutionCtx; +use crate::kernel::encodings::DecodedVarBin; + +// The kernels fix the dictionary row stride at 16 bytes (two `uint2` reads). +const _: () = assert!(MAX_TOKEN_SIZE == 16); + +/// Tokens per decode batch: one decode-kernel warp emits 128 tokens (4 per +/// thread). Must match `ONPAIR_TOKENS_PER_BATCH` in `cub/kernels/onpair.cu`. +const TOKENS_PER_BATCH: usize = 128; +/// Threads per block for the warp-per-batch kernels (16 warps). +const BLOCK_THREADS: u32 = 512; +const WARPS_PER_BLOCK: usize = (BLOCK_THREADS / 32) as usize; + +/// CUDA decoder for OnPair. +#[derive(Debug)] +pub(crate) struct OnPairExecutor; + +#[async_trait] +impl CudaExecute for OnPairExecutor { + #[instrument(level = "trace", skip_all, fields(executor = ?self))] + async fn execute( + &self, + array: ArrayRef, + ctx: &mut CudaExecutionCtx, + ) -> VortexResult { + let onpair = array + .as_typed::() + .ok_or_else(|| vortex_err!("Expected OnPairArray"))?; + decode_onpair(onpair, ctx).await + } +} + +async fn decode_onpair( + onpair: ArrayView<'_, OnPair>, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + let dtype = onpair.dtype().clone(); + let validity = onpair.array_validity(); + let num_rows = onpair.len(); + + if onpair.is_empty() { + return Ok(Canonical::empty(&dtype)); + } + + let Some(decoded) = decode_onpair_bytes(onpair, ctx).await? else { + return empty_views(num_rows, dtype, validity, ctx).await; + }; + let OnPairDecoded { + bytes, + total_size, + lengths, + } = decoded; + + // Fast path: the decoded window fits a single BinaryView backing buffer + // (`MAX_BUFFER_LEN`, i32::MAX), so the per-row offsets fit Arrow's i32 + // range and build on device from the device-resident lengths — nothing + // touches the host. + if total_size <= MAX_BUFFER_LEN { + let I32Offsets { + buffer: row_offsets, + total, + } = i32_offsets_from_lengths(lengths, ctx).await?; + let row_total = u64::try_from(total)?; + vortex_ensure!( + row_total == total_size as u64, + "OnPair codes decode to {total_size} bytes but uncompressed_lengths records {row_total}" + ); + let row_offsets_view = row_offsets.cuda_view::()?; + let bytes_view = bytes.cuda_view::()?; + let mut device_views = ctx.device_alloc::(num_rows)?; + let num_rows_u64 = u64::try_from(num_rows)?; + let build_views_fn = ctx.load_function_with_suffixes("onpair", &["build_views"])?; + ctx.launch_kernel(&build_views_fn, num_rows, |args| { + args.arg(&row_offsets_view) + .arg(&bytes_view) + .arg(&mut device_views) + .arg(&num_rows_u64); + })?; + + let views = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(device_views))); + return Ok(Canonical::VarBinView(unsafe { + VarBinViewArray::new_handle_unchecked(views, Arc::from([bytes]), dtype, validity) + })); + } + + // BinaryView offsets are u32. Windows that need multiple backing buffers + // roll the decoded bytes over on host, mirroring the CPU canonical path; + // only here do the lengths leave the device. The host views index the + // copied window, so validate the lengths first. + let lengths = Canonical::Primitive(lengths) + .into_host() + .await? + .into_primitive(); + let row_total = sum_lengths(&lengths)?; + vortex_ensure!( + row_total == total_size as u64, + "OnPair codes decode to {total_size} bytes but uncompressed_lengths records {row_total}" + ); + let host_bytes = bytes.try_to_host()?.await?; + + let (buffers, views) = match_each_integer_ptype!(lengths.ptype(), |P| { + build_views(0, MAX_BUFFER_LEN, host_bytes, lengths.as_slice::

()) + }); + + Ok(Canonical::VarBinView(unsafe { + VarBinViewArray::new_unchecked(views, Arc::from(buffers), dtype, validity) + })) +} + +/// Decode OnPair directly into Arrow-compatible i32 offsets and contiguous +/// values on device, mirroring the FSST varbin path. +pub(crate) async fn decode_onpair_varbin( + onpair: OnPairArray, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + let dtype = onpair.dtype().clone(); + let validity = onpair.array_validity(); + let len = onpair.len(); + + let Some(decoded) = decode_onpair_bytes(onpair.as_view(), ctx).await? else { + // Zero decoded bytes: all-zero offsets and an empty values heap. + let offsets = ctx.copy_to_device(vec![0i32; len + 1])?.await?; + let allocation = CudaDeviceBuffer::new(ctx.device_alloc::(1)?); + let values = BufferHandle::new_device(allocation.slice(0..0)); + return Ok(DecodedVarBin { + dtype, + len, + offsets, + values, + validity, + }); + }; + + let OnPairDecoded { + bytes, + total_size, + lengths, + } = decoded; + + // Build the Arrow i32 offsets from the lengths on device; this also + // rejects windows beyond Arrow's i32 offset range. + let I32Offsets { + buffer: offsets, + total, + } = i32_offsets_from_lengths(lengths, ctx).await?; + let row_total = u64::try_from(total)?; + vortex_ensure!( + row_total == total_size as u64, + "OnPair codes decode to {total_size} bytes but uncompressed_lengths records {row_total}" + ); + + Ok(DecodedVarBin { + dtype, + len, + offsets, + values: bytes, + validity, + }) +} + +/// The shared result of the OnPair GPU decode pipeline. +struct OnPairDecoded { + /// This array's rows' decoded bytes. + bytes: BufferHandle, + /// Byte size of the window, computed on device. + total_size: usize, + /// Per-row lengths. The varbin path and the canonical fast path build + /// their row offsets from them on device; only the host rollover path + /// materialises them. + lengths: PrimitiveArray, +} + +/// Run the OnPair decode pipeline: stage the codes and dictionary, regenerate +/// per-batch output offsets on the device, resolve and validate the visible +/// token window, then decode only the batches intersecting that window. The +/// retained full code stream never round-trips through the host. Returns +/// `Ok(None)` when there is nothing to decode: the array is empty, every row +/// is null, or the code window is empty. +async fn decode_onpair_bytes( + onpair: ArrayView<'_, OnPair>, + ctx: &mut CudaExecutionCtx, +) -> VortexResult> { + let num_rows = onpair.len(); + + if num_rows == 0 { + return Ok(None); + } + + // Every row null (cheap metadata check): nothing to decode. A sliced + // all-null window usually carries a validity child instead of the + // `AllInvalid` marker and is caught when the token window resolves empty. + if onpair.array_validity().definitely_all_null() { + return Ok(None); + } + + // No codes at all (e.g. every row empty): the child's length is host + // metadata, so this early-out costs no device read. + if onpair.codes().is_empty() { + let lengths = decode_primitive_child(onpair.uncompressed_lengths().clone(), ctx).await?; + ensure_zero_lengths(lengths).await?; + return Ok(None); + } + + // The three children and the dictionary staging all run on `ctx`'s + // stream: the per-row lengths (consumed only by the output paths), the + // per-row code boundaries (the token window is resolved from them by a + // kernel, never by host scalar reads), the codes themselves, and the + // split dictionary staging. + let lengths = decode_primitive_child(onpair.uncompressed_lengths().clone(), ctx).await?; + let codes_offsets = decode_primitive_child(onpair.codes_offsets().clone(), ctx).await?; + let codes = decode_primitive_child(onpair.codes().clone(), ctx).await?; + let dict = stage_dict(onpair, ctx).await?; + + // The kernels are instantiated for the two widths OnPair stores — u16 + // natively, u8 when the compressor narrowed the codes — so no widening + // pass is needed. + match codes.ptype() { + PType::U8 => decode_window::(codes, codes_offsets, lengths, dict, ctx).await, + PType::U16 => decode_window::(codes, codes_offsets, lengths, dict, ctx).await, + other => vortex_bail!("OnPair codes must decompress to u8 or u16, got {other}"), + } +} + +/// The dictionary staged device-resident in the decode kernel's split layout. +struct StagedDict { + dict_s8: BufferHandle, + dict_padded: BufferHandle, + lens: BufferHandle, + dict_size: u32, +} + +/// Stage the dictionary in the decode kernel's split layout: fixed 16-byte +/// rows (`dict_padded`, the rare `len > 8` read), the first 8 bytes of every +/// row (`dict_s8`, the common-case read), and the per-code lengths. +async fn stage_dict( + onpair: ArrayView<'_, OnPair>, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + // Both dictionary parts are read here by the host staging loop below, so both must be + // host-resident first. `dict_view` reads them through `BufferHandle::as_host` and a CPU + // cast kernel, which panic once the array's buffers live on the device. + let dict_bytes = onpair.dict_bytes_handle().clone().try_into_host()?.await?; + let dict_offsets = host_dict_offsets(onpair, ctx).await?; + let dict = OnPairDictionary::try_new(dict_bytes, dict_offsets)?; + let dict = dict.as_view(); + let dict_size = dict.num_tokens(); + let dict_size_u32 = u32::try_from(dict_size)?; + let mut dict_padded = vec![0u8; dict_size * MAX_TOKEN_SIZE]; + let mut dict_s8 = vec![0u8; dict_size * 8]; + let mut lens = vec![0u8; dict_size]; + for code in 0..dict_size { + let token = + dict.token(u16::try_from(code).vortex_expect("dictionary has at most 2^16 tokens")); + let len = token.len(); + lens[code] = u8::try_from(len).vortex_expect("token length is at most MAX_TOKEN_SIZE"); + dict_padded[code * MAX_TOKEN_SIZE..code * MAX_TOKEN_SIZE + len].copy_from_slice(token); + let head = len.min(8); + dict_s8[code * 8..code * 8 + head].copy_from_slice(&token[..head]); + } + + let (s8_dev, padded_dev, lens_dev) = futures::try_join!( + ctx.copy_to_device(dict_s8)?, + ctx.copy_to_device(dict_padded)?, + ctx.copy_to_device(lens)?, + )?; + + Ok(StagedDict { + dict_s8: s8_dev, + dict_padded: padded_dev, + lens: lens_dev, + dict_size: dict_size_u32, + }) +} + +/// Stage the codes at their native width `C`, resolve the window bounds on +/// device, validate the stream, and decode the visible token window. Returns +/// `Ok(None)` when the token window turns out to be empty. +async fn decode_window( + codes: PrimitiveArray, + codes_offsets: PrimitiveArray, + lengths: PrimitiveArray, + dict: StagedDict, + ctx: &mut CudaExecutionCtx, +) -> VortexResult> +where + C: NativePType + DeviceRepr + Send + Sync + 'static, +{ + // Corruption flag raised by the batch-sizes kernel for a code outside the + // dictionary; checked before the unchecked decode kernel is allowed to run. + let mut status = ctx.device_alloc::(1)?; + ctx.stream() + .memset_zeros(&mut status) + .map_err(|e| vortex_err!("Failed to zero OnPair status flag: {e}"))?; + + let staged = stage_codes(codes, &dict, &mut status, ctx).await?; + + let ptype = C::PTYPE.to_string(); + let num_tokens_u64 = u64::try_from(staged.num_tokens)?; + let codes_view = staged.codes.cuda_view::()?; + let s8_view = dict.dict_s8.cuda_view::()?; + let padded_view = dict.dict_padded.cuda_view::()?; + let lens_view = dict.lens.cuda_view::()?; + + // The readback scratch: the fused window-offsets kernel writes the token + // window into slots 0..2, the byte window into slots 2..4, and packs the + // sweep's outputs — the decoded heap size and the corruption status flag + // — into slots 4..6, so one readback gates the decode kernel. + let mut scratch = ctx.device_alloc::(6)?; + + // Token and byte bounds of this array's rows, resolved on device by one + // fused launch — one warp per boundary. Each warp reads its token + // boundary from the (possibly slice-narrowed) `codes_offsets` child (the + // offsets are nondecreasing, so the window's min and max are its first + // and last elements) and resolves its byte position inside the decoded + // stream: the whole-batch prefix from `chunk_offsets` plus a + // partial-batch reduction over the boundary batch's head. The kernel is + // instantiated per (code width, offsets ptype) pair. + let offsets_ptype = codes_offsets.ptype(); + let last_boundary = u64::try_from(codes_offsets.len().saturating_sub(1))?; + let PrimitiveDataParts { + buffer: offsets_buffer, + .. + } = codes_offsets.into_data_parts(); + let offsets_dev = ctx.ensure_on_device(offsets_buffer).await?; + let window_fn = ctx.load_function_with_suffixes( + "onpair", + &["window_offsets", &ptype, &offsets_ptype.to_string()], + )?; + match_each_integer_ptype!(offsets_ptype, |O| { + let offsets_view = offsets_dev.cuda_view::()?; + ctx.launch_kernel_config( + &window_fn, + LaunchConfig { + grid_dim: (2, 1, 1), + block_dim: (32, 1, 1), + shared_mem_bytes: 0, + }, + 2, + |args| { + args.arg(&codes_view) + .arg(&lens_view) + .arg(&dict.dict_size) + .arg(&staged.chunk_offsets) + .arg(&num_tokens_u64) + .arg(&offsets_view) + .arg(&last_boundary) + .arg(&status) + .arg(&mut scratch); + }, + )?; + }); + + // One synchronizing readback gates the decode kernel — whose dictionary + // gathers and output scatters are unchecked — and yields the GPU-computed + // heap size, token window, byte window, and corruption status in a single + // round trip. The lengths child is validated against the window size by + // whichever output path materialises row offsets. + let scratch = ctx + .stream() + .clone_dtoh(&scratch) + .map_err(|e| vortex_err!("Failed to copy OnPair window scratch to host: {e}"))?; + let [ + token_start, + token_end, + byte_start, + byte_end, + chunk_total, + status, + ] = scratch[..] + else { + vortex_bail!("OnPair window resolution returned no bounds"); + }; + if status != 0 { + vortex_bail!("OnPair code out of dictionary range"); + } + let heap_size = usize::try_from(chunk_total)?; + vortex_ensure!( + token_start <= token_end, + "OnPair codes_offsets must be nondecreasing" + ); + vortex_ensure!( + token_end <= num_tokens_u64, + "OnPair codes_offsets end {token_end} exceeds codes len {num_tokens_u64}" + ); + if token_start == token_end { + // No codes in the window (e.g. a slice covering only null rows). + ensure_zero_lengths(lengths).await?; + return Ok(None); + } + let byte_start_u64 = byte_start; + let byte_start = usize::try_from(byte_start_u64)?; + let byte_end = usize::try_from(byte_end)?; + vortex_ensure!( + byte_start <= byte_end && byte_end <= heap_size, + "OnPair window bounds [{byte_start}, {byte_end}) exceed decoded heap size {heap_size}" + ); + let total_size = byte_end - byte_start; + // A conformant dictionary has no zero-length tokens, so a non-empty code window decodes to at + // least one byte. + vortex_ensure!(total_size > 0, "OnPair has codes but decodes to zero bytes"); + + // Decode only batches intersecting the visible token window. The kernel's drain gates 16-byte + // stores on `out_start % 16` relative to the buffer base, so the base must be 16-aligned. + let mut bytes = ctx.device_alloc::(total_size)?; + let (bytes_base_ptr, _) = bytes.device_ptr(ctx.stream()); + assert_eq!( + bytes_base_ptr % 16, + 0, + "output base not 16-aligned: {bytes_base_ptr:#x}", + ); + + let first_batch = token_start / TOKENS_PER_BATCH as u64; + let end_batch = token_end.div_ceil(TOKENS_PER_BATCH as u64); + let window_batches = usize::try_from(end_batch - first_batch)?; + let launch_config = batch_launch_config(window_batches)?; + let decode_fn = ctx.load_function_with_suffixes("onpair_shmem_4tpt_split8read", &[&ptype])?; + ctx.launch_kernel_config( + &decode_fn, + launch_config, + usize::try_from(token_end - token_start)?, + |args| { + args.arg(&codes_view) + .arg(&staged.chunk_offsets) + .arg(&s8_view) + .arg(&padded_view) + .arg(&lens_view) + .arg(&mut bytes) + .arg(&token_start) + .arg(&token_end) + .arg(&first_batch) + .arg(&byte_start_u64); + }, + )?; + + Ok(Some(OnPairDecoded { + bytes: BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(bytes))), + total_size, + lengths, + })) +} + +/// The device-staged compressed token stream: the full code stream at its +/// native width and the regenerated per-batch output offsets. +struct StagedCodes { + codes: BufferHandle, + /// Exclusive per-batch output offsets, `num_batches + 1` entries; the last + /// is the total decoded byte count of the full code stream. + chunk_offsets: CudaSlice, + num_tokens: usize, +} + +/// Stage this array's device-decompressed codes and regenerate the decode +/// kernel's per-batch output offsets from them and the staged dictionary in +/// one fused sweep (see [`onpair_batch_offsets`]). The caller has validated +/// that the codes are u8 or u16; the sweep reads them at their native width. +async fn stage_codes( + codes: PrimitiveArray, + dict: &StagedDict, + status: &mut CudaSlice, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + let num_tokens = codes.len(); + let code_width = u32::try_from(codes.ptype().byte_width())?; + let PrimitiveDataParts { + buffer: codes_buffer, + .. + } = codes.into_data_parts(); + let codes_dev = ctx.ensure_on_device(codes_buffer).await?; + + let num_batches = num_tokens.div_ceil(TOKENS_PER_BATCH); + let chunk_offsets = onpair_batch_offsets( + &codes_dev, + code_width, + &dict.lens, + dict.dict_size, + num_tokens, + num_batches, + status, + ctx, + )?; + + Ok(StagedCodes { + codes: codes_dev, + chunk_offsets, + num_tokens, + }) +} + +/// Launch config for the warp-per-batch kernels: one warp per 128-token batch. +fn batch_launch_config(num_batches: usize) -> VortexResult { + let grid_dim = u32::try_from(num_batches.div_ceil(WARPS_PER_BLOCK))?; + Ok(LaunchConfig { + grid_dim: (grid_dim, 1, 1), + block_dim: (BLOCK_THREADS, 1, 1), + shared_mem_bytes: 0, + }) +} + +/// All-empty output: `num_rows` inline empty views and no backing buffers. +async fn empty_views( + num_rows: usize, + dtype: DType, + validity: Validity, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + let views = ctx.copy_to_device(vec![0i128; num_rows])?.await?; + Ok(Canonical::VarBinView(unsafe { + VarBinViewArray::new_handle_unchecked(views, Arc::from([]), dtype, validity) + })) +} + +/// Decompress a child array to a canonical primitive on `ctx`'s stream. +async fn decode_primitive_child( + child: ArrayRef, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + Ok(child.execute_cuda(ctx).await?.into_primitive()) +} + +/// The `dict_offsets` child, decoded and widened to the `u32` the dictionary stores. +/// +/// Decoded on the GPU like the other children, then copied back: the dictionary is staged +/// on the host, and the child is `dict_size + 1` elements, so the round trip is bounded by +/// the dictionary size rather than the array length. +async fn host_dict_offsets( + onpair: ArrayView<'_, OnPair>, + ctx: &mut CudaExecutionCtx, +) -> VortexResult> { + let offsets = decode_primitive_child(onpair.dict_offsets().clone(), ctx).await?; + let offsets = Canonical::Primitive(offsets) + .into_host() + .await? + .into_primitive(); + let nullability = offsets.dtype().nullability(); + Ok(offsets + .into_array() + .cast(DType::Primitive(PType::U32, nullability))? + .execute::(ctx.execution_ctx())? + .into_buffer::()) +} + +/// Cold path: the window has no codes, so the rows must decode to zero bytes. +async fn ensure_zero_lengths(lengths: PrimitiveArray) -> VortexResult<()> { + let lengths = Canonical::Primitive(lengths) + .into_host() + .await? + .into_primitive(); + let total = sum_lengths(&lengths)?; + vortex_ensure!( + total == 0, + "OnPair records {total} decoded bytes but has no codes" + ); + Ok(()) +} + +/// Checked host sum of the per-row decoded lengths. A negative length +/// sign-extends and surfaces as overflow here or as a mismatch against the +/// GPU-computed window size. +fn sum_lengths(lengths: &PrimitiveArray) -> VortexResult { + match_each_integer_ptype!(lengths.ptype(), |P| { + let mut acc = 0u64; + for &length in lengths.as_slice::

() { + acc = acc + .checked_add(AsPrimitive::::as_(length)) + .ok_or_else(|| vortex_err!("OnPair decoded size overflow"))?; + } + Ok(acc) + }) +} + +#[cfg(test)] +mod tests { + use arrow_schema::DataType; + use arrow_schema::Field; + use rstest::rstest; + use vortex::array::IntoArray; + use vortex::array::arrays::VarBinArray; + use vortex::array::assert_arrays_eq; + use vortex::buffer::Buffer; + use vortex::dtype::Nullability; + use vortex::error::VortexExpect; + use vortex_array::VortexSessionExecute; + use vortex_onpair::DEFAULT_CONFIG; + use vortex_onpair::onpair_compress; + + use super::*; + use crate::CanonicalCudaExt; + use crate::arrow::DeviceArrayExt; + use crate::arrow::release_device_array; + use crate::arrow::release_schema; + use crate::device_buffer::cuda_backing_allocation; + use crate::session::CudaSession; + use crate::session::VarBinExportLayout; + + fn cuda_ctx_with_varbin_layout(layout: VarBinExportLayout) -> VortexResult { + let session = vortex::array::array_session() + .with_some(CudaSession::try_default()?.with_varbin_export_layout(layout)); + CudaSession::create_execution_ctx(&session) + } + + fn assert_device_resident(canonical: &Canonical) { + let varbinview = canonical.as_varbinview(); + assert!(varbinview.views_handle().is_on_device()); + assert!( + varbinview + .data_buffers() + .iter() + .all(BufferHandle::is_on_device) + ); + } + + fn compress_onpair( + strings: Vec>, + dtype: DType, + ctx: &mut CudaExecutionCtx, + ) -> VortexResult { + let varbin = VarBinArray::from_iter(strings, dtype).into_array(); + let onpair = onpair_compress(&varbin, DEFAULT_CONFIG, ctx.execution_ctx())?; + vortex_ensure!( + onpair.as_opt::().is_some(), + "expected OnPair array, got {}", + onpair.encoding_id() + ); + Ok(onpair) + } + + #[rstest] + #[case::binary_non_null( + vec![Some(&b"the quick brown fox"[..]), + Some(&b"jumps over the lazy dog"[..]), + Some(&b"hello world"[..]), + Some(&b"vortex onpair test string"[..])], + DType::Binary(Nullability::NonNullable), + )] + #[case::utf8_non_null( + vec![Some(&b"the quick brown fox"[..]), + Some(&b"jumps over the lazy dog"[..]), + Some(&b"hello world"[..]), + Some(&b"vortex onpair test string"[..])], + DType::Utf8(Nullability::NonNullable), + )] + #[case::utf8_inline_boundary( + vec![Some(&b""[..]), + Some(&b"123456789012"[..]), + Some(&b"1234567890123"[..]), + Some(&b"this is another outlined value"[..])], + DType::Utf8(Nullability::NonNullable), + )] + #[case::utf8_partial_nulls( + vec![Some(&b"alpha"[..]), None, Some(&b"gamma"[..]), None, Some(&b"epsilon"[..])], + DType::Utf8(Nullability::Nullable), + )] + #[case::binary_all_empty( + vec![Some(&b""[..]), Some(&b""[..]), Some(&b""[..])], + DType::Binary(Nullability::NonNullable), + )] + #[crate::test] + async fn test_cuda_onpair_decompression_roundtrip( + #[case] strings: Vec>, + #[case] dtype: DType, + ) -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let onpair = compress_onpair(strings, dtype.clone(), &mut cuda_ctx)?; + + let gpu_result = OnPairExecutor + .execute(onpair.clone(), &mut cuda_ctx) + .await + .vortex_expect("GPU decompression failed"); + assert_eq!(gpu_result.dtype(), &dtype); + assert_device_resident(&gpu_result); + + let host_result = gpu_result.into_host().await?.into_array(); + assert_arrays_eq!(onpair, host_result, &mut ctx); + Ok(()) + } + + /// Arrays read through the CUDA flat layout arrive with their buffers already on the + /// device. Staging the dictionary must copy it back rather than assuming a host buffer, + /// which is what panicked with "expected host buffer" on `taxi` and TPC-H `l_comment`. + #[crate::test] + async fn test_cuda_onpair_device_resident_dictionary() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let strings = vec![ + Some(&b"the quick brown fox"[..]), + Some(&b"jumps over the lazy dog"[..]), + Some(&b"hello world"[..]), + Some(&b"vortex onpair test string"[..]), + ]; + let dtype = DType::Utf8(Nullability::NonNullable); + let onpair = compress_onpair(strings, dtype.clone(), &mut cuda_ctx)?; + let view = onpair + .as_typed::() + .vortex_expect("expected OnPair array"); + + // Rebuild the array with its dictionary blob resident on the device. + let dict_bytes = view.dict_bytes_handle().as_host().to_vec(); + let dict_device = cuda_ctx.copy_to_device(dict_bytes)?.await?; + let device_onpair = OnPair::try_new( + dtype.clone(), + dict_device, + view.dict_offsets().clone(), + view.codes().clone(), + view.codes_offsets().clone(), + view.uncompressed_lengths().clone(), + view.array_validity(), + )? + .into_array(); + + let gpu_result = OnPairExecutor + .execute(device_onpair, &mut cuda_ctx) + .await + .vortex_expect("GPU decompression failed"); + assert_eq!(gpu_result.dtype(), &dtype); + + let host_result = gpu_result.into_host().await?.into_array(); + assert_arrays_eq!(onpair, host_result, &mut ctx); + Ok(()) + } + + /// A slice keeps the whole `codes` child and narrows only `codes_offsets`, + /// so this exercises the nonzero `code_start` window. + #[crate::test] + async fn test_cuda_onpair_decompression_sliced() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + let values = vec![ + Some(&b"before the window"[..]), + None, + Some(&b"the quick brown fox"[..]), + None, + Some(&b"after the window"[..]), + ]; + let onpair = compress_onpair(values, DType::Utf8(Nullability::Nullable), &mut cuda_ctx)?; + let sliced = onpair.slice(1..4)?; + + let gpu_result = OnPairExecutor + .execute(sliced.clone(), &mut cuda_ctx) + .await?; + assert_device_resident(&gpu_result); + let host_result = gpu_result.into_host().await?.into_array(); + assert_arrays_eq!(sliced, host_result, &mut ctx); + Ok(()) + } + + /// A slice covering only null rows decodes zero bytes and takes the + /// empty-views path. + #[crate::test] + async fn test_cuda_onpair_decompression_null_slice() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + let values = vec![Some(&b"alpha"[..]), None, None, Some(&b"omega"[..])]; + let onpair = compress_onpair(values, DType::Utf8(Nullability::Nullable), &mut cuda_ctx)?; + let sliced = onpair.slice(1..3)?; + + let gpu_result = OnPairExecutor + .execute(sliced.clone(), &mut cuda_ctx) + .await?; + assert_device_resident(&gpu_result); + let host_result = gpu_result.into_host().await?.into_array(); + assert_arrays_eq!(sliced, host_result, &mut ctx); + Ok(()) + } + + /// Exercises many 128-token batches and the multi-block decode grid. + #[crate::test] + async fn test_cuda_onpair_decompression_roundtrip_large() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let strings: Vec = (0..100_000) + .map(|i| format!("https://www.example.com/path/{i}/segment?q={}", i % 97)) + .collect(); + let varbin = VarBinArray::from_iter( + strings.iter().map(|s| Some(s.as_str())), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + let onpair = onpair_compress(&varbin, DEFAULT_CONFIG, cuda_ctx.execution_ctx())?; + + let gpu_result = OnPairExecutor + .execute(onpair.clone(), &mut cuda_ctx) + .await + .vortex_expect("GPU decompression failed"); + assert_device_resident(&gpu_result); + + let host_result = gpu_result.into_host().await?.into_array(); + assert_arrays_eq!(onpair, host_result, &mut ctx); + Ok(()) + } + + /// A slice deep into a large array: both code-window boundaries land + /// mid-batch in non-zero batches, exercising the on-device window-bounds + /// resolution (whole-batch prefix plus partial-batch reduction) and the + /// exact-sized decoded window allocation. + #[crate::test] + async fn test_cuda_onpair_decompression_sliced_large() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let strings: Vec = (0..40_000) + .map(|i| format!("https://www.example.com/path/{i}/segment?q={}", i % 97)) + .collect(); + let varbin = VarBinArray::from_iter( + strings.iter().map(|s| Some(s.as_str())), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + let onpair = onpair_compress(&varbin, DEFAULT_CONFIG, cuda_ctx.execution_ctx())?; + let sliced = onpair.slice(19_997..20_101)?; + + let gpu_result = OnPairExecutor + .execute(sliced.clone(), &mut cuda_ctx) + .await?; + assert_device_resident(&gpu_result); + let expected_bytes = strings[19_997..20_101] + .iter() + .map(String::len) + .sum::(); + let values = &gpu_result.as_varbinview().data_buffers()[0]; + assert_eq!(cuda_backing_allocation(values)?.len(), expected_bytes); + let host_result = gpu_result.into_host().await?.into_array(); + assert_arrays_eq!(sliced, host_result, &mut ctx); + Ok(()) + } + + /// Codes narrowed to u8 dispatch the u8 kernel instantiations end to end. + /// A trained dictionary always holds the 256 single-byte tokens sorted + /// among its merges, so real merge codes never fit u8; the u8-addressable + /// case is the minimal alphabet-only dictionary, where token id `b` is + /// exactly the byte `b` and every row is coded byte per byte. + #[crate::test] + async fn test_cuda_onpair_decompression_u8_codes() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let strings: Vec<&[u8]> = [ + &b"tokenized token stream"[..], + b"tokenized", + b"token stream", + b"stream of tokens", + ] + .into_iter() + .cycle() + .take(800) + .collect(); + + // The alphabet-only compact dictionary: the 256 single-byte tokens + // (sorted by construction) plus the trailing read padding. + let mut dict_bytes: Vec = (0..=u8::MAX).collect(); + dict_bytes.resize(255 + MAX_TOKEN_SIZE, 0); + let dict_offsets: Vec = (0..=256).collect(); + + let codes: Vec = strings.concat(); + let mut codes_offsets = vec![0u32]; + let mut lengths = Vec::with_capacity(strings.len()); + let mut acc = 0u32; + for s in &strings { + let len = u32::try_from(s.len())?; + lengths.push(len); + acc += len; + codes_offsets.push(acc); + } + + let onpair = OnPair::try_new( + DType::Utf8(Nullability::NonNullable), + BufferHandle::new_host(Buffer::from(dict_bytes).into_byte_buffer()), + Buffer::from(dict_offsets).into_array(), + Buffer::from(codes).into_array(), + Buffer::from(codes_offsets).into_array(), + Buffer::from(lengths).into_array(), + Validity::NonNullable, + )?; + + let expected = VarBinArray::from_iter( + strings.iter().map(|s| Some(*s)), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + + let gpu_result = OnPairExecutor + .execute(onpair.into_array(), &mut cuda_ctx) + .await?; + assert_device_resident(&gpu_result); + let host_result = gpu_result.into_host().await?.into_array(); + assert_arrays_eq!(expected, host_result, &mut ctx); + Ok(()) + } + + #[crate::test] + async fn test_cuda_onpair_direct_varbin_output() -> VortexResult<()> { + let mut cuda_ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBin)?; + let values: [&[u8]; 3] = [ + b"", + b"short", + b"this value is stored directly in the values buffer", + ]; + let onpair = compress_onpair( + values.iter().map(|v| Some(*v)).collect(), + DType::Utf8(Nullability::NonNullable), + &mut cuda_ctx, + )? + .try_downcast::() + .map_err(|array| vortex_err!("expected OnPair array, got {}", array.encoding_id()))?; + + let output = decode_onpair_varbin(onpair, &mut cuda_ctx).await?; + assert_eq!(output.dtype, DType::Utf8(Nullability::NonNullable)); + assert_eq!(output.len, values.len()); + assert!(output.offsets.is_on_device()); + assert!(output.values.is_on_device()); + + let offsets = Buffer::::from_byte_buffer(output.offsets.try_to_host()?.await?); + assert_eq!( + offsets.as_slice(), + &[0, 0, 5, i32::try_from(5 + values[2].len())?,] + ); + assert_eq!( + output.values.try_to_host()?.await?.as_ref(), + values.concat() + ); + Ok(()) + } + + #[rstest] + #[case::binary( + DType::Binary(Nullability::NonNullable), + VarBinExportLayout::VarBin, + DataType::Binary, + 3 + )] + #[case::utf8( + DType::Utf8(Nullability::NonNullable), + VarBinExportLayout::VarBin, + DataType::Utf8, + 3 + )] + #[case::binary_view( + DType::Binary(Nullability::NonNullable), + VarBinExportLayout::VarBinView, + DataType::BinaryView, + 4 + )] + #[case::utf8_view( + DType::Utf8(Nullability::NonNullable), + VarBinExportLayout::VarBinView, + DataType::Utf8View, + 4 + )] + #[crate::test] + async fn test_cuda_onpair_arrow_export_uses_dtype_layout( + #[case] dtype: DType, + #[case] layout: VarBinExportLayout, + #[case] expected_data_type: DataType, + #[case] expected_n_buffers: i64, + ) -> VortexResult<()> { + let mut cuda_ctx = cuda_ctx_with_varbin_layout(layout)?; + let values = vec![ + Some(&b"short"[..]), + Some(&b"this value is stored out of line"[..]), + ]; + let onpair = compress_onpair(values, dtype, &mut cuda_ctx)?; + + let mut exported = onpair + .export_device_array_with_schema(&mut cuda_ctx) + .await?; + assert_eq!( + Field::try_from(&exported.schema)?, + Field::new("", expected_data_type, false) + ); + assert_eq!(exported.array.array.n_buffers, expected_n_buffers); + + release_device_array(&mut exported.array); + release_schema(&mut exported.schema); + Ok(()) + } +} diff --git a/vortex-cuda/src/kernel/encodings/runend.rs b/vortex-cuda/src/kernel/encodings/runend.rs index 4ac042e2972..da6973ff11c 100644 --- a/vortex-cuda/src/kernel/encodings/runend.rs +++ b/vortex-cuda/src/kernel/encodings/runend.rs @@ -10,8 +10,10 @@ use tracing::instrument; use vortex::array::ArrayRef; use vortex::array::Canonical; use vortex::array::IntoArray; +use vortex::array::arrays::BoolArray; use vortex::array::arrays::ConstantArray; use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::bool::BoolDataParts; use vortex::array::arrays::primitive::PrimitiveDataParts; use vortex::array::buffer::BufferHandle; use vortex::array::match_each_native_ptype; @@ -22,6 +24,7 @@ use vortex::dtype::PType; use vortex::encodings::runend::RunEnd; use vortex::encodings::runend::RunEndArray; use vortex::encodings::runend::RunEndArrayExt; +use vortex::encodings::runend::RunEndArraySlotsExt; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::error::vortex_ensure; @@ -122,23 +125,24 @@ async fn decode_runend_typed(output_len)?; - let output_device = CudaDeviceBuffer::new(output_slice); + let mut output_slice = ctx.device_alloc::(output_len)?; let ends_view = ends_device.cuda_view::()?; let values_view = values_device.cuda_view::()?; - let output_view = output_device.as_view::(); // Load kernel function let cuda_function = ctx.load_function("runend", &[value_ptype, E::PTYPE])?; + let num_runs_u64 = num_runs as u64; + let offset_u64 = offset as u64; + let output_len_u64 = output_len as u64; ctx.launch_kernel(&cuda_function, output_len, |args| { args.arg(&ends_view) - .arg(&num_runs) + .arg(&num_runs_u64) .arg(&values_view) - .arg(&offset) - .arg(&output_len) - .arg(&output_view); + .arg(&offset_u64) + .arg(&output_len_u64) + .arg(&mut output_slice); })?; let output_validity = match values_validity { @@ -147,13 +151,45 @@ async fn decode_runend_typed { unreachable!("AllInvalid should be handled by RunEndExecutor::execute") } - Validity::Array(_) => { - vortex_bail!( - "RunEnd GPU decoding does not yet support per-element validity in values; falling back to CPU" - ); + Validity::Array(validity) => { + // Expand the per-run validity bitmap through the same run mapping as the values. + let validity_bools = validity.execute_cuda(ctx).await?.into_bool(); + let validity_len = validity_bools.len(); + let BoolDataParts { + bits: validity_bits, + meta: validity_meta, + } = validity_bools.into_data().into_parts(validity_len); + let validity_device = ctx.ensure_on_device(validity_bits).await?; + let validity_view = validity_device.cuda_view::()?; + + // Each thread owns a whole output byte, so threads never race on bits in one byte. + let output_bytes = output_len.div_ceil(8); + let mut validity_out = ctx.device_alloc::(output_bytes)?; + let validity_offset_u64 = validity_meta.offset() as u64; + + let ends_ptype = E::PTYPE.to_string(); + let validity_function = + ctx.load_function_with_suffixes("runend", &["bool", &ends_ptype])?; + ctx.launch_kernel(&validity_function, output_bytes, |args| { + args.arg(&ends_view) + .arg(&num_runs_u64) + .arg(&validity_view) + .arg(&validity_offset_u64) + .arg(&offset_u64) + .arg(&output_len_u64) + .arg(&mut validity_out); + })?; + + let validity_buffer = + BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(validity_out))); + Validity::Array( + BoolArray::new_handle(validity_buffer, 0, output_len, Validity::NonNullable) + .into_array(), + ) } }; + let output_device = CudaDeviceBuffer::new(output_slice); Ok(Canonical::Primitive(PrimitiveArray::from_buffer_handle( BufferHandle::new_device(Arc::new(output_device)), value_ptype, @@ -179,7 +215,6 @@ mod tests { use super::*; use crate::CanonicalCudaExt; - use crate::executor::CudaArrayExt; use crate::session::CudaSession; fn make_runend_array(ends: Vec, values: Vec, ctx: &mut ExecutionCtx) -> RunEndArray @@ -300,7 +335,7 @@ mod tests { } #[crate::test] - async fn test_cuda_runend_nullable_values_falls_back_to_cpu() -> VortexResult<()> { + async fn test_cuda_runend_nullable_values() -> VortexResult<()> { let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -315,13 +350,48 @@ mod tests { PrimitiveArray::new(Buffer::from(vec![10i32, 0, 30]), validity).into_array(); let runend_array = RunEnd::new(ends_array, values_array, cuda_ctx.execution_ctx()); - // execute_cuda should fall back to CPU and still produce the correct result. - let gpu_result = runend_array - .clone() - .into_array() - .execute_cuda(&mut cuda_ctx) + // The GPU expands the per-run validity bitmap through the run mapping. + // Call the executor directly: `execute_cuda` would silently fall back to CPU for a + // host-resident array, hiding a GPU failure. + let gpu_result = RunEndExecutor + .execute(runend_array.clone().into_array(), &mut cuda_ctx) + .await + .vortex_expect("GPU decompression failed") + .into_host() + .await? + .into_array(); + + assert_arrays_eq!(runend_array, gpu_result, &mut ctx); + + Ok(()) + } + + /// Validity expansion packs bits a byte at a time, so exercise a run layout whose runs + /// straddle output byte boundaries. + #[crate::test] + async fn test_cuda_runend_nullable_values_across_byte_boundaries() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let num_runs = 301; + let ends: Vec = (1..=num_runs).map(|run| run * 3).collect(); + let values: Vec = (0..num_runs as i32).collect(); + let validity = Validity::Array( + BoolArray::from_iter((0..num_runs).map(|run| run % 3 != 0)).into_array(), + ); + + let ends_array = + PrimitiveArray::new(Buffer::from(ends), Validity::NonNullable).into_array(); + let values_array = PrimitiveArray::new(Buffer::from(values), validity).into_array(); + let runend_array = RunEnd::new(ends_array, values_array, cuda_ctx.execution_ctx()); + + // Call the executor directly: `execute_cuda` would silently fall back to CPU for a + // host-resident array, hiding a GPU failure. + let gpu_result = RunEndExecutor + .execute(runend_array.clone().into_array(), &mut cuda_ctx) .await - .vortex_expect("GPU/CPU fallback should succeed") + .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); diff --git a/vortex-cuda/src/kernel/encodings/sequence.rs b/vortex-cuda/src/kernel/encodings/sequence.rs index e52a706c8b7..0543e43f3f5 100644 --- a/vortex-cuda/src/kernel/encodings/sequence.rs +++ b/vortex-cuda/src/kernel/encodings/sequence.rs @@ -11,11 +11,11 @@ use vortex::array::ArrayRef; use vortex::array::Canonical; use vortex::array::arrays::PrimitiveArray; use vortex::array::buffer::BufferHandle; -use vortex::array::match_each_native_ptype; +use vortex::array::match_each_unsigned_integer_ptype; use vortex::dtype::NativePType; use vortex::dtype::Nullability; +use vortex::dtype::PType; use vortex::encodings::sequence::Sequence; -use vortex::encodings::sequence::SequenceDataParts; use vortex::error::VortexResult; use vortex::error::vortex_err; @@ -29,6 +29,8 @@ pub(crate) struct SequenceExecutor; #[async_trait] impl CudaExecute for SequenceExecutor { + // The kernel intentionally truncates both operands to the output width. + #[expect(clippy::cast_possible_truncation)] #[instrument(level = "trace", skip_all, fields(executor = ?self))] async fn execute( &self, @@ -41,17 +43,21 @@ impl CudaExecute for SequenceExecutor { let len = array.len(); let nullability = array.dtype().nullability(); - - let SequenceDataParts { - base, - multiplier, - ptype, - } = array.into_data().into_parts(); - - match_each_native_ptype!(ptype, |P| { - let base = base.cast::

()?; - let multiplier = multiplier.cast::

()?; - execute_typed::

(base, multiplier, len, nullability, ctx).await + let output_ptype = PType::try_from(array.dtype())?; + + // Unsigned wrapping arithmetic handles every signedness pair at the output width. + let (base, multiplier) = array.wrapping_bits()?; + + match_each_unsigned_integer_ptype!(output_ptype.to_unsigned(), |U| { + execute_typed::( + base as U, + multiplier as U, + len, + output_ptype, + nullability, + ctx, + ) + .await }) } } @@ -60,24 +66,28 @@ async fn execute_typed( base: T, multiplier: T, len: usize, + output_ptype: PType, nullability: Nullability, ctx: &mut CudaExecutionCtx, ) -> VortexResult { - let buffer = ctx.device_alloc::(len)?; + let mut buffer = ctx.device_alloc::(len)?; let len_u64 = len as u64; let kernel_func = ctx.load_function("sequence", &[T::PTYPE])?; ctx.launch_kernel(&kernel_func, len, |args| { - args.arg(&buffer).arg(&base).arg(&multiplier).arg(&len_u64); + args.arg(&mut buffer) + .arg(&base) + .arg(&multiplier) + .arg(&len_u64); })?; let output_buf = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(buffer))); Ok(Canonical::Primitive(PrimitiveArray::from_buffer_handle( output_buf, - T::PTYPE, + output_ptype, nullability.into(), ))) } @@ -88,9 +98,13 @@ mod tests { use rstest::rstest; use vortex::array::IntoArray; use vortex::array::assert_arrays_eq; + use vortex::array::builtins::ArrayBuiltins; + use vortex::dtype::DType; use vortex::dtype::NativePType; use vortex::dtype::Nullability; + use vortex::dtype::PType; use vortex::encodings::sequence::Sequence; + use vortex::error::VortexResult; use vortex::scalar::PValue; use vortex_array::VortexSessionExecute; @@ -143,4 +157,25 @@ mod tests { assert_arrays_eq!(array, gpu_result, &mut ctx); } + + #[crate::test] + async fn test_sequence_arithmetic_ptype_differs_from_output() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).unwrap(); + + let array = Sequence::try_new_typed(100i32, -10i32, Nullability::NonNullable, 5)? + .into_array() + .cast(DType::Primitive(PType::U8, Nullability::NonNullable))?; + + let gpu_result = SequenceExecutor + .execute(array.clone(), &mut cuda_ctx) + .await? + .into_host() + .await? + .into_array(); + + assert_arrays_eq!(array, gpu_result, &mut ctx); + + Ok(()) + } } diff --git a/vortex-cuda/src/kernel/encodings/zigzag.rs b/vortex-cuda/src/kernel/encodings/zigzag.rs index 107ad2b36a1..212e958ef9f 100644 --- a/vortex-cuda/src/kernel/encodings/zigzag.rs +++ b/vortex-cuda/src/kernel/encodings/zigzag.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::fmt::Debug; +use std::sync::Arc; use async_trait::async_trait; use cudarc::driver::DeviceRepr; @@ -11,17 +12,20 @@ use vortex::array::ArrayRef; use vortex::array::Canonical; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::primitive::PrimitiveDataParts; +use vortex::array::buffer::BufferHandle; use vortex::array::match_each_unsigned_integer_ptype; use vortex::dtype::NativePType; use vortex::dtype::PType; use vortex::encodings::zigzag::ZigZag; use vortex::encodings::zigzag::ZigZagArray; -use vortex::encodings::zigzag::ZigZagArrayExt; +use vortex::encodings::zigzag::ZigZagArraySlotsExt; use vortex::error::VortexResult; use vortex::error::vortex_ensure; use vortex::error::vortex_err; use crate::CudaBufferExt; +use crate::CudaDeviceBuffer; +use crate::device_buffer::with_cuda_view_mut; use crate::executor::CudaArrayExt; use crate::executor::CudaExecute; use crate::executor::CudaExecutionCtx; @@ -50,6 +54,10 @@ impl CudaExecute for ZigZagExecutor { // The encoded array is unsigned, we decode to signed of the same width. let encoded_ptype = array.encoded().dtype().as_ptype(); let output_ptype = PType::try_from(array.dtype())?; + vortex_ensure!( + output_ptype == encoded_ptype.to_signed(), + "ZigZag output type {output_ptype} must be the signed equivalent of {encoded_ptype}" + ); match_each_unsigned_integer_ptype!(encoded_ptype, |U| { decode_zigzag::(array, output_ptype, ctx).await @@ -75,22 +83,39 @@ where buffer, validity, .. } = primitive.into_data_parts(); - let device_buffer = ctx.ensure_on_device(buffer).await?; - - // Get CUDA view of the buffer - let cuda_view = device_buffer.cuda_view::()?; + let mut device_buffer = ctx.ensure_on_device(buffer).await?; let array_len_u64 = array_len as u64; - // Load kernel function let cuda_function = ctx.load_function("zigzag", &[U::PTYPE])?; + let (next, launch) = with_cuda_view_mut::(device_buffer, |view| { + ctx.launch_kernel(&cuda_function, array_len, |args| { + args.arg(view).arg(&array_len_u64); + }) + })?; + device_buffer = next; + if let Some(launch) = launch { + launch?; + return Ok(Canonical::Primitive(PrimitiveArray::from_buffer_handle( + device_buffer, + output_ptype, + validity, + ))); + } + // Preserve aliased inputs by copying into a unique allocation, then use the same in-place + // kernel as the zero-copy path. + let input_view = device_buffer.cuda_view::()?; + let mut output = ctx.device_alloc::(array_len)?; + ctx.stream() + .memcpy_dtod(&input_view, &mut output) + .map_err(|err| vortex_err!("Failed to copy shared ZigZag input: {err}"))?; ctx.launch_kernel(&cuda_function, array_len, |args| { - args.arg(&cuda_view).arg(&array_len_u64); + args.arg(&mut output).arg(&array_len_u64); })?; - // Build result - in-place, reinterpret as signed + let output = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(output))); Ok(Canonical::Primitive(PrimitiveArray::from_buffer_handle( - device_buffer, + output, output_ptype, validity, ))) diff --git a/vortex-cuda/src/kernel/encodings/zstd.rs b/vortex-cuda/src/kernel/encodings/zstd.rs index 6d14cc46e86..01f8f964615 100644 --- a/vortex-cuda/src/kernel/encodings/zstd.rs +++ b/vortex-cuda/src/kernel/encodings/zstd.rs @@ -331,7 +331,7 @@ async fn decode_zstd(array: ZstdArray, ctx: &mut CudaExecutionCtx) -> VortexResu .indices() { AllOr::All => { - let (buffers, all_views) = reconstruct_views(&host_buffer, MAX_BUFFER_LEN); + let (buffers, all_views) = reconstruct_views(&host_buffer, 0, MAX_BUFFER_LEN); let sliced_views = all_views.slice(slice_value_idx_start..slice_value_idx_stop); Ok(Canonical::VarBinView(unsafe { diff --git a/vortex-cuda/src/kernel/filter/mod.rs b/vortex-cuda/src/kernel/filter/mod.rs index 3c504113b11..39d897f3fcb 100644 --- a/vortex-cuda/src/kernel/filter/mod.rs +++ b/vortex-cuda/src/kernel/filter/mod.rs @@ -19,7 +19,7 @@ use tracing::instrument; use vortex::array::ArrayRef; use vortex::array::Canonical; use vortex::array::arrays::Filter; -use vortex::array::arrays::filter::FilterArrayExt; +use vortex::array::arrays::filter::FilterArraySlotsExt; use vortex::array::buffer::BufferHandle; use vortex::array::match_each_decimal_value_type; use vortex::array::match_each_native_simd_ptype; @@ -109,7 +109,7 @@ async fn filter_sized(temp_bytes.max(1))?; + let mut d_temp = ctx.device_alloc::(temp_bytes.max(1))?; let mut d_output = ctx.device_alloc::(output_len)?; let mut d_num_selected = ctx.device_alloc::(1)?; // Get raw pointers for FFI. @@ -134,7 +134,7 @@ async fn filter_sized(); let (d_packed_ptr, record_d_packed) = d_packed_view.device_ptr(stream); - let (d_temp_ptr, record_d_temp) = d_temp.device_ptr(stream); + let (d_temp_ptr, record_d_temp) = d_temp.device_ptr_mut(stream); let (d_output_ptr, record_d_output) = d_output.device_ptr_mut(stream); let (d_num_selected_ptr, record_d_num_selected) = d_num_selected.device_ptr_mut(stream); diff --git a/vortex-cuda/src/kernel/mod.rs b/vortex-cuda/src/kernel/mod.rs index 323458db2d8..fa27c7479a5 100644 --- a/vortex-cuda/src/kernel/mod.rs +++ b/vortex-cuda/src/kernel/mod.rs @@ -20,9 +20,9 @@ use vortex::error::vortex_err; use vortex::utils::aliases::dash_map::DashMap; mod arrays; -// PTX kernels embedded at build time; generated by build.rs's generate_embedded_ptx. -mod embedded_ptx { - include!(concat!(env!("OUT_DIR"), "/embedded_ptx.rs")); +// Fatbins let the CUDA driver select among the configured real and virtual architectures. +mod embedded_kernels { + include!(concat!(env!("OUT_DIR"), "/embedded_kernels.rs")); } mod encodings; mod filter; @@ -31,6 +31,8 @@ mod slice; pub(crate) use arrays::ConstantNumericExecutor; pub(crate) use arrays::DictExecutor; +pub(crate) use arrays::ListExecutor; +pub(crate) use arrays::MaskedExecutor; pub(crate) use arrays::SharedExecutor; pub use encodings::ZstdKernelPrep; pub use encodings::zstd_kernel_prepare; @@ -193,9 +195,7 @@ pub(crate) fn launch_cuda_kernel_with_config( } } -/// Loader for CUDA kernels with PTX caching. -/// -/// Handles loading PTX files, compiling modules, and loading functions. +/// Loader for embedded CUDA kernels with module caching. #[derive(Debug)] pub(crate) struct KernelLoader { /// Cache of loaded CUDA modules, keyed by module name @@ -217,7 +217,7 @@ impl KernelLoader { /// /// # Arguments /// - /// * `module_name` - Name of the module (`kernels/{module_name}.ptx`) + /// * `module_name` - Kernel source name without the `.cu` extension /// * `type_suffixes` - List of type suffix strings for the kernel name (`kernel_i128`) /// * `cuda_context` - CUDA context for loading the module pub fn load_function( @@ -249,17 +249,17 @@ impl KernelLoader { .map_err(|e| vortex_err!("Failed to load kernel function '{}': {}", kernel_name, e)) } - /// Loads a CUDA module from PTX embedded in the binary. + /// Loads a CUDA module from a fatbin embedded in the binary. fn load_module( module_name: &str, cuda_context: &Arc, ) -> VortexResult> { - let ptx = embedded_ptx::embedded_ptx(module_name).ok_or_else(|| { + let fatbin = embedded_kernels::embedded_kernel(module_name).ok_or_else(|| { vortex_err!("CUDA module {module_name} was not embedded at build time") })?; cuda_context - .load_module(Ptx::from_src(ptx)) + .load_module(Ptx::from_binary(fatbin.to_vec())) .map_err(|e| vortex_err!("Failed to load embedded CUDA module {module_name}: {e}")) } } @@ -290,7 +290,7 @@ mod tests { // Allocate output buffer for 3 u32 values // SAFETY: Allocating uninitialized memory that will be written by kernel - let output = unsafe { + let mut output = unsafe { stream .alloc::(3) .expect("failed to allocate output buffer") @@ -309,7 +309,7 @@ mod tests { }; let mut launch_args = stream.launch_builder(&function); - launch_args.arg(&output); + launch_args.arg(&mut output); // SAFETY: kernel only writes to output buffer unsafe { diff --git a/vortex-cuda/src/kernel/patches/mod.rs b/vortex-cuda/src/kernel/patches/mod.rs index bd6d2b737db..7c651b7e507 100644 --- a/vortex-cuda/src/kernel/patches/mod.rs +++ b/vortex-cuda/src/kernel/patches/mod.rs @@ -21,6 +21,7 @@ use vortex::dtype::PType; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::error::vortex_ensure; +use vortex::error::vortex_err; use crate::CudaBufferExt; use crate::CudaDeviceBuffer; @@ -91,7 +92,10 @@ pub(crate) fn build_gpu_patches( } } -/// Apply a set of patches in-place onto a [`CudaDeviceBuffer`] holding `ValuesT`. +/// Apply a set of patches to a [`CudaDeviceBuffer`] holding `ValuesT`. +/// +/// The target is updated in place when its allocation is uniquely owned. Shared allocations are +/// copied before patching so aliases remain unchanged. /// /// Naive scatter kernel. Kept as a reusable fallback for encoders that cannot /// use the chunk-based fused patching path (e.g., where `chunk_offsets` are @@ -103,7 +107,7 @@ pub(crate) async fn execute_patches< IndicesT: NativePType + DeviceRepr, >( patches: Patches, - target: CudaDeviceBuffer, + mut target: CudaDeviceBuffer, ctx: &mut CudaExecutionCtx, ) -> VortexResult { let indices = patches.indices().clone(); @@ -151,21 +155,35 @@ pub(crate) async fn execute_patches< let d_patch_indices = ctx.ensure_on_device(indices_buffer).await?; let d_patch_values = ctx.ensure_on_device(values_buffer).await?; - - let d_target_view = target.as_view::(); let d_patch_indices_view = d_patch_indices.cuda_view::()?; let d_patch_values_view = d_patch_values.cuda_view::()?; - let kernel_func = ctx.load_function("patches", &[ValuesT::PTYPE, IndicesT::PTYPE])?; + if let Some(launch) = target.with_view_mut::(|view| { + ctx.launch_kernel(&kernel_func, patches_len, |args| { + args.arg(view) + .arg(&d_patch_indices_view) + .arg(&d_patch_values_view) + .arg(&patches_len_u64); + }) + }) { + launch?; + return Ok(target); + } + + let target_view = target.as_view::(); + let mut output = ctx.device_alloc::(target_view.len())?; + ctx.stream() + .memcpy_dtod(&target_view, &mut output) + .map_err(|err| vortex_err!("Failed to copy CUDA patch target: {err}"))?; ctx.launch_kernel(&kernel_func, patches_len, |args| { - args.arg(&d_target_view) + args.arg(&mut output) .arg(&d_patch_indices_view) .arg(&d_patch_values_view) .arg(&patches_len_u64); })?; - Ok(target) + Ok(CudaDeviceBuffer::new(output)) } #[cfg(test)] @@ -246,10 +264,15 @@ mod tests { .downcast_ref::() .unwrap() .clone(); + let input_ptr = device_buf.offset_ptr(); + // Release the handle's allocation reference so `execute_patches` can + // prove unique ownership and patch the existing device buffer in place. + drop(handle); let patched_buf = execute_patches::(patches, device_buf, &mut cuda_ctx) .await .unwrap(); + assert_eq!(patched_buf.offset_ptr(), input_ptr); let gpu_result = PrimitiveArray::from_buffer_handle( BufferHandle::new_device(Arc::new(patched_buf)), diff --git a/vortex-cuda/src/kernel/slice/mod.rs b/vortex-cuda/src/kernel/slice/mod.rs index 30b1fdaf640..9ac4b4c70cf 100644 --- a/vortex-cuda/src/kernel/slice/mod.rs +++ b/vortex-cuda/src/kernel/slice/mod.rs @@ -7,7 +7,7 @@ use vortex::array::ArrayRef; use vortex::array::Canonical; use vortex::array::IntoArray; use vortex::array::arrays::Slice; -use vortex::array::arrays::slice::SliceArrayExt; +use vortex::array::arrays::slice::SliceArraySlotsExt; use vortex::error::VortexResult; use vortex::error::vortex_err; diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index 1b1d37cd14d..216b3369bd1 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -13,21 +13,16 @@ use async_trait::async_trait; use futures::FutureExt; use futures::StreamExt; use futures::future::BoxFuture; -use vortex::array::ArrayContext; -use vortex::array::ArrayId; use vortex::array::ArrayRef; use vortex::array::ArrayVTable; -use vortex::array::DeserializeMetadata; use vortex::array::MaskFuture; use vortex::array::ProstMetadata; use vortex::array::VortexSessionExecute; use vortex::array::arrays::Constant; -use vortex::array::expr::Expression; +use vortex::array::expr::BoundExpression; use vortex::array::expr::stats::Precision; use vortex::array::expr::stats::Stat; use vortex::array::expr::stats::StatsProvider; -use vortex::array::normalize::NormalizeOptions; -use vortex::array::normalize::Operation; use vortex::array::serde::SerializeOptions; use vortex::array::serde::SerializedArray; use vortex::array::stats::StatsSetRef; @@ -39,26 +34,27 @@ use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::error::vortex_panic; -use vortex::layout::IntoLayout; -use vortex::layout::LayoutBuildContext; +use vortex::layout::Layout; use vortex::layout::LayoutChildType; -use vortex::layout::LayoutChildren; +use vortex::layout::LayoutDeserializeArgs; use vortex::layout::LayoutEncodingRef; use vortex::layout::LayoutId; +use vortex::layout::LayoutParts; use vortex::layout::LayoutReader; use vortex::layout::LayoutReaderRef; use vortex::layout::LayoutRef; use vortex::layout::LayoutStrategy; +use vortex::layout::LayoutWriterContext; use vortex::layout::RowSplits; use vortex::layout::SplitRange; use vortex::layout::VTable; +use vortex::layout::layout_children; use vortex::layout::layouts::SharedArrayFuture; use vortex::layout::segments::SegmentId; use vortex::layout::segments::SegmentSinkRef; use vortex::layout::segments::SegmentSource; use vortex::layout::sequence::SendableSequentialStream; use vortex::layout::sequence::SequencePointer; -use vortex::layout::vtable; use vortex::mask::Mask; use vortex::scalar::Scalar; use vortex::scalar::ScalarTruncation; @@ -68,7 +64,6 @@ use vortex::session::VortexSession; use vortex::session::registry::CachedId; use vortex::session::registry::ReadContext; use vortex::utils::aliases::hash_map::HashMap; -use vortex::utils::aliases::hash_set::HashSet; /// A buffer inlined into layout metadata for host-side access. #[derive(Clone, prost::Message)] @@ -88,15 +83,16 @@ pub struct CudaFlatLayoutMetadata { pub host_buffers: Vec, } -vtable!(CudaFlat); +/// CUDA flat layout vtable. +#[derive(Clone, Debug)] +pub struct CudaFlat; -#[derive(Debug)] -pub struct CudaFlatLayoutEncoding; +/// Backwards-compatible plugin name. +pub use CudaFlat as CudaFlatLayoutEncoding; +/// CUDA-flat-specific data. #[derive(Clone, Debug)] -pub struct CudaFlatLayout { - row_count: u64, - dtype: DType, +pub struct CudaFlatData { segment_id: SegmentId, ctx: ReadContext, array_tree: ByteBuffer, @@ -104,7 +100,10 @@ pub struct CudaFlatLayout { host_buffers: Arc>, } -impl CudaFlatLayout { +/// A CUDA-optimized terminal layout. +pub type CudaFlatLayout = Layout; + +impl CudaFlatData { #[inline] pub fn segment_id(&self) -> SegmentId { self.segment_id @@ -127,28 +126,15 @@ impl CudaFlatLayout { } impl VTable for CudaFlat { - type Layout = CudaFlatLayout; - type Encoding = CudaFlatLayoutEncoding; + type LayoutData = CudaFlatData; type Metadata = ProstMetadata; - fn id(_encoding: &Self::Encoding) -> LayoutId { + fn id(&self) -> LayoutId { static ID: CachedId = CachedId::new("vortex.cuda_flat"); *ID } - fn encoding(_layout: &Self::Layout) -> LayoutEncodingRef { - LayoutEncodingRef::new_ref(CudaFlatLayoutEncoding.as_ref()) - } - - fn row_count(layout: &Self::Layout) -> u64 { - layout.row_count - } - - fn dtype(layout: &Self::Layout) -> &DType { - &layout.dtype - } - - fn metadata(layout: &Self::Layout) -> Self::Metadata { + fn metadata(layout: &Layout) -> Self::Metadata { ProstMetadata(CudaFlatLayoutMetadata { array_encoding_tree: layout.array_tree.to_vec(), host_buffers: layout @@ -162,24 +148,40 @@ impl VTable for CudaFlat { }) } - fn segment_ids(layout: &Self::Layout) -> Vec { - vec![layout.segment_id] - } - - fn nchildren(_layout: &Self::Layout) -> usize { - 0 + fn deserialize( + &self, + args: &LayoutDeserializeArgs<'_>, + metadata: &CudaFlatLayoutMetadata, + ) -> VortexResult { + if args.segment_ids.len() != 1 { + vortex_bail!("CudaFlatLayout must have exactly one segment ID"); + } + if args.children.nchildren() != 0 { + vortex_bail!("CudaFlatLayout must not have children"); + } + let host_buffers = metadata + .host_buffers + .iter() + .map(|hb| (hb.buffer_index, ByteBuffer::from(hb.data.clone()))) + .collect(); + Ok(CudaFlatData { + segment_id: args.segment_ids[0], + ctx: args.array_read_ctx.clone(), + array_tree: ByteBuffer::from(metadata.array_encoding_tree.clone()), + host_buffers: Arc::new(host_buffers), + }) } - fn child(_layout: &Self::Layout, _idx: usize) -> VortexResult { - vortex_bail!("CudaFlatLayout has no children"); + fn child_dtype(_layout: &Layout, idx: usize) -> VortexResult { + vortex_bail!("CudaFlatLayout has no child {idx}"); } - fn child_type(_layout: &Self::Layout, _idx: usize) -> LayoutChildType { + fn child_type(_layout: &Layout, _idx: usize) -> LayoutChildType { vortex_panic!("CudaFlatLayout has no children"); } fn new_reader( - layout: &Self::Layout, + layout: &Layout, name: Arc, segment_source: Arc, session: &VortexSession, @@ -193,40 +195,6 @@ impl VTable for CudaFlat { array: Default::default(), })) } - - fn build( - _encoding: &Self::Encoding, - dtype: &DType, - row_count: u64, - metadata: &::Output, - segment_ids: Vec, - _children: &dyn LayoutChildren, - build_ctx: &LayoutBuildContext<'_>, - ) -> VortexResult { - if segment_ids.len() != 1 { - vortex_bail!("CudaFlatLayout must have exactly one segment ID"); - } - let host_buffers: HashMap = metadata - .host_buffers - .iter() - .map(|hb| (hb.buffer_index, ByteBuffer::from(hb.data.clone()))) - .collect(); - Ok(CudaFlatLayout { - row_count, - dtype: dtype.clone(), - segment_id: segment_ids[0], - ctx: build_ctx.array_read_ctx.clone(), - array_tree: ByteBuffer::from(metadata.array_encoding_tree.clone()), - host_buffers: Arc::new(host_buffers), - }) - } - - fn with_children(_layout: &mut Self::Layout, children: Vec) -> VortexResult<()> { - if !children.is_empty() { - vortex_bail!("CudaFlatLayout has no children, got {}", children.len()); - } - Ok(()) - } } // Threshold to order filter and apply expression, copied from FlatLayout. @@ -244,14 +212,14 @@ impl CudaFlatReader { fn array_future(&self) -> SharedArrayFuture { self.array .get_or_init(|| { - let row_count = usize::try_from(self.layout.row_count) + let row_count = usize::try_from(self.layout.row_count()) .vortex_expect("row count must fit in usize"); let segment_fut = self.segment_source.request(self.layout.segment_id); let ctx = self.layout.ctx.clone(); let session = self.session.clone(); - let dtype = self.layout.dtype.clone(); + let dtype = self.layout.dtype().clone(); let array_tree = self.layout.array_tree.clone(); let host_buffers = Arc::clone(&self.layout.host_buffers); @@ -279,11 +247,11 @@ impl LayoutReader for CudaFlatReader { } fn dtype(&self) -> &DType { - &self.layout.dtype + self.layout.dtype() } fn row_count(&self) -> u64 { - self.layout.row_count + self.layout.row_count() } fn register_splits( @@ -292,7 +260,7 @@ impl LayoutReader for CudaFlatReader { split_range: &SplitRange, splits: &mut RowSplits, ) -> VortexResult<()> { - split_range.check_bounds(self.layout.row_count)?; + split_range.check_bounds(self.layout.row_count())?; splits.push(split_range.root_row_range().end); Ok(()) } @@ -300,7 +268,7 @@ impl LayoutReader for CudaFlatReader { fn pruning_evaluation( &self, _row_range: &Range, - _expr: &Expression, + _expr: &BoundExpression, mask: Mask, ) -> VortexResult { Ok(MaskFuture::ready(mask)) @@ -309,7 +277,7 @@ impl LayoutReader for CudaFlatReader { fn filter_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult { let row_range = usize::try_from(row_range.start) @@ -331,13 +299,13 @@ impl LayoutReader for CudaFlatReader { let mask_density = mask.density(); let array_mask = if mask_density < EXPR_EVAL_THRESHOLD { - let array = array.apply(&expr)?; + let array = array.apply_bound(&expr)?; let array = array.filter(mask.clone())?; let mut ctx = session.create_execution_ctx(); let array_mask = array.null_as_false().execute(&mut ctx)?; mask.intersect_by_rank(&array_mask) } else { - let array = array.apply(&expr)?; + let array = array.apply_bound(&expr)?; let mut ctx = session.create_execution_ctx(); let array_mask = array.null_as_false().execute(&mut ctx)?; mask.bitand(&array_mask) @@ -358,7 +326,7 @@ impl LayoutReader for CudaFlatReader { fn projection_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult>> { let row_range = usize::try_from(row_range.start) @@ -383,7 +351,7 @@ impl LayoutReader for CudaFlatReader { array = array.filter(mask)?; } - array = array.apply(&expr)?; + array = array.apply_bound(&expr)?; Ok(array) } @@ -403,8 +371,6 @@ pub struct CudaFlatLayoutStrategy { pub include_padding: bool, /// Maximum length of variable length statistics. pub max_variable_length_statistics_size: usize, - /// Optional set of allowed array encodings for normalization. - pub allowed_encodings: Option>, } impl Default for CudaFlatLayoutStrategy { @@ -412,7 +378,6 @@ impl Default for CudaFlatLayoutStrategy { Self { include_padding: true, max_variable_length_statistics_size: 64, - allowed_encodings: None, } } } @@ -427,11 +392,6 @@ impl CudaFlatLayoutStrategy { self.max_variable_length_statistics_size = size; self } - - pub fn with_allow_encodings(mut self, allow_encodings: HashSet) -> Self { - self.allowed_encodings = Some(allow_encodings); - self - } } fn truncate_scalar_stat Option<(Scalar, bool)>>( @@ -454,13 +414,12 @@ fn truncate_scalar_stat Option<(Scalar, bool)>>( impl LayoutStrategy for CudaFlatLayoutStrategy { async fn write_stream( &self, - ctx: ArrayContext, + ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, mut stream: SendableSequentialStream, _eof: SequencePointer, session: &VortexSession, ) -> VortexResult { - let ctx = ctx.clone(); let options = self.clone(); let Some(chunk) = stream.next().await else { vortex_bail!("CudaFlatLayoutStrategy needs a single chunk"); @@ -508,20 +467,11 @@ impl LayoutStrategy for CudaFlatLayoutStrategy { _ => {} } - let chunk = if let Some(allowed) = &options.allowed_encodings { - chunk.normalize(&mut NormalizeOptions { - allowed, - operation: Operation::Error, - })? - } else { - chunk - }; - // Scan for constant array buffers before serialization (while data is still on host). let host_buffers = extract_constant_buffers(&chunk); let buffers = chunk.serialize( - &ctx, + ctx.array_ctx(), session, &SerializeOptions { offset: 0, @@ -544,14 +494,19 @@ impl LayoutStrategy for CudaFlatLayoutStrategy { .map(|hb| (hb.buffer_index, ByteBuffer::from(hb.data.clone()))) .collect(); - Ok(CudaFlatLayout { + Ok(LayoutParts::new( + CudaFlat, + stream.dtype().clone(), row_count, - dtype: stream.dtype().clone(), - segment_id, - ctx: ReadContext::new(ctx.to_ids()), - array_tree, - host_buffers: Arc::new(host_buffer_map), - } + vec![segment_id], + layout_children(Vec::new()), + CudaFlatData { + segment_id, + ctx: ReadContext::new(ctx.array_ctx().to_ids()), + array_tree, + host_buffers: Arc::new(host_buffer_map), + }, + ) .into_layout()) } } @@ -586,5 +541,5 @@ pub fn register_cuda_layout(session: &VortexSession) { use vortex::layout::session::LayoutSessionExt; session .layouts() - .register(LayoutEncodingRef::new_ref(CudaFlatLayoutEncoding.as_ref())); + .register(LayoutEncodingRef::new_ref(&CudaFlat)); } diff --git a/vortex-cuda/src/lib.rs b/vortex-cuda/src/lib.rs index 5817e15436b..2447dfdbad5 100644 --- a/vortex-cuda/src/lib.rs +++ b/vortex-cuda/src/lib.rs @@ -14,6 +14,7 @@ mod device_buffer; mod device_read_at; pub mod dynamic_dispatch; pub mod executor; +mod file; pub mod hybrid_dispatch; mod kernel; pub mod layout; @@ -34,22 +35,27 @@ pub use device_read_at::CopyDeviceReadAt; pub use executor::CudaDispatchMode; pub use executor::CudaExecutionCtx; pub use executor::CudaKernelEvents; +pub use file::CudaOpenOptions; +pub use file::CudaOpenOptionsExt; use kernel::ALPExecutor; use kernel::BitPackedExecutor; use kernel::ConstantNumericExecutor; use kernel::DateTimePartsExecutor; use kernel::DecimalBytePartsExecutor; pub use kernel::DefaultLaunchStrategy; +use kernel::DeltaExecutor; use kernel::DictExecutor; use kernel::FSSTExecutor; use kernel::FilterExecutor; use kernel::FoRExecutor; pub use kernel::LaunchStrategy; +use kernel::ListExecutor; +use kernel::MaskedExecutor; +use kernel::OnPairExecutor; use kernel::RunEndExecutor; use kernel::SharedExecutor; pub use kernel::TracingLaunchStrategy; use kernel::ZigZagExecutor; -#[cfg(feature = "unstable_encodings")] use kernel::ZstdBuffersExecutor; use kernel::ZstdExecutor; pub use kernel::ZstdKernelPrep; @@ -59,6 +65,7 @@ pub use pinned::PinnedPoolStats; pub use pinned::PooledPinnedBuffer; pub use pooled_read_at::PooledByteBufferReadAt; pub use pooled_read_at::PooledFileReadAt; +pub use pooled_read_at::PooledFileReadAtOptions; pub use pooled_read_at::PooledObjectStoreReadAt; pub use session::CudaSession; pub use session::CudaSessionExt; @@ -69,23 +76,26 @@ use vortex::array::ArrayVTable; use vortex::array::arrays::Constant; use vortex::array::arrays::Dict; use vortex::array::arrays::Filter; +use vortex::array::arrays::List; +use vortex::array::arrays::Masked; use vortex::array::arrays::Shared; use vortex::array::arrays::Slice; use vortex::encodings::alp::ALP; use vortex::encodings::datetime_parts::DateTimeParts; use vortex::encodings::decimal_byte_parts::DecimalByteParts; use vortex::encodings::fastlanes::BitPacked; +use vortex::encodings::fastlanes::Delta; use vortex::encodings::fastlanes::FoR; use vortex::encodings::fsst::FSST; use vortex::encodings::runend::RunEnd; use vortex::encodings::sequence::Sequence; use vortex::encodings::zigzag::ZigZag; use vortex::encodings::zstd::Zstd; -#[cfg(feature = "unstable_encodings")] use vortex::encodings::zstd::ZstdBuffers; #[cfg(test)] use vortex_cuda_macros::test; pub use vortex_nvcomp as nvcomp; +use vortex_onpair::OnPair; use crate::kernel::SequenceExecutor; use crate::kernel::SliceExecutor; @@ -111,14 +121,17 @@ pub fn initialize_cuda(session: &CudaSession) { session.register_kernel(DateTimeParts.id(), &DateTimePartsExecutor); session.register_kernel(DecimalByteParts.id(), &DecimalBytePartsExecutor); session.register_kernel(Dict.id(), &DictExecutor); + session.register_kernel(List.id(), &ListExecutor); + session.register_kernel(Masked.id(), &MaskedExecutor); session.register_kernel(Shared.id(), &SharedExecutor); + session.register_kernel(Delta.id(), &DeltaExecutor); session.register_kernel(FoR.id(), &FoRExecutor); session.register_kernel(FSST.id(), &FSSTExecutor); + session.register_kernel(OnPair.id(), &OnPairExecutor); session.register_kernel(RunEnd.id(), &RunEndExecutor); session.register_kernel(Sequence.id(), &SequenceExecutor); session.register_kernel(ZigZag.id(), &ZigZagExecutor); session.register_kernel(Zstd.id(), &ZstdExecutor); - #[cfg(feature = "unstable_encodings")] session.register_kernel(ZstdBuffers.id(), &ZstdBuffersExecutor); // Operation kernels diff --git a/vortex-cuda/src/pinned.rs b/vortex-cuda/src/pinned.rs index 52b1aa46e2a..11796c5899a 100644 --- a/vortex-cuda/src/pinned.rs +++ b/vortex-cuda/src/pinned.rs @@ -309,6 +309,19 @@ impl PooledPinnedBuffer { .unwrap_or_else(|e| vortex_panic!("failed to access pinned host buffer: {e}")) } + /// Shortens the logical length without reallocating the pinned buffer. + /// + /// This is an O(1) metadata update; the allocation capacity is unchanged. + #[cfg(target_os = "linux")] + pub(crate) fn truncate(&mut self, len: usize) { + let inner = self + .inner + .as_mut() + .unwrap_or_else(|| vortex_panic!("buffer already consumed")); + assert!(len <= inner.len()); + inner.set_logical_len(len); + } + /// Submits a non-blocking H2D DMA transfer and returns a device buffer. /// /// The pinned buffer is placed in the pool's inflight queue, gated on a `CudaEvent` marking diff --git a/vortex-cuda/src/pooled_read_at.rs b/vortex-cuda/src/pooled_read_at.rs index 1b5ec4fb31a..f5352a7c06a 100644 --- a/vortex-cuda/src/pooled_read_at.rs +++ b/vortex-cuda/src/pooled_read_at.rs @@ -1,9 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::fs::File; use std::io; -use std::path::Path; use std::sync::Arc; use futures::FutureExt; @@ -30,95 +28,13 @@ use vortex::io::std_file::read_exact_at; use crate::pinned::PinnedByteBufferPool; use crate::stream::VortexCudaStream; -/// Default number of concurrent requests to allow for local file I/O. -pub const DEFAULT_FILE_CONCURRENCY: usize = 32; -/// Default number of concurrent requests to allow for object store I/O. -pub const DEFAULT_OBJECT_STORE_CONCURRENCY: usize = 192; - -/// File reader that uses CUDA pinned host memory for I/O buffers and transfers -/// directly to the GPU. -/// -/// Reads into a pooled pinned (page-locked) buffer, then submits a non-blocking -/// H2D DMA transfer and returns a device `BufferHandle`. -#[derive(Clone)] -pub struct PooledFileReadAt { - uri: Arc, - file: Arc, - handle: Handle, - pool: Arc, - stream: VortexCudaStream, -} - -impl PooledFileReadAt { - /// Open a file for pooled reading with direct device transfer. - pub fn open( - path: impl AsRef, - handle: Handle, - pool: Arc, - stream: VortexCudaStream, - ) -> VortexResult { - let path = path.as_ref(); - let uri = Arc::from(path.to_string_lossy().to_string()); - let file = Arc::new(File::open(path)?); - Ok(Self { - uri, - file, - handle, - pool, - stream, - }) - } -} +mod file; -impl VortexReadAt for PooledFileReadAt { - fn uri(&self) -> Option<&Arc> { - Some(&self.uri) - } +pub use file::PooledFileReadAt; +pub use file::PooledFileReadAtOptions; - fn coalesce_config(&self) -> Option { - Some(CoalesceConfig::file()) - } - - fn concurrency(&self) -> usize { - DEFAULT_FILE_CONCURRENCY - } - - fn size(&self) -> BoxFuture<'static, VortexResult> { - let file = Arc::clone(&self.file); - async move { - let metadata = file.metadata()?; - Ok(metadata.len()) - } - .boxed() - } - - fn read_at( - &self, - offset: u64, - length: usize, - _alignment: Alignment, - ) -> BoxFuture<'static, VortexResult> { - let file = Arc::clone(&self.file); - let handle = self.handle.clone(); - let stream = self.stream.clone(); - let pool = Arc::clone(&self.pool); - - async move { - let mut target = pool.get(length)?; - let target = handle - .spawn_blocking(move || { - read_exact_at(&file, target.as_mut_slice(), offset)?; - Ok::<_, io::Error>(target) - }) - .await - .map_err(VortexError::from)?; - - let cuda_buf = target.transfer_to_device(&stream)?; - Ok(BufferHandle::new_device(Arc::new(cuda_buf))) - } - .boxed() - } -} +/// Default number of concurrent requests to allow for object store I/O. +pub const DEFAULT_OBJECT_STORE_CONCURRENCY: usize = 192; /// Object store reader that uses CUDA pinned host memory for I/O buffers and /// transfers directly to the GPU. diff --git a/vortex-cuda/src/pooled_read_at/file.rs b/vortex-cuda/src/pooled_read_at/file.rs new file mode 100644 index 00000000000..b2b44fed2bf --- /dev/null +++ b/vortex-cuda/src/pooled_read_at/file.rs @@ -0,0 +1,231 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#[cfg(target_os = "linux")] +mod direct; + +use std::fs::File; +use std::ops::Range; +use std::path::Path; +use std::sync::Arc; + +use futures::FutureExt; +use futures::future::BoxFuture; +use vortex::array::buffer::BufferHandle; +use vortex::buffer::Alignment; +use vortex::error::VortexResult; +use vortex::io::CoalesceConfig; +use vortex::io::VortexReadAt; +use vortex::io::runtime::Handle; +use vortex::io::std_file::read_exact_at; + +#[cfg(target_os = "linux")] +use self::direct::DirectFileReadBackend; +use crate::pinned::PinnedByteBufferPool; +use crate::pinned::PooledPinnedBuffer; +use crate::stream::VortexCudaStream; + +/// Default number of concurrent requests to allow for local file I/O. +pub const DEFAULT_FILE_CONCURRENCY: usize = 32; + +/// Options controlling how [`PooledFileReadAt`] opens and reads a local file. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct PooledFileReadAtOptions { + direct_io: bool, +} + +impl PooledFileReadAtOptions { + /// Bypass the operating system page cache for pooled file reads. + /// + /// This option is available only on Linux. Unaligned logical reads are widened to satisfy the + /// filesystem's direct-I/O requirements and sliced back to the requested range after transfer + /// to the device. + #[cfg(target_os = "linux")] + pub fn with_direct_io(mut self) -> Self { + self.direct_io = true; + self + } +} + +struct PooledHostRead { + buffer: PooledPinnedBuffer, + requested_range: Range, +} + +trait FileReadBackend: Send + Sync { + fn size(&self) -> VortexResult; + + fn read( + &self, + pool: &Arc, + offset: u64, + length: usize, + ) -> VortexResult; +} + +struct BufferedFileReadBackend { + file: File, +} + +impl BufferedFileReadBackend { + fn open(path: &Path) -> VortexResult { + Ok(Self { + file: File::open(path)?, + }) + } +} + +impl FileReadBackend for BufferedFileReadBackend { + fn size(&self) -> VortexResult { + Ok(self.file.metadata()?.len()) + } + + fn read( + &self, + pool: &Arc, + offset: u64, + length: usize, + ) -> VortexResult { + let mut buffer = pool.get(length)?; + read_exact_at(&self.file, buffer.as_mut_slice(), offset)?; + Ok(PooledHostRead { + buffer, + requested_range: 0..length, + }) + } +} + +#[cfg(target_os = "linux")] +fn open_backend( + path: &Path, + options: PooledFileReadAtOptions, +) -> VortexResult> { + if options.direct_io { + Ok(Arc::new(DirectFileReadBackend::open(path)?)) + } else { + Ok(Arc::new(BufferedFileReadBackend::open(path)?)) + } +} + +#[cfg(not(target_os = "linux"))] +fn open_backend( + path: &Path, + _options: PooledFileReadAtOptions, +) -> VortexResult> { + Ok(Arc::new(BufferedFileReadBackend::open(path)?)) +} + +/// File reader that uses CUDA pinned host memory for I/O buffers and transfers +/// directly to the GPU. +/// +/// Reads into a pooled pinned (page-locked) buffer, then submits a non-blocking +/// H2D DMA transfer and returns a device `BufferHandle`. +/// +/// This is a data-plane reader. To open a complete local Vortex file, prefer +/// [`crate::CudaOpenOptionsExt::with_cuda`], which keeps the footer and zone maps on the host. +#[derive(Clone)] +pub struct PooledFileReadAt { + uri: Arc, + backend: Arc, + handle: Handle, + pool: Arc, + stream: VortexCudaStream, +} + +impl PooledFileReadAt { + /// Open a file for pooled reading with direct device transfer. + pub fn open( + path: impl AsRef, + handle: Handle, + pool: Arc, + stream: VortexCudaStream, + ) -> VortexResult { + Self::open_with_options( + path, + handle, + pool, + stream, + PooledFileReadAtOptions::default(), + ) + } + + /// Open a file for pooled reading with explicit options. + pub fn open_with_options( + path: impl AsRef, + handle: Handle, + pool: Arc, + stream: VortexCudaStream, + options: PooledFileReadAtOptions, + ) -> VortexResult { + let path = path.as_ref(); + let uri = Arc::from(path.to_string_lossy().to_string()); + let backend = open_backend(path, options)?; + Ok(Self { + uri, + backend, + handle, + pool, + stream, + }) + } +} + +impl VortexReadAt for PooledFileReadAt { + fn uri(&self) -> Option<&Arc> { + Some(&self.uri) + } + + fn coalesce_config(&self) -> Option { + Some(CoalesceConfig::file()) + } + + fn concurrency(&self) -> usize { + DEFAULT_FILE_CONCURRENCY + } + + fn size(&self) -> BoxFuture<'static, VortexResult> { + let backend = Arc::clone(&self.backend); + async move { backend.size() }.boxed() + } + + fn read_at( + &self, + offset: u64, + length: usize, + _alignment: Alignment, + ) -> BoxFuture<'static, VortexResult> { + let backend = Arc::clone(&self.backend); + let handle = self.handle.clone(); + let stream = self.stream.clone(); + let pool = Arc::clone(&self.pool); + + async move { + let read = handle + .spawn_blocking(move || backend.read(&pool, offset, length)) + .await?; + let cuda_buf = read.buffer.transfer_to_device(&stream)?; + Ok(BufferHandle::new_device(Arc::new(cuda_buf)).slice(read.requested_range)) + } + .boxed() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn pooled_file_read_options_default_to_buffered_io() { + assert!(!PooledFileReadAtOptions::default().direct_io); + } + + #[cfg(target_os = "linux")] + #[test] + fn pooled_file_read_options_enable_direct_io() { + assert!( + PooledFileReadAtOptions::default() + .with_direct_io() + .direct_io + ); + } +} diff --git a/vortex-cuda/src/pooled_read_at/file/direct.rs b/vortex-cuda/src/pooled_read_at/file/direct.rs new file mode 100644 index 00000000000..53366c36915 --- /dev/null +++ b/vortex-cuda/src/pooled_read_at/file/direct.rs @@ -0,0 +1,288 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fs::File; +use std::io; +use std::ops::Range; +use std::os::unix::fs::FileExt; +use std::path::Path; +use std::sync::Arc; + +use rustix::fs::AtFlags; +use rustix::fs::Mode; +use rustix::fs::OFlags; +use rustix::fs::StatxFlags; +use vortex::error::VortexResult; +use vortex::error::vortex_ensure; +use vortex::error::vortex_err; + +use super::FileReadBackend; +use super::PooledHostRead; +use crate::pinned::PinnedByteBufferPool; + +/// Conservative direct-I/O alignment used when Linux cannot report the filesystem constraints. +/// +/// A page-sized fallback is accepted by common block devices and filesystems. If the actual +/// requirement is stricter, the read fails with the underlying `EINVAL`. +const FALLBACK_DIRECT_IO_ALIGNMENT: usize = 4096; + +#[derive(Clone, Copy)] +struct DirectIoConstraints { + /// Required alignment of the address of the userspace I/O buffer. + memory_alignment: usize, + /// Required alignment of both the file offset and the I/O length. + offset_alignment: usize, +} + +impl Default for DirectIoConstraints { + fn default() -> Self { + Self { + memory_alignment: FALLBACK_DIRECT_IO_ALIGNMENT, + offset_alignment: FALLBACK_DIRECT_IO_ALIGNMENT, + } + } +} + +#[derive(Debug, PartialEq, Eq)] +struct DirectIoRange { + read_offset: u64, + read_length: usize, + requested_range: Range, +} + +pub(super) struct DirectFileReadBackend { + file: File, + constraints: DirectIoConstraints, +} + +impl DirectFileReadBackend { + pub(super) fn open(path: &Path) -> VortexResult { + let file = File::from( + rustix::fs::open( + path, + OFlags::RDONLY | OFlags::CLOEXEC | OFlags::DIRECT, + Mode::empty(), + ) + .map_err(io::Error::from)?, + ); + let constraints = direct_io_constraints(&file)?; + Ok(Self { file, constraints }) + } +} + +impl FileReadBackend for DirectFileReadBackend { + fn size(&self) -> VortexResult { + Ok(self.file.metadata()?.len()) + } + + fn read( + &self, + pool: &Arc, + offset: u64, + length: usize, + ) -> VortexResult { + let direct_range = direct_io_range(offset, length, self.constraints.offset_alignment)?; + let mut buffer = pool.get(direct_range.read_length)?; + let address = buffer.as_mut_slice().as_ptr() as usize; + vortex_ensure!( + address.is_multiple_of(self.constraints.memory_alignment), + "pinned buffer address {address:#x} is not aligned to {} bytes", + self.constraints.memory_alignment + ); + + let bytes_read = read_direct_at( + &self.file, + buffer.as_mut_slice(), + direct_range.read_offset, + direct_range.requested_range.end, + self.constraints, + )?; + buffer.truncate(bytes_read); + Ok(PooledHostRead { + buffer, + requested_range: direct_range.requested_range, + }) + } +} + +fn direct_io_range(offset: u64, length: usize, alignment: usize) -> VortexResult { + vortex_ensure!(alignment > 0, "direct I/O alignment must be non-zero"); + if length == 0 { + return Ok(DirectIoRange { + read_offset: offset, + read_length: 0, + requested_range: 0..0, + }); + } + + let alignment_u64 = u64::try_from(alignment)?; + let length_u64 = u64::try_from(length)?; + let requested_end = offset.checked_add(length_u64).ok_or_else(|| { + vortex_err!("direct I/O range overflow: offset={offset}, length={length}") + })?; + let read_offset = offset - offset % alignment_u64; + let read_end = requested_end + .checked_next_multiple_of(alignment_u64) + .ok_or_else(|| vortex_err!("direct I/O aligned end overflow"))?; + let read_length = usize::try_from(read_end - read_offset)?; + let slice_start = usize::try_from(offset - read_offset)?; + let slice_end = slice_start.checked_add(length).ok_or_else(|| { + vortex_err!("direct I/O range overflow: offset={offset}, length={length}") + })?; + + Ok(DirectIoRange { + read_offset, + read_length, + requested_range: slice_start..slice_end, + }) +} + +fn read_direct_at( + file: &File, + buffer: &mut [u8], + offset: u64, + required_bytes: usize, + constraints: DirectIoConstraints, +) -> io::Result { + let mut initialized = 0; + while initialized < required_bytes { + let initialized_u64 = u64::try_from(initialized) + .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "read offset overflow"))?; + let read_offset = offset + .checked_add(initialized_u64) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "read offset overflow"))?; + let bytes_read = match file.read_at(&mut buffer[initialized..], read_offset) { + Err(error) if error.kind() == io::ErrorKind::Interrupted => continue, + result => result?, + }; + if bytes_read == 0 { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + format!( + "direct read returned {initialized} bytes, but {required_bytes} bytes were required" + ), + )); + } + initialized += bytes_read; + if initialized < required_bytes + && (!initialized.is_multiple_of(constraints.offset_alignment) + || !initialized.is_multiple_of(constraints.memory_alignment)) + { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + format!( + "direct read returned an unaligned short read of {initialized} bytes before the required {required_bytes} bytes" + ), + )); + } + } + + Ok(initialized) +} + +fn direct_io_constraints(file: &File) -> VortexResult { + let Ok(stat) = rustix::fs::statx( + file, + c"", + AtFlags::EMPTY_PATH | AtFlags::STATX_DONT_SYNC, + StatxFlags::DIOALIGN, + ) else { + return Ok(DirectIoConstraints::default()); + }; + if stat.stx_mask & StatxFlags::DIOALIGN.bits() == 0 { + return Ok(DirectIoConstraints::default()); + } + + let Ok(memory_alignment) = usize::try_from(stat.stx_dio_mem_align) else { + return Ok(DirectIoConstraints::default()); + }; + let Ok(offset_alignment) = usize::try_from(stat.stx_dio_offset_align) else { + return Ok(DirectIoConstraints::default()); + }; + if memory_alignment == 0 || offset_alignment == 0 { + return Ok(DirectIoConstraints::default()); + } + vortex_ensure!( + memory_alignment.is_power_of_two(), + "direct I/O memory alignment must be a power of two, got {memory_alignment}" + ); + vortex_ensure!( + offset_alignment.is_power_of_two(), + "direct I/O offset alignment must be a power of two, got {offset_alignment}" + ); + + Ok(DirectIoConstraints { + memory_alignment, + offset_alignment, + }) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case(0, 0, 4096, 0, 0, 0)] + #[case(5, 0, 4096, 5, 0, 0)] + #[case(5, 10, 4096, 0, 4096, 5)] + #[case(4090, 20, 4096, 0, 8192, 4090)] + #[case(4096, 4096, 4096, 4096, 4096, 0)] + #[case(513, 1, 512, 512, 512, 1)] + #[case(4096, 8193, 4096, 4096, 12288, 0)] + fn widens_direct_read_to_block_boundaries( + #[case] offset: u64, + #[case] length: usize, + #[case] alignment: usize, + #[case] expected_offset: u64, + #[case] expected_length: usize, + #[case] expected_prefix: usize, + ) -> VortexResult<()> { + assert_eq!( + direct_io_range(offset, length, alignment)?, + DirectIoRange { + read_offset: expected_offset, + read_length: expected_length, + requested_range: expected_prefix..expected_prefix + length, + } + ); + Ok(()) + } + + #[rstest] + #[case(u64::MAX, 2, 4096)] + #[case(0, 1, 0)] + fn rejects_invalid_direct_read_range( + #[case] offset: u64, + #[case] length: usize, + #[case] alignment: usize, + ) { + assert!(direct_io_range(offset, length, alignment).is_err()); + } + + #[test] + fn aligned_ranges_cover_requested_bytes() -> VortexResult<()> { + for alignment in [512, 4096] { + for offset in 0..alignment * 2 { + for length in [0, 1, alignment - 1, alignment, alignment + 1] { + let range = direct_io_range(offset as u64, length, alignment)?; + if length == 0 { + assert_eq!(range.read_length, 0); + continue; + } + + assert_eq!(range.read_offset % alignment as u64, 0); + assert_eq!(range.read_length % alignment, 0); + assert_eq!(range.requested_range.len(), length); + assert!(range.requested_range.end <= range.read_length); + assert_eq!( + range.read_offset + range.requested_range.start as u64, + offset as u64 + ); + } + } + } + Ok(()) + } +} diff --git a/vortex-cuda/src/session.rs b/vortex-cuda/src/session.rs index a5410db0c99..0924854405f 100644 --- a/vortex-cuda/src/session.rs +++ b/vortex-cuda/src/session.rs @@ -169,12 +169,12 @@ impl CudaSession { /// /// # Arguments /// - /// * `module_name` - Name of the module (`kernels/{module_name}.ptx`) + /// * `module_name` - Kernel source name without the `.cu` extension /// * `type_suffixes` - List of type suffix strings to generate kernel name /// /// # Errors /// - /// Returns an error if PTX file cannot be read or kernel cannot be loaded. + /// Returns an error if the module was not embedded or the kernel cannot be loaded. pub fn load_function_with_suffixes( &self, module_name: &str, diff --git a/vortex-cuda/src/stream.rs b/vortex-cuda/src/stream.rs index 4d31a7937f9..d54f8d0940b 100644 --- a/vortex-cuda/src/stream.rs +++ b/vortex-cuda/src/stream.rs @@ -259,6 +259,9 @@ fn register_stream_callback(stream: &CudaStream) -> VortexResult VortexResult<()> { + let ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let bits = ctx + .stream() + .copy_to_device(vec![0b1010_1100_u8, 0b0110_1001, 0b1100_0011])? + .await?; + let array = BoolArray::new_handle(bits, 3, 18, Validity::NonNullable).into_array(); + + let sliced = array.slice(7..16)?; + + assert_eq!(sliced.len(), 9); + assert!(sliced.buffer_handles()[0].is_on_device()); + Ok(()) + } } diff --git a/vortex-datafusion/examples/vortex_table.rs b/vortex-datafusion/examples/vortex_table.rs index 2570fc47afd..90a7ba5e1e4 100644 --- a/vortex-datafusion/examples/vortex_table.rs +++ b/vortex-datafusion/examples/vortex_table.rs @@ -70,9 +70,7 @@ async fn main() -> anyhow::Result<()> { .ok_or_else(|| vortex_err!("Path is not valid UTF-8"))?, )?; let config = ListingTableConfig::new(table_url) - .with_listing_options( - ListingOptions::new(format).with_session_config_options(ctx.state().config()), - ) + .with_listing_options(ListingOptions::new(format)) .infer_schema(&ctx.state()) .await?; diff --git a/vortex-datafusion/src/convert/exprs.rs b/vortex-datafusion/src/convert/exprs.rs index 701b4c995fa..9b0d7090507 100644 --- a/vortex-datafusion/src/convert/exprs.rs +++ b/vortex-datafusion/src/convert/exprs.rs @@ -15,12 +15,12 @@ use datafusion_expr::Operator as DFOperator; use datafusion_functions::core::getfield::GetFieldFunc; use datafusion_functions::string::octet_length::OctetLengthFunc; use datafusion_functions_nested::length::ArrayLength; +use datafusion_physical_expr::DynamicFilterTracking; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::ScalarFunctionExpr; use datafusion_physical_expr::projection::ProjectionExpr; use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_expr::utils::collect_columns; -use datafusion_physical_expr_common::physical_expr::is_dynamic_physical_expr; use datafusion_physical_plan::expressions as df_expr; use itertools::Itertools; use vortex::VortexSessionDefault; @@ -48,11 +48,13 @@ use vortex::scalar_fn::fns::operators::Operator; use vortex::session::VortexSession; use vortex_arrow::ArrowSessionExt; -use crate::convert::FromDataFusion; +use crate::convert::scalar_from_df; /// Result of splitting a projection into Vortex expressions and leftover DataFusion projections. pub struct ProcessedProjection { + /// Projection evaluated by the Vortex scan. pub scan_projection: Expression, + /// Projection evaluated by DataFusion after the Vortex scan. pub leftover_projection: ProjectionExprs, } @@ -73,6 +75,47 @@ pub(crate) fn make_vortex_predicate( } /// Trait for converting DataFusion expressions to Vortex ones. +/// +/// # Implementing a custom convertor +/// +/// ``` +/// use std::sync::Arc; +/// +/// use arrow_schema::Schema; +/// use datafusion_common::Result as DFResult; +/// use datafusion_physical_expr::PhysicalExpr; +/// use datafusion_physical_expr::projection::ProjectionExprs; +/// use vortex::expr::Expression; +/// use vortex_datafusion::convert::DefaultExpressionConvertor; +/// use vortex_datafusion::convert::ExpressionConvertor; +/// use vortex_datafusion::convert::ProcessedProjection; +/// +/// struct CustomExpressionConvertor(DefaultExpressionConvertor); +/// +/// impl ExpressionConvertor for CustomExpressionConvertor { +/// fn can_be_pushed_down(&self, expr: &Arc, schema: &Schema) -> bool { +/// self.0.can_be_pushed_down(expr, schema) +/// } +/// +/// fn convert(&self, expr: &dyn PhysicalExpr) -> DFResult { +/// self.0.convert(expr) +/// } +/// +/// fn split_projection( +/// &self, +/// source_projection: ProjectionExprs, +/// input_schema: &Schema, +/// output_schema: &Schema, +/// ) -> DFResult { +/// self.0 +/// .split_projection(source_projection, input_schema, output_schema) +/// } +/// } +/// +/// let _convertor: Arc = Arc::new(CustomExpressionConvertor( +/// DefaultExpressionConvertor::default(), +/// )); +/// ``` pub trait ExpressionConvertor: Send + Sync { /// Can an expression be pushed down given a specific schema fn can_be_pushed_down(&self, expr: &Arc, schema: &Schema) -> bool; @@ -302,7 +345,7 @@ impl ExpressionConvertor for DefaultExpressionConvertor { } if let Some(literal) = df.downcast_ref::() { - let value = Scalar::from_df(literal.value()); + let value = scalar_from_df(literal.value(), &self.session); return Ok(lit(value)); } @@ -333,7 +376,7 @@ impl ExpressionConvertor for DefaultExpressionConvertor { .iter() .map(|e| { if let Some(lit) = e.downcast_ref::() { - Ok(Scalar::from_df(lit.value())) + Ok(scalar_from_df(lit.value(), &self.session)) } else { Err(exec_datafusion_err!("Failed to cast sub-expression")) } @@ -487,7 +530,7 @@ fn try_operator_from_df(value: &DFOperator) -> DFResult { fn can_be_pushed_down_impl(expr: &Arc, schema: &Schema) -> bool { // We currently do not support pushdown of dynamic expressions in DF. // See issue: https://github.com/vortex-data/vortex/issues/4034 - if is_dynamic_physical_expr(expr) { + if DynamicFilterTracking::classify(expr).contains_dynamic_filter() { return false; } @@ -496,8 +539,7 @@ fn can_be_pushed_down_impl(expr: &Arc, schema: &Schema) -> boo } else if let Some(col) = expr.downcast_ref::() { schema .field_with_name(col.name()) - .ok() - .is_some_and(|field| supported_data_types(field.data_type())) + .is_ok_and(|field| supported_data_types(field.data_type())) } else if let Some(like) = expr.downcast_ref::() { can_be_pushed_down_impl(like.expr(), schema) && can_be_pushed_down_impl(like.pattern(), schema) @@ -1244,7 +1286,6 @@ mod tests { use vortex::array::Canonical; use vortex::array::VortexSessionExecute as _; use vortex::session::VortexSession; - use vortex_arrow::FromArrowArray; // Create test data let values = Arc::new(Int32Array::from(vec![1, 5, 10, 15, 20])); @@ -1291,10 +1332,13 @@ mod tests { let vortex_expr = expr_convertor.try_convert_case_expr(&case_expr).unwrap(); // Convert batch to Vortex array - let vortex_array: ArrayRef = ArrayRef::from_arrow(&batch, false).unwrap(); + let session = VortexSession::default(); + let vortex_array: ArrayRef = session + .arrow() + .from_arrow_record_batch(batch.clone(), &batch.schema()) + .unwrap(); // Apply Vortex expression - let session = VortexSession::default(); let mut ctx = session.create_execution_ctx(); let vortex_result = vortex_array .apply(&vortex_expr) diff --git a/vortex-datafusion/src/convert/mod.rs b/vortex-datafusion/src/convert/mod.rs index 6a6fb8aef08..476f32afa49 100644 --- a/vortex-datafusion/src/convert/mod.rs +++ b/vortex-datafusion/src/convert/mod.rs @@ -18,6 +18,8 @@ pub(crate) mod stats; pub use exprs::DefaultExpressionConvertor; pub use exprs::ExpressionConvertor; +pub use exprs::ProcessedProjection; +pub use scalars::scalar_from_df; /// First-party trait for implementing conversion from DataFusion types to Vortex types. pub trait FromDataFusion: Sized { diff --git a/vortex-datafusion/src/convert/scalars.rs b/vortex-datafusion/src/convert/scalars.rs index ec262b479d9..c2b1dfc79b2 100644 --- a/vortex-datafusion/src/convert/scalars.rs +++ b/vortex-datafusion/src/convert/scalars.rs @@ -4,10 +4,12 @@ use std::sync::Arc; use arrow_array::Array; +use arrow_array::ArrayRef as ArrowArrayRef; use arrow_array::StructArray; use arrow_schema::Field; use arrow_schema::Fields; use datafusion_common::ScalarValue; +use vortex::array::VortexSessionExecute; use vortex::buffer::ByteBuffer; use vortex::dtype::DType; use vortex::dtype::DecimalDType; @@ -26,9 +28,9 @@ use vortex::extension::datetime::TemporalMetadata; use vortex::extension::datetime::TimeUnit; use vortex::scalar::DecimalValue; use vortex::scalar::Scalar; -use vortex_arrow::FromArrowType; +use vortex::session::VortexSession; +use vortex_arrow::ArrowSessionExt; -use crate::convert::FromDataFusion; use crate::convert::TryToDataFusion; impl TryToDataFusion for Scalar { @@ -111,13 +113,13 @@ impl TryToDataFusion for Scalar { } // SAFETY: By construction Utf8 scalar values are utf8 DType::Utf8(_) => ScalarValue::Utf8(self.as_utf8().value().cloned().map(|s| unsafe { - String::from_utf8_unchecked(Vec::::from(s.into_inner().into_inner())) + String::from_utf8_unchecked(Vec::::from(s.into_inner().into_bytes())) })), DType::Binary(_) => ScalarValue::Binary( self.as_binary() .value() .cloned() - .map(|b| Vec::::from(b.into_inner())), + .map(|b| Vec::::from(b.into_bytes())), ), dtype @ DType::List(..) => vortex_bail!( "cannot convert Vortex scalar dtype {dtype} to DataFusion ScalarValue: unsupported scalar type" @@ -125,6 +127,9 @@ impl TryToDataFusion for Scalar { dtype @ DType::FixedSizeList(..) => vortex_bail!( "cannot convert Vortex scalar dtype {dtype} to DataFusion ScalarValue: unsupported scalar type" ), + dtype @ DType::Map(..) => vortex_bail!( + "cannot convert Vortex scalar dtype {dtype} to DataFusion ScalarValue: unsupported scalar type" + ), DType::Struct(..) => struct_to_df(self)?, dtype @ DType::Union(..) => vortex_bail!( "cannot convert Vortex scalar dtype {dtype} to DataFusion ScalarValue: unsupported scalar type" @@ -180,131 +185,136 @@ impl TryToDataFusion for Scalar { } } -impl FromDataFusion for Scalar { - fn from_df(value: &ScalarValue) -> Scalar { - match value { - ScalarValue::Null => Scalar::null(DType::Null), - ScalarValue::Boolean(b) => b - .map(Scalar::from) - .unwrap_or_else(|| Scalar::null(DType::Bool(Nullability::Nullable))), - ScalarValue::Float16(f) => f.map(Scalar::from).unwrap_or_else(|| { - Scalar::null(DType::Primitive(PType::F16, Nullability::Nullable)) - }), - ScalarValue::Float32(f) => f.map(Scalar::from).unwrap_or_else(|| { - Scalar::null(DType::Primitive(PType::F32, Nullability::Nullable)) - }), - ScalarValue::Float64(f) => f.map(Scalar::from).unwrap_or_else(|| { - Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable)) - }), - ScalarValue::Int8(i) => i.map(Scalar::from).unwrap_or_else(|| { - Scalar::null(DType::Primitive(PType::I8, Nullability::Nullable)) - }), - ScalarValue::Int16(i) => i.map(Scalar::from).unwrap_or_else(|| { - Scalar::null(DType::Primitive(PType::I16, Nullability::Nullable)) - }), - ScalarValue::Int32(i) => i.map(Scalar::from).unwrap_or_else(|| { - Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)) - }), - ScalarValue::Int64(i) => i.map(Scalar::from).unwrap_or_else(|| { - Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable)) - }), - ScalarValue::UInt8(i) => i.map(Scalar::from).unwrap_or_else(|| { - Scalar::null(DType::Primitive(PType::U8, Nullability::Nullable)) - }), - ScalarValue::UInt16(i) => i.map(Scalar::from).unwrap_or_else(|| { - Scalar::null(DType::Primitive(PType::U16, Nullability::Nullable)) - }), - ScalarValue::UInt32(i) => i.map(Scalar::from).unwrap_or_else(|| { - Scalar::null(DType::Primitive(PType::U32, Nullability::Nullable)) - }), - ScalarValue::UInt64(i) => i.map(Scalar::from).unwrap_or_else(|| { - Scalar::null(DType::Primitive(PType::U64, Nullability::Nullable)) - }), - ScalarValue::Utf8(s) | ScalarValue::Utf8View(s) | ScalarValue::LargeUtf8(s) => s - .as_ref() - .map(|s| Scalar::from(s.as_str())) - .unwrap_or_else(|| Scalar::null(DType::Utf8(Nullability::Nullable))), - ScalarValue::Binary(b) - | ScalarValue::BinaryView(b) - | ScalarValue::LargeBinary(b) - | ScalarValue::FixedSizeBinary(_, b) => b - .as_ref() - .map(|b| Scalar::binary(ByteBuffer::from(b.clone()), Nullability::Nullable)) - .unwrap_or_else(|| Scalar::null(DType::Binary(Nullability::Nullable))), - ScalarValue::Date32(v) - | ScalarValue::Time32Second(v) - | ScalarValue::Time32Millisecond(v) => { - let dtype = DType::from_arrow((&value.data_type(), Nullability::Nullable)); - Scalar::try_new(dtype, v.map(vortex::scalar::ScalarValue::from)) - .vortex_expect("unable to create a time `Scalar`") - } - ScalarValue::Date64(v) - | ScalarValue::Time64Microsecond(v) - | ScalarValue::Time64Nanosecond(v) - | ScalarValue::TimestampSecond(v, _) - | ScalarValue::TimestampMillisecond(v, _) - | ScalarValue::TimestampMicrosecond(v, _) - | ScalarValue::TimestampNanosecond(v, _) => { - let dtype = DType::from_arrow((&value.data_type(), Nullability::Nullable)); - Scalar::try_new(dtype, v.map(vortex::scalar::ScalarValue::from)) - .vortex_expect("unable to create a time `Scalar`") - } - ScalarValue::Decimal32(decimal, precision, scale) => { - let decimal_dtype = DecimalDType::new(*precision, *scale); - let nullable = Nullability::Nullable; - if let Some(value) = decimal { - Scalar::decimal( - DecimalValue::I32(*value), - decimal_dtype, - Nullability::Nullable, - ) - } else { - Scalar::null(DType::Decimal(decimal_dtype, nullable)) - } +/// Converts a DataFusion [`ScalarValue`] to a Vortex [`Scalar`], resolving Arrow types through +/// `session`'s [`ArrowSession`](vortex_arrow::ArrowSession). +pub fn scalar_from_df(value: &ScalarValue, session: &VortexSession) -> Scalar { + let arrow = session.arrow(); + match value { + ScalarValue::Null => Scalar::null(DType::Null), + ScalarValue::Boolean(b) => b + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Bool(Nullability::Nullable))), + ScalarValue::Float16(f) => f + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::F16, Nullability::Nullable))), + ScalarValue::Float32(f) => f + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::F32, Nullability::Nullable))), + ScalarValue::Float64(f) => f + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable))), + ScalarValue::Int8(i) => i + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::I8, Nullability::Nullable))), + ScalarValue::Int16(i) => i + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::I16, Nullability::Nullable))), + ScalarValue::Int32(i) => i + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable))), + ScalarValue::Int64(i) => i + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable))), + ScalarValue::UInt8(i) => i + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::U8, Nullability::Nullable))), + ScalarValue::UInt16(i) => i + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::U16, Nullability::Nullable))), + ScalarValue::UInt32(i) => i + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::U32, Nullability::Nullable))), + ScalarValue::UInt64(i) => i + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::U64, Nullability::Nullable))), + ScalarValue::Utf8(s) | ScalarValue::Utf8View(s) | ScalarValue::LargeUtf8(s) => s + .as_ref() + .map(|s| Scalar::from(s.as_str())) + .unwrap_or_else(|| Scalar::null(DType::Utf8(Nullability::Nullable))), + ScalarValue::Binary(b) + | ScalarValue::BinaryView(b) + | ScalarValue::LargeBinary(b) + | ScalarValue::FixedSizeBinary(_, b) => b + .as_ref() + .map(|b| Scalar::binary(ByteBuffer::from(b.clone()), Nullability::Nullable)) + .unwrap_or_else(|| Scalar::null(DType::Binary(Nullability::Nullable))), + ScalarValue::Date32(v) + | ScalarValue::Time32Second(v) + | ScalarValue::Time32Millisecond(v) => { + let dtype = arrow + .from_arrow_datatype(&value.data_type(), Nullability::Nullable) + .vortex_expect("arrow data type to dtype"); + Scalar::try_new(dtype, v.map(vortex::scalar::ScalarValue::from)) + .vortex_expect("unable to create a time `Scalar`") + } + ScalarValue::Date64(v) + | ScalarValue::Time64Microsecond(v) + | ScalarValue::Time64Nanosecond(v) + | ScalarValue::TimestampSecond(v, _) + | ScalarValue::TimestampMillisecond(v, _) + | ScalarValue::TimestampMicrosecond(v, _) + | ScalarValue::TimestampNanosecond(v, _) => { + let dtype = arrow + .from_arrow_datatype(&value.data_type(), Nullability::Nullable) + .vortex_expect("arrow data type to dtype"); + Scalar::try_new(dtype, v.map(vortex::scalar::ScalarValue::from)) + .vortex_expect("unable to create a time `Scalar`") + } + ScalarValue::Decimal32(decimal, precision, scale) => { + let decimal_dtype = DecimalDType::new(*precision, *scale); + let nullable = Nullability::Nullable; + if let Some(value) = decimal { + Scalar::decimal( + DecimalValue::I32(*value), + decimal_dtype, + Nullability::Nullable, + ) + } else { + Scalar::null(DType::Decimal(decimal_dtype, nullable)) } - ScalarValue::Decimal64(decimal, precision, scale) => { - let decimal_dtype = DecimalDType::new(*precision, *scale); - let nullable = Nullability::Nullable; - if let Some(value) = decimal { - Scalar::decimal( - DecimalValue::I64(*value), - decimal_dtype, - Nullability::Nullable, - ) - } else { - Scalar::null(DType::Decimal(decimal_dtype, nullable)) - } + } + ScalarValue::Decimal64(decimal, precision, scale) => { + let decimal_dtype = DecimalDType::new(*precision, *scale); + let nullable = Nullability::Nullable; + if let Some(value) = decimal { + Scalar::decimal( + DecimalValue::I64(*value), + decimal_dtype, + Nullability::Nullable, + ) + } else { + Scalar::null(DType::Decimal(decimal_dtype, nullable)) } - ScalarValue::Decimal128(decimal, precision, scale) => { - let decimal_dtype = DecimalDType::new(*precision, *scale); - let nullable = Nullability::Nullable; - if let Some(value) = decimal { - Scalar::decimal( - DecimalValue::I128(*value), - decimal_dtype, - Nullability::Nullable, - ) - } else { - Scalar::null(DType::Decimal(decimal_dtype, nullable)) - } + } + ScalarValue::Decimal128(decimal, precision, scale) => { + let decimal_dtype = DecimalDType::new(*precision, *scale); + let nullable = Nullability::Nullable; + if let Some(value) = decimal { + Scalar::decimal( + DecimalValue::I128(*value), + decimal_dtype, + Nullability::Nullable, + ) + } else { + Scalar::null(DType::Decimal(decimal_dtype, nullable)) } - ScalarValue::Decimal256(decimal, precision, scale) => { - let decimal_dtype = DecimalDType::new(*precision, *scale); - let nullable = Nullability::Nullable; - if let Some(value) = decimal { - Scalar::decimal( - DecimalValue::I256(i256::from_le_bytes(value.to_le_bytes())), - decimal_dtype, - Nullability::Nullable, - ) - } else { - Scalar::null(DType::Decimal(decimal_dtype, nullable)) - } + } + ScalarValue::Decimal256(decimal, precision, scale) => { + let decimal_dtype = DecimalDType::new(*precision, *scale); + let nullable = Nullability::Nullable; + if let Some(value) = decimal { + Scalar::decimal( + DecimalValue::I256(i256::from_le_bytes(value.to_le_bytes())), + decimal_dtype, + Nullability::Nullable, + ) + } else { + Scalar::null(DType::Decimal(decimal_dtype, nullable)) } - ScalarValue::Dictionary(_, v) => Scalar::from_df(v.as_ref()), - ScalarValue::Struct(array) => struct_from_df(array), - _ => unimplemented!("Can't convert {value:?} value to a Vortex scalar"), } + ScalarValue::Dictionary(_, v) => scalar_from_df(v.as_ref(), session), + ScalarValue::Struct(array) => struct_from_df(array, session), + _ => unimplemented!("Can't convert {value:?} value to a Vortex scalar"), } } @@ -350,20 +360,31 @@ fn struct_to_df(scalar: &Scalar) -> VortexResult { } /// Converts a DataFusion `ScalarValue::Struct` (a one-row struct array) to a Vortex struct scalar. -fn struct_from_df(array: &StructArray) -> Scalar { - let dtype = DType::from_arrow((array.data_type(), Nullability::Nullable)); +/// +/// The struct dtype comes from the Arrow `Fields`, so the children must be converted from those +/// same fields. Going through `ScalarValue` instead would drop each field's `ARROW:extension:name` +/// (and its declared nullability), yielding storage-typed children that +/// [`Scalar::struct_`] rejects against an extension-typed struct dtype. +fn struct_from_df(array: &StructArray, session: &VortexSession) -> Scalar { + let arrow = session.arrow(); + let dtype = arrow + .from_arrow_datatype(array.data_type(), Nullability::Nullable) + .vortex_expect("arrow data type to dtype"); if array.is_null(0) { Scalar::null(dtype) } else { + let mut ctx = session.create_execution_ctx(); let children = array .columns() .iter() - .map(|column| { - Scalar::from_df( - &ScalarValue::try_from_array(column.as_ref(), 0).unwrap_or_else(|e| { + .zip(array.fields().iter()) + .map(|(column, field)| { + arrow + .from_arrow_array(ArrowArrayRef::clone(column), field) + .and_then(|column| column.execute_scalar(0, &mut ctx)) + .unwrap_or_else(|e| { vortex_panic!("cannot convert struct field to a Vortex scalar: {e}") - }), - ) + }) }) .collect::>(); Scalar::struct_(dtype, children) @@ -375,6 +396,7 @@ mod tests { use datafusion_common::ScalarValue; use datafusion_common::arrow::datatypes::i256 as arrow_i256; use rstest::rstest; + use vortex::VortexSessionDefault; use vortex::buffer::ByteBuffer; use vortex::dtype::DType; use vortex::dtype::DecimalDType; @@ -389,6 +411,11 @@ mod tests { use super::*; + /// Test shim: convert with a default `VortexSession` passed explicitly. + fn from_df(value: &ScalarValue) -> Scalar { + scalar_from_df(value, &VortexSession::default()) + } + #[rstest] #[case::u8_some(Scalar::from(42u8), ScalarValue::UInt8(Some(42)))] #[case::u8_null( @@ -574,7 +601,7 @@ mod tests { #[case] df_scalar: ScalarValue, #[case] expected_vortex: Scalar, ) { - let result = Scalar::from_df(&df_scalar); + let result = from_df(&df_scalar); assert_eq!(result.dtype(), expected_vortex.dtype()); assert_eq!(result.is_null(), expected_vortex.is_null()); @@ -592,7 +619,7 @@ mod tests { #[case::decimal256_some(ScalarValue::Decimal256(Some(arrow_i256::from_i128(12345)), 50, 10))] #[case::decimal256_null(ScalarValue::Decimal256(None, 50, 10))] fn test_from_datafusion_decimals(#[case] df_scalar: ScalarValue) { - let result = Scalar::from_df(&df_scalar); + let result = from_df(&df_scalar); match &df_scalar { ScalarValue::Decimal128(value, precision, scale) => { if let DType::Decimal(decimal_type, _) = result.dtype() { @@ -639,7 +666,7 @@ mod tests { None ))] fn test_from_datafusion_temporals(#[case] df_scalar: ScalarValue) { - let result = Scalar::from_df(&df_scalar); + let result = from_df(&df_scalar); // All temporal types should convert to extension types if let DType::Extension(_) = result.dtype() { @@ -668,7 +695,7 @@ mod tests { #[case::binary(Scalar::binary(ByteBuffer::from(vec![1u8, 2, 3, 4, 5]), Nullability::NonNullable))] fn test_round_trip_conversions(#[case] original: Scalar) { let df_scalar = original.try_to_df().unwrap(); - let round_trip = Scalar::from_df(&df_scalar); + let round_trip = from_df(&df_scalar); // Check that core types match (ignoring nullability differences that can occur in round-trip) assert!( @@ -737,7 +764,7 @@ mod tests { assert_eq!(df_result, expected_df_null); // Test DataFusion -> Vortex - let vortex_result = Scalar::from_df(&expected_df_null); + let vortex_result = from_df(&expected_df_null); assert!(vortex_result.is_null()); assert!( vortex_result @@ -751,7 +778,7 @@ mod tests { #[case::utf8_view(ScalarValue::Utf8View(Some("test string".to_string())))] #[case::large_utf8(ScalarValue::LargeUtf8(Some("test string".to_string())))] fn test_utf8_variants(#[case] variant: ScalarValue) { - let result = Scalar::from_df(&variant); + let result = from_df(&variant); assert_eq!(result.as_utf8().value().unwrap().as_str(), "test string"); } @@ -761,17 +788,48 @@ mod tests { #[case::large_binary(ScalarValue::LargeBinary(Some(vec![1u8, 2, 3, 4, 5])))] #[case::fixed_size_binary(ScalarValue::FixedSizeBinary(5, Some(vec![1u8, 2, 3, 4, 5])))] fn test_binary_variants(#[case] variant: ScalarValue) { - let result = Scalar::from_df(&variant); + let result = from_df(&variant); let result_bytes: Vec = result .as_binary() .value() .cloned() .unwrap() - .into_inner() + .into_bytes() .into(); assert_eq!(result_bytes, vec![1u8, 2, 3, 4, 5]); } + /// A DataFusion struct whose child field carries Arrow extension metadata: the struct dtype + /// derived from the Arrow type keeps the extension, so the child scalars must keep it too or + /// `Scalar::struct_` rejects them. + #[test] + fn struct_from_df_preserves_extension_child() -> VortexResult<()> { + use arrow_array::FixedSizeBinaryArray; + use arrow_schema::DataType; + use arrow_schema::extension::Uuid as ArrowUuid; + use vortex::extension::uuid::Uuid; + + let mut id_field = Field::new("id", DataType::FixedSizeBinary(16), false); + id_field.try_with_extension_type(ArrowUuid)?; + let ids = FixedSizeBinaryArray::try_from_iter([*b"0123456789abcdef"].into_iter())?; + let struct_array = StructArray::try_new( + Fields::from(vec![Arc::new(id_field)]), + vec![Arc::new(ids) as ArrowArrayRef], + None, + )?; + + let scalar = from_df(&ScalarValue::Struct(Arc::new(struct_array))); + let DType::Struct(fields, _) = scalar.dtype() else { + panic!("expected a struct dtype, got {}", scalar.dtype()); + }; + let id_dtype = fields.field_by_index(0).vortex_expect("one field"); + assert!( + id_dtype.as_extension().is::(), + "expected a Uuid extension field, got {id_dtype}" + ); + Ok(()) + } + #[test] fn struct_scalar_round_trips() -> VortexResult<()> { let dtype = DType::Struct( @@ -793,7 +851,7 @@ mod tests { assert!(matches!(df, ScalarValue::Struct(_))); // Back through `from_df` and out again yields the identical DataFusion struct value. - let back = Scalar::from_df(&df); + let back = from_df(&df); assert_eq!(back.try_to_df()?, df); Ok(()) } @@ -813,7 +871,7 @@ mod tests { let df = Scalar::null(dtype).try_to_df()?; assert!(matches!(df, ScalarValue::Struct(_))); - assert!(Scalar::from_df(&df).is_null()); + assert!(from_df(&df).is_null()); Ok(()) } diff --git a/vortex-datafusion/src/convert/schema.rs b/vortex-datafusion/src/convert/schema.rs index 2ec17758593..36f3b5f4cd6 100644 --- a/vortex-datafusion/src/convert/schema.rs +++ b/vortex-datafusion/src/convert/schema.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; + use arrow_schema::DataType; use arrow_schema::Field; use arrow_schema::Schema; @@ -45,7 +47,7 @@ pub fn calculate_physical_schema( arrow_session, )?; Ok( - Field::new(name.to_string(), arrow_type, field_dtype.is_nullable()) + Field::new(name.as_ref(), arrow_type, field_dtype.is_nullable()) .with_metadata(logical_field.metadata().clone()), ) } @@ -56,7 +58,10 @@ pub fn calculate_physical_schema( }) .collect::>>()?; - Ok(Schema::new(fields)) + Ok(Schema::new_with_metadata( + fields, + reference_logical_schema.metadata().clone(), + )) } /// Calculate the physical Arrow type for a field, preferring the logical type when the @@ -105,12 +110,14 @@ fn calculate_physical_field_type( logical_field.data_type(), arrow_session, )?; - Ok(Field::new( - name.to_string(), - arrow_type, - field_dtype.is_nullable(), + Ok( + Field::new( + name.as_ref(), + arrow_type, + field_dtype.is_nullable(), + ) + .with_metadata(logical_field.metadata().clone()), ) - .with_metadata(logical_field.metadata().clone())) } None => arrow_session .to_arrow_field(name.as_ref(), &field_dtype) @@ -175,6 +182,52 @@ fn calculate_physical_field_type( } } + // Map field names and child metadata come from the reference schema, while child + // types are recursively reconciled so nested extension metadata is preserved. + DataType::Map(logical_entries, keys_sorted) => { + let DType::Map(map_dtype, _) = dtype else { + return Err(exec_datafusion_err!( + "Failed to convert dtype to arrow: Vortex DType is {dtype} which is not compatible with {logical_type}" + )); + }; + let DataType::Struct(logical_fields) = logical_entries.data_type() else { + return Err(exec_datafusion_err!( + "Failed to convert dtype to arrow: Arrow Map entries must be a Struct, got {:?}", + logical_entries.data_type() + )); + }; + if logical_fields.len() != 2 { + return Err(exec_datafusion_err!( + "Failed to convert dtype to arrow: Arrow Map entries must contain exactly two fields" + )); + } + + let key = Field::new( + logical_fields[0].name(), + calculate_physical_field_type( + &map_dtype.key_dtype(), + logical_fields[0].data_type(), + arrow_session, + )?, + false, + ) + .with_metadata(logical_fields[0].metadata().clone()); + let value = Field::new( + logical_fields[1].name(), + calculate_physical_field_type( + &map_dtype.value_dtype(), + logical_fields[1].data_type(), + arrow_session, + )?, + logical_fields[1].is_nullable(), + ) + .with_metadata(logical_fields[1].metadata().clone()); + let entries = Field::new_struct(logical_entries.name(), vec![key, value], false) + .with_metadata(logical_entries.metadata().clone()); + + DataType::Map(Arc::new(entries), *keys_sorted) + } + // For list view types, recursively check the element type DataType::ListView(logical_elem) | DataType::LargeListView(logical_elem) => { if let DType::List(elem_dtype, _) = dtype { @@ -246,6 +299,32 @@ mod tests { ); } + #[test] + fn test_schema_metadata_preserved() -> DFResult<()> { + let logical_schema = Schema::new_with_metadata( + vec![Field::new("col", DataType::Int32, false)], + [("table".to_string(), "metadata".to_string())] + .into_iter() + .collect(), + ); + let dtype = DType::Struct( + StructFields::from_iter([( + "col", + DType::Primitive(PType::I32, Nullability::NonNullable), + )]), + Nullability::NonNullable, + ); + + let physical_schema = + calculate_physical_schema(&dtype, &logical_schema, &ArrowSession::default())?; + + assert_eq!( + physical_schema.metadata().get("table"), + Some(&"metadata".to_string()) + ); + Ok(()) + } + #[test] fn test_utf8_variants_preserved() { // Non-view string types become view types after roundtrip through DType, @@ -409,6 +488,65 @@ mod tests { } } + #[test] + fn test_map_schema_conversion_preserves_reference_fields() { + let key = Field::new("custom_key", DataType::Int32, false) + .with_metadata([("key_metadata".to_owned(), "key_value".to_owned())].into()); + let value = Field::new("custom_value", DataType::Utf8, true) + .with_metadata([("value_metadata".to_owned(), "value_value".to_owned())].into()); + let entries = Field::new_struct("custom_entries", vec![key, value], false) + .with_metadata([("entries_metadata".to_owned(), "entries_value".to_owned())].into()); + let logical_schema = Schema::new(vec![Field::new( + "map_col", + DataType::Map(Arc::new(entries), true), + true, + )]); + let dtype = DType::Struct( + StructFields::from_iter([( + "map_col", + DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + Nullability::Nullable, + ) + .unwrap(), + )]), + Nullability::NonNullable, + ); + + let physical_schema = + calculate_physical_schema(&dtype, &logical_schema, &ArrowSession::default()).unwrap(); + let field = physical_schema.field(0); + assert!(field.is_nullable()); + let DataType::Map(entries, keys_sorted) = field.data_type() else { + panic!("expected Map type, got {:?}", field.data_type()); + }; + assert!(*keys_sorted); + assert_eq!(entries.name(), "custom_entries"); + assert_eq!( + entries.metadata().get("entries_metadata"), + Some(&"entries_value".to_owned()) + ); + let DataType::Struct(fields) = entries.data_type() else { + panic!("expected map entries struct, got {:?}", entries.data_type()); + }; + assert_eq!(fields[0].name(), "custom_key"); + assert_eq!(fields[0].data_type(), &DataType::Int32); + assert!(!fields[0].is_nullable()); + assert_eq!( + fields[0].metadata().get("key_metadata"), + Some(&"key_value".to_owned()) + ); + assert_eq!(fields[1].name(), "custom_value"); + assert_eq!(fields[1].data_type(), &DataType::Utf8); + assert!(fields[1].is_nullable()); + assert_eq!( + fields[1].metadata().get("value_metadata"), + Some(&"value_value".to_owned()) + ); + } + #[test] fn test_non_struct_dtype_error() { // Test that non-struct DType produces an error @@ -425,4 +563,87 @@ mod tests { .contains("Expected struct dtype") ); } + + /// Names carrying raw control bytes must reach Arrow byte-for-byte. + /// `FieldName`'s `Display` escapes via `escape_debug`, so building the + /// field with `to_string` renames `\x08` to the literal five characters + /// `\u{8}`. The scanned batch then disagrees with the table schema the + /// scan was planned against, and the query fails with "column types must + /// match schema types". + #[test] + fn test_control_byte_column_names_are_not_escaped() { + let column_names = ["plain", "\u{8}", "check_id\u{10}"]; + let logical_schema = Schema::new( + column_names + .iter() + .map(|n| Field::new(*n, DataType::Utf8, true)) + .collect::(), + ); + + let dtype = DType::Struct( + StructFields::from_iter( + column_names + .iter() + .map(|n| (*n, DType::Utf8(Nullability::Nullable))), + ), + Nullability::NonNullable, + ); + + let physical_schema = + calculate_physical_schema(&dtype, &logical_schema, &ArrowSession::default()).unwrap(); + + let names: Vec<&str> = physical_schema + .fields() + .iter() + .map(|f| f.name().as_str()) + .collect(); + assert_eq!(names, column_names); + } + + /// Same, one level down: struct children are built at a separate call + /// site in `calculate_physical_field_type`. The children here are + /// run-end-encoded dictionaries, so this also covers the branches that + /// take the type from the reference schema instead of the `DType`. + #[test] + fn test_control_byte_struct_field_names_are_not_escaped() { + let label_names = ["app", "\u{8}", "check_id\u{10}"]; + let ree = DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int32, false)), + Arc::new(Field::new( + "values", + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), + true, + )), + ); + let logical_schema = Schema::new(vec![Field::new_struct( + "labels", + label_names + .iter() + .map(|n| Field::new(*n, ree.clone(), true)) + .collect::(), + false, + )]); + + let labels_dtype = DType::Struct( + StructFields::from_iter( + label_names + .iter() + .map(|n| (*n, DType::Utf8(Nullability::Nullable))), + ), + Nullability::NonNullable, + ); + let dtype = DType::Struct( + StructFields::from_iter([("labels", labels_dtype)]), + Nullability::NonNullable, + ); + + let physical_schema = + calculate_physical_schema(&dtype, &logical_schema, &ArrowSession::default()).unwrap(); + + let DataType::Struct(labels) = physical_schema.field(0).data_type() else { + panic!("expected labels to be a struct"); + }; + let names: Vec<&str> = labels.iter().map(|f| f.name().as_str()).collect(); + assert_eq!(names, label_names); + } } diff --git a/vortex-datafusion/src/lib.rs b/vortex-datafusion/src/lib.rs index 65bb021d258..fe20d191e5d 100644 --- a/vortex-datafusion/src/lib.rs +++ b/vortex-datafusion/src/lib.rs @@ -142,12 +142,11 @@ mod common_tests { use object_store::memory::InMemory; use url::Url; use vortex::VortexSessionDefault; - use vortex::array::ArrayRef; use vortex::file::WriteOptionsSessionExt; use vortex::io::VortexWrite; use vortex::io::object_store::ObjectStoreWrite; use vortex::session::VortexSession; - use vortex_arrow::FromArrowArray; + use vortex_arrow::ArrowSessionExt; use crate::VortexFormatFactory; use crate::VortexTableOptions; @@ -168,12 +167,18 @@ mod common_tests { impl TestSessionContext { /// Create a new test session context with the given projection pushdown setting. pub fn new(projection_pushdown: bool) -> Self { - let store = Arc::new(InMemory::new()); let opts = VortexTableOptions { projection_pushdown, ..Default::default() }; let factory = Arc::new(VortexFormatFactory::new().with_options(opts)); + + Self::new_with_factory(factory) + } + + /// Create a new test session context with the given Vortex format factory. + pub fn new_with_factory(factory: Arc) -> Self { + let store = Arc::new(InMemory::new()); let mut session_state_builder = SessionStateBuilder::new() .with_default_features() .with_table_factory( @@ -200,7 +205,10 @@ mod common_tests { where P: Into, { - let array = ArrayRef::from_arrow(batch, false)?; + let schema = batch.schema(); + let array = VX_SESSION + .arrow() + .from_arrow_record_batch(batch.clone(), &schema)?; let mut write = ObjectStoreWrite::new(Arc::clone(&self.store), &path.into()).await?; VX_SESSION .write_options() diff --git a/vortex-datafusion/src/persistent/cache.rs b/vortex-datafusion/src/persistent/cache.rs index a76574031b6..21a32dfda7c 100644 --- a/vortex-datafusion/src/persistent/cache.rs +++ b/vortex-datafusion/src/persistent/cache.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use datafusion_common::HashMap; use datafusion_execution::cache::cache_manager::FileMetadata; use vortex::file::Footer; use vortex::file::VortexFile; @@ -40,8 +41,7 @@ impl FileMetadata for CachedVortexMetadata { .unwrap_or(1024 * 64) } - #[expect(clippy::disallowed_types)] - fn extra_info(&self) -> std::collections::HashMap { + fn extra_info(&self) -> HashMap { Default::default() } } diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs index d936f27d366..8d375ceef0a 100644 --- a/vortex-datafusion/src/persistent/format.rs +++ b/vortex-datafusion/src/persistent/format.rs @@ -5,6 +5,7 @@ use std::fmt::Debug; use std::fmt::Formatter; use std::sync::Arc; +use arrow_schema::DataType; use arrow_schema::Schema; use arrow_schema::SchemaRef; use async_trait::async_trait; @@ -70,6 +71,7 @@ use super::cache::CachedVortexMetadata; use super::sink::VortexSink; use super::source::VortexSource; use crate::PrecisionExt as _; +use crate::convert::ExpressionConvertor; use crate::convert::TryToDataFusion; use crate::convert::stats::is_constant_to_distinct_count; @@ -109,7 +111,7 @@ const DEFAULT_TARGET_FILE_SIZE_MB: usize = 128; /// let table_url = ListingTableUrl::parse(dir.path().to_str().unwrap())?; /// let config = ListingTableConfig::new(table_url) /// .with_listing_options( -/// ListingOptions::new(format).with_session_config_options(ctx.state().config()), +/// ListingOptions::new(format), /// ) /// .infer_schema(&ctx.state()) /// .await?; @@ -126,12 +128,17 @@ const DEFAULT_TARGET_FILE_SIZE_MB: usize = 128; pub struct VortexFormat { session: VortexSession, opts: VortexTableOptions, + expression_convertor: Option>, } impl Debug for VortexFormat { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("VortexFormat") .field("opts", &self.opts) + .field( + "has_expression_convertor", + &self.expression_convertor.is_some(), + ) .finish() } } @@ -279,10 +286,23 @@ impl ConfigExtension for VortexTableOptions { /// ``` /// /// [`ListingTable`]: https://docs.rs/datafusion/latest/datafusion/datasource/listing/struct.ListingTable.html -#[derive(Debug)] pub struct VortexFormatFactory { session: VortexSession, options: Option, + expression_convertor: Option>, +} + +impl Debug for VortexFormatFactory { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("VortexFormatFactory") + .field("session", &self.session) + .field("options", &self.options) + .field( + "has_expression_convertor", + &self.expression_convertor.is_some(), + ) + .finish() + } } impl GetExt for VortexFormatFactory { @@ -306,6 +326,20 @@ impl VortexFormatFactory { Self { session: VortexSession::default(), options: None, + expression_convertor: None, + } + } + + /// Creates a factory with an explicit session and session-driven table options. + /// + /// Formats created by this factory start from the DataFusion session's `vortex` options, + /// falling back to [`VortexTableOptions::default`]. Table-level `OPTIONS(...)` are still + /// applied last. + pub fn new_with_session(session: VortexSession) -> Self { + Self { + session, + options: None, + expression_convertor: None, } } @@ -320,6 +354,7 @@ impl VortexFormatFactory { Self { session, options: Some(options), + expression_convertor: None, } } @@ -346,6 +381,15 @@ impl VortexFormatFactory { self.options = Some(options); self } + + /// Sets the [`ExpressionConvertor`] used by formats and sources created by this factory. + pub fn with_expression_convertor( + mut self, + expression_convertor: Arc, + ) -> Self { + self.expression_convertor = Some(expression_convertor); + self + } } impl FileFormatFactory for VortexFormatFactory { @@ -383,14 +427,19 @@ impl FileFormatFactory for VortexFormatFactory { } } - Ok(Arc::new(VortexFormat::new_with_options( - self.session.clone(), - opts, - ))) + let mut format = VortexFormat::new_with_options(self.session.clone(), opts); + if let Some(expression_convertor) = &self.expression_convertor { + format = format.with_expression_convertor(Arc::clone(expression_convertor)); + } + Ok(Arc::new(format)) } fn default(&self) -> Arc { - Arc::new(VortexFormat::new(self.session.clone())) + let mut format = VortexFormat::new(self.session.clone()); + if let Some(expression_convertor) = &self.expression_convertor { + format = format.with_expression_convertor(Arc::clone(expression_convertor)); + } + Arc::new(format) } } @@ -407,7 +456,11 @@ impl VortexFormat { /// Creates a format with explicit [`VortexTableOptions`]. pub fn new_with_options(session: VortexSession, opts: VortexTableOptions) -> Self { - Self { session, opts } + Self { + session, + opts, + expression_convertor: None, + } } /// Returns the format-specific configuration that will be copied into the @@ -415,6 +468,15 @@ impl VortexFormat { pub fn options(&self) -> &VortexTableOptions { &self.opts } + + /// Sets the [`ExpressionConvertor`] used by every [`VortexSource`] created by this format. + pub fn with_expression_convertor( + mut self, + expression_convertor: Arc, + ) -> Self { + self.expression_convertor = Some(expression_convertor); + self + } } #[async_trait] @@ -494,7 +556,13 @@ impl FileFormat for VortexFormat { }) .map(|f| f.vortex_expect("Failed to spawn infer_schema")) }) - .buffer_unordered(state.config_options().execution.meta_fetch_concurrency) + .buffer_unordered( + state + .config_options() + .execution + .meta_fetch_concurrency + .get(), + ) .try_collect::>() .await .map_err(|e| DataFusionError::Execution(format!("Failed to infer schema: {e}")))?; @@ -613,36 +681,26 @@ impl FileFormat for VortexFormat { let column_size = stats_set.get_as::(Stat::UncompressedSizeInBytes, &PType::U64.into()); - let target_dtype = - session - .arrow() - .from_arrow_field(field.as_ref()) - .map_err(|e| { - DataFusionError::Execution(format!( - "Failed to derive Vortex DType for field {}: {e}", - field.name() - )) - })?; let min = scalar_stat_to_df( Stat::Min, stats_set.get(Stat::Min), stats_dtype, - &target_dtype, + field.data_type(), ); let max = scalar_stat_to_df( Stat::Max, stats_set.get(Stat::Max), stats_dtype, - &target_dtype, + field.data_type(), ); let null_count = stats_set.get_as::(Stat::NullCount, &PType::U64.into()); column_statistics.push(ColumnStatistics { null_count: null_count.to_df(), - min_value: min.to_df(), - max_value: max.to_df(), + min_value: min, + max_value: max, sum_value: DFPrecision::Absent, distinct_count: is_constant_to_distinct_count( stats_set.get_as::( @@ -704,18 +762,6 @@ impl FileFormat for VortexFormat { return not_impl_err!("Overwrites are not implemented yet for Vortex"); } - let target_file_size = (self.opts.target_file_size_mb > 0) - .then(|| { - u64::try_from(self.opts.target_file_size_mb) - .map_err(|e| { - internal_datafusion_err!( - "target_file_size_mb cannot be represented as u64: {e}" - ) - }) - .map(|v| v.saturating_mul(1024 * 1024).max(1)) - }) - .transpose()?; - // For non-partitioned writes, force a single input stream so VortexSink // performs one coordinated write per statement instead of one // independent write per CPU/input partition. @@ -734,20 +780,18 @@ impl FileFormat for VortexFormat { }; let schema = Arc::clone(conf.output_schema()); - let sink = Arc::new(VortexSink::new( - conf, - schema, - self.session.clone(), - target_file_size, - )); + let sink = Arc::new(VortexSink::new(conf, schema, self.session.clone())); Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _) } fn file_source(&self, table_schema: TableSchema) -> Arc { - Arc::new( - VortexSource::new(table_schema, self.session.clone()).with_options(self.opts.clone()), - ) as _ + let mut source = + VortexSource::new(table_schema, self.session.clone()).with_options(self.opts.clone()); + if let Some(expression_convertor) = &self.expression_convertor { + source = source.with_expression_convertor(Arc::clone(expression_convertor)); + } + Arc::new(source) as _ } } @@ -755,27 +799,226 @@ fn scalar_stat_to_df( stat: Stat, value: Precision, stats_dtype: &DType, - target_dtype: &DType, -) -> Precision { + target_dtype: &DataType, +) -> DFPrecision { let Some(stat_dtype) = stat.dtype(stats_dtype) else { - return Precision::Absent; + return DFPrecision::Absent; }; value - .map(|stat_value| { - Scalar::try_new(stat_dtype, Some(stat_value))? - .cast(target_dtype)? + .and_then(|stat_value| { + let scalar = Scalar::try_new(stat_dtype, Some(stat_value)) + .ok()? .try_to_df() + .ok()?; + scalar.cast_to(target_dtype).ok() }) - .transpose() - .unwrap_or(Precision::Absent) + .to_df() } #[cfg(test)] mod tests { + use std::sync::atomic::AtomicBool; + use std::sync::atomic::Ordering; + + use arrow_array::Int32Array; + use arrow_schema::DataType; + use arrow_schema::Field; + use arrow_schema::TimeUnit; + use datafusion_common::ScalarValue; + use datafusion_common::config::ConfigOptions; + use datafusion_expr::Operator; + use datafusion_physical_expr::PhysicalExpr; + use datafusion_physical_expr::expressions as df_expr; + use datafusion_physical_expr::projection::ProjectionExprs; + use datafusion_physical_plan::filter_pushdown::PushedDown; + use rstest::rstest; + use vortex::expr::Expression; use super::*; use crate::common_tests::TestSessionContext; + use crate::convert::DefaultExpressionConvertor; + use crate::convert::ProcessedProjection; + use crate::convert::scalar_from_df; + + #[rstest] + #[case::timestamp_upscale( + ScalarValue::TimestampMillisecond(Some(1_234), None), + ScalarValue::TimestampMicrosecond(Some(1_234_000), None) + )] + #[case::timestamp_downscale( + ScalarValue::TimestampNanosecond(Some(-1_234_567), None), + ScalarValue::TimestampMicrosecond(Some(-1_234), None) + )] + #[case::timestamp_timezone( + ScalarValue::TimestampSecond(Some(123), Some("UTC".into())), + ScalarValue::TimestampMillisecond(Some(123_000), Some("UTC".into())) + )] + #[case::date_upscale(ScalarValue::Date32(Some(2)), ScalarValue::Date64(Some(172_800_000)))] + #[case::date_downscale(ScalarValue::Date64(Some(-172_800_000)), ScalarValue::Date32(Some(-2)))] + #[case::time_upscale( + ScalarValue::Time32Second(Some(123)), + ScalarValue::Time64Nanosecond(Some(123_000_000_000)) + )] + #[case::time_downscale( + ScalarValue::Time64Microsecond(Some(1_234_567)), + ScalarValue::Time32Millisecond(Some(1_234)) + )] + #[case::decimal_scale( + ScalarValue::Decimal32(Some(123), 5, 2), + ScalarValue::Decimal32(Some(1_230), 6, 3) + )] + #[case::integer_widening(ScalarValue::Int32(Some(123)), ScalarValue::Int64(Some(123)))] + #[case::string_view(ScalarValue::Utf8(Some("value".into())), ScalarValue::Utf8View(Some("value".into())))] + #[case::decimal_width( + ScalarValue::Decimal32(Some(123), 5, 2), + ScalarValue::Decimal128(Some(123), 5, 2) + )] + fn test_scalar_stat_to_df_cast( + #[case] value: ScalarValue, + #[case] expected: ScalarValue, + #[values(Stat::Min, Stat::Max)] stat: Stat, + #[values(true, false)] exact: bool, + ) -> VortexResult<()> { + let session = VortexSession::default(); + let scalar = scalar_from_df(&value, &session); + let value = scalar + .value() + .cloned() + .ok_or_else(|| vortex_err!("expected non-null scalar"))?; + let target_dtype = expected.data_type(); + let (value, expected) = if exact { + (Precision::Exact(value), DFPrecision::Exact(expected)) + } else { + (Precision::Inexact(value), DFPrecision::Inexact(expected)) + }; + + assert_eq!( + scalar_stat_to_df(stat, value, scalar.dtype(), &target_dtype), + expected + ); + Ok(()) + } + + #[rstest] + #[case::timestamp_overflow( + ScalarValue::TimestampSecond(Some(10_000_000_000), None), + DataType::Timestamp(TimeUnit::Nanosecond, None) + )] + #[case::timestamp_underflow( + ScalarValue::TimestampSecond(Some(-10_000_000_000), None), + DataType::Timestamp(TimeUnit::Nanosecond, None) + )] + #[case::integer_overflow(ScalarValue::Int64(Some(i64::MAX)), DataType::Int32)] + #[case::unsupported(ScalarValue::Boolean(Some(true)), DataType::Date32)] + fn test_scalar_stat_to_df_failed_cast( + #[case] value: ScalarValue, + #[case] target_dtype: DataType, + #[values(Stat::Min, Stat::Max)] stat: Stat, + ) -> VortexResult<()> { + let session = VortexSession::default(); + let scalar = scalar_from_df(&value, &session); + let value = scalar + .value() + .cloned() + .ok_or_else(|| vortex_err!("expected non-null scalar"))?; + assert_eq!( + scalar_stat_to_df(stat, Precision::Exact(value), scalar.dtype(), &target_dtype), + DFPrecision::Absent + ); + Ok(()) + } + + #[derive(Clone, Copy)] + enum PushdownMode { + Reject, + Delegate, + } + + #[derive(Default)] + struct ExpressionConvertorCalls { + can_be_pushed_down: AtomicBool, + convert: AtomicBool, + } + + impl ExpressionConvertorCalls { + fn reset(&self) { + self.can_be_pushed_down.store(false, Ordering::Relaxed); + self.convert.store(false, Ordering::Relaxed); + } + } + + struct TestExpressionConvertor { + inner: DefaultExpressionConvertor, + pushdown_mode: PushdownMode, + calls: Arc, + } + + impl TestExpressionConvertor { + fn new( + session: VortexSession, + pushdown_mode: PushdownMode, + calls: Arc, + ) -> Self { + Self { + inner: DefaultExpressionConvertor::new(session), + pushdown_mode, + calls, + } + } + } + + impl ExpressionConvertor for TestExpressionConvertor { + fn can_be_pushed_down(&self, expr: &Arc, schema: &Schema) -> bool { + self.calls.can_be_pushed_down.store(true, Ordering::Relaxed); + match self.pushdown_mode { + PushdownMode::Reject => false, + PushdownMode::Delegate => self.inner.can_be_pushed_down(expr, schema), + } + } + + fn convert(&self, expr: &dyn PhysicalExpr) -> DFResult { + self.calls.convert.store(true, Ordering::Relaxed); + self.inner.convert(expr) + } + + fn split_projection( + &self, + source_projection: ProjectionExprs, + input_schema: &Schema, + output_schema: &Schema, + ) -> DFResult { + self.inner + .split_projection(source_projection, input_schema, output_schema) + } + } + + fn expression_convertor_test_schema() -> Arc { + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])) + } + + fn expression_convertor_test_filter() -> Arc { + let column = Arc::new(df_expr::Column::new("a", 0)) as Arc; + let literal = + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(1)))) as Arc; + Arc::new(df_expr::BinaryExpr::new(column, Operator::Gt, literal)) + } + + fn assert_rejects_pushdown_with_expression_convertor( + format: &dyn FileFormat, + calls: &ExpressionConvertorCalls, + ) -> anyhow::Result<()> { + let source = format.file_source(TableSchema::from(expression_convertor_test_schema())); + let result = source.try_pushdown_filters( + vec![expression_convertor_test_filter()], + &ConfigOptions::new(), + )?; + + assert!(calls.can_be_pushed_down.load(Ordering::Relaxed)); + assert!(!calls.convert.load(Ordering::Relaxed)); + assert!(matches!(result.filters.as_slice(), [PushedDown::No])); + Ok(()) + } #[tokio::test] async fn create_table() -> anyhow::Result<()> { @@ -847,7 +1090,7 @@ mod tests { ..Default::default() }; let format = VortexFormat::new_with_options(VortexSession::default(), opts.clone()); - let table_schema = TableSchema::from_file_schema(Arc::new(Schema::empty())); + let table_schema = TableSchema::from(Arc::new(Schema::empty())); let source = format.file_source(table_schema); let source = source @@ -862,4 +1105,83 @@ mod tests { assert_eq!(source.options().scan_concurrency, opts.scan_concurrency); Ok(()) } + + #[test] + fn format_plumbs_expression_convertor() -> anyhow::Result<()> { + let session = VortexSession::default(); + let calls = Arc::new(ExpressionConvertorCalls::default()); + let convertor = Arc::new(TestExpressionConvertor::new( + session.clone(), + PushdownMode::Reject, + Arc::clone(&calls), + )); + let format = VortexFormat::new(session).with_expression_convertor(convertor); + + assert_rejects_pushdown_with_expression_convertor(&format, &calls) + } + + #[test] + fn factory_plumbs_expression_convertor() -> anyhow::Result<()> { + let calls = Arc::new(ExpressionConvertorCalls::default()); + let convertor = Arc::new(TestExpressionConvertor::new( + VortexSession::default(), + PushdownMode::Reject, + Arc::clone(&calls), + )); + let factory = VortexFormatFactory::new().with_expression_convertor(convertor); + let ctx = TestSessionContext::default(); + + let format = factory.create(&ctx.session.state(), &Default::default())?; + assert_rejects_pushdown_with_expression_convertor(format.as_ref(), &calls)?; + + calls.reset(); + let format = FileFormatFactory::default(&factory); + assert_rejects_pushdown_with_expression_convertor(format.as_ref(), &calls) + } + + #[tokio::test] + async fn external_table_query_uses_factory_expression_convertor() -> anyhow::Result<()> { + let calls = Arc::new(ExpressionConvertorCalls::default()); + let convertor = Arc::new(TestExpressionConvertor::new( + VortexSession::default(), + PushdownMode::Delegate, + Arc::clone(&calls), + )); + let factory = Arc::new(VortexFormatFactory::new().with_expression_convertor(convertor)); + let ctx = TestSessionContext::new_with_factory(factory); + + ctx.session + .sql( + "CREATE EXTERNAL TABLE numbers (a INT NOT NULL) \ + STORED AS vortex LOCATION '/expression-convertor/'", + ) + .await?; + ctx.session + .sql("INSERT INTO numbers VALUES (1), (2), (3)") + .await? + .collect() + .await?; + + calls.reset(); + let batches = ctx + .session + .sql("SELECT a FROM numbers WHERE a > 1 ORDER BY a") + .await? + .collect() + .await?; + + assert!(calls.can_be_pushed_down.load(Ordering::Relaxed)); + assert!(calls.convert.load(Ordering::Relaxed)); + let mut values = Vec::new(); + for batch in batches { + let array = batch + .column(0) + .as_any() + .downcast_ref::() + .ok_or_else(|| anyhow::anyhow!("expected Int32 result column"))?; + values.extend(array.values().iter().copied()); + } + assert_eq!(values, vec![2, 3]); + Ok(()) + } } diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 823a47af1ff..b768b9b1f50 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -5,11 +5,13 @@ use std::ops::Range; use std::sync::Arc; use std::sync::Weak; +use arrow_array::RecordBatchOptions; use arrow_schema::Field; use arrow_schema::Schema; use datafusion_common::DataFusionError; use datafusion_common::Result as DFResult; use datafusion_common::ScalarValue; +use datafusion_common::Statistics; use datafusion_common::arrow::array::AsArray; use datafusion_common::arrow::array::RecordBatch; use datafusion_common::exec_datafusion_err; @@ -23,28 +25,27 @@ use datafusion_physical_expr::PhysicalExprRef; use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_expr::simplifier::PhysicalExprSimplifier; use datafusion_physical_expr::split_conjunction; +use datafusion_physical_expr::utils::collect_columns; use datafusion_physical_expr::utils::reassign_expr_columns; use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory; use datafusion_physical_expr_adapter::replace_columns_with_literals; -use datafusion_physical_expr_common::physical_expr::is_dynamic_physical_expr; -use datafusion_physical_plan::metrics::Count; +use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; +use datafusion_physical_plan::metrics::MetricBuilder; +use datafusion_physical_plan::metrics::MetricCategory; use datafusion_pruning::FilePruner; use futures::FutureExt; use futures::StreamExt; use futures::TryStreamExt; use futures::stream; -use itertools::Itertools; use object_store::path::Path; use tracing::Instrument; use vortex::array::VortexSessionExecute; -use vortex::dtype::FieldMask; use vortex::error::VortexError; use vortex::error::VortexExpect; use vortex::file::OpenOptionsSessionExt; use vortex::io::InstrumentedReadAt; use vortex::layout::LayoutReader; use vortex::layout::scan::scan_builder::ScanBuilder; -use vortex::layout::scan::split_by::SplitBy; use vortex::metrics::Label; use vortex::metrics::MetricsRegistry; use vortex::session::VortexSession; @@ -83,24 +84,24 @@ pub(crate) struct VortexOpener { /// This is the table's schema without partition columns. It may contain fields which do /// not exist in the file. Missing columns are null-filled during batch remapping. pub table_schema: TableSchema, - /// A hint for the desired row count of record batches returned from the scan. - pub batch_size: usize, /// If provided, the scan will not return more than this many rows. pub limit: Option, /// A metrics object for tracking performance of the scan. pub metrics_registry: Arc, + /// DataFusion-native metrics exposed through `DataSourceExec`. + pub df_metrics: ExecutionPlanMetricsSet, /// A shared cache of file readers. /// /// To save on the overhead of reparsing FlatBuffers and rebuilding the layout tree, we cache /// a file reader the first time we read a file. pub layout_readers: Arc>>, - /// Shared full-file natural split ranges keyed by file path. - pub natural_split_ranges: Arc]>>>, + /// Shared full-file natural splits keyed by file path. + pub natural_splits: Arc>>, /// Whether the query has output ordering specified pub has_output_ordering: bool, pub expression_convertor: Arc, - pub file_metadata_cache: Option>, + pub file_metadata_cache: Option>, /// Whether to enable expression pushdown into the underlying Vortex scan. pub projection_pushdown: bool, pub scan_concurrency: Option, @@ -108,6 +109,12 @@ pub(crate) struct VortexOpener { impl FileOpener for VortexOpener { fn open(&self, file: PartitionedFile) -> DFResult { + // Calculate the output schema before replacing partition columns with literals so it + // retains the table and partition-field metadata declared by the plan. + let output_schema = Arc::new( + self.projection + .project_schema(self.table_schema.table_schema())?, + ); let session = self.session.clone(); let metrics_registry = Arc::clone(&self.metrics_registry); let labels = vec![ @@ -123,21 +130,24 @@ impl FileOpener for VortexOpener { let reader = InstrumentedReadAt::new_with_labels(reader, metrics_registry.as_ref(), labels.clone()); - let file_pruning_predicate = self.file_pruning_predicate.clone(); + let mut file_pruning_predicate = self.file_pruning_predicate.clone(); let expr_adapter_factory = Arc::clone(&self.expr_adapter_factory); let file_metadata_cache = self.file_metadata_cache.clone(); let unified_file_schema = Arc::clone(self.table_schema.file_schema()); - let batch_size = self.batch_size; let limit = self.limit; let layout_readers = Arc::clone(&self.layout_readers); - let natural_split_ranges = Arc::clone(&self.natural_split_ranges); + let natural_splits = Arc::clone(&self.natural_splits); let has_output_ordering = self.has_output_ordering; let scan_concurrency = self.scan_concurrency; let expr_convertor = Arc::clone(&self.expression_convertor); let projection_pushdown = self.projection_pushdown; + let predicate_creation_errors = MetricBuilder::new(&self.df_metrics) + .with_category(MetricCategory::Rows) + .global_counter("num_predicate_creation_errors"); + // Replace column access for partition columns with literals #[expect(clippy::disallowed_types)] let literal_value_cols = self @@ -149,6 +159,13 @@ impl FileOpener for VortexOpener { .zip(file.partition_values.clone()) .collect::>(); + let predicate_uses_partition_columns = + file_pruning_predicate.as_ref().is_some_and(|predicate| { + collect_columns(predicate) + .iter() + .any(|column| literal_value_cols.contains_key(column.name())) + }); + if !literal_value_cols.is_empty() { projection = projection.try_map_exprs(|expr| { replace_columns_with_literals(Arc::clone(&expr), &literal_value_cols) @@ -156,24 +173,30 @@ impl FileOpener for VortexOpener { filter = filter .map(|p| replace_columns_with_literals(p, &literal_value_cols)) .transpose()?; + file_pruning_predicate = file_pruning_predicate + .map(|p| replace_columns_with_literals(p, &literal_value_cols)) + .transpose()?; } Ok(async move { - // Create FilePruner when we have a predicate and either dynamic expressions - // or file statistics available. The pruner can eliminate files without - // opening them based on File-level statistics (min/max values per column) + // FilePruner requires a statistics object even when the rewritten predicate + // only contains partition literals. Supply unknown file-column statistics in + // that case so static and dynamic partition predicates can still prune. + let synthetic_statistics = (!file.has_statistics() && predicate_uses_partition_columns) + .then(|| { + file.clone() + .with_statistics(Arc::new(Statistics::new_unknown(&unified_file_schema))) + }); + let pruning_file = synthetic_statistics.as_ref().unwrap_or(&file); + let mut file_pruner = file_pruning_predicate - .filter(|p| { - // Only create pruner if we have dynamic expressions or file statistics - // to work with. Static predicates without stats won't benefit from pruning. - is_dynamic_physical_expr(p) || file.has_statistics() - }) + .filter(|_| file.has_statistics() || predicate_uses_partition_columns) .and_then(|predicate| { FilePruner::try_new( Arc::clone(&predicate), &unified_file_schema, - &file, - Count::default(), + pruning_file, + predicate_creation_errors, ) }); @@ -240,8 +263,6 @@ impl FileOpener for VortexOpener { &session.arrow(), )?); - let projected_physical_schema = projection.project_schema(&unified_file_schema)?; - let expr_adapter = expr_adapter_factory.create( Arc::clone(&unified_file_schema), Arc::clone(&this_file_schema), @@ -268,7 +289,7 @@ impl FileOpener for VortexOpener { expr_convertor.split_projection( projection.clone(), &this_file_schema, - &projected_physical_schema, + output_schema.as_ref(), )? } else { // When projection pushdown is disabled, read only the required columns @@ -278,14 +299,18 @@ impl FileOpener for VortexOpener { // The schema of the stream returned from the vortex scan. // We use a reference schema for types that don't roundtrip (Dictionary, Utf8, etc.). - let scan_dtype = scan_projection.return_dtype(vxf.dtype()).map_err(|_e| { - exec_datafusion_err!("Couldn't get the dtype for the underlying Vortex scan") - })?; + let scan_projection = scan_projection + .optimize_recursive(vxf.dtype()) + .and_then(|projection| projection.bind(vxf.dtype())) + .map_err(|_e| { + exec_datafusion_err!("Couldn't get the dtype for the underlying Vortex scan") + })?; + let scan_dtype = scan_projection.dtype().clone(); // When projection pushdown is enabled, the scan outputs the projected columns. // When disabled, the scan outputs raw columns and the projection is applied after. let scan_reference_schema = if projection_pushdown { - projected_physical_schema + (*output_schema).clone() } else { // Build schema from the raw columns being read let column_indices = projection.column_indices(); @@ -293,7 +318,7 @@ impl FileOpener for VortexOpener { .into_iter() .map(|idx| this_file_schema.field(idx).clone()) .collect(); - Schema::new(fields) + Schema::new_with_metadata(fields, this_file_schema.metadata().clone()) }; let stream_schema = calculate_physical_schema(&scan_dtype, &scan_reference_schema, &session.arrow())?; @@ -336,34 +361,6 @@ impl FileOpener for VortexOpener { scan_builder = vortex_plan.apply_to_builder(scan_builder); } - if let Some(file_range) = file.range { - let byte_range = Range { - start: u64::try_from(file_range.start) - .map_err(|_| exec_datafusion_err!("Vortex file range start is negative"))?, - end: u64::try_from(file_range.end) - .map_err(|_| exec_datafusion_err!("Vortex file range end is negative"))?, - }; - if byte_range.start != 0 || byte_range.end != file.object_meta.size { - // Full-file scans already cover every natural split. Only translate the - // byte range back into row boundaries when DataFusion has trimmed the file. - let natural_split_ranges = natural_split_ranges_for_file( - natural_split_ranges.as_ref(), - &file.object_meta.location, - &layout_reader, - )?; - - let Some(row_range) = split_aligned_row_range( - byte_range, - file.object_meta.size, - natural_split_ranges.as_ref(), - ) else { - return Ok(stream::empty().boxed()); - }; - - scan_builder = scan_builder.with_row_range(row_range); - } - } - let filter = filter .and_then(|f| { // Verify that all filters we've accepted from DataFusion get pushed down. @@ -395,6 +392,10 @@ impl FileOpener for VortexOpener { make_vortex_predicate(expr_convertor.as_ref(), &pushed).transpose() }) .transpose()?; + let filter = filter + .map(|filter| filter.optimize_recursive(vxf.dtype())?.bind(vxf.dtype())) + .transpose() + .map_err(|e| exec_datafusion_err!("Couldn't bind Vortex scan filter: {e}"))?; if let Some(limit) = limit && filter.is_none() @@ -406,11 +407,46 @@ impl FileOpener for VortexOpener { scan_builder = scan_builder.with_concurrency(concurrency); } + // Set before the byte-range translation below, which computes natural splits for + // the fields the scan's projection and filter reference. + scan_builder = scan_builder + .with_projection(scan_projection) + .with_some_filter(filter); + + if let Some(file_range) = file.range { + let byte_range = Range { + start: u64::try_from(file_range.start) + .map_err(|_| exec_datafusion_err!("Vortex file range start is negative"))?, + end: u64::try_from(file_range.end) + .map_err(|_| exec_datafusion_err!("Vortex file range end is negative"))?, + }; + if byte_range.start != 0 || byte_range.end != file.object_meta.size { + // Full-file scans already cover every natural split. Only translate the + // byte range back into row boundaries when DataFusion has trimmed the file. + let natural_splits = natural_splits_for_file( + natural_splits.as_ref(), + &file.object_meta.location, + &scan_builder, + file.object_meta.size, + )?; + + let Some(row_range) = + split_aligned_row_range(byte_range, natural_splits.as_ref()) + else { + return Ok(stream::empty().boxed()); + }; + + scan_builder = scan_builder + .with_row_range(row_range) + // Hand the shared full-file boundaries back to the scan so prepare() + // skips its own layout walk. + .with_natural_splits(Arc::clone(&natural_splits.row_boundaries)); + } + } + let stream_target_field = Field::new_struct("", stream_schema.fields().clone(), false); let stream = scan_builder .with_metrics_registry(metrics_registry) - .with_projection(scan_projection) - .with_some_filter(filter) .with_ordered(has_output_ordering) .map(move |chunk| { let mut ctx = session.create_execution_ctx(); @@ -424,44 +460,32 @@ impl FileOpener for VortexOpener { }) .into_stream() .map_err(|e| exec_datafusion_err!("Failed to create Vortex stream: {e}"))? - .map_ok(move |rb| { - // We try and slice the stream into respecting datafusion's configured batch size. - stream::iter( - (0..rb.num_rows().div_ceil(batch_size * 2)) - .flat_map(move |block_idx| { - let offset = block_idx * batch_size * 2; - - // If we have less than two batches worth of rows left, we keep them together as a single batch. - if rb.num_rows() - offset < 2 * batch_size { - let length = rb.num_rows() - offset; - [Some(rb.slice(offset, length)), None].into_iter() - } else { - let first = rb.slice(offset, batch_size); - let second = rb.slice(offset + batch_size, batch_size); - [Some(first), Some(second)].into_iter() - } - }) - .flatten() - .map(Ok), - ) - }) .map_err(move |e: VortexError| { DataFusionError::External(Box::new(e.with_context(format!( "Failed to read Vortex file: {}", file.object_meta.location )))) }) - .try_flatten() .map(move |batch| { - if projector.projection().as_ref().is_empty() { + let batch = if projector.projection().as_ref().is_empty() { batch } else { batch.and_then(|b| projector.project_batch(&b)) - } + }?; + + let (_, columns, row_count) = batch.into_parts(); + RecordBatch::try_new_with_options( + Arc::clone(&output_schema), + columns, + &RecordBatchOptions::new().with_row_count(Some(row_count)), + ) + .map_err(Into::into) }) .boxed(); - if let Some(file_pruner) = file_pruner { + if let Some(file_pruner) = file_pruner + && file_pruner.is_watching() + { Ok(PrunableStream::new(file_pruner, stream).boxed()) } else { Ok(stream) @@ -472,38 +496,98 @@ impl FileOpener for VortexOpener { } } -fn natural_split_ranges_for_file( - natural_split_ranges: &DashMap]>>, - path: &Path, - layout_reader: &Arc, -) -> DFResult]>> { - if let Some(split_ranges) = natural_split_ranges.get(path) { - return Ok(Arc::clone(split_ranges.value())); +/// A file's natural split boundaries plus the precomputed byte each split is assigned to, +/// enabling [`split_aligned_row_range`] to translate a DataFusion byte range into row +/// boundaries with a binary search instead of re-projecting every split per partition. +/// +/// The boundaries are computed for the fields referenced by the scan's projection and filter. +/// All partitions translate through the first opener's cached entry (the cache lives on the +/// source, so projection and filter are fixed for its lifetime), which keeps the byte ranges +/// tiling the file's rows exactly once. +#[derive(Debug)] +pub(crate) struct NaturalSplits { + /// Sorted row boundaries of the natural splits; split `i` covers + /// `row_boundaries[i]..row_boundaries[i + 1]`. Shared so partitions can hand the + /// boundaries back to the scan via [`ScanBuilder::with_natural_splits`], skipping the + /// per-partition layout walk in `prepare`. + row_boundaries: Arc<[u64]>, + /// For each split, the byte a DataFusion byte range must contain to own it (see + /// [`split_assignment_byte`]); one entry per split, sorted because split midpoints + /// increase monotonically under the row-to-byte projection. + assignment_bytes: Box<[u64]>, +} + +impl NaturalSplits { + fn new(row_boundaries: Arc<[u64]>, total_size: u64) -> Self { + let row_count = row_boundaries.last().copied().unwrap_or_default(); + let assignment_bytes = if row_count == 0 { + Box::default() + } else { + row_boundaries + .windows(2) + .enumerate() + .map(|(idx, boundaries)| { + split_assignment_byte( + idx, + &(boundaries[0]..boundaries[1]), + row_count, + total_size, + ) + }) + .collect() + }; + + debug_assert!(assignment_bytes.is_sorted()); + debug_assert_eq!( + assignment_bytes.len() + usize::from(!row_boundaries.is_empty()), + row_boundaries.len() + ); + + Self { + row_boundaries, + assignment_bytes, + } } +} - let split_ranges = compute_natural_split_ranges(layout_reader.as_ref())?; +/// Return the cached [`NaturalSplits`] for `path`, computing and caching them on first use. +fn natural_splits_for_file( + natural_splits: &DashMap>, + path: &Path, + scan_builder: &ScanBuilder, + total_size: u64, +) -> DFResult> { + if let Some(splits) = natural_splits.get(path) { + return Ok(Arc::clone(splits.value())); + } - match natural_split_ranges.entry(path.clone()) { + // Compute while holding the entry so concurrent partitions opening the same file wait + // for the winner instead of all walking the layout tree; the redundant walks contend on + // the lazily-initialized layout children and dominate the cost of the computation itself. + match natural_splits.entry(path.clone()) { Entry::Occupied(entry) => Ok(Arc::clone(entry.get())), Entry::Vacant(entry) => { - entry.insert(Arc::clone(&split_ranges)); - Ok(split_ranges) + let splits = compute_natural_splits(scan_builder, total_size)?; + entry.insert(Arc::clone(&splits)); + Ok(splits) } } } -fn compute_natural_split_ranges(layout_reader: &dyn LayoutReader) -> DFResult]>> { - let row_count = layout_reader.row_count(); - let row_range = 0..row_count; - let split_points: Vec<_> = SplitBy::Layout - .splits(layout_reader, &row_range, &[FieldMask::All]) - .map_err(|e| exec_datafusion_err!("Failed to compute Vortex natural splits: {e}"))? - .into_iter() - .tuple_windows() - .map(|(s, e)| s..e) - .collect::>(); - - Ok(split_points.into()) +/// Walk the layout tree to compute the file's full natural split boundaries for the fields +/// referenced by the scan's projection and filter. +fn compute_natural_splits( + scan_builder: &ScanBuilder, + total_size: u64, +) -> DFResult> { + let row_boundaries = scan_builder + .full_file_splits() + .map_err(|e| exec_datafusion_err!("Failed to compute Vortex natural splits: {e}"))?; + + Ok(Arc::new(NaturalSplits::new( + row_boundaries.into(), + total_size, + ))) } /// Translate a DataFusion byte range to the contiguous natural split ranges it owns. @@ -511,33 +595,25 @@ fn compute_natural_split_ranges(layout_reader: &dyn LayoutReader) -> DFResult, - total_size: u64, - split_ranges: &[Range], + natural_splits: &NaturalSplits, ) -> Option> { if byte_range.start >= byte_range.end { return None; } - let row_count = split_ranges.last().map(|split| split.end)?; - if row_count == 0 { + let first_split = natural_splits + .assignment_bytes + .partition_point(|&assignment_byte| assignment_byte < byte_range.start); + let after_last_split = natural_splits + .assignment_bytes + .partition_point(|&assignment_byte| assignment_byte < byte_range.end); + if first_split == after_last_split { return None; } - let mut owned_splits = split_ranges - .iter() - .enumerate() - .filter_map(|(idx, split_range)| { - let assignment_byte = split_assignment_byte(idx, split_range, row_count, total_size); - byte_range.contains(&assignment_byte).then_some(split_range) - }); - - let first_split = owned_splits.next()?; - let mut row_range = first_split.start..first_split.end; - for split_range in owned_splits { - row_range.end = split_range.end; - } - - Some(row_range) + Some( + natural_splits.row_boundaries[first_split]..natural_splits.row_boundaries[after_last_split], + ) } fn split_assignment_byte( @@ -564,13 +640,16 @@ fn split_midpoint_to_byte(split_range: &Range, row_count: u64, total_size: #[cfg(test)] mod tests { + use std::fmt; use std::sync::Arc; use std::sync::LazyLock; + use arrow_array::record_batch; use arrow_schema::Field; use arrow_schema::Fields; use arrow_schema::SchemaRef; use datafusion::arrow::array::DictionaryArray; + use datafusion::arrow::array::Int32Array; use datafusion::arrow::array::RecordBatch; use datafusion::arrow::array::StringArray; use datafusion::arrow::array::StructArray; @@ -579,15 +658,17 @@ mod tests { use datafusion::arrow::datatypes::UInt32Type; use datafusion::arrow::util::display::FormatOptions; use datafusion::arrow::util::pretty::pretty_format_batches_with_options; - use datafusion::common::record_batch; use datafusion::logical_expr::col; use datafusion::logical_expr::lit; use datafusion::physical_expr::planner::logical2physical; use datafusion::physical_expr_adapter::DefaultPhysicalExprAdapterFactory; use datafusion::scalar::ScalarValue; - use datafusion_execution::cache::DefaultFilesMetadataCache; + use datafusion_common::stats::Precision; + use datafusion_execution::cache::default_cache::DefaultCache; use datafusion_expr::Operator; + use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::expressions as df_expr; + use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr; use datafusion_physical_expr::projection::ProjectionExpr; use insta::assert_snapshot; use itertools::Itertools; @@ -595,15 +676,14 @@ mod tests { use object_store::memory::InMemory; use rstest::rstest; use vortex::VortexSessionDefault; - use vortex::array::ArrayRef; use vortex::buffer::Buffer; use vortex::file::WriteOptionsSessionExt; use vortex::io::VortexWrite; use vortex::io::object_store::ObjectStoreWrite; use vortex::metrics::DefaultMetricsRegistry; use vortex::scan::selection::Selection; + use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; use vortex::session::VortexSession; - use vortex_arrow::FromArrowArray; use super::*; use crate::VortexAccessPlan; @@ -612,12 +692,70 @@ mod tests { static SESSION: LazyLock = LazyLock::new(VortexSession::default); + /// Test-only expr used to test error reporting. + #[derive(Debug, Eq, Hash, PartialEq)] + struct SnapshotErrorExpr; + + impl fmt::Display for SnapshotErrorExpr { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "snapshot_error") + } + } + + impl PhysicalExpr for SnapshotErrorExpr { + fn data_type(&self, _input_schema: &Schema) -> DFResult { + Ok(DataType::Boolean) + } + + fn nullable(&self, _input_schema: &Schema) -> DFResult { + Ok(false) + } + + fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Display::fmt(self, f) + } + + fn evaluate(&self, _batch: &RecordBatch) -> DFResult { + Err(DataFusionError::Internal( + "intentional snapshot error".to_owned(), + )) + } + + fn children(&self) -> Vec<&PhysicalExprRef> { + Vec::new() + } + + fn with_new_children( + self: Arc, + children: Vec, + ) -> DFResult { + assert!(children.is_empty()); + Ok(self) + } + + fn snapshot(&self) -> DFResult> { + Err(DataFusionError::Internal( + "intentional snapshot error".to_owned(), + )) + } + } + + fn natural_splits(total_size: u64, split_ranges: &[Range]) -> NaturalSplits { + let mut row_boundaries = Vec::with_capacity(split_ranges.len() + 1); + if let Some(first) = split_ranges.first() { + row_boundaries.push(first.start); + row_boundaries.extend(split_ranges.iter().map(|range| range.end)); + } + NaturalSplits::new(row_boundaries.into(), total_size) + } + #[rstest] #[case(0..3, 10, vec![0..2, 2..5, 5..10], Some(0..2))] #[case(3..7, 10, vec![0..2, 2..5, 5..10], Some(2..5))] #[case(1..8, 10, vec![0..1, 1..9, 9..10], Some(1..9))] #[case(1..4, 16, vec![0..1, 1..2, 2..3, 3..4], None)] #[case(0..1, 10, vec![0..2, 2..10], Some(0..2))] + #[case(0..2, 2, vec![], None)] fn test_split_aligned_row_range( #[case] byte_range: Range, #[case] total_size: u64, @@ -625,7 +763,7 @@ mod tests { #[case] expected: Option>, ) { assert_eq!( - split_aligned_row_range(byte_range, total_size, &split_ranges), + split_aligned_row_range(byte_range, &natural_splits(total_size, &split_ranges)), expected ); } @@ -634,10 +772,11 @@ mod tests { fn test_split_aligned_ranges_cover_splits_exactly_once() { let split_ranges = vec![0..1, 1..4, 4..10, 10..13]; let byte_ranges = [0..4, 4..8, 8..12, 12..16]; + let natural_splits = natural_splits(16, &split_ranges); let assigned = byte_ranges .into_iter() - .filter_map(|byte_range| split_aligned_row_range(byte_range, 16, &split_ranges)) + .filter_map(|byte_range| split_aligned_row_range(byte_range, &natural_splits)) .collect::>(); assert_eq!(assigned, vec![0..4, 4..10, 10..13]); @@ -668,12 +807,35 @@ mod tests { } } + #[rstest] + #[case(vec![], 10)] + #[case(vec![0], 10)] + #[case(vec![], 0)] + #[case(vec![0], 0)] + fn test_natural_splits_empty_file(#[case] row_boundaries: Vec, #[case] total_size: u64) { + let splits = NaturalSplits::new(row_boundaries.clone().into(), total_size); + + assert!(splits.assignment_bytes.is_empty()); + assert_eq!(splits.row_boundaries.as_ref(), row_boundaries.as_slice()); + assert_eq!(split_aligned_row_range(0..u64::MAX, &splits), None); + } + + #[test] + fn test_split_aligned_row_range_keeps_colliding_assignments_together() { + let natural_splits = natural_splits(2, &[0..1, 1..2, 2..3, 3..4]); + + assert_eq!(natural_splits.assignment_bytes.as_ref(), [0, 0, 1, 1]); + assert_eq!(split_aligned_row_range(0..1, &natural_splits), Some(0..2)); + assert_eq!(split_aligned_row_range(1..2, &natural_splits), Some(2..4)); + } + async fn write_arrow_to_vortex( object_store: Arc, path: &str, rb: RecordBatch, ) -> anyhow::Result { - let array = ArrayRef::from_arrow(rb, false)?; + let schema = rb.schema(); + let array = SESSION.arrow().from_arrow_record_batch(rb, &schema)?; let path = Path::parse(path)?; let mut write = ObjectStoreWrite::new(object_store, &path).await?; @@ -700,11 +862,11 @@ mod tests { file_pruning_predicate: None, expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), table_schema, - batch_size: 100, limit: None, metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), - natural_split_ranges: Default::default(), + natural_splits: Default::default(), has_output_ordering: false, expression_convertor: Arc::new(DefaultExpressionConvertor::default()), file_metadata_cache: None, @@ -725,10 +887,9 @@ mod tests { let mut file = PartitionedFile::new(file_path.to_string(), data_size); file.partition_values = vec![ScalarValue::Int32(Some(1))]; - let table_schema = TableSchema::new( - Arc::clone(&file_schema), - vec![Arc::new(Field::new("part", DataType::Int32, false))], - ); + let table_schema = TableSchema::builder(Arc::clone(&file_schema)) + .with_table_partition_cols(vec![Arc::new(Field::new("part", DataType::Int32, false))]) + .build(); // filter matches partition value let filter = col("part").eq(lit(1)); @@ -766,6 +927,159 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_open_preserves_declared_schema_metadata() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "part=1/file.vortex"; + let batch = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)]))?; + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + + let file_schema = Arc::new( + batch.schema().as_ref().clone().with_metadata( + [("table".to_string(), "metadata".to_string())] + .into_iter() + .collect(), + ), + ); + let table_schema = TableSchema::builder(file_schema) + .with_table_partition_cols(vec![Arc::new( + Field::new("part", DataType::Int32, false).with_metadata( + [("partition".to_string(), "metadata".to_string())] + .into_iter() + .collect(), + ), + )]) + .build(); + let projection = ProjectionExprs::from_indices(&[0, 1], table_schema.table_schema()); + let expected_schema = Arc::new(projection.project_schema(table_schema.table_schema())?); + + assert_eq!( + expected_schema.metadata().get("table"), + Some(&"metadata".to_string()) + ); + assert_eq!( + expected_schema.field(1).metadata().get("partition"), + Some(&"metadata".to_string()) + ); + + for projection_pushdown in [false, true] { + let mut opener = make_opener(Arc::clone(&object_store), table_schema.clone(), None); + opener.projection = projection.clone(); + opener.projection_pushdown = projection_pushdown; + + let mut file = PartitionedFile::new(file_path.to_string(), data_size); + file.partition_values = vec![ScalarValue::Int32(Some(1))]; + let batches = opener.open(file)?.await?.try_collect::>().await?; + + assert!(!batches.is_empty()); + for batch in batches { + assert_eq!(batch.schema().as_ref(), expected_schema.as_ref()); + } + } + + Ok(()) + } + + #[tokio::test] + async fn test_open_all_valid_nullable_columns_with_nonnullable_table_schema() + -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "nullable/file.vortex"; + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])), + vec![Arc::new(Int32Array::from(vec![Some(1), Some(2), Some(3)]))], + )?; + let data_size = write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch).await?; + + let expected_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let table_schema = TableSchema::from(Arc::clone(&expected_schema)); + + for projection_pushdown in [false, true] { + let mut opener = make_opener(Arc::clone(&object_store), table_schema.clone(), None); + opener.projection_pushdown = projection_pushdown; + + let file = PartitionedFile::new(file_path.to_string(), data_size); + let batches = opener.open(file)?.await?.try_collect::>().await?; + + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].schema().as_ref(), expected_schema.as_ref()); + } + + Ok(()) + } + + #[tokio::test] + async fn test_file_pruning_replaces_partition_columns_without_file_statistics() + -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let table_schema = TableSchema::builder(Arc::clone(&file_schema)) + .with_table_partition_cols(vec![Arc::new(Field::new("part", DataType::Int32, false))]) + .build(); + + let partition_column = Arc::new(df_expr::Column::new("part", 1)) as PhysicalExprRef; + let predicate = Arc::new(df_expr::BinaryExpr::new( + Arc::clone(&partition_column), + Operator::Gt, + df_expr::lit(ScalarValue::Int32(Some(1))), + )) as PhysicalExprRef; + let dynamic_predicate = Arc::new(DynamicFilterPhysicalExpr::new( + vec![partition_column], + predicate, + )) as PhysicalExprRef; + + let mut opener = make_opener(object_store, table_schema, None); + opener.file_pruning_predicate = Some(dynamic_predicate); + let df_metrics = opener.df_metrics.clone(); + + // The file does not exist and has no statistics. Replacing `part` with 1 + // makes the predicate false, so pruning must happen before any file I/O. + let mut file = PartitionedFile::new("missing.vortex", 1); + file.partition_values = vec![ScalarValue::Int32(Some(1))]; + let batches = opener.open(file)?.await?.try_collect::>().await?; + + assert!(batches.is_empty()); + assert_eq!( + df_metrics + .clone_inner() + .sum_by_name("num_predicate_creation_errors") + .map(|metric| metric.as_usize()), + Some(0) + ); + + Ok(()) + } + + #[tokio::test] + async fn test_file_pruning_creation_errors_are_reported() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "metrics/file.vortex"; + let batch = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + let mut statistics = Statistics::new_unknown(batch.schema().as_ref()); + statistics.column_statistics[0].null_count = Precision::Exact(0); + let file = PartitionedFile::new(file_path, data_size).with_statistics(Arc::new(statistics)); + + let mut opener = make_opener(object_store, TableSchema::from(batch.schema()), None); + opener.file_pruning_predicate = Some(Arc::new(SnapshotErrorExpr)); + let df_metrics = opener.df_metrics.clone(); + + let batches = opener.open(file)?.await?.try_collect::>().await?; + + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 3); + assert_eq!( + df_metrics + .clone_inner() + .sum_by_name("num_predicate_creation_errors") + .map(|metric| metric.as_usize()), + Some(1) + ); + + Ok(()) + } + #[tokio::test] async fn test_open_empty_file() -> anyhow::Result<()> { use futures::TryStreamExt; @@ -782,7 +1096,7 @@ mod tests { let file = PartitionedFile::new_with_range(file_path.to_string(), file_size, 0, file_size as i64); - let table_schema = TableSchema::from_file_schema(Arc::clone(&file_schema)); + let table_schema = TableSchema::from(Arc::clone(&file_schema)); let opener = make_opener(object_store, table_schema, None); let stream = opener.open(file)?.await?; @@ -802,10 +1116,11 @@ mod tests { write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; let file = PartitionedFile::new(file_path.to_string(), data_size); - let table_schema = TableSchema::from_file_schema(batch.schema()); + let table_schema = TableSchema::from(batch.schema()); - let cache: Arc = - Arc::new(DefaultFilesMetadataCache::new(64 * 1024 * 1024)); + let cache: Arc = Arc::new( + DefaultCache::::new(64 * 1024 * 1024), + ); let mut opener = make_opener(Arc::clone(&object_store), table_schema, None); opener.file_metadata_cache = Some(Arc::clone(&cache)); @@ -855,7 +1170,7 @@ mod tests { }; // Table schema has can accommodate both files - let table_schema = TableSchema::from_file_schema(Arc::new(Schema::new(vec![Field::new( + let table_schema = TableSchema::from(Arc::new(Schema::new(vec![Field::new( "a", DataType::Int32, true, @@ -872,11 +1187,11 @@ mod tests { file_pruning_predicate: None, expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), table_schema: table_schema.clone(), - batch_size: 100, limit: None, metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), - natural_split_ranges: Default::default(), + natural_splits: Default::default(), has_output_ordering: false, expression_convertor: Arc::new(DefaultExpressionConvertor::default()), file_metadata_cache: None, @@ -958,12 +1273,12 @@ mod tests { filter: None, file_pruning_predicate: None, expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), - table_schema: TableSchema::from_file_schema(Arc::clone(&table_schema)), - batch_size: 100, + table_schema: TableSchema::from(Arc::clone(&table_schema)), limit: None, metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), - natural_split_ranges: Default::default(), + natural_splits: Default::default(), has_output_ordering: false, expression_convertor: Arc::new(DefaultExpressionConvertor::default()), file_metadata_cache: None, @@ -1022,7 +1337,7 @@ mod tests { let data_size = write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch).await?; // Table schema has an extra utf8 field. - let table_schema = TableSchema::from_file_schema(Arc::new(Schema::new(vec![Field::new( + let table_schema = TableSchema::from(Arc::new(Schema::new(vec![Field::new( "my_struct", DataType::Struct(Fields::from(vec![ Field::new( @@ -1084,18 +1399,15 @@ mod tests { // Table schema has columns in DIFFERENT order: c, a, b // and different types that require casting (Utf8 -> Dictionary) - let table_schema = TableSchema::new( - Arc::new(Schema::new(vec![ - Field::new("c", DataType::Int32, true), - Field::new("a", DataType::Int32, true), - Field::new( - "b", - DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), - true, - ), - ])), - vec![], - ); + let table_schema = TableSchema::from(Arc::new(Schema::new(vec![ + Field::new("c", DataType::Int32, true), + Field::new("a", DataType::Int32, true), + Field::new( + "b", + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), + true, + ), + ]))); // Project columns [0, 2] from table schema, which should give us: c, b // Before the fix, the schema adapter would get confused about which fields @@ -1116,11 +1428,11 @@ mod tests { file_pruning_predicate: None, expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), table_schema: table_schema.clone(), - batch_size: 100, limit: None, metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), - natural_split_ranges: Default::default(), + natural_splits: Default::default(), has_output_ordering: false, expression_convertor: Arc::new(DefaultExpressionConvertor::default()), file_metadata_cache: None, @@ -1175,12 +1487,12 @@ mod tests { filter: None, file_pruning_predicate: None, expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), - table_schema: TableSchema::from_file_schema(schema), - batch_size: 100, + table_schema: TableSchema::from(schema), limit: None, metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), - natural_split_ranges: Default::default(), + natural_splits: Default::default(), has_output_ordering: false, expression_convertor: Arc::new(DefaultExpressionConvertor::default()), file_metadata_cache: None, @@ -1193,8 +1505,6 @@ mod tests { // Test that Selection::IncludeByIndex filters to specific row indices. async fn test_selection_include_by_index() -> anyhow::Result<()> { use datafusion::arrow::util::pretty::pretty_format_batches_with_options; - use vortex::buffer::Buffer; - use vortex::scan::selection::Selection; let object_store = Arc::new(InMemory::new()) as Arc; let file_path = "/path/file.vortex"; @@ -1208,7 +1518,7 @@ mod tests { file.extensions .insert( VortexAccessPlan::default().with_selection(Selection::IncludeByIndex( - Buffer::from_iter(vec![1, 3, 5, 7]), + StrictSortedBuffer::try_new(Buffer::from_iter(vec![1, 3, 5, 7]))?, )), ); @@ -1252,7 +1562,7 @@ mod tests { file.extensions .insert( VortexAccessPlan::default().with_selection(Selection::ExcludeByIndex( - Buffer::from_iter(vec![0, 2, 4, 6, 8]), + StrictSortedBuffer::try_new(Buffer::from_iter(vec![0, 2, 4, 6, 8]))?, )), ); @@ -1357,7 +1667,7 @@ mod tests { write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; let file_schema = batch.schema(); - let table_schema = TableSchema::from_file_schema(Arc::clone(&file_schema)); + let table_schema = TableSchema::from(Arc::clone(&file_schema)); // Create a projection that includes an arithmetic expression: a + b * 2 let col_a = df_expr::col("a", &file_schema)?; @@ -1385,11 +1695,11 @@ mod tests { file_pruning_predicate: None, expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), table_schema, - batch_size: 100, limit: None, metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), - natural_split_ranges: Default::default(), + natural_splits: Default::default(), has_output_ordering: false, expression_convertor: Arc::new(DefaultExpressionConvertor::default()), file_metadata_cache: None, diff --git a/vortex-datafusion/src/persistent/sink.rs b/vortex-datafusion/src/persistent/sink.rs index d36d57e994b..37074e5fe44 100644 --- a/vortex-datafusion/src/persistent/sink.rs +++ b/vortex-datafusion/src/persistent/sink.rs @@ -3,17 +3,13 @@ use std::sync::Arc; -use arrow_schema::Schema; use arrow_schema::SchemaRef; use async_trait::async_trait; use datafusion_common::DataFusionError; use datafusion_common::Result as DFResult; -use datafusion_common::arrow::array::RecordBatch; -use datafusion_common::arrow::array::RecordBatchOptions; use datafusion_common::exec_datafusion_err; use datafusion_common_runtime::JoinSet; use datafusion_common_runtime::SpawnedTask; -use datafusion_datasource::ListingTableUrl; use datafusion_datasource::display::FileGroupDisplay; use datafusion_datasource::file_sink_config::FileSink; use datafusion_datasource::file_sink_config::FileSinkConfig; @@ -24,148 +20,53 @@ use datafusion_execution::SendableRecordBatchStream; use datafusion_execution::TaskContext; use datafusion_physical_plan::DisplayAs; use datafusion_physical_plan::DisplayFormatType; +use datafusion_physical_plan::metrics::Count; +use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; +use datafusion_physical_plan::metrics::MetricBuilder; +use datafusion_physical_plan::metrics::MetricCategory; use datafusion_physical_plan::metrics::MetricsSet; -use futures::SinkExt; -use futures::Stream; use futures::StreamExt; -use futures::channel::mpsc; use object_store::ObjectStore; -use object_store::ObjectStoreExt; use object_store::path::Path; -use tokio::task::JoinHandle; use tokio_stream::wrappers::ReceiverStream; -use uuid::Uuid; use vortex::array::stream::ArrayStreamAdapter; -use vortex::dtype::DType; use vortex::file::WriteOptionsSessionExt; use vortex::file::WriteSummary; use vortex::io::VortexWrite; use vortex::io::object_store::ObjectStoreWrite; use vortex::session::VortexSession; use vortex_arrow::ArrowSessionExt; -use vortex_utils::aliases::hash_set::HashSet; - -struct WriteOutputOptions<'a> { - base_output_path: &'a ListingTableUrl, - target_file_size: Option, - extension: &'a str, - write_id: &'a str, - partition_column_names: &'a [String], - keep_partition_by_columns: bool, - single_file_output: bool, -} - -#[derive(Clone, Copy)] -struct CompressionEstimate { - prev_compressed_bytes: u64, - prev_uncompressed_bytes: u64, -} - -impl CompressionEstimate { - fn identity() -> Self { - Self { - prev_compressed_bytes: 1, - prev_uncompressed_bytes: 1, - } - } - - fn from_file_sizes(compressed_bytes: u64, uncompressed_bytes: u64) -> DFResult { - if uncompressed_bytes == 0 { - return Err(exec_datafusion_err!( - "Cannot derive compression estimate from zero uncompressed bytes" - )); - } - - Ok(Self { - prev_compressed_bytes: compressed_bytes, - prev_uncompressed_bytes: uncompressed_bytes, - }) - } - - fn estimate_compressed_size(self, uncompressed_bytes: u64) -> DFResult { - if self.prev_uncompressed_bytes == 0 { - return Err(exec_datafusion_err!( - "Compression estimate denominator must be non-zero" - )); - } - - let estimated = u128::from(uncompressed_bytes) - .checked_mul(u128::from(self.prev_compressed_bytes)) - .ok_or_else(|| { - exec_datafusion_err!( - "Compressed size estimate overflow for {} * {}", - uncompressed_bytes, - self.prev_compressed_bytes - ) - })? - / u128::from(self.prev_uncompressed_bytes); - - u64::try_from(estimated).map_err(|_| { - exec_datafusion_err!("Compressed size estimate does not fit in u64: {estimated}") - }) - } -} - -struct ActiveFileWriter { - path: Path, - sender: mpsc::Sender, - /// `None` once the task has been joined, or once it has been taken to be - /// aborted. Joining leaves the handle in place so that a cancelled join - /// still aborts the task. - task: Option>>, -} - -impl Drop for ActiveFileWriter { - /// Aborts a writer that was neither finished nor cleaned up. A bare - /// [`JoinHandle`] keeps running when dropped, so a cancelled write would - /// otherwise let the task complete its multipart upload and publish a file - /// that nothing accounts for. - fn drop(&mut self) { - if let Some(task) = self.task.take() { - task.abort(); - } - } -} /// Implements [`DataSink`] for writing Vortex files. pub struct VortexSink { config: FileSinkConfig, schema: SchemaRef, session: VortexSession, - target_file_size: Option, + metrics: ExecutionPlanMetricsSet, + rows_written: Count, + bytes_written: Count, } impl VortexSink { /// Creates a new [`VortexSink`] instance. - pub fn new( - config: FileSinkConfig, - schema: SchemaRef, - session: VortexSession, - target_file_size: Option, - ) -> Self { + pub fn new(config: FileSinkConfig, schema: SchemaRef, session: VortexSession) -> Self { + let metrics = ExecutionPlanMetricsSet::new(); + let rows_written = MetricBuilder::new(&metrics) + .with_category(MetricCategory::Rows) + .global_counter("rows_written"); + let bytes_written = MetricBuilder::new(&metrics) + .with_category(MetricCategory::Bytes) + .global_counter("bytes_written"); + Self { config, schema, session, - target_file_size, + metrics, + rows_written, + bytes_written, } } - - fn base_output_path(&self) -> DFResult<&ListingTableUrl> { - self.config - .table_paths - .first() - .ok_or_else(|| exec_datafusion_err!("Vortex sink requires at least one table path")) - } - - fn writer_dtype(&self, writer_schema: &SchemaRef) -> DFResult { - self.session - .arrow() - .from_arrow_schema(writer_schema) - .map_err(|e| { - exec_datafusion_err!("Failed to derive Vortex DType from writer schema: {e}") - }) - } } impl std::fmt::Debug for VortexSink { @@ -192,7 +93,7 @@ impl DisplayAs for VortexSink { #[async_trait] impl DataSink for VortexSink { fn metrics(&self) -> Option { - None + Some(self.metrics.clone_inner()) } /// Returns the sink schema @@ -205,46 +106,7 @@ impl DataSink for VortexSink { data: SendableRecordBatchStream, context: &Arc, ) -> DFResult { - if !self.config.table_partition_cols.is_empty() { - return FileSink::write_all(self, data, context).await; - } - - let object_store = context - .runtime_env() - .object_store(&self.config.object_store_url)?; - let writer_schema = get_writer_schema(&self.config); - let dtype = self.writer_dtype(&writer_schema)?; - let write_id = Uuid::now_v7().simple().to_string(); - let base_output_path = self.base_output_path()?; - let partition_column_names = self - .config - .table_partition_cols - .iter() - .map(|(name, _)| name.clone()) - .collect::>(); - - let summaries = write_record_batch_stream_to_files( - self.session.clone(), - object_store, - dtype, - data, - Arc::clone(&writer_schema), - &WriteOutputOptions { - base_output_path, - target_file_size: self.target_file_size, - extension: &self.config.file_extension, - write_id: &write_id, - partition_column_names: &partition_column_names, - keep_partition_by_columns: self.config.keep_partition_by_columns, - single_file_output: self - .config - .file_output_mode - .single_file_output(base_output_path), - }, - ) - .await?; - - aggregate_write_summaries(summaries) + FileSink::write_all(self, data, context).await } } @@ -261,50 +123,70 @@ impl FileSink for VortexSink { mut file_stream_rx: DemuxedStreamReceiver, object_store: Arc, ) -> DFResult { + let mut file_write_tasks: JoinSet> = JoinSet::new(); let writer_schema = get_writer_schema(&self.config); - let dtype = self.writer_dtype(&writer_schema)?; - let mut file_write_tasks: JoinSet>> = JoinSet::new(); + let dtype = self + .session + .arrow() + .from_arrow_schema(&writer_schema) + .map_err(|e| { + exec_datafusion_err!("Failed to derive Vortex DType from writer schema: {e}") + })?; + // TODO(adamg): + // 1. We can probably be better at signaling how much memory we're consuming (potentially when reading too), see ParquetSink::spawn_writer_tasks_and_join. while let Some((path, rx)) = file_stream_rx.recv().await { let session = self.session.clone(); let object_store = Arc::clone(&object_store); - let dtype = dtype.clone(); - let import_schema = Arc::clone(&writer_schema); - let target_file_size = self.target_file_size; - let extension = self.config.file_extension.clone(); + // We need to spawn work because there's a dependency between the different files. If one file has too many batches buffered, + // the demux task might deadlock itself. + let arrow_session = session.clone(); + let import_schema = Arc::clone(&writer_schema); + let dtype = dtype.clone(); file_write_tasks.spawn(async move { - let stream = ReceiverStream::new(rx).map(Ok); - write_record_batch_stream_to_paths( - session, - object_store, - dtype, - stream, - import_schema, - target_file_size, - |file_index| { - if file_index == 0 { - path.clone() - } else { - numbered_path(&path, file_index, &extension) - } - }, - ) - .await + let stream = ReceiverStream::new(rx).map(move |rb| { + arrow_session + .arrow() + .from_arrow_record_batch(rb, &import_schema) + }); + + let stream_adapter = ArrayStreamAdapter::new(dtype, stream); + + let mut object_writer = ObjectStoreWrite::new(object_store, &path) + .await + .map_err(|e| exec_datafusion_err!("Failed to create ObjectStoreWrite: {e}"))?; + + let summary = session + .write_options() + .write(&mut object_writer, stream_adapter) + .await + .map_err(|e| exec_datafusion_err!("Failed to write Vortex file: {e}"))?; + + object_writer + .shutdown() + .await + .map_err(|e| exec_datafusion_err!("Failed to shutdown Vortex writer: {e}"))?; + + Ok((path, summary)) }); } - let mut row_count = 0_u64; + let mut row_count = 0; + while let Some(result) = file_write_tasks.join_next().await { match result { - Ok(summaries) => { - row_count = row_count - .checked_add(aggregate_write_summaries(summaries?)?) - .ok_or_else(|| { - exec_datafusion_err!( - "Row count overflow while aggregating file writer tasks" - ) - })?; + Ok(r) => { + let (path, summary) = r?; + + let rows = summary.row_count(); + row_count += rows; + self.rows_written + .add(usize::try_from(rows).unwrap_or(usize::MAX)); + self.bytes_written + .add(usize::try_from(summary.size()).unwrap_or(usize::MAX)); + + tracing::info!(path = %path, "Successfully written file"); } Err(e) => { if e.is_panic() { @@ -325,359 +207,6 @@ impl FileSink for VortexSink { } } -fn aggregate_write_summaries(summaries: Vec<(Path, WriteSummary)>) -> DFResult { - let mut row_count = 0_u64; - for (path, summary) in summaries { - row_count = row_count.checked_add(summary.row_count()).ok_or_else(|| { - exec_datafusion_err!( - "Row count overflow while aggregating sink summaries (current={}, file={})", - row_count, - summary.row_count() - ) - })?; - tracing::debug!(path = %path, "Successfully written file"); - } - - Ok(row_count) -} - -/// Write batches from a single input stream to one or more output files. -/// -/// For collection paths, files are emitted using the `{write_id}_{file_index:05}.{extension}` -/// naming scheme as produced by the underlying writer implementation. -/// For a single-file path, the original target path is used unless rolling is needed, -/// in which case additional files follow the same naming scheme. -async fn write_record_batch_stream_to_files( - session: VortexSession, - object_store: Arc, - dtype: DType, - data: SendableRecordBatchStream, - import_schema: SchemaRef, - output_options: &WriteOutputOptions<'_>, -) -> DFResult> { - let stream = data.map(|batch| { - let batch = batch?; - if output_options.keep_partition_by_columns - || output_options.partition_column_names.is_empty() - { - Ok(batch) - } else { - remove_partition_columns(&batch, output_options.partition_column_names) - } - }); - - write_record_batch_stream_to_paths( - session, - object_store, - dtype, - stream, - import_schema, - output_options.target_file_size, - |file_index| { - output_file_path( - output_options.base_output_path, - file_index, - output_options.extension, - output_options.single_file_output, - output_options.write_id, - ) - }, - ) - .await -} - -async fn write_record_batch_stream_to_paths( - session: VortexSession, - object_store: Arc, - dtype: DType, - data: S, - import_schema: SchemaRef, - target_file_size: Option, - mut output_path: F, -) -> DFResult> -where - S: Stream>, - F: FnMut(usize) -> Path, -{ - let target = target_file_size.map(|t| t.max(1)); - let mut results: Vec<(Path, WriteSummary)> = Vec::new(); - let mut active_writer: Option = None; - // Every path a writer was opened for. A file that fails while being - // finalized is already out of `active_writer` and never reaches `results`, - // so cleanup has to work from this rather than from either of those. - let mut attempted_paths: Vec = Vec::new(); - let mut uncompressed_bytes_in_file = 0_u64; - let mut file_index = 0_usize; - let mut compression_estimate = CompressionEstimate::identity(); - - futures::pin_mut!(data); - - let write_result: DFResult<()> = async { - while let Some(batch) = data.next().await.transpose()? { - if active_writer.is_none() { - let path = output_path(file_index); - attempted_paths.push(path.clone()); - active_writer = Some(start_file_writer( - &session, - Arc::clone(&object_store), - path, - dtype.clone(), - Arc::clone(&import_schema), - )); - } - - let batch_bytes = batch_uncompressed_bytes(&batch)?; - let writer = active_writer.as_mut().ok_or_else(|| { - exec_datafusion_err!("Missing active file writer for sink output") - })?; - send_batch_to_active_writer(writer, batch).await?; - uncompressed_bytes_in_file = uncompressed_bytes_in_file - .checked_add(batch_bytes) - .ok_or_else(|| { - exec_datafusion_err!( - "Uncompressed byte counter overflow for output file {}", - writer.path - ) - })?; - - if let Some(target) = target { - let estimated_compressed = - compression_estimate.estimate_compressed_size(uncompressed_bytes_in_file)?; - if estimated_compressed >= target { - let writer = active_writer.take().ok_or_else(|| { - exec_datafusion_err!( - "Missing active file writer while finalizing rotated output file" - ) - })?; - let file_path = writer.path.clone(); - let summary = finish_file_writer(writer).await?; - if uncompressed_bytes_in_file > 0 { - compression_estimate = CompressionEstimate::from_file_sizes( - summary.size(), - uncompressed_bytes_in_file, - )?; - } - - results.push((file_path, summary)); - uncompressed_bytes_in_file = 0; - file_index += 1; - } - } - } - - if let Some(writer) = active_writer.take() { - let file_path = writer.path.clone(); - let summary = finish_file_writer(writer).await?; - results.push((file_path, summary)); - } - - Ok(()) - } - .await; - - if let Err(err) = write_result { - cleanup_failed_write( - object_store, - active_writer.into_iter().collect(), - attempted_paths, - ) - .await; - return Err(err); - } - - Ok(results) -} - -/// Generate a numbered file path from an existing path for size-based splitting. -/// -/// Given `base/file.vortex`, produces `base/file_00000.vortex`. -/// If the path has no recognized extension, appends `_00000.{extension}`. -fn numbered_path(original: &Path, index: usize, extension: &str) -> Path { - let s = original.to_string(); - let suffix = format!(".{extension}"); - if let Some(stem) = s.strip_suffix(&suffix) { - Path::from(format!("{stem}_{index:05}{suffix}")) - } else { - Path::from(format!("{s}_{index:05}.{extension}")) - } -} - -fn start_file_writer( - session: &VortexSession, - object_store: Arc, - path: Path, - dtype: DType, - import_schema: SchemaRef, -) -> ActiveFileWriter { - let (sender, receiver) = mpsc::channel::(1); - let session = session.clone(); - let path_for_task = path.clone(); - - let task = tokio::spawn(async move { - let mut object_writer = ObjectStoreWrite::new(object_store, &path_for_task) - .await - .map_err(|e| { - exec_datafusion_err!( - "Failed to create ObjectStoreWrite for '{}': {e}", - path_for_task - ) - })?; - - let arrow_session = session.clone(); - let stream = receiver.map(move |rb| { - arrow_session - .arrow() - .from_arrow_record_batch(rb, &import_schema) - }); - let stream_adapter = ArrayStreamAdapter::new(dtype, stream); - - let summary = session - .write_options() - .write(&mut object_writer, stream_adapter) - .await - .map_err(|e| { - exec_datafusion_err!("Failed to write Vortex file '{}': {e}", path_for_task) - })?; - - object_writer.shutdown().await.map_err(|e| { - exec_datafusion_err!("Failed to shutdown Vortex writer '{}': {e}", path_for_task) - })?; - - Ok(summary) - }); - - ActiveFileWriter { - path, - sender, - task: Some(task), - } -} - -async fn cleanup_failed_write( - object_store: Arc, - active_writers: Vec, - // Every path a writer was opened for, not just the ones that finished: a - // file that failed while being finalized is in neither `active_writers` - // nor the caller's results. - attempted_paths: Vec, -) { - let mut cleanup_paths = HashSet::new(); - - for mut writer in active_writers { - cleanup_paths.insert(writer.path.clone()); - if let Some(task) = writer.task.take() { - task.abort(); - drop(task.await); - } - } - - for path in attempted_paths { - cleanup_paths.insert(path); - } - - for path in cleanup_paths { - if let Err(e) = object_store.delete(&path).await { - tracing::warn!(path = %path, error = %e, "Failed to delete sink output during error cleanup"); - } - } -} - -async fn send_batch_to_active_writer( - writer: &mut ActiveFileWriter, - batch: RecordBatch, -) -> DFResult<()> { - writer.sender.send(batch).await.map_err(|e| { - exec_datafusion_err!( - "Failed to send batch to active writer for '{}': {e}", - writer.path - ) - }) -} - -async fn finish_file_writer(mut writer: ActiveFileWriter) -> DFResult { - writer.sender.close_channel(); - // Join through the handle still owned by `writer` so that cancelling this - // future drops `writer` and aborts the task. Taking the handle out first - // would leave a bare JoinHandle, which detaches instead of aborting, and - // this is the window in which the file is published. - let joined = { - let task = writer.task.as_mut().ok_or_else(|| { - exec_datafusion_err!( - "Vortex writer task for '{}' was already consumed", - writer.path - ) - })?; - task.await - }; - // Completed, so there is nothing left for the drop guard to abort. - writer.task = None; - - match joined { - Ok(result) => result, - Err(e) => Err(exec_datafusion_err!( - "Vortex writer task for '{}' failed to join: {e}", - writer.path - )), - } -} - -fn batch_uncompressed_bytes(batch: &RecordBatch) -> DFResult { - u64::try_from(batch.get_array_memory_size()).map_err(|_| { - exec_datafusion_err!( - "RecordBatch memory size does not fit in u64: {}", - batch.get_array_memory_size() - ) - }) -} - -fn remove_partition_columns( - batch: &RecordBatch, - partition_column_names: &[String], -) -> DFResult { - let partition_names: Vec<_> = partition_column_names.iter().map(String::as_str).collect(); - let (columns, fields): (Vec<_>, Vec<_>) = batch - .columns() - .iter() - .zip(batch.schema().fields()) - .filter(|(_, field)| !partition_names.contains(&field.name().as_str())) - .map(|(array, field)| (Arc::clone(array), (**field).clone())) - .unzip(); - - let schema = Schema::new(fields); - if columns.is_empty() { - let options = RecordBatchOptions::default().with_row_count(Some(batch.num_rows())); - return Ok(RecordBatch::try_new_with_options( - Arc::new(schema), - columns, - &options, - )?); - } - - Ok(RecordBatch::try_new(Arc::new(schema), columns)?) -} - -/// Build the output path for a rolling write. -fn output_file_path( - base_output_path: &ListingTableUrl, - file_index: usize, - extension: &str, - single_file_output: bool, - write_id: &str, -) -> Path { - if single_file_output { - if file_index == 0 { - return base_output_path.prefix().clone(); - } - return numbered_path(base_output_path.prefix(), file_index, extension); - } - - let mut base = base_output_path.prefix().to_string(); - if !base.ends_with('/') { - base.push('/'); - } - Path::from(format!("{base}{write_id}_{file_index:05}.{extension}")) -} - #[cfg(test)] mod tests { use std::sync::Arc; @@ -696,8 +225,6 @@ mod tests { use datafusion::logical_expr::Values; use datafusion::logical_expr::dml::InsertOp; use datafusion_common::ScalarValue; - use datafusion_common::exec_datafusion_err; - use datafusion_datasource::ListingTableUrl; use datafusion_datasource::TableSchema; use datafusion_datasource::file_format::format_as_file_type; use datafusion_datasource::file_groups::FileGroup; @@ -710,26 +237,12 @@ mod tests { use datafusion_physical_plan::empty::EmptyExec; use futures::TryStreamExt; use rstest::rstest; - use tokio::sync::oneshot; - use tokio::time::Duration; use vortex::file::VORTEX_FILE_EXTENSION; use vortex::session::VortexSession; use super::*; use crate::common_tests::TestSessionContext; use crate::persistent::VortexFormatFactory; - use crate::persistent::VortexTableOptions; - use crate::persistent::sink::ActiveFileWriter; - use crate::persistent::sink::finish_file_writer; - - fn split_path(base_path: &Path, file_index: usize, extension: &str) -> Path { - let mut base = base_path.to_string(); - if !base.ends_with('/') { - base.push('/'); - } - let filename = format!("part-{file_index:05}.{extension}"); - Path::from(format!("{base}{filename}")) - } #[tokio::test] async fn test_insert_into_sql() -> anyhow::Result<()> { @@ -744,11 +257,28 @@ mod tests { ) .await?; - ctx.session + let insert = ctx + .session .sql("INSERT INTO my_tbl VALUES ('hello', 1), ('world', 2);") - .await? - .collect() .await?; + let physical_plan = insert.create_physical_plan().await?; + datafusion_physical_plan::collect(Arc::clone(&physical_plan), ctx.session.task_ctx()) + .await?; + + let metrics = physical_plan + .metrics() + .ok_or_else(|| anyhow::anyhow!("Vortex insert did not expose sink metrics"))?; + assert_eq!( + metrics + .sum_by_name("rows_written") + .map(|metric| metric.as_usize()), + Some(2) + ); + assert!( + metrics + .sum_by_name("bytes_written") + .is_some_and(|metric| metric.as_usize() > 0) + ); let batches = ctx .session @@ -829,11 +359,10 @@ mod tests { } /// Reproduction by . - /// Uses a 1MB target file size to exercise file splitting behavior. #[rstest] - #[case(1_000, 1)] - #[case(5_000_000, 6)] - #[case(10_000_000, 10)] + #[case(1000, 1)] + #[case(40_961, 4)] + #[case(1_000_000, 4)] #[tokio::test] async fn test_write_large_batch( #[case] entries: usize, @@ -841,26 +370,15 @@ mod tests { ) -> anyhow::Result<()> { let ctx = TestSessionContext::default(); - let opts = VortexTableOptions { - target_file_size_mb: 1, - ..Default::default() - }; - - let factory = VortexFormatFactory::new().with_options(opts); - - let values: Vec = (0..entries) - .map(|i| i8::try_from(i % 127)) - .collect::>()?; - let data = ctx.session.read_batch(RecordBatch::try_new( Arc::new(Schema::new(vec![Field::new("a", DataType::Int8, false)])), - vec![Arc::new(Int8Array::from(values))], + vec![Arc::new(Int8Array::from(vec![0i8; entries]))], )?)?; let logical_plan = LogicalPlanBuilder::copy_to( data.logical_plan().clone(), "/table/".to_string(), - format_as_file_type(Arc::new(factory)), + format_as_file_type(Arc::new(VortexFormatFactory::new())), Default::default(), vec![], )? @@ -896,57 +414,50 @@ mod tests { entries, count_value ); - let file_metas = ctx - .store - .list(Some(&"table".into())) - .try_collect::>() - .await?; - - assert!( - file_metas.len() >= expected_files, - "Expected at least {expected_files} files for {entries} values, got {}", - file_metas.len() - ); - - Ok(()) - } - - #[tokio::test] - async fn test_write_large_batch_default_target_is_128mb() -> anyhow::Result<()> { - let ctx = TestSessionContext::default(); - - let entries = 1_000_000; - let values: Vec = (0..entries) - .map(|i| i8::try_from(i % 127)) - .collect::>()?; - - let data = ctx.session.read_batch(RecordBatch::try_new( - Arc::new(Schema::new(vec![Field::new("a", DataType::Int8, false)])), - vec![Arc::new(Int8Array::from(values))], - )?)?; - - let logical_plan = LogicalPlanBuilder::copy_to( - data.logical_plan().clone(), - "/table/".to_string(), - format_as_file_type(Arc::new(VortexFormatFactory::new())), - Default::default(), - vec![], - )? - .build()?; - - ctx.session - .execute_logical_plan(logical_plan) + let all_data = ctx + .session + .sql("SELECT a FROM '/table/'") .await? .collect() .await?; + let mut total_rows = 0; + for batch in all_data { + let col = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + + for i in 0..batch.num_rows() { + assert_eq!( + col.value(i), + 0i8, + "Expected value 0 at row {}, but found {}", + total_rows + i, + col.value(i) + ); + } + total_rows += batch.num_rows(); + } + + assert_eq!( + total_rows, entries, + "Total rows read ({}) doesn't match expected entries ({})", + total_rows, entries + ); + let file_metas = ctx .store .list(Some(&"/table".into())) .try_collect::>() .await?; - assert_eq!(file_metas.len(), 1); + assert_eq!( + file_metas.len(), + expected_files, + "Expected {expected_files} files for {entries} values" + ); Ok(()) } @@ -994,1095 +505,29 @@ mod tests { } #[test] - fn test_split_path_basic() { - let path = Path::from("data/output"); - assert_eq!( - split_path(&path, 0, "vortex").to_string(), - "data/output/part-00000.vortex" - ); - assert_eq!( - split_path(&path, 12, "vortex").to_string(), - "data/output/part-00012.vortex" - ); - } - - #[test] - fn test_split_path_preserves_trailing_slash() { - let path = Path::from("nested/path/"); - assert_eq!( - split_path(&path, 3, "vx").to_string(), - "nested/path/part-00003.vx" - ); - } - - #[test] - fn test_numbered_path() { - use super::numbered_path; - - let path = Path::from("table/c1=alpha/abc123.vortex"); - assert_eq!( - numbered_path(&path, 0, "vortex").to_string(), - "table/c1=alpha/abc123_00000.vortex" - ); - assert_eq!( - numbered_path(&path, 5, "vortex").to_string(), - "table/c1=alpha/abc123_00005.vortex" - ); - } - - #[test] - fn test_numbered_path_no_extension() { - use super::numbered_path; - - let path = Path::from("table/output"); - assert_eq!( - numbered_path(&path, 0, "vortex").to_string(), - "table/output_00000.vortex" - ); - } - - #[test] - fn test_output_file_path_single_file_and_collection() { - use super::output_file_path; - - let single = ListingTableUrl::parse("file:///tmp/output.vortex").unwrap(); - assert_eq!( - output_file_path(&single, 0, "vortex", true, "wid").to_string(), - "tmp/output.vortex" - ); - assert_eq!( - output_file_path(&single, 2, "vortex", true, "wid").to_string(), - "tmp/output_00002.vortex" - ); - - let collection = ListingTableUrl::parse("file:///tmp/table/").unwrap(); - assert_eq!( - output_file_path(&collection, 3, "vortex", false, "wid").to_string(), - "tmp/table/wid_00003.vortex" - ); - } - - #[tokio::test] - async fn test_dropping_active_writer_aborts_its_task() -> anyhow::Result<()> { - use std::sync::atomic::AtomicBool; - use std::sync::atomic::Ordering; - - struct AbortSignal(Arc); - - impl Drop for AbortSignal { - fn drop(&mut self) { - self.0.store(true, Ordering::SeqCst); - } - } - - let dropped = Arc::new(AtomicBool::new(false)); - let (sender, receiver) = mpsc::channel::(1); - drop(receiver); - let (_gate_tx, gate_rx) = oneshot::channel::<()>(); - - let signal = AbortSignal(Arc::clone(&dropped)); - let writer = ActiveFileWriter { - path: Path::from("table/cancelled.vortex"), - sender, - task: Some(tokio::spawn(async move { - // Held across the await so aborting drops it. - let _signal = signal; - let _ = gate_rx.await; - Err(exec_datafusion_err!("writer task should have been aborted")) - })), - }; - - // Let the task reach its await before cancelling. - tokio::task::yield_now().await; - drop(writer); - - tokio::time::timeout(Duration::from_secs(5), async { - while !dropped.load(Ordering::SeqCst) { - tokio::task::yield_now().await; - } - }) - .await - .map_err(|_| anyhow::anyhow!("dropping ActiveFileWriter did not abort its writer task"))?; - - Ok(()) - } - - #[tokio::test] - async fn test_finish_file_writer_waits_for_task_completion() -> anyhow::Result<()> { - let (sender, receiver) = mpsc::channel::(1); - drop(receiver); - - let (gate_tx, gate_rx) = oneshot::channel::<()>(); - - let writer = ActiveFileWriter { - path: Path::from("table/pending.vortex"), - sender, - task: Some(tokio::spawn(async move { - let _ = gate_rx.await; - Err(exec_datafusion_err!( - "synthetic writer failure after completion gate" - )) - })), - }; - - let mut finish_fut = Box::pin(finish_file_writer(writer)); - assert!( - tokio::time::timeout(Duration::from_millis(50), &mut finish_fut) - .await - .is_err(), - "finish_file_writer returned before writer task completed" - ); - - gate_tx - .send(()) - .map_err(|_| anyhow::anyhow!("failed to release writer completion gate in test"))?; + fn test_display_as() { + let session = VortexSession::empty(); + let table_schema = TableSchema::from(Arc::new(Schema::empty())); - let err = match finish_fut.await { - Ok(_) => { - return Err(anyhow::anyhow!( - "finish_file_writer unexpectedly succeeded after gate release" - )); - } - Err(err) => err, + let config = FileSinkConfig { + original_url: "".to_owned(), + object_store_url: ObjectStoreUrl::local_filesystem(), + file_group: FileGroup::new(Vec::new()), + table_paths: Vec::new(), + output_schema: Arc::new(Schema::empty()), + table_partition_cols: Vec::new(), + insert_op: InsertOp::Overwrite, + keep_partition_by_columns: false, + file_extension: VORTEX_FILE_EXTENSION.to_owned(), + file_output_mode: FileOutputMode::SingleFile, }; - assert!( - err.to_string().contains("synthetic writer failure"), - "unexpected error: {err}" - ); - - Ok(()) - } - #[test] - fn test_remove_partition_columns() -> anyhow::Result<()> { - use datafusion::arrow::array::StringArray; - - use super::remove_partition_columns; - - let batch = RecordBatch::try_new( - Arc::new(Schema::new(vec![ - Field::new("part", DataType::Utf8, false), - Field::new("val", DataType::Int64, false), - ])), - vec![ - Arc::new(StringArray::from(vec!["x", "y"])), - Arc::new(Int64Array::from(vec![1, 2])), - ], - )?; - - let out = remove_partition_columns(&batch, &["part".to_string()])?; - assert_eq!(out.num_columns(), 1); - assert_eq!(out.schema().field(0).name(), "val"); - assert_eq!(out.num_rows(), 2); - - Ok(()) - } - - #[test] - fn test_remove_partition_columns_all_columns_partitioned() -> anyhow::Result<()> { - use datafusion::arrow::array::StringArray; - - use super::remove_partition_columns; - - let batch = RecordBatch::try_new( - Arc::new(Schema::new(vec![Field::new("part", DataType::Utf8, false)])), - vec![Arc::new(StringArray::from(vec!["x", "y", "z"]))], - )?; - - let out = remove_partition_columns(&batch, &["part".to_string()])?; - assert_eq!(out.num_columns(), 0); - assert_eq!(out.num_rows(), 3); - - Ok(()) - } - - /// Generate `count` pseudo-random i64 values using a simple LCG. - /// These values resist compression (unlike sequential or modular data), - /// giving more realistic compressed file sizes. - fn pseudo_random_i64s(count: usize, seed: i64) -> Vec { - let mut v = seed; - (0..count) - .map(|_| { - v = v - .wrapping_mul(6_364_136_223_846_793_005) - .wrapping_add(1_442_695_040_888_963_407); - v - }) - .collect() - } - - /// Tests file splitting through the full DataFusion pipeline. - /// - /// Writes ~62MB of pseudo-random Int64 data (near 1:1 compression ratio) - /// via COPY TO with a 16MB target file size. Verifies that exactly 4 files - /// are produced and each file's compressed size is approximately 16MB. - /// - /// This exercises the complete COPY TO write path through DataFusion and - /// VortexSink, unlike a direct `write_stream_to_files` call. - #[tokio::test] - async fn test_file_splitting_62mb_into_4_files() -> anyhow::Result<()> { - use datafusion::datasource::MemTable; - use datafusion_datasource::file_format::format_as_file_type; - - let ctx = TestSessionContext::default(); - - let target_mb = 16_usize; - let opts = VortexTableOptions { - target_file_size_mb: target_mb, - ..Default::default() - }; - let factory = VortexFormatFactory::new().with_options(opts); - - let batch_rows = 8192_usize; - let total_elements = 62 * 1024 * 1024 / 8; // ~8,126,464 i64 values ≈ 62MB Arrow memory - let num_batches = total_elements / batch_rows; - let expected_total_rows = (num_batches * batch_rows) as i64; - - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - - let mut batches = Vec::new(); - for i in 0..num_batches { - let values = pseudo_random_i64s(batch_rows, (i * batch_rows) as i64); - batches.push(RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int64Array::from(values))], - )?); - } - - let table = MemTable::try_new(Arc::clone(&schema), vec![batches])?; - ctx.session.register_table("source", Arc::new(table))?; - - let source = ctx.session.table("source").await?; - let logical_plan = LogicalPlanBuilder::copy_to( - source.logical_plan().clone(), - "/table/".to_string(), - format_as_file_type(Arc::new(factory)), - Default::default(), - vec![], - )? - .build()?; - - ctx.session - .execute_logical_plan(logical_plan) - .await? - .collect() - .await?; - - let file_metas = ctx - .store - .list(Some(&"/table".into())) - .try_collect::>() - .await?; - - assert_eq!( - file_metas.len(), - 4, - "Expected 4 files for ~62MB data with {target_mb}MB target, got {} (sizes: {:?})", - file_metas.len(), - file_metas.iter().map(|m| m.size).collect::>() - ); - - let target_bytes = u64::try_from(target_mb * 1024 * 1024)?; - for meta in &file_metas { - assert!( - meta.size > target_bytes / 2, - "File {} is {}B, expected at least {}B (target/2)", - meta.location, - meta.size, - target_bytes / 2 - ); - } - - // Verify total row count. - let result = ctx - .session - .sql("SELECT COUNT(*) as cnt FROM '/table/'") - .await? - .collect() - .await?; - - let count = result[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - - assert_eq!(count, expected_total_rows, "Total row count mismatch"); - - Ok(()) - } - - /// Tests file splitting with compressible data through the full pipeline. - /// - /// Uses low-entropy Int64 values (repeating 0..255) which compress ~8:1 in - /// Vortex. With the current code that compares Arrow memory size against - /// `target_file_size`, files are split far too early, producing many tiny - /// compressed files instead of files that are close to the target. - /// - /// For ~32MB of Arrow data (~4MB compressed at 8:1) with a 1MB target: - /// - **Correct**: 4 files of ~1MB compressed each - /// - **Bug**: 32 files of ~0.125MB compressed each - #[tokio::test] - async fn test_file_splitting_compressible_data() -> anyhow::Result<()> { - use datafusion::datasource::MemTable; - use datafusion_datasource::file_format::format_as_file_type; - - let ctx = TestSessionContext::default(); - - let target_mb = 1_usize; - let opts = VortexTableOptions { - target_file_size_mb: target_mb, - ..Default::default() - }; - let factory = VortexFormatFactory::new().with_options(opts); - - // Generate low-entropy Int64 values: repeating 0..255. - // Arrow memory: 4M × 8 bytes = 32MB. - // Vortex compressed: each value only needs ~1 byte → ~4MB total. - let total_elements = 4_000_000_usize; - let batch_rows = 8192_usize; - let num_batches = total_elements / batch_rows; - - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - - let mut batches = Vec::new(); - for i in 0..num_batches { - let values: Vec = (0..batch_rows) - .map(|j| ((i * batch_rows + j) % 256) as i64) - .collect(); - batches.push(RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int64Array::from(values))], - )?); - } - - let table = MemTable::try_new(Arc::clone(&schema), vec![batches])?; - ctx.session.register_table("source", Arc::new(table))?; - - let source = ctx.session.table("source").await?; - let logical_plan = LogicalPlanBuilder::copy_to( - source.logical_plan().clone(), - "/table/".to_string(), - format_as_file_type(Arc::new(factory)), - Default::default(), - vec![], - )? - .build()?; - - ctx.session - .execute_logical_plan(logical_plan) - .await? - .collect() - .await?; - - let file_metas = ctx - .store - .list(Some(&"/table".into())) - .try_collect::>() - .await?; - - // With compressible data, there should be few files (not > 10). - // The buggy code produces many tiny files because it splits on Arrow - // memory (32MB / 1MB = 32 files) instead of compressed size (~4MB / 1MB = 4 files). - let total_compressed: u64 = file_metas.iter().map(|m| m.size).sum(); - let target_bytes = u64::try_from(target_mb * 1024 * 1024)?; - - // We should have at most ~(total_compressed / target) + 1 files, not - // ~(arrow_memory / target) files. - let max_expected = usize::try_from(total_compressed / target_bytes + 2)?; - assert!( - file_metas.len() <= max_expected, - "Too many files: got {} but total compressed is {}B with {}B target \ - (expected at most {max_expected}). Files are being split on Arrow memory \ - instead of compressed size. Sizes: {:?}", - file_metas.len(), - total_compressed, - target_bytes, - file_metas.iter().map(|m| m.size).collect::>() - ); - - // Files in the middle should be reasonably sized. The first may be smaller - // because the compression ratio is unknown until the first write completes, - // and the last holds whatever remained when the input ended, so its size - // depends on how the input divides rather than on the target. - let mut sizes: Vec<(String, u64)> = file_metas - .iter() - .map(|m| (m.location.to_string(), m.size)) - .collect(); - sizes.sort(); - let without_last = sizes.len().saturating_sub(1); - for (location, size) in sizes.iter().take(without_last).skip(1) { - assert!( - *size > target_bytes / 4, - "File {location} is {size}B, far below target {target_bytes}B — \ - splitting on Arrow memory, not compressed size. Sizes: {sizes:?}" - ); - } - - Ok(()) - } - - #[tokio::test] - async fn test_write_large_batch_target_file_size_disabled() -> anyhow::Result<()> { - use datafusion::datasource::MemTable; - use datafusion_datasource::file_format::format_as_file_type; - - let ctx = TestSessionContext::default(); - - let opts = VortexTableOptions { - // Disable sink-side rolling/splitting. - target_file_size_mb: 0, - ..Default::default() - }; - let factory = VortexFormatFactory::new().with_options(opts); - - let rows_per_partition = 300_000_usize; - let num_partitions = 8_usize; - let expected_total_rows = (rows_per_partition * num_partitions) as i64; - - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - - let mut partitions: Vec> = Vec::new(); - for p in 0..num_partitions { - let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); - partitions.push(vec![RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int64Array::from(values))], - )?]); - } - - let table = MemTable::try_new(schema, partitions)?; - ctx.session.register_table("source", Arc::new(table))?; - - let source = ctx.session.table("source").await?; - let logical_plan = LogicalPlanBuilder::copy_to( - source.logical_plan().clone(), - "/table/".to_string(), - format_as_file_type(Arc::new(factory)), - Default::default(), - vec![], - )? - .build()?; - - ctx.session - .execute_logical_plan(logical_plan) - .await? - .collect() - .await?; - - let file_metas = ctx - .store - .list(Some(&"/table".into())) - .try_collect::>() - .await?; - - let unique_write_ids: HashSet<_> = file_metas - .iter() - .filter_map(|m| { - m.location - .filename() - .and_then(|name| name.split_once('_')) - .map(|(prefix, _)| prefix.to_string()) - }) - .collect(); - - assert_eq!( - unique_write_ids.len(), - 1, - "Expected one write_id with target size disabled; got {:?} from files: {:?}", - unique_write_ids, - file_metas - .iter() - .map(|m| format!("{}: {}B", m.location, m.size)) - .collect::>() - ); - - assert_eq!( - file_metas.len(), - 1, - "Expected exactly one output file when target size is disabled, got {}", - file_metas.len() - ); - - let result = ctx - .session - .sql("SELECT COUNT(*) as cnt FROM '/table/'") - .await? - .collect() - .await?; - - let count = result[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - - assert_eq!(count, expected_total_rows, "Total row count mismatch"); - - Ok(()) - } - - #[tokio::test] - async fn test_target_file_size_uses_single_sink_input_partition() -> anyhow::Result<()> { - use datafusion::datasource::MemTable; - use datafusion_datasource::file_format::format_as_file_type; - - let ctx = TestSessionContext::default(); - - let opts = VortexTableOptions { - // Enable sink-side sizing, but make the threshold large enough - // that all input data should fit in a single file. - target_file_size_mb: 512, - ..Default::default() - }; - let factory = VortexFormatFactory::new().with_options(opts); - - let rows_per_partition = 300_000_usize; - let num_partitions = 8_usize; - let expected_total_rows = (rows_per_partition * num_partitions) as i64; - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - - // Build a MemTable with multiple physical input partitions to mimic - // DataFusion's parallel writer inputs. - let mut partitions: Vec> = Vec::new(); - for p in 0..num_partitions { - let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); - partitions.push(vec![RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int64Array::from(values))], - )?]); - } - - let table = MemTable::try_new(schema, partitions)?; - ctx.session.register_table("source", Arc::new(table))?; - - let source = ctx.session.table("source").await?; - let logical_plan = LogicalPlanBuilder::copy_to( - source.logical_plan().clone(), - "/table/".to_string(), - format_as_file_type(Arc::new(factory)), - Default::default(), - vec![], - )? - .build()?; - - ctx.session - .execute_logical_plan(logical_plan) - .await? - .collect() - .await?; - - let file_metas = ctx - .store - .list(Some(&"/table".into())) - .try_collect::>() - .await?; - - let unique_write_ids: HashSet<_> = file_metas - .iter() - .filter_map(|m| { - m.location - .filename() - .and_then(|name| name.split_once('_')) - .map(|(prefix, _)| prefix.to_string()) - }) - .collect(); - - assert_eq!( - unique_write_ids.len(), - 1, - "Expected one write_id (single sink stream), got {:?} from files: {:?}", - unique_write_ids, - file_metas - .iter() - .map(|m| format!("{}: {}B", m.location, m.size)) - .collect::>() - ); - - assert!( - file_metas.len() < num_partitions, - "Expected fewer output files than input partitions after coalescing; got {} files for {num_partitions} input partitions", - file_metas.len() - ); - - let result = ctx - .session - .sql("SELECT COUNT(*) AS cnt FROM '/table/'") - .await? - .collect() - .await?; - - let count = result[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - assert_eq!(count, expected_total_rows, "Total row count mismatch"); - - Ok(()) - } - - #[tokio::test] - async fn test_insert_sql_target_size_multi_partition_source_single_write_id() - -> anyhow::Result<()> { - use datafusion::datasource::MemTable; - - let ctx = TestSessionContext::default(); - - ctx.session - .sql( - "CREATE EXTERNAL TABLE my_tbl \ - (a BIGINT NOT NULL) \ - STORED AS vortex \ - LOCATION 'table/' \ - OPTIONS(target_file_size_mb '64');", - ) - .await?; - - let rows_per_partition = 300_000_usize; - let num_partitions = 8_usize; - let expected_total_rows = (rows_per_partition * num_partitions) as i64; - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - - let mut partitions: Vec> = Vec::new(); - for p in 0..num_partitions { - let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); - partitions.push(vec![RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int64Array::from(values))], - )?]); - } - - let source = MemTable::try_new(schema, partitions)?; - ctx.session.register_table("source", Arc::new(source))?; - - ctx.session - .sql("INSERT INTO my_tbl SELECT a FROM source") - .await? - .collect() - .await?; - - let all_files = ctx.store.list(None).try_collect::>().await?; - - let unique_write_ids: HashSet<_> = all_files - .iter() - .filter_map(|m| { - m.location - .filename() - .and_then(|name| name.split_once('_')) - .map(|(prefix, _)| prefix.to_string()) - }) - .collect(); - - assert_eq!( - unique_write_ids.len(), - 1, - "Expected one write_id, got {:?} from files: {:?}", - unique_write_ids, - all_files - .iter() - .map(|m| format!("{}: {}B", m.location, m.size)) - .collect::>() - ); - assert!( - all_files.len() < num_partitions, - "Expected fewer files than input partitions; got {} files for {num_partitions} input partitions", - all_files.len() - ); - - let result = ctx - .session - .sql("SELECT COUNT(*) AS cnt FROM my_tbl") - .await? - .collect() - .await?; - let count = result[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - assert_eq!(count, expected_total_rows, "Total row count mismatch"); - - Ok(()) - } - - #[tokio::test] - async fn test_insert_sql_streaming_source_single_write_id() -> anyhow::Result<()> { - use arrow_schema::SchemaRef; - use datafusion::catalog::streaming::StreamingTable; - use datafusion::physical_plan::stream::RecordBatchStreamAdapter; - use datafusion::physical_plan::streaming::PartitionStream; - use futures::stream; - - #[derive(Debug)] - struct StaticPartitionStream { - schema: SchemaRef, - batch: RecordBatch, - } - - impl PartitionStream for StaticPartitionStream { - fn schema(&self) -> &SchemaRef { - &self.schema - } - - fn execute(&self, _ctx: Arc) -> SendableRecordBatchStream { - let schema = Arc::clone(&self.schema); - let batch = self.batch.clone(); - Box::pin(RecordBatchStreamAdapter::new( - schema, - stream::iter(vec![Ok(batch)]), - )) - } - } - - let ctx = TestSessionContext::default(); - - ctx.session - .sql( - "CREATE EXTERNAL TABLE my_tbl \ - (a BIGINT NOT NULL) \ - STORED AS vortex \ - LOCATION 'table/' \ - OPTIONS(target_file_size_mb '64');", - ) - .await?; - - let rows_per_partition = 300_000_usize; - let num_partitions = 8_usize; - let expected_total_rows = (rows_per_partition * num_partitions) as i64; - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - - let mut partitions: Vec> = Vec::new(); - for p in 0..num_partitions { - let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int64Array::from(values))], - )?; - - partitions.push(Arc::new(StaticPartitionStream { - schema: Arc::clone(&schema), - batch, - })); - } - - let source = StreamingTable::try_new(schema, partitions)?; - ctx.session - .register_table("source_stream", Arc::new(source))?; - - ctx.session - .sql("INSERT INTO my_tbl SELECT a FROM source_stream") - .await? - .collect() - .await?; - - let all_files = ctx.store.list(None).try_collect::>().await?; - - let unique_write_ids: HashSet<_> = all_files - .iter() - .filter_map(|m| { - m.location - .filename() - .and_then(|name| name.split_once('_')) - .map(|(prefix, _)| prefix.to_string()) - }) - .collect(); - - assert_eq!( - unique_write_ids.len(), - 1, - "Expected one write_id for streaming source insert, got {:?} from files: {:?}", - unique_write_ids, - all_files - .iter() - .map(|m| format!("{}: {}B", m.location, m.size)) - .collect::>() - ); - - let result = ctx - .session - .sql("SELECT COUNT(*) AS cnt FROM my_tbl") - .await? - .collect() - .await?; - let count = result[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - assert_eq!(count, expected_total_rows, "Total row count mismatch"); - - Ok(()) - } - - #[tokio::test] - async fn test_listing_table_direct_insert_into_streaming_exec_single_write_id() - -> anyhow::Result<()> { - use arrow_schema::SchemaRef; - use datafusion::physical_plan::ExecutionPlan; - use datafusion::physical_plan::collect; - use datafusion::physical_plan::stream::RecordBatchStreamAdapter; - use datafusion::physical_plan::streaming::PartitionStream; - use datafusion::physical_plan::streaming::StreamingTableExec; - use datafusion_expr::dml::InsertOp; - use futures::stream; - - #[derive(Debug)] - struct StaticPartitionStream { - schema: SchemaRef, - batch: RecordBatch, - } - - impl PartitionStream for StaticPartitionStream { - fn schema(&self) -> &SchemaRef { - &self.schema - } - - fn execute(&self, _ctx: Arc) -> SendableRecordBatchStream { - let schema = Arc::clone(&self.schema); - let batch = self.batch.clone(); - Box::pin(RecordBatchStreamAdapter::new( - schema, - stream::iter(vec![Ok(batch)]), - )) - } - } - - let ctx = TestSessionContext::default(); - - ctx.session - .sql( - "CREATE EXTERNAL TABLE my_tbl \ - (a BIGINT NOT NULL) \ - STORED AS vortex \ - LOCATION 'table/' \ - OPTIONS(target_file_size_mb '64');", + let get_sink = || { + VortexSink::new( + config.clone(), + Arc::clone(table_schema.file_schema()), + session.clone(), ) - .await?; - - let table_provider = ctx.session.table_provider("my_tbl").await?; - - let rows_per_partition = 300_000_usize; - let num_partitions = 8_usize; - let expected_total_rows = (rows_per_partition * num_partitions) as i64; - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - - let mut partitions: Vec> = Vec::new(); - for p in 0..num_partitions { - let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int64Array::from(values))], - )?; - - partitions.push(Arc::new(StaticPartitionStream { - schema: Arc::clone(&schema), - batch, - })); - } - - let input = Arc::new(StreamingTableExec::try_new( - schema, - partitions, - None, - Vec::new(), - false, - None, - )?) as Arc; - - let plan = table_provider - .insert_into(&ctx.session.state(), input, InsertOp::Append) - .await?; - let _count_batches = collect(plan, ctx.session.task_ctx()).await?; - - let all_files = ctx.store.list(None).try_collect::>().await?; - - let unique_write_ids: HashSet<_> = all_files - .iter() - .filter_map(|m| { - m.location - .filename() - .and_then(|name| name.split_once('_')) - .map(|(prefix, _)| prefix.to_string()) - }) - .collect(); - - assert_eq!( - unique_write_ids.len(), - 1, - "Expected one write_id for direct insert_into streaming exec, got {:?} from files: {:?}", - unique_write_ids, - all_files - .iter() - .map(|m| format!("{}: {}B", m.location, m.size)) - .collect::>() - ); - - let result = ctx - .session - .sql("SELECT COUNT(*) AS cnt FROM my_tbl") - .await? - .collect() - .await?; - let count = result[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - assert_eq!(count, expected_total_rows, "Total row count mismatch"); - - Ok(()) - } - - #[tokio::test] - async fn test_listing_table_direct_insert_into_unbounded_streaming_exec_single_write_id() - -> anyhow::Result<()> { - use arrow_schema::SchemaRef; - use datafusion::physical_plan::ExecutionPlan; - use datafusion::physical_plan::collect; - use datafusion::physical_plan::stream::RecordBatchStreamAdapter; - use datafusion::physical_plan::streaming::PartitionStream; - use datafusion::physical_plan::streaming::StreamingTableExec; - use datafusion_expr::dml::InsertOp; - use futures::stream; - - #[derive(Debug)] - struct StaticPartitionStream { - schema: SchemaRef, - batch: RecordBatch, - } - - impl PartitionStream for StaticPartitionStream { - fn schema(&self) -> &SchemaRef { - &self.schema - } - - fn execute(&self, _ctx: Arc) -> SendableRecordBatchStream { - let schema = Arc::clone(&self.schema); - let batch = self.batch.clone(); - Box::pin(RecordBatchStreamAdapter::new( - schema, - stream::iter(vec![Ok(batch)]), - )) - } - } - - let ctx = TestSessionContext::default(); - - ctx.session - .sql( - "CREATE EXTERNAL TABLE my_tbl \ - (a BIGINT NOT NULL) \ - STORED AS vortex \ - LOCATION 'table/' \ - OPTIONS(target_file_size_mb '64');", - ) - .await?; - - let table_provider = ctx.session.table_provider("my_tbl").await?; - - let rows_per_partition = 100_000_usize; - let num_partitions = 8_usize; - let expected_total_rows = (rows_per_partition * num_partitions) as i64; - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - - let mut partitions: Vec> = Vec::new(); - for p in 0..num_partitions { - let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int64Array::from(values))], - )?; - - partitions.push(Arc::new(StaticPartitionStream { - schema: Arc::clone(&schema), - batch, - })); - } - - let input = Arc::new(StreamingTableExec::try_new( - schema, - partitions, - None, - Vec::new(), - true, - None, - )?) as Arc; - - let plan = table_provider - .insert_into(&ctx.session.state(), input, InsertOp::Append) - .await?; - let _count_batches = collect(plan, ctx.session.task_ctx()).await?; - - let all_files = ctx.store.list(None).try_collect::>().await?; - - let unique_write_ids: HashSet<_> = all_files - .iter() - .filter_map(|m| { - m.location - .filename() - .and_then(|name| name.split_once('_')) - .map(|(prefix, _)| prefix.to_string()) - }) - .collect(); - - assert_eq!( - unique_write_ids.len(), - 1, - "Expected one write_id for unbounded streaming insert, got {:?} from files: {:?}", - unique_write_ids, - all_files - .iter() - .map(|m| format!("{}: {}B", m.location, m.size)) - .collect::>() - ); - - let result = ctx - .session - .sql("SELECT COUNT(*) AS cnt FROM my_tbl") - .await? - .collect() - .await?; - let count = result[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - assert_eq!(count, expected_total_rows, "Total row count mismatch"); - - Ok(()) - } - - #[test] - fn test_display_as() { - let session = VortexSession::empty(); - let table_schema = TableSchema::new(Arc::new(Schema::empty()), Vec::new()); - - let config = FileSinkConfig { - original_url: "".to_owned(), - object_store_url: ObjectStoreUrl::local_filesystem(), - file_group: FileGroup::new(Vec::new()), - table_paths: Vec::new(), - output_schema: Arc::new(Schema::empty()), - table_partition_cols: Vec::new(), - insert_op: InsertOp::Overwrite, - keep_partition_by_columns: false, - file_extension: VORTEX_FILE_EXTENSION.to_owned(), - file_output_mode: FileOutputMode::SingleFile, - }; - - let get_sink = || VortexSink { - config: config.clone(), - schema: Arc::clone(table_schema.file_schema()), - session: session.clone(), - target_file_size: None, }; insta::assert_snapshot!(DefaultDisplay(get_sink()).to_string(), @"VortexSink(file_groups=[])"); diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index 469759527f6..fbd37f04ded 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -2,12 +2,12 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::fmt::Formatter; -use std::ops::Range; use std::sync::Arc; use std::sync::Weak; use datafusion_common::Result as DFResult; use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_datasource::TableSchema; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_scan_config::FileScanConfig; @@ -29,7 +29,6 @@ use datafusion_physical_plan::filter_pushdown::PushedDownPredicate; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use object_store::ObjectStore; use object_store::path::Path; -use vortex::error::VortexExpect; use vortex::file::VORTEX_FILE_EXTENSION; use vortex::layout::LayoutReader; use vortex::metrics::DefaultMetricsRegistry; @@ -37,6 +36,7 @@ use vortex::metrics::MetricsRegistry; use vortex::session::VortexSession; use vortex_utils::aliases::dash_map::DashMap; +use super::opener::NaturalSplits; use super::opener::VortexOpener; use crate::VortexTableOptions; use crate::convert::exprs::DefaultExpressionConvertor; @@ -190,19 +190,19 @@ pub struct VortexSource { /// Subset of predicates that can be pushed down into Vortex scan operations. /// These are expressions that Vortex can efficiently evaluate during scanning. pub(crate) vortex_predicate: Option, - pub(crate) batch_size: Option, - _unused_df_metrics: ExecutionPlanMetricsSet, + /// DataFusion-native metrics exposed through `DataSourceExec`. + df_metrics: ExecutionPlanMetricsSet, /// Shared layout readers, the source only lives as long as one scan. /// /// Sharing the readers allows us to only read every layout once from the file, even across partitions. layout_readers: Arc>>, - /// Shared full-file natural split ranges keyed by path. - natural_split_ranges: Arc]>>>, + /// Shared full-file natural splits keyed by path. + natural_splits: Arc>>, expression_convertor: Arc, pub(crate) vortex_reader_factory: Option>, pub(crate) ordered: bool, vx_metrics_registry: Arc, - file_metadata_cache: Option>, + file_metadata_cache: Option>, /// Options controlling scan planning and execution behavior. options: VortexTableOptions, } @@ -228,10 +228,9 @@ impl VortexSource { projection, full_predicate: None, vortex_predicate: None, - batch_size: None, - _unused_df_metrics: Default::default(), + df_metrics: Default::default(), layout_readers: Arc::new(DashMap::default()), - natural_split_ranges: Arc::new(DashMap::default()), + natural_splits: Arc::new(DashMap::default()), expression_convertor, vortex_reader_factory: None, vx_metrics_registry: Arc::new(DefaultMetricsRegistry::default()), @@ -297,10 +296,7 @@ impl VortexSource { } /// Overrides the metadata cache used to reuse Vortex footers across scans. - pub fn with_file_metadata_cache( - mut self, - file_metadata_cache: Arc, - ) -> Self { + pub fn with_file_metadata_cache(mut self, file_metadata_cache: Arc) -> Self { self.file_metadata_cache = Some(file_metadata_cache); self } @@ -335,10 +331,6 @@ impl VortexSource { base_config: &FileScanConfig, partition: usize, ) -> DFResult { - let batch_size = self - .batch_size - .vortex_expect("batch_size must be supplied to VortexSource"); - let expr_adapter_factory = base_config .expr_adapter_factory .clone() @@ -358,11 +350,11 @@ impl VortexSource { file_pruning_predicate: self.full_predicate.clone(), expr_adapter_factory, table_schema: self.table_schema.clone(), - batch_size, limit: base_config.limit.map(|l| l as u64), metrics_registry: Arc::clone(&self.vx_metrics_registry), + df_metrics: self.df_metrics.clone(), layout_readers: Arc::clone(&self.layout_readers), - natural_split_ranges: Arc::clone(&self.natural_split_ranges), + natural_splits: Arc::clone(&self.natural_splits), has_output_ordering: !base_config.output_ordering.is_empty() || self.ordered, expression_convertor: Arc::clone(&self.expression_convertor), file_metadata_cache: self.file_metadata_cache.clone(), @@ -388,10 +380,9 @@ impl FileSource for VortexSource { )?)) } - fn with_batch_size(&self, batch_size: usize) -> Arc { - let mut source = self.clone(); - source.batch_size = Some(batch_size); - Arc::new(source) + fn with_batch_size(&self, _batch_size: usize) -> Arc { + // DataSourceExec applies BatchSplitStream after the FileSource stream. + Arc::new(self.clone()) } fn filter(&self) -> Option> { @@ -399,7 +390,7 @@ impl FileSource for VortexSource { } fn metrics(&self) -> &ExecutionPlanMetricsSet { - &self._unused_df_metrics + &self.df_metrics } fn file_type(&self) -> &str { @@ -541,6 +532,19 @@ impl FileSource for VortexSource { fn table_schema(&self) -> &TableSchema { &self.table_schema } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> DFResult, + ) -> DFResult { + datafusion_physical_plan::apply_expression_roots( + self.full_predicate + .iter() + .chain(self.vortex_predicate.iter()) + .chain(self.projection.iter().map(|expr| &expr.expr)), + f, + ) + } } #[cfg(test)] @@ -610,10 +614,7 @@ mod tests { } fn sort_test_source(schema: Arc) -> VortexSource { - VortexSource::new( - TableSchema::from_file_schema(schema), - VortexSession::default(), - ) + VortexSource::new(TableSchema::from(schema), VortexSession::default()) } fn octet_length_filter(schema: &Schema) -> PhysicalExprRef { @@ -667,12 +668,8 @@ mod tests { inner: DefaultExpressionConvertor::default(), }) as Arc; - let mut source = VortexSource::new( - TableSchema::from_file_schema(file_schema), - VortexSession::default(), - ) - .with_expression_convertor(Arc::clone(&expression_convertor)); - source.batch_size = Some(100); + let source = VortexSource::new(TableSchema::from(file_schema), VortexSession::default()) + .with_expression_convertor(Arc::clone(&expression_convertor)); let config = FileScanConfigBuilder::new( ObjectStoreUrl::local_filesystem(), diff --git a/vortex-datafusion/src/persistent/tests.rs b/vortex-datafusion/src/persistent/tests.rs index 41d8141165f..35dd745461d 100644 --- a/vortex-datafusion/src/persistent/tests.rs +++ b/vortex-datafusion/src/persistent/tests.rs @@ -26,6 +26,8 @@ use vortex::array::arrays::VarBinArray; use vortex::array::validity::Validity; use vortex::buffer::Buffer; use vortex::buffer::buffer; +use vortex::editions::CORE_2026_08_3; +use vortex::editions::EditionSessionExt; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; use vortex::io::VortexWrite; @@ -497,8 +499,8 @@ async fn arrow_uuid_extension_roundtrip() -> anyhow::Result<()> { use vortex_arrow::ArrowSessionExt; let ctx = TestSessionContext::default(); - // Default vortex session has importer/exporter for Arrow UUID let session = VortexSession::default(); + session.enable_edition(CORE_2026_08_3)?; let mut uuid_field = Field::new("id", DataType::FixedSizeBinary(16), false); uuid_field.try_with_extension_type(Uuid)?; @@ -565,6 +567,7 @@ async fn arrow_uuid_extension_roundtrip_nested_struct() -> anyhow::Result<()> { let ctx = TestSessionContext::default(); let session = VortexSession::default(); + session.enable_edition(CORE_2026_08_3)?; let mut inner_uuid_field = Field::new("id", DataType::FixedSizeBinary(16), false); inner_uuid_field.try_with_extension_type(Uuid)?; diff --git a/vortex-datafusion/src/tests/schema_evolution.rs b/vortex-datafusion/src/tests/schema_evolution.rs index 093fed6c148..bdc8f822af6 100644 --- a/vortex-datafusion/src/tests/schema_evolution.rs +++ b/vortex-datafusion/src/tests/schema_evolution.rs @@ -5,6 +5,7 @@ use std::sync::Arc; +use arrow_array::record_batch; use arrow_schema::DataType; use arrow_schema::Field; use arrow_schema::Fields; @@ -19,7 +20,6 @@ use datafusion::arrow::datatypes::UInt32Type; use datafusion::assert_batches_sorted_eq; use datafusion_common::assert_batches_eq; use datafusion_common::create_array; -use datafusion_common::record_batch; use datafusion_expr::col; use datafusion_expr::lit; use datafusion_functions::expr_fn::get_field; diff --git a/vortex-datafusion/src/v2/source.rs b/vortex-datafusion/src/v2/source.rs index 325f8eae92f..6437cb24dc5 100644 --- a/vortex-datafusion/src/v2/source.rs +++ b/vortex-datafusion/src/v2/source.rs @@ -81,6 +81,7 @@ use datafusion_common::Statistics; use datafusion_common::arrow::array::AsArray; use datafusion_common::arrow::array::RecordBatch; use datafusion_common::stats::Precision as DFPrecision; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_datasource::source::DataSource; use datafusion_execution::SendableRecordBatchStream; use datafusion_execution::TaskContext; @@ -362,7 +363,7 @@ impl fmt::Debug for VortexDataSource { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { f.debug_struct("VortexScanSource") .field("schema", &self.leftover_schema) - .field("projection", &format!("{}", &self.projected_projection)) + .field("projection", &format!("{}", self.projected_projection)) .field("filter", &self.filter.as_ref().map(|e| format!("{}", e))) .field("limit", &self.limit) .finish() @@ -652,6 +653,19 @@ impl DataSource for VortexDataSource { .with_updated_node(Arc::new(this) as _), ) } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> DFResult, + ) -> DFResult { + datafusion_physical_plan::apply_expression_roots( + self.leftover_projection + .iter() + .flat_map(|projection| projection.iter()) + .map(|expr| &expr.expr), + f, + ) + } } /// Convert a Vortex [`Option`] to a DataFusion diff --git a/vortex-duckdb/Cargo.toml b/vortex-duckdb/Cargo.toml index 59dbe8d4003..fd611af030d 100644 --- a/vortex-duckdb/Cargo.toml +++ b/vortex-duckdb/Cargo.toml @@ -38,10 +38,19 @@ static_assertions = { workspace = true } tracing = { workspace = true } tracing-subscriber = { workspace = true } url = { workspace = true } -vortex = { workspace = true, features = ["files", "tokio", "object_store"] } -vortex-geo = { workspace = true } +vortex = { workspace = true, features = [ + "files", + "hf", + "tokio", + "object_store", + "object_store_registry", +] } +vortex-spatial = { workspace = true } vortex-utils = { workspace = true, features = ["dashmap"] } +[features] +opendal = ["vortex/opendal"] + [dev-dependencies] anyhow = { workspace = true } divan = { workspace = true } diff --git a/vortex-duckdb/README.md b/vortex-duckdb/README.md index 5d93d1011a0..3efda0af58a 100644 --- a/vortex-duckdb/README.md +++ b/vortex-duckdb/README.md @@ -1,80 +1,74 @@ -# Vortex DuckDB - -Rust bindings for DuckDB. Supports DuckDB precompiled libraries for fast builds and from source builds for debugging. - -## Prerequisites - -- **Ninja**: `brew install ninja` (macOS) | `apt-get install ninja-build` (Ubuntu) -- **CMake**: `brew install cmake` (macOS) | `apt-get install cmake` (Ubuntu) -- **C++20 compatible compiler**: GCC or Clang - -## Build Modes - -### Default (Release) - -Link against the precompiled DuckDB release build. - -```bash -cargo build -p vortex-duckdb -``` - -### Debug Build - -Opt into DuckDB debug build: `VX_DUCKDB_DEBUG=1`. - -```bash -VX_DUCKDB_DEBUG=1 cargo build -p vortex-duckdb -``` - -### AddressSanitizer & ThreadSanitizer - -Enable both ASAN & TSAN: `VX_DUCKDB_SAN=1`. - -```bash -VX_DUCKDB_DEBUG=1 VX_DUCKDB_SAN=1 cargo build -p vortex-duckdb -``` - -## Environment Variables - -| Variable | Effect | -| ----------------- | ------------------------------- | -| `VX_DUCKDB_DEBUG` | Build from source in debug mode | -| `VX_DUCKDB_ASAN` | Enable AddressSanitizer | - -## Running Tests - -```bash -# By default, link against the precompiled DuckDB release build. -cargo test -p vortex-duckdb - -# Link against the DuckDB debug build from source. -VX_DUCKDB_DEBUG=1 cargo test -p vortex-duckdb - -# Link against the DuckDB debug build from source with ASAN & TSAN. -ASAN_OPTIONS=detect_container_overflow=0 VX_DUCKDB_DEBUG=1 VX_DUCKDB_SAN=1 cargo test -p vortex-duckdb -``` - -## Testing the extension with DuckDB - -By default, our tests use a precompiled build which means you don't get an -.extension which you can load in DuckDB. If you want to test a full setup, - -1. Clone [duckdb-vortex](https://github.com/vortex-data/duckdb-vortex) - repository. - -2. If there is an api difference between duckdb-vortex's duckdb submodule and - vortex's vortex-duckdb/duckdb submodule, checkout duckdb-vortex to previous - commit. For example, if duckdb-vortex's HEAD uses 1.6 API but vortex's HEAD - uses 1.5.2, checkout duckdb-vortex at 8a41ee6ebd9. - -3. Update duckdb-vortex's submodules. Replace vortex/ submodule by a softlink to - your local vortex repository. -4. Inside duckdb-vortex, run make -j. - -./target/release/duckdb will be a duckdb instance with vortex-duckdb already -loaded. - -## Testing a custom DuckDB tag - -Change `DUCKDB_VERSION` environment variable value to a preferred hash or commit -(local build), or change build.rs (for testing in CI). +# DuckDB extension for Vortex + +User documentation is on the [docs website](https://docs.vortex.dev/user-guide/duckdb). + +## Build + +Install `ninja`, `cmake`, and a C++20 compiler. + +- For debug builds, launch `VX_DUCKDB_DEBUG=1 cargo build -p vortex-duckdb`. +- For ASAN/TSAN builds, launch `VX_DUCKDB_DEBUG=1 VX_DUCKDB_SAN=1 cargo build -p + vortex-duckdb`. +- For builds against a custom DuckDB tag or commit, change `build.rs`'s + `DEFAULT_DUCKDB_VERSION`. + +## Test + +Our sqllogic tests use a build which links against `libduckdb.so`. This means +you don't get an extension file to load in DuckDB. If you want to test just +sqllogic, run `cargo test -p vortex-sqllogictest --test sqllogictests`. If you +want to test a full setup, + +1. Clone [duckdb-vortex](https://github.com/vortex-data/duckdb-vortex). +2. Update duckdb-vortex's submodules. Update vortex reference in duckdb-vortex's + `vortex-extension/Cargo.toml` with this Vortex checkout. +4. Inside duckdb-vortex, run `make reldebug -j`. + +`./build/reldebug/duckdb` will be a DuckDB shell with Vortex linked statically. + +## Update extension with same DuckDB version + +1. Merge the PR with your changes, if any. +2. Make a new Vortex release. +3. Get the commit of the release, i.e. + `afb005379dd9a1b7dc7ae2cb49f4f465d91cc2b3` is the commit for + [0.85](https://github.com/vortex-data/vortex/releases/tag/0.85.0). +4. In duckdb-vortex, update vortex reference in `vortex-extension/Cargo.toml` to + the release commit hash. +5. Open a PR in duckdb-vortex, check the tests and merge it. +6. Ensure the PR has been backported to the current release branch, i.e. for + duckdb 1.5.* the release branch is + [v1.5-variegata](https://github.com/vortex-data/duckdb-vortex/tree/v1.5-variegata). + This should happen automatically. +7. Notify DuckDB maintainers in the Discord `#spiraldb` channel of the new + commit and ask to update the extension version. This will likely take some + days since they need to bump the extension version, trigger a deploy job and + update the artifacts on their website. Provide the commit in duckdb-vortex's + release branch, i.e. + [`5390e62`](https://github.com/vortex-data/duckdb-vortex/commit/5390e6270da964f04afbd6df4a5c9bc83a2e7880). +8. Check `UPDATE EXTENSIONS; INSTALL vortex; LOAD vortex` in your local duckdb + loads the latest version. + +## Update duckdb version + +- Update the version in `build.rs`'s `DEFAULT_DUCKDB_VERSION`. +- Create a PR with your changes and ensure it builds. Merge the PR. If your PR + is for a pre-release build, CI will build DuckDB from source and upload the + artifacts to R2 so you can reuse them locally. +- Create a new Vortex release. +- In duckdb-vortex, change all occurrences of old version to the new version. See + [1.5.5](https://github.com/vortex-data/duckdb-vortex/pull/105/changes) upgrade + as an example. +- In duckdb-vortex, upgrade submodules to the new version. +- Follow from step (3) in extension update plan. + +## Build artifacts storage + +CI uses duckdb prebuilt (`.so/.a/.dylib`) artifacts from Github releases. +However, release downloads from Github are very unstable, and we're observed a +~25% build rate failure when doing so. Thus we mirror release archives to R2. + +A job is triggered on every PR which checks artifacts are present in R2, and if +they aren't, tries to mirror them from DuckDB Github releases via a deployment +request. If `DEFAULT_DUCKDB_VERSION` is a hash and files weren't present, the +job builds DuckDB from source and uploads the artifacts to R2 as well. diff --git a/vortex-duckdb/build.rs b/vortex-duckdb/build.rs index b4b6c0057ff..5630312e2ed 100644 --- a/vortex-duckdb/build.rs +++ b/vortex-duckdb/build.rs @@ -26,8 +26,11 @@ const DUCKDB_SOURCE_COMMIT_URL: &str = "https://github.com/duckdb/duckdb/archive const DEFAULT_DUCKDB_VERSION: &str = "1.5.5"; const BUILD_ARTIFACTS: [&str; 3] = ["libduckdb.dylib", "libduckdb.so", "libduckdb_static.a"]; +const BUILD_MARKER: &str = ".vx-build-complete"; +const DUCKDB_CACHE_DIR: &str = "vortex-duckdb-cache"; +const EXTRACT_MARKER: &str = ".vx-extract-complete"; -const SOURCE_FILES: [&str; 11] = [ +const SOURCE_FILES: [&str; 12] = [ "cpp/vortex_duckdb.cpp", "cpp/copy_function.cpp", "cpp/expr.cpp", @@ -37,13 +40,14 @@ const SOURCE_FILES: [&str; 11] = [ "cpp/cast_pushdown.cpp", "cpp/aggregate_fn_pushdown.cpp", "cpp/table_filter.cpp", + "cpp/multi_file_reader.cpp", "cpp/table_function.cpp", "cpp/vector.cpp", ]; // Duckdb C API function we use. // This lowers codegen'd src/cpp.rs by four times. -const DUCKDB_C_API_FUNCTIONS: [&str; 133] = [ +const DUCKDB_C_API_FUNCTIONS: [&str; 135] = [ "duckdb_array_type_array_size", "duckdb_array_type_child_type", "duckdb_array_vector_get_child", @@ -65,6 +69,7 @@ const DUCKDB_C_API_FUNCTIONS: [&str; 133] = [ "duckdb_create_decimal_type", "duckdb_create_double", "duckdb_create_float", + "duckdb_create_hugeint", "duckdb_create_int16", "duckdb_create_int32", "duckdb_create_int64", @@ -76,6 +81,7 @@ const DUCKDB_C_API_FUNCTIONS: [&str; 133] = [ "duckdb_create_selection_vector", "duckdb_create_struct_type", "duckdb_create_time", + "duckdb_create_time_ns", "duckdb_create_timestamp", "duckdb_create_timestamp_ms", "duckdb_create_timestamp_ns", @@ -315,6 +321,29 @@ fn download_url(url: &str, path: &Path) { } } +fn env_true(key: &str) -> bool { + env::var(key).is_ok_and(|v| matches!(v.as_str(), "1" | "true")) +} + +fn duckdb_cache_root(out_dir: &Path) -> PathBuf { + out_dir + .parent() + .unwrap() + .parent() + .unwrap() + .join(DUCKDB_CACHE_DIR) +} + +fn clear_dir(dir: &Path) { + if let Err(err) = fs::remove_dir_all(dir) + && err.kind() != io::ErrorKind::NotFound + { + println!("cargo:error=Failed to clear {}: {err}", dir.display()); + exit(1); + } + fs::create_dir_all(dir).unwrap(); +} + fn extract(archive: &Path, dest: &Path) { println!( "cargo:info=Extracting {} to {}", @@ -325,11 +354,52 @@ fn extract(archive: &Path, dest: &Path) { zip::ZipArchive::new(file).unwrap().extract(dest).unwrap(); } +fn git_apply(repo_dir: &Path, patch: &Path, args: &[&str]) -> bool { + let output = Command::new("git") + .current_dir(repo_dir) + .args(["apply", "-p1"]) + .args(args) + .arg(patch) + .output(); + match output { + Ok(out) => out.status.success(), + Err(e) => { + println!("cargo:error=git is required to patch DuckDB sources: {e}"); + exit(1); + } + } +} + +fn apply_source_patches(crate_dir: &Path, repo_dir: &Path) { + let mut patches: Vec = fs::read_dir(crate_dir.join("patches")) + .unwrap() + .map(|entry| entry.unwrap().path()) + .filter(|path| path.extension().is_some_and(|ext| ext == "diff")) + .collect(); + patches.sort(); + + for patch in patches { + // A successful reverse dry-run means the patch is already applied. + if git_apply(repo_dir, &patch, &["--check", "--reverse"]) { + continue; + } + if !git_apply(repo_dir, &patch, &[]) { + println!( + "cargo:error=Failed to apply {} to {}; delete that directory to re-extract \ + DuckDB sources", + patch.display(), + repo_dir.display() + ); + exit(1); + } + println!("cargo:info=Applied {}", patch.display()); + } +} + /// Download DuckDB library archive from R2 and extract it. /// Return false if archive is not available or download failed -fn download(version: &DuckDBVersion, library_dir: &Path) -> bool { - let target = env::var("TARGET").unwrap(); - let (platform, arch) = match target.as_str() { +fn download_prebuilt(version: &DuckDBVersion, library_dir: &Path, target: &str) -> bool { + let (platform, arch) = match target { "aarch64-apple-darwin" | "x86_64-apple-darwin" => ("osx", "universal"), "x86_64-unknown-linux-gnu" => ("linux", "amd64"), "aarch64-unknown-linux-gnu" => ("linux", "arm64"), @@ -342,19 +412,21 @@ fn download(version: &DuckDBVersion, library_dir: &Path) -> bool { let archive_name = format!("libduckdb-{platform}-{arch}.zip"); let url = format!("{DUCKDB_RELEASES_URL}/{version}/{archive_name}"); let archive_path = library_dir.join(&archive_name); + let extract_marker = library_dir.join(EXTRACT_MARKER); - fs::create_dir_all(library_dir).unwrap(); + if extract_marker.exists() { + drop(fs::remove_file(&archive_path)); + return true; + } + + clear_dir(library_dir); if !try_download_url(&url, &archive_path) { return false; } - let duckdb_lib_dir = archive_path.parent().unwrap().to_path_buf(); - for artifact in BUILD_ARTIFACTS { - if duckdb_lib_dir.join(artifact).exists() { - return true; - } - } - extract(&archive_path, &duckdb_lib_dir); + extract(&archive_path, library_dir); + fs::remove_file(&archive_path).unwrap(); + fs::write(&extract_marker, format!("{version}\n{target}\n")).unwrap(); true } @@ -369,12 +441,11 @@ fn build_duckdb(version: &DuckDBVersion, duckdb_repo_dir: &Path) { } println!("cargo:info=Building DuckDB from source (this may take a while)..."); - let (asan_option, tsan_option) = - if env::var("VX_DUCKDB_SAN").is_ok_and(|v| matches!(v.as_str(), "1" | "true")) { - ("0", "1") // DISABLE_SANITIZER=0 enables ASAN, THREADSAN=1 enables TSAN - } else { - ("1", "0") - }; + let (asan_option, tsan_option) = if env_true("VX_DUCKDB_SAN") { + ("0", "1") // DISABLE_SANITIZER=0 enables ASAN, THREADSAN=1 enables TSAN + } else { + ("1", "0") + }; // If we're building from a commit we need to build benchmark // extensions statically, otherwise DuckDB tries to load them from an http @@ -415,34 +486,16 @@ fn try_build_duckdb( version: &DuckDBVersion, build_type: &str, ) { - let inner_dir_name = version.archive_inner_dir_name(); - let repo_dir = source_dir.join(&inner_dir_name); - let build_dir = repo_dir.join("build").join(build_type); - let build_src_dir = build_dir.join("src"); - - let mut build = true; - for artifact in BUILD_ARTIFACTS { - let path = build_src_dir.join(artifact); - if path.exists() { - println!("cargo:info=Found {artifact} in {}", path.display()); - build = false; - break; - } - } - - if build { - build_duckdb(version, &repo_dir); + let repo_dir = source_dir.join(version.archive_inner_dir_name()); + let library_marker = library_dir.join(BUILD_MARKER); + if library_marker.exists() { + return; } - let library_dir_str = library_dir.display(); - if let Err(err) = fs::remove_dir_all(library_dir) - && err.kind() != io::ErrorKind::NotFound - { - println!("cargo:error=Failed to remove {library_dir_str}: {err}"); - exit(1); - }; - fs::create_dir_all(library_dir).unwrap(); + build_duckdb(version, &repo_dir); + clear_dir(library_dir); + let build_src_dir = repo_dir.join("build").join(build_type).join("src"); let mut found_artifact = false; for artifact in BUILD_ARTIFACTS { let src = build_src_dir.join(artifact); @@ -459,6 +512,7 @@ fn try_build_duckdb( println!("cargo:error=Failed to find any of {artifacts} after build"); exit(1); } + fs::write(&library_marker, format!("{version}\n{build_type}\n")).unwrap(); } /// Generate rust functions with bindgen from C sources. @@ -515,10 +569,25 @@ fn bindgen_c2rust(crate_dir: &Path, duckdb_include_dir: &Path) { /// Generate libvortex_duckdb.* fn compile_cpp(duckdb_include_dir: &Path) { - cc::Build::new() + let mut build = cc::Build::new(); + let has_debuginfo = env::var("DEBUG") + .map(|v| !matches!(v.as_str(), "false" | "0" | "none" | "")) + .unwrap_or(false); + if env_true("VX_DUCKDB_DEBUG") || has_debuginfo { + build.define("DEBUG", None); + } else { + build.define("NDEBUG", None); + } + build .std("c++20") .flags(["-Wall", "-Wextra", "-Wpedantic", "-Werror"]) .cpp(true) + // Duckdb 1.5.5 uses C++11. spatial_overrides.o uses + // duckdb::ScalarFunctionCatalogEntry::Name which is constexpr but not + // inline. Our code uses C++20 where constexpr implies inline. GCC + // emits this symbol with STB_GNU_UNIQUE and this conflicts on link stage + // in duckdb-vortex where libvortex_duckdb.a is linked statically + .flag_if_supported("-fno-gnu-unique") // We don't want compiler warnings inside duckdb headers, pass as flags .flag("-isystem") .flag(duckdb_include_dir) @@ -558,10 +627,41 @@ fn cbindgen_rust2c(crate_dir: &Path) { } } +fn git(crate_dir: &Path, args: &[&str]) -> Option { + let output = Command::new("git") + .arg("-C") + .arg(crate_dir) + .args(args) + .output() + .ok()?; + if !output.status.success() { + return None; + } + let stdout = String::from_utf8(output.stdout).ok()?; + let stdout = stdout.trim(); + (!stdout.is_empty()).then(|| stdout.to_owned()) +} + +fn vortex_version(crate_dir: &Path) { + println!("cargo:rerun-if-env-changed=VORTEX_VERSION"); + let version = env::var("VORTEX_VERSION") + .ok() + .filter(|version| !version.is_empty()) + // If this commit belongs to a tag + .or_else(|| git(crate_dir, &["describe", "--tags", "--exact-match", "HEAD"])) + // If this commit doesn't belong to a tag + .or_else(|| git(crate_dir, &["rev-parse", "HEAD"])) + // No version, can't build + .unwrap(); + println!("cargo:rustc-env=VORTEX_VERSION={version}"); +} + fn main() { println!("cargo:rerun-if-changed=cpp/include"); + println!("cargo:rerun-if-changed=patches"); println!("cargo:rerun-if-env-changed=VX_DUCKDB_DEBUG"); println!("cargo:rerun-if-env-changed=VX_DUCKDB_SAN"); + println!("cargo:rerun-if-env-changed=DEBUG"); println!("cargo:rerun-if-env-changed=CARGO_HTTP_TIMEOUT"); println!("cargo:rerun-if-env-changed=HTTP_TIMEOUT"); println!("cargo:rerun-if-env-changed=TARGET"); @@ -582,6 +682,7 @@ fn main() { // in vortex's CI. let crate_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + vortex_version(&crate_dir); if let Some(source_dir) = env::var_os("DUCKDB_SOURCE_DIR") { let source_dir = PathBuf::from(source_dir); let duckdb_include_dir = source_dir.join("src").join("include"); @@ -589,6 +690,7 @@ fn main() { "cargo:info=Using DuckDB source from DUCKDB_SOURCE_DIR={}", source_dir.display() ); + apply_source_patches(&crate_dir, &source_dir); bindgen_c2rust(&crate_dir, &duckdb_include_dir); cbindgen_rust2c(&crate_dir); compile_cpp(&duckdb_include_dir); @@ -609,84 +711,93 @@ fn main() { let duckdb_dir = crate_dir.join("duckdb"); let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap()); - let library_dir = out_dir.join(format!("duckdb-lib-{version}")); - - let library_dir_str = library_dir.display(); - println!("cargo:rustc-link-search=native={library_dir_str}"); - println!("cargo:rustc-link-lib=dylib=duckdb"); - - // Set rpath for binaries built directly from this crate. This is not - // inherited by downstream crates. - println!("cargo:rustc-link-arg=-Wl,-rpath,{library_dir_str}"); + let target = env::var("TARGET").unwrap(); + let cache_root = duckdb_cache_root(&out_dir); + fs::create_dir_all(&cache_root).unwrap(); - // Export the library path for downstream crates via the `links` manifest key. - // Downstream crates can access this via `env::var("DEP_DUCKDB_LIB_DIR")` in their build.rs - // and add their own rpath: - // - // if let Ok(duckdb_lib) = env::var("DEP_DUCKDB_LIB_DIR") { - // println!("cargo:rustc-link-arg=-Wl,-rpath,{duckdb_lib}"); - // } - // - // Alternatively, set LD_LIBRARY_PATH (Linux) or DYLD_LIBRARY_PATH (macOS) at runtime. - println!("cargo:lib_dir={library_dir_str}"); + let debug_duckdb = env_true("VX_DUCKDB_DEBUG"); + let build_type = if debug_duckdb { "debug" } else { "release" }; + let sanitizer_suffix = if env_true("VX_DUCKDB_SAN") { + "-san" + } else { + "" + }; - let source_dir = out_dir.join(format!("duckdb-source-{version}")); + let source_dir = cache_root.join(format!("duckdb-source-{version}")); let source_archive_url = match &version { DuckDBVersion::Release(v) => format!("{DUCKDB_SOURCE_RELEASE_URL}/v{v}.zip"), DuckDBVersion::Commit(c) => format!("{DUCKDB_SOURCE_COMMIT_URL}/{c}.zip"), }; - let source_archive_path = source_dir.with_extension("zip"); - download_url(&source_archive_url, &source_archive_path); - let inner_dir = source_dir.join(version.archive_inner_dir_name()); - let extract_marker = source_dir.join(".vx-extract-complete"); - if !extract_marker.exists() { - if let Err(err) = fs::remove_dir_all(&source_dir) - && err.kind() != io::ErrorKind::NotFound - { - println!( - "cargo:error=Failed to clear {}: {err}", - source_dir.display() - ); - exit(1); - } - fs::create_dir_all(&source_dir).unwrap(); + let extract_marker = source_dir.join(EXTRACT_MARKER); + let source_archive_path = cache_root.join(format!("duckdb-source-{version}.zip")); + if extract_marker.exists() { + drop(fs::remove_file(&source_archive_path)); + } else { + download_url(&source_archive_url, &source_archive_path); + clear_dir(&source_dir); extract(&source_archive_path, &source_dir); + fs::remove_file(&source_archive_path).unwrap(); fs::write(&extract_marker, version.to_string()).unwrap(); } + apply_source_patches(&crate_dir, &inner_dir); + drop(fs::remove_file(&duckdb_dir)); drop(fs::remove_dir_all(&duckdb_dir)); symlink(&source_dir, &duckdb_dir).unwrap(); - let has_debug_env = - env::var("VX_DUCKDB_DEBUG").is_ok_and(|v| matches!(v.as_str(), "1" | "true")); - let build_type = match has_debug_env { - true => "debug", - false => "release", - }; println!("cargo:info=building DuckDB in {build_type} mode"); - if has_debug_env { - try_build_duckdb(&source_dir, &library_dir, &version, build_type); + let prebuilt_library_dir = cache_root.join(format!("duckdb-lib-{version}-{target}-prebuilt")); + let source_library_dir = cache_root.join(format!( + "duckdb-lib-{version}-{target}-{build_type}{sanitizer_suffix}" + )); + + let library_dir = if debug_duckdb { + try_build_duckdb(&source_dir, &source_library_dir, &version, build_type); + source_library_dir } else { match &version { DuckDBVersion::Release(_) => { - if !download(&version, &library_dir) { + if !download_prebuilt(&version, &prebuilt_library_dir, &target) { println!("cargo:error=DuckDB release {version} not available in R2"); exit(1); } + prebuilt_library_dir } DuckDBVersion::Commit(_) => { - if !download(&version, &library_dir) { + if download_prebuilt(&version, &prebuilt_library_dir, &target) { + prebuilt_library_dir + } else { println!("cargo:info=DuckDB commit {version} not in R2, building from source"); - try_build_duckdb(&source_dir, &library_dir, &version, build_type); + try_build_duckdb(&source_dir, &source_library_dir, &version, build_type); + source_library_dir } } } }; + let library_dir_str = library_dir.display(); + println!("cargo:rustc-link-search=native={library_dir_str}"); + println!("cargo:rustc-link-lib=dylib=duckdb"); + + // Set rpath for binaries built directly from this crate. This is not + // inherited by downstream crates. + println!("cargo:rustc-link-arg=-Wl,-rpath,{library_dir_str}"); + + // Export the library path for downstream crates via the `links` manifest key. + // Downstream crates can access this via `env::var("DEP_DUCKDB_LIB_DIR")` in their build.rs + // and add their own rpath: + // + // if let Ok(duckdb_lib) = env::var("DEP_DUCKDB_LIB_DIR") { + // println!("cargo:rustc-link-arg=-Wl,-rpath,{duckdb_lib}"); + // } + // + // Alternatively, set LD_LIBRARY_PATH (Linux) or DYLD_LIBRARY_PATH (macOS) at runtime. + println!("cargo:lib_dir={library_dir_str}"); + let duckdb_include_dir = inner_dir.join("src").join("include"); bindgen_c2rust(&crate_dir, &duckdb_include_dir); cbindgen_rust2c(&crate_dir); diff --git a/vortex-duckdb/cpp/aggregate_fn_pushdown.cpp b/vortex-duckdb/cpp/aggregate_fn_pushdown.cpp index 2caa3463d53..4163c4db898 100644 --- a/vortex-duckdb/cpp/aggregate_fn_pushdown.cpp +++ b/vortex-duckdb/cpp/aggregate_fn_pushdown.cpp @@ -87,22 +87,26 @@ LogicalOperatorPtr TryReplaceAggregate(ClientContext &context, return op; } - // GET now returns one column per aggregate. Expand existing columns - auto &column_ids = get->GetMutableColumnIds(); - get->types.resize(aggregates_len); - get->returned_types.resize(aggregates_len); - column_ids.resize(aggregates_len); - vector names(aggregates_len); // need a copy because we reference original names + const vector original_column_ids = get->GetColumnIds(); for (idx_t i = 0; i < aggregates_len; i++) { const auto &[column_index, expr] = input[i]; if (column_index == COUNT_STAR_PROJ_IDX) { names[i] = "count_star()"; } else { - const TableColumnStorageIndex storage_index = get->GetColumnIds()[column_index].GetPrimaryIndex(); + const TableColumnStorageIndex storage_index = original_column_ids[column_index].GetPrimaryIndex(); names[i] = get->names[storage_index]; } + } + + // GET now returns one column per aggregate. + auto &column_ids = get->GetMutableColumnIds(); + get->types.resize(aggregates_len); + get->returned_types.resize(aggregates_len); + column_ids.resize(aggregates_len); + for (idx_t i = 0; i < aggregates_len; i++) { + const auto &[_, expr] = input[i]; get->types[i] = expr.return_type; get->returned_types[i] = expr.return_type; column_ids[i] = ColumnIndex {i}; @@ -136,5 +140,5 @@ LogicalGet *GetChildGet(const LogicalAggregate &agg) { return nullptr; } LogicalGet &get = op->Cast(); - return get.function.bind == duckdb_vx_table_function_bind ? &get : nullptr; + return is_vortex_scan(get.function) ? &get : nullptr; } diff --git a/vortex-duckdb/cpp/cast_pushdown.cpp b/vortex-duckdb/cpp/cast_pushdown.cpp index 3acd5b04161..d1a3b849c00 100644 --- a/vortex-duckdb/cpp/cast_pushdown.cpp +++ b/vortex-duckdb/cpp/cast_pushdown.cpp @@ -25,7 +25,7 @@ static bool ReachesPushdownGet(const LogicalOperator &op) { cur = cur->children[0].get(); switch (cur->type) { case LogicalOperatorType::LOGICAL_GET: - return cur->Cast().function.bind == duckdb_vx_table_function_bind; + return is_vortex_scan(cur->Cast().function); case LogicalOperatorType::LOGICAL_PROJECTION: continue; default: diff --git a/vortex-duckdb/cpp/copy_function.cpp b/vortex-duckdb/cpp/copy_function.cpp index 4740b92bbc3..840d159f266 100644 --- a/vortex-duckdb/cpp/copy_function.cpp +++ b/vortex-duckdb/cpp/copy_function.cpp @@ -1,64 +1,52 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "copy_function.hpp" #include "data.hpp" #include "error.hpp" #include "vortex_duckdb.h" #include "table_function.h" #include "vortex.h" -#include "duckdb/function/copy_function.hpp" +#include "duckdb/common/types/column/column_data_collection.hpp" #include "duckdb/main/capi/capi_internal.hpp" #include "duckdb/main/client_context.hpp" #include "duckdb/main/connection.hpp" +#include "duckdb/parser/keyword_helper.hpp" #include "duckdb/parser/parsed_data/create_copy_function_info.hpp" -using namespace duckdb; - -struct CopyBindData final : TableFunctionData { - CopyBindData(unique_ptr ffi_data) : ffi_data(std::move(ffi_data)) { - } - unique_ptr ffi_data; -}; - -struct CopyGlobalData final : GlobalFunctionData { - CopyGlobalData(unique_ptr ffi_data) : ffi_data(std::move(ffi_data)) { - } - - unique_ptr ffi_data; -}; - unique_ptr copy_to_bind(ClientContext &, CopyFunctionBindInput &, - const vector &column_names, - const vector &column_types) { - vector ffi_column_names(column_names.size()); - for (size_t i = 0; i < column_names.size(); ++i) { - ffi_column_names[i] = column_names[i].c_str(); + const vector &names, + const vector &types) { + vector ffi_names(names.size()); + for (size_t i = 0; i < names.size(); ++i) { + ffi_names[i] = names[i].c_str(); } - vector ffi_column_types(column_types.size()); - for (size_t i = 0; i < column_types.size(); ++i) { + vector ffi_types(types.size()); + for (size_t i = 0; i < types.size(); ++i) { // duckdb C api doesn't allow passing const LogicalTypes. We never // modify input in copy function. - ffi_column_types[i] = - reinterpret_cast(const_cast(&column_types[i])); + ffi_types[i] = reinterpret_cast(const_cast(&types[i])); } duckdb_vx_error error_out = nullptr; - const duckdb_vx_data ffi_bind_data = duckdb_copy_function_copy_to_bind(ffi_column_names.data(), - ffi_column_names.size(), - ffi_column_types.data(), - ffi_column_types.size(), + const duckdb_vx_data ffi_bind_data = duckdb_copy_function_copy_to_bind(ffi_names.data(), + ffi_names.size(), + ffi_types.data(), + ffi_types.size(), &error_out); if (error_out) { throw BinderException(IntoErrString(error_out)); } auto cdata = unique_ptr(reinterpret_cast(ffi_bind_data)); - return make_uniq(std::move(cdata)); + return make_uniq(std::move(cdata), names); } unique_ptr copy_to_initialize_global(ClientContext &, FunctionData &bind_data, const string &file_path) { - void *const ffi_bind = bind_data.Cast().ffi_data->DataPtr(); + const VortexCopyBindData &bind = bind_data.Cast(); + const void *const ffi_bind = bind.ffi_bind->DataPtr(); duckdb_vx_error error_out = nullptr; const duckdb_vx_data ffi_global = @@ -68,7 +56,7 @@ copy_to_initialize_global(ClientContext &, FunctionData &bind_data, const string } auto cdata = unique_ptr(reinterpret_cast(ffi_global)); - return make_uniq(std::move(cdata)); + return make_uniq(std::move(cdata)); } void copy_to_sink(ExecutionContext &, @@ -76,9 +64,13 @@ void copy_to_sink(ExecutionContext &, GlobalFunctionData &gstate, LocalFunctionData &, DataChunk &input) { - void *const ffi_bind = bind_data.Cast().ffi_data->DataPtr(); - void *const ffi_global = gstate.Cast().ffi_data->DataPtr(); - auto ffi_chunk = reinterpret_cast(&input); + const VortexCopyBindData &bind = bind_data.Cast(); + const VortexCopyGlobalState &global = gstate.Cast(); + + const void *const ffi_bind = bind.ffi_bind->DataPtr(); + const void *const ffi_global = global.ffi_global->DataPtr(); + + duckdb_data_chunk ffi_chunk = reinterpret_cast(&input); duckdb_vx_error error_out = nullptr; duckdb_copy_function_copy_to_sink(ffi_bind, ffi_global, ffi_chunk, &error_out); if (error_out) { @@ -86,13 +78,114 @@ void copy_to_sink(ExecutionContext &, } } -void copy_to_finalize(ClientContext &, FunctionData &, GlobalFunctionData &gstate) { - void *const ffi_global = gstate.Cast().ffi_data->DataPtr(); +// CopyToFileInfo::file_stats is owned by the operator's sink state, which outlives gstate. +void copy_to_get_written_statistics(ClientContext &, + FunctionData &, + GlobalFunctionData &gstate, + CopyFunctionFileStatistics &statistics) { + gstate.Cast().written_stats = &statistics; +} + +void copy_to_finalize(ClientContext &, FunctionData &bind_data, GlobalFunctionData &gstate) { + auto &global = gstate.Cast(); + void *const ffi_global = global.ffi_global->DataPtr(); duckdb_vx_error error_out = nullptr; duckdb_copy_function_copy_to_finalize(ffi_global, &error_out); if (error_out) { throw ExecutorException(IntoErrString(error_out)); } + + if (!global.written_stats) { + return; + } + auto &names = bind_data.Cast().column_names; + duckdb_vx_written_file_statistics file_stats; + if (!duckdb_copy_function_get_written_file_statistics(ffi_global, &file_stats)) { + // Statistics were requested (written_stats is set) but the finished write produced none; + // that is an internal inconsistency, not a silently empty result. + throw InternalException("vortex COPY: written statistics were requested but not produced"); + } + if (file_stats.num_columns != names.size()) { + throw InternalException("vortex COPY: %llu statistics columns for %llu written columns", + file_stats.num_columns, + names.size()); + } + D_ASSERT(global.written_stats != nullptr); + global.written_stats->row_count = file_stats.row_count; + global.written_stats->file_size_bytes = file_stats.file_size_bytes; + global.written_stats->footer_size_bytes = Value::UBIGINT(file_stats.footer_size_bytes); + // Keyed by top-level column name only. The vortex footer reports one statistics set per + // top-level field, so nested struct/list leaf columns get no statistics here (unlike parquet, + // which recurses to leaf paths). Flat tables are fully covered. + for (idx_t i = 0; i < file_stats.num_columns; i++) { + duckdb_vx_written_column_statistics col_stats {}; + duckdb_vx_error col_error = nullptr; + if (!duckdb_copy_function_get_written_column_statistics(ffi_global, i, &col_stats, &col_error)) { + if (col_error) { + throw ExecutorException(IntoErrString(col_error)); + } + throw InternalException("vortex COPY: no statistics for column %llu after finalize", i); + } + case_insensitive_map_t column; + column["num_values"] = Value::UBIGINT(col_stats.num_values); + if (col_stats.has_column_size) { + column["column_size_bytes"] = Value::UBIGINT(col_stats.column_size_bytes); + } + if (col_stats.has_null_count) { + column["null_count"] = Value::UBIGINT(col_stats.null_count); + } + if (col_stats.min) { + column["min"] = Value(reinterpret_cast(col_stats.min)->ToString()); + duckdb_destroy_value(&col_stats.min); + } + if (col_stats.max) { + column["max"] = Value(reinterpret_cast(col_stats.max)->ToString()); + duckdb_destroy_value(&col_stats.max); + } + if (col_stats.has_nan_stat) { + column["has_nan"] = Value::BOOLEAN(col_stats.contains_nan); + } + // DuckLake keys column statistics by a quoted, dot-separated path (see + // DuckLakeUtil::ParseQuotedList); match the parquet writer, which quotes each name. + global.written_stats->column_statistics.emplace(KeywordHelper::WriteQuoted(names[i], '"'), + std::move(column)); + } +} + +unique_ptr copy_to_prepare_batch(ClientContext &, + FunctionData &bind_data, + GlobalFunctionData &, + unique_ptr collection) { + const VortexCopyBindData &bind = bind_data.Cast(); + + const void *const ffi_bind = bind.ffi_bind->DataPtr(); + auto ffi_batch = unique_ptr(reinterpret_cast(duckdb_copy_function_prepare_batch_new())); + duckdb_vx_error error_out = nullptr; + + for (DataChunk &chunk : collection->Chunks()) { + duckdb_data_chunk ffi_chunk = reinterpret_cast(&chunk); + duckdb_copy_function_prepare_batch_push(ffi_bind, ffi_batch->DataPtr(), ffi_chunk, &error_out); + if (error_out) { + throw ExecutorException(IntoErrString(error_out)); + } + } + return make_uniq(std::move(ffi_batch)); +} + +void copy_to_flush_batch(ClientContext &, + FunctionData &, + GlobalFunctionData &gstate, + PreparedBatchData &batch_data) { + const VortexCopyGlobalState &global = gstate.Cast(); + const VortexCopyPreparedBatchData &batch = batch_data.Cast(); + + const void *const ffi_global = global.ffi_global->DataPtr(); + const void *const ffi_batch = batch.ffi_copy_prepared->DataPtr(); + duckdb_vx_error error_out = nullptr; + duckdb_copy_function_flush_batch(ffi_global, ffi_batch, &error_out); + if (error_out) { + throw ExecutorException(IntoErrString(error_out)); + } } extern "C" duckdb_state duckdb_vx_register_copy_function(duckdb_database ffi_db) { @@ -103,18 +196,30 @@ extern "C" duckdb_state duckdb_vx_register_copy_function(duckdb_database ffi_db) CopyFunction fn("vortex"); fn.copy_to_bind = copy_to_bind; fn.copy_to_initialize_global = copy_to_initialize_global; + // required by duckdb fn.copy_to_initialize_local = [](auto &, auto &) { return make_uniq(); }; fn.copy_to_sink = copy_to_sink; + // required by duckdb for PARTITION_BY + fn.copy_to_combine = [](ExecutionContext &, FunctionData &, GlobalFunctionData &, LocalFunctionData &) { + }; fn.copy_to_finalize = copy_to_finalize; + fn.prepare_batch = copy_to_prepare_batch; + fn.flush_batch = copy_to_flush_batch; + fn.copy_to_get_written_statistics = copy_to_get_written_statistics; fn.extension = "vortex"; - // TODO(joe): expose this via c our api - fn.execution_mode = [](bool, bool) { - return CopyFunctionExecutionMode::REGULAR_COPY_TO_FILE; + fn.execution_mode = [](bool preserve_insertion_order, bool supports_batch_index) { + using enum CopyFunctionExecutionMode; + if (!preserve_insertion_order) { + return PARALLEL_COPY_TO_FILE; + } + if (supports_batch_index) { + return BATCH_COPY_TO_FILE; + } + return REGULAR_COPY_TO_FILE; }; - // TODO(joe): handle parameters as in table_function try { Catalog &system_catalog = Catalog::GetSystemCatalog(db); diff --git a/vortex-duckdb/cpp/include/copy_function.hpp b/vortex-duckdb/cpp/include/copy_function.hpp new file mode 100644 index 00000000000..ec8cccd7048 --- /dev/null +++ b/vortex-duckdb/cpp/include/copy_function.hpp @@ -0,0 +1,32 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#pragma once +#include "data.hpp" +#include "duckdb/function/copy_function.hpp" + +using namespace duckdb; + +struct VortexCopyBindData final : TableFunctionData { + VortexCopyBindData(unique_ptr ffi_bind, vector column_names) + : ffi_bind(std::move(ffi_bind)), column_names(std::move(column_names)) { + } + unique_ptr ffi_bind; + // Column names in write order, used to key WRITTEN_FILE_STATISTICS. + vector column_names; +}; + +struct VortexCopyGlobalState final : GlobalFunctionData { + VortexCopyGlobalState(unique_ptr ffi_global) : ffi_global(std::move(ffi_global)) { + } + unique_ptr ffi_global; + // Non-owning; null when the plan does not request statistics. + CopyFunctionFileStatistics *written_stats = nullptr; +}; + +struct VortexCopyPreparedBatchData final : PreparedBatchData { + VortexCopyPreparedBatchData(unique_ptr ffi_copy_prepared) + : ffi_copy_prepared(std::move(ffi_copy_prepared)) { + } + unique_ptr ffi_copy_prepared; +}; diff --git a/vortex-duckdb/cpp/include/multi_file_reader.hpp b/vortex-duckdb/cpp/include/multi_file_reader.hpp new file mode 100644 index 00000000000..6a62bfbf25c --- /dev/null +++ b/vortex-duckdb/cpp/include/multi_file_reader.hpp @@ -0,0 +1,196 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once + +#include "data.hpp" +#include "table_function.h" +#include "duckdb/common/multi_file/multi_file_function.hpp" + +using namespace duckdb; + +// Convert Vortex statistics to DuckDB statistics +unique_ptr to_duckdb_statistics(duckdb_column_statistics &statistics); + +struct VortexBindData final : TableFunctionData { + VortexBindData() = default; + unique_ptr Copy() const override; + bool Equals(const FunctionData &other) const override; + + unique_ptr ffi_bind_data; +}; + +struct VortexBindResult { + vector &return_types; + vector &names; +}; + +struct VortexGlobalState final : GlobalTableFunctionState { + VortexGlobalState() = default; + ~VortexGlobalState() override = default; + + const void *ffi_bind_data = nullptr; // needed for local state partial accumulation + unique_ptr ffi_global_state; +}; + +struct VortexLocalState final : LocalTableFunctionState { + VortexLocalState() = default; + unique_ptr ffi_local_state; +}; + +struct VortexMultiFileReader final : MultiFileReader { + inline unique_ptr Copy() const override { + return make_uniq(); + } + + /* + * Called after InitializeGlobalState but before TryInitializeScan under + * file-local lock. Used to avoid initializing file splits if footer + * statistics prove false for pushed filter or file index is not present in + * file selection. + * + * TODO(myrrc): + * InitializeReader prunes unused files e.g. ones filtered by "file_index". + * These files are opened but not initialized. Duckdb 2.0 allows filtering + * on virtual columns to skip opening files, but until it's released, + * filtered files (say, second pass in late materialization queries like + * clickbench 23) will be opened, and we will read their footers. + * + * ComplexFilterPushdown/DynamicFilterPushdown callbacks technically allow + * filtering file list via filter, but with file_index specifically it + * doesn't work. + * Say we have a filter "file_index IN (1, 2)" and files 0,1,2. We filter + * first file and return a new list of two files. Duckdb 1.5 then + * renumbers files 1->0, 2->1 (it uses file index). Then it evaluates + * the old filter with new file index and file 1 (with new index 0) is + * filtered out which produces invalid results. + */ + ReaderInitializeType InitializeReader(MultiFileReaderData &reader_data, + const MultiFileBindData &bind_data, + const vector &global_columns, + const vector &global_column_ids, + optional_ptr table_filters, + ClientContext &context, + MultiFileGlobalState &gstate) override; +}; + +struct VortexReaderInterface final : MultiFileReaderInterface { + static unique_ptr CreateInterface(ClientContext &) { + return make_uniq(); + } + + inline unique_ptr InitializeOptions(ClientContext &, + optional_ptr) override { + return make_uniq(); + } + + inline bool ParseCopyOption(ClientContext &, + const string &, + const vector &, + BaseFileReaderOptions &, + vector &, + vector &) override { + return false; + }; + + inline bool ParseOption(ClientContext &, + const string &, + const Value &, + MultiFileOptions &, + BaseFileReaderOptions &) override { + return false; + } + + inline unique_ptr InitializeBindData(MultiFileBindData &, + unique_ptr) override { + return make_uniq(); + } + + void BindReader(ClientContext &context, + vector &types, + vector &names, + MultiFileBindData &bind_data) override; + + unique_ptr InitializeGlobalState(ClientContext &context, + MultiFileBindData &bind_data, + MultiFileGlobalState &global_state) override; + + unique_ptr InitializeLocalState(ExecutionContext &context, + GlobalTableFunctionState &global_state) override; + + inline shared_ptr CreateReader(ClientContext &, + GlobalTableFunctionState &, + BaseUnionData &, + const MultiFileBindData &) override { + throw BinderException("UNION BY NAME for Vortex files is not supported"); + } + + shared_ptr CreateReader(ClientContext &context, + GlobalTableFunctionState &gstate, + const OpenFileInfo &file, + idx_t file_idx, + const MultiFileBindData &bind_data) override; + + shared_ptr CreateReader(ClientContext &context, + const OpenFileInfo &file, + BaseFileReaderOptions &options, + const MultiFileOptions &file_options) override; + + unique_ptr GetCardinality(const MultiFileBindData &bind_data, idx_t file_count) override; + + inline FileGlobInput GetGlobInput() override { + return {FileGlobOptions::FALLBACK_GLOB, "vortex"}; + } + + inline unique_ptr Copy() override { + return make_uniq(); + } + + void GetVirtualColumns(ClientContext &, MultiFileBindData &, virtual_column_map_t &result) override; + bool FinalizeScan(ClientContext &, GlobalTableFunctionState &gstate, DataChunk &output) override; + void FinishReading(ClientContext &, + GlobalTableFunctionState &gstate, + LocalTableFunctionState &lstate) override; +}; + +struct VortexBaseReader final : BaseFileReader { + VortexBaseReader(OpenFileInfo file, unique_ptr ffi_file) + : BaseFileReader(file), ffi_file(std::move(ffi_file)) { + } + + unique_ptr ffi_file; + /* + * Populated only for first file reader in scan when BindReader() is + * called on it. Used in GetStatistics() which is called only for first + * file after BindReader(). + */ + const void *ffi_bind = nullptr; + + inline void AddVirtualColumn(column_t) override { + } + + /* + * Called by all threads on current file under global lock. Once + * TryInitializeScan returns false, first thread to receive it advances + * to next file and calls TryInitializeScan on it. + */ + bool TryInitializeScan(ClientContext &, + GlobalTableFunctionState &, + LocalTableFunctionState &local_state) override; + + AsyncResult Scan(ClientContext &, + GlobalTableFunctionState &global_state, + LocalTableFunctionState &local_state, + DataChunk &chunk) override; + + inline void FinishFile(ClientContext &, GlobalTableFunctionState &) override { + } + + double GetProgressInFile(ClientContext &) override; + + // Called on first file in scan after BindReader() has been called on it + unique_ptr GetStatistics(ClientContext &context, const string &name) override; + + inline string GetReaderType() const override { + return "Vortex"; + } +}; diff --git a/vortex-duckdb/cpp/include/table_function.h b/vortex-duckdb/cpp/include/table_function.h index 65a1a3f1dd3..82090c26387 100644 --- a/vortex-duckdb/cpp/include/table_function.h +++ b/vortex-duckdb/cpp/include/table_function.h @@ -10,16 +10,10 @@ extern "C" { #endif -// Info passed into the bind callback. The callback should set error or else add result columns. -typedef struct duckdb_vx_tfunc_bind_input_ *duckdb_vx_tfunc_bind_input; -typedef struct duckdb_vx_tfunc_bind_result_ *duckdb_vx_tfunc_bind_result; - -// Fetch a parameter from the bind info. -// The caller is responsible for freeing the value using duckdb_value_free. -duckdb_value duckdb_vx_tfunc_bind_input_get_parameter(duckdb_vx_tfunc_bind_input ffi_input, size_t index); +typedef struct duckdb_bind_result_ *duckdb_bind_result; // Add a result column to the bind info. -void duckdb_vx_tfunc_bind_result_add_column(duckdb_vx_tfunc_bind_result ffi_result, +void duckdb_vx_tfunc_bind_result_add_column(duckdb_bind_result ffi_result, const char *name_str, size_t name_len, duckdb_logical_type ffi_type); @@ -31,29 +25,8 @@ void duckdb_vx_string_map_insert(duckdb_vx_string_map map, const char *key, cons // Input data passed into the init_global and init_local callbacks. typedef struct { const void *bind_data; - - /** - * Projected columns that are requested to be read. These are not - * all columns, only the ones DuckDB optimizer thinks we should read. - */ idx_t *column_ids; size_t column_ids_count; - - /** - * Post filter projected columns. Our table function implements filter - * pushdown so this list is a subset of columns referenced in column_ids - * after filter pushdown and filter pruning. May be empty, in which case - * column_ids should be used. - * Indices in this list reference values from column_ids. I.e. if - * column_ids=[1,5,6], projection_ids=[1], output column should be - * column_ids[1] = 5 - * - * Example usage: - * https://github.com/duckdb/duckdb/blob/dc11eadd8f0a7c600f0034810706605ebe10d5b9/src/include/duckdb/function/table_function.hpp#L147 - */ - const idx_t *projection_ids; - size_t projection_ids_count; - duckdb_vx_table_filter_set filters; duckdb_client_context client_context; } duckdb_vx_tfunc_init_input; @@ -74,19 +47,36 @@ typedef struct { // set only for strings uint64_t max_string_length; bool has_null; + // owned column type + duckdb_logical_type type; } duckdb_column_statistics; -const idx_t INVALID_IDX = UINT64_MAX; +// File-level statistics of a written Vortex file, for the DuckLake +// WRITTEN_FILE_STATISTICS return path. Filled by Rust from the WriteSummary. +typedef struct { + uint64_t row_count; + uint64_t file_size_bytes; + uint64_t footer_size_bytes; + uint64_t num_columns; +} duckdb_vx_written_file_statistics; +// Per-column statistics of a written Vortex file. typedef struct { - idx_t partition_index; - // Either INVALID_IDX or position of column in output for file_index column - size_t file_index_column_pos; - // File index for the exported partition. - size_t file_index; -} duckdb_vx_partition_data; + // Owned values, null if absent; the caller must destroy them. + duckdb_value min; + duckdb_value max; + bool has_null_count; + uint64_t null_count; + uint64_t num_values; + bool has_column_size; + uint64_t column_size_bytes; + // Whether a NaN-count statistic was available (float columns), and whether it saw any NaN. + bool has_nan_stat; + bool contains_nan; +} duckdb_vx_written_column_statistics; duckdb_state duckdb_vx_register_table_functions(duckdb_database ffi_db); +duckdb_state duckdb_vx_register_version_function(duckdb_database ffi_db, const char *version); typedef struct duckdb_vx_agg_input_ *duckdb_vx_agg_input; idx_t duckdb_vx_aggregate_len(duckdb_vx_agg_input ffi); diff --git a/vortex-duckdb/cpp/include/table_function.hpp b/vortex-duckdb/cpp/include/table_function.hpp index 07d035917be..9af3a27c2c5 100644 --- a/vortex-duckdb/cpp/include/table_function.hpp +++ b/vortex-duckdb/cpp/include/table_function.hpp @@ -3,33 +3,56 @@ #pragma once +#include "data.hpp" #include "duckdb.h" #include "duckdb/function/function.hpp" #include "duckdb/function/table_function.hpp" +using namespace duckdb; + static_assert(sizeof(idx_t) == 8); -// We need this exposed to compare function addresses in optimizer.cpp -duckdb::unique_ptr -duckdb_vx_table_function_bind(duckdb::ClientContext &context, - duckdb::TableFunctionBindInput &input, - duckdb::vector &return_types, - duckdb::vector &names); +bool is_vortex_scan(const TableFunction &function); struct TableFunctionProjectionExpressionInput { - const duckdb::LogicalGet &get; - const duckdb::Expression &expression; + const LogicalGet &get; + const Expression &expression; idx_t projection_idx; }; // true if we can push down the expression, false otherwise -bool projection_expression_pushdown(duckdb::ClientContext &context, +bool projection_expression_pushdown(ClientContext &context, const TableFunctionProjectionExpressionInput &input); struct TableFunctionUngroupedAggregateInput { - const duckdb::LogicalGet &get; + const LogicalGet &get; // Column scan index -> aggregate expression - const duckdb::vector> &projections; + const vector> &projections; }; -bool aggregate_pushdown(duckdb::ClientContext &context, const TableFunctionUngroupedAggregateInput &input); +bool aggregate_pushdown(ClientContext &context, const TableFunctionUngroupedAggregateInput &input); + +/* + * DuckDB uses partition row groups for two purposes: + * + * 1. If optimizer proves query (e.g. SELECT min(col)) can be answered from + * metadata, it replaces a real scan with a call to GetColumnStatistics. + * 2. If optimizer proves a row group can be pruned because of statistics, it + * removes row group's read. We don't use this as we implement own prunung. + * + * For (1) we care about providing statistics fast, so we report a file as + * a "row group". + */ +struct VortexRowGroup final : PartitionRowGroup { + explicit VortexRowGroup(unique_ptr ffi_footer) : ffi_footer(std::move(ffi_footer)) { + } + + unique_ptr ffi_footer; + + unique_ptr GetColumnStatistics(const StorageIndex &storage_index) override; + bool MinMaxIsExact(const BaseStatistics &, const StorageIndex &) override { + // TODO(myrrc): in duckdb 2.0 we should report false for strings and + // also add TRUNCATED_STATS type for them + return true; + } +}; diff --git a/vortex-duckdb/cpp/include/vector.h b/vortex-duckdb/cpp/include/vector.h index 24d01748666..983a7ecc479 100644 --- a/vortex-duckdb/cpp/include/vector.h +++ b/vortex-duckdb/cpp/include/vector.h @@ -22,8 +22,6 @@ void duckdb_vx_vector_dictionary(duckdb_vector ffi_vector, duckdb_selection_vector ffi_sel_vec, idx_t count); -void duckdb_vx_set_dictionary_vector_length(duckdb_vector dict, unsigned int len); - // Reset vector's validity mask to nullptr, making all vector's elements valid. // vector must not be a DictionaryVector or a SequenceVector void duckdb_vx_vector_set_all_valid(duckdb_vector ffi_vector); diff --git a/vortex-duckdb/cpp/multi_file_reader.cpp b/vortex-duckdb/cpp/multi_file_reader.cpp new file mode 100644 index 00000000000..02d4da3caaa --- /dev/null +++ b/vortex-duckdb/cpp/multi_file_reader.cpp @@ -0,0 +1,412 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "multi_file_reader.hpp" +#include "error.hpp" +#include "table_function.h" +#include "vortex_duckdb.h" +#include "vortex.h" + +unique_ptr VortexBindData::Copy() const { + auto result = make_uniq(); + if (ffi_bind_data) { + const duckdb_vx_data copy = duckdb_table_function_bind_data_clone(ffi_bind_data->DataPtr()); + result->ffi_bind_data = unique_ptr(reinterpret_cast(copy)); + } + return result; +} + +bool VortexBindData::Equals(const FunctionData &other_base) const { + const VortexBindData &other = other_base.Cast(); + return ffi_bind_data.get() == other.ffi_bind_data.get(); +} + +bool SkipMultiFileReaderColumn(size_t i, + optional_idx filename_idx, + const std::vector &hive_indices) { + if (filename_idx.IsValid() && i == filename_idx.GetIndex()) { + return true; + } + for (const auto &idx : hive_indices) { + if (idx.index == i) { + return true; + } + } + return false; +} + +ReaderInitializeType +VortexMultiFileReader::InitializeReader(MultiFileReaderData &reader_data, + const MultiFileBindData &bind_data, + const vector &global_columns, + const vector &global_column_ids, + optional_ptr table_filters, + ClientContext &context, + MultiFileGlobalState &gstate) { + D_ASSERT(reader_data.reader != nullptr); + D_ASSERT(gstate.global_state != nullptr); + + VortexBaseReader &reader = reader_data.reader->Cast(); + const VortexBindData &bind = bind_data.bind_data->Cast(); + + const void *const ffi_bind = bind.ffi_bind_data->DataPtr(); + + reader.columns = global_columns; + + const optional_idx filename_idx = bind_data.reader_bind.filename_idx; + const auto &hive_indices = bind_data.reader_bind.hive_partitioning_indexes; + + // Aggregate scan columns are aggregate results so base column mapping is + // wrong. + if (!duckdb_reader_is_aggregate(ffi_bind)) { + // Projection expression pushdown changes types of columns + for (size_t i = 0; i < global_columns.size(); i++) { + if (SkipMultiFileReaderColumn(i, filename_idx, hive_indices)) { + continue; + } + + const duckdb_logical_type type = duckdb_reader_bind_column_type(ffi_bind, i); + reader.columns[i].type = *reinterpret_cast(type); + } + + // this prunes files for virtual columns like "file_index" + const ReaderInitializeType base_skip = MultiFileReader::InitializeReader(reader_data, + bind_data, + reader.columns, + global_column_ids, + table_filters, + context, + gstate); + if (base_skip == ReaderInitializeType::SKIP_READING_FILE) { + return base_skip; + } + } + + const VortexGlobalState &global = gstate.global_state->Cast(); + + duckdb_vx_error error = nullptr; + const void *const ffi_global = global.ffi_global_state->DataPtr(); + void *const ffi_file = reader.ffi_file->DataPtr(); + const bool skip = duckdb_reader_initialize(ffi_global, ffi_file, &error); + if (error) { + throw InvalidInputException(IntoErrString(error)); + } + + return skip ? ReaderInitializeType::SKIP_READING_FILE : ReaderInitializeType::INITIALIZED; +} + +void VortexReaderInterface::BindReader(ClientContext &context, + vector &types, + vector &names, + MultiFileBindData &bind_data) { + BaseFileReaderOptions options; + + VortexBindResult result = {types, names}; + + VortexBindData &bind = bind_data.bind_data->Cast(); + const OpenFileInfo first_file = bind_data.file_list->GetFirstFile(); + bind_data.initial_reader = CreateReader(context, first_file, options, bind_data.file_options); + VortexBaseReader &initial_reader = bind_data.initial_reader->Cast(); + + duckdb_vx_error error = nullptr; + const void *const ffi_file = initial_reader.ffi_file->DataPtr(); + duckdb_bind_result ffi_result = reinterpret_cast(&result); + + duckdb_vx_data ffi_bind_data = duckdb_reader_bind(ffi_file, ffi_result, &error); + if (error) { + throw BinderException(IntoErrString(error)); + } + + bind.ffi_bind_data = unique_ptr(reinterpret_cast(ffi_bind_data)); + initial_reader.ffi_bind = bind.ffi_bind_data->DataPtr(); + + // Fills bind_data.file_options which are used for hive partitioning and + // "filename" column. + bind_data.multi_file_reader->BindOptions(bind_data.file_options, + *bind_data.file_list, + types, + names, + bind_data.reader_bind); +} + +unique_ptr +VortexReaderInterface::InitializeGlobalState(ClientContext &context, + MultiFileBindData &bind_data, + MultiFileGlobalState &input) { + const VortexBindData &bind = bind_data.bind_data->Cast(); + + const optional_idx filename_idx = bind_data.reader_bind.filename_idx; + const auto &hive_indices = bind_data.reader_bind.hive_partitioning_indexes; + + vector column_ids(input.column_indexes.size()); + for (size_t i = 0; i < input.column_indexes.size(); ++i) { + idx_t storage_index = input.column_indexes[i].GetPrimaryIndex(); + + if (SkipMultiFileReaderColumn(storage_index, filename_idx, hive_indices)) { + // replace with a virtual column which Vortex filters in Rust + storage_index = COLUMN_IDENTIFIER_EMPTY; + } + + column_ids[i] = storage_index; + } + + void *const ffi_bind = bind.ffi_bind_data->DataPtr(); + duckdb_vx_tfunc_init_input ffi_input = { + .bind_data = ffi_bind, + .column_ids = column_ids.data(), + .column_ids_count = column_ids.size(), + .filters = reinterpret_cast(input.filters.get()), + .client_context = reinterpret_cast(&context), + }; + + duckdb_vx_error error_out = nullptr; + duckdb_vx_data ffi_global_state = duckdb_table_function_init_global(&ffi_input, &error_out); + if (error_out) { + throw BinderException(IntoErrString(error_out)); + } + + auto result = make_uniq(); + result->ffi_bind_data = ffi_bind; + result->ffi_global_state = unique_ptr(reinterpret_cast(ffi_global_state)); + return result; +} + +unique_ptr +VortexReaderInterface::InitializeLocalState(ExecutionContext &, GlobalTableFunctionState &global_state) { + auto &global = global_state.Cast(); + const void *const ffi_global = global.ffi_global_state->DataPtr(); + duckdb_vx_data ffi_local_state = duckdb_table_function_init_local(global.ffi_bind_data, ffi_global); + + auto result = make_uniq(); + result->ffi_local_state = unique_ptr(reinterpret_cast(ffi_local_state)); + return result; +} + +static shared_ptr OpenReader(const OpenFileInfo &file) { + duckdb_vx_error error = nullptr; + + const char *const ffi_file_path = file.path.c_str(); + const size_t ffi_file_size = file.path.size(); + + duckdb_vx_data ffi_file = duckdb_reader_open(ffi_file_path, ffi_file_size, &error); + if (error) { + throw IOException(IntoErrString(error)); + } + + auto cdata = unique_ptr(reinterpret_cast(ffi_file)); + return make_shared_ptr(file, std::move(cdata)); +} + +shared_ptr VortexReaderInterface::CreateReader(ClientContext &, + GlobalTableFunctionState &, + const OpenFileInfo &file, + idx_t, + const MultiFileBindData &) { + return OpenReader(file); +} + +shared_ptr VortexReaderInterface::CreateReader(ClientContext &, + const OpenFileInfo &file, + BaseFileReaderOptions &, + const MultiFileOptions &) { + return OpenReader(file); +} + +unique_ptr VortexReaderInterface::GetCardinality(const MultiFileBindData &data, + idx_t file_count) { + const VortexBindData &bind_data = data.bind_data->Cast(); + const void *const ffi_bind = bind_data.ffi_bind_data->DataPtr(); + + duckdb_vx_node_statistics stats = {}; + duckdb_table_function_cardinality(ffi_bind, file_count, &stats); + + auto out = make_uniq(); + out->has_estimated_cardinality = stats.has_estimated_cardinality; + out->estimated_cardinality = stats.estimated_cardinality; + out->has_max_cardinality = stats.has_max_cardinality; + out->max_cardinality = stats.max_cardinality; + return out; +} + +bool VortexBaseReader::TryInitializeScan(ClientContext &, + GlobalTableFunctionState &, + LocalTableFunctionState &local_state) { + VortexLocalState &local = local_state.Cast(); + + void *const ffi_local = local.ffi_local_state->DataPtr(); + void *const ffi_file_ptr = ffi_file->DataPtr(); + return duckdb_reader_try_initialize_scan(ffi_local, ffi_file_ptr); +} + +AsyncResult VortexBaseReader::Scan(ClientContext &, + GlobalTableFunctionState &global_state, + LocalTableFunctionState &local_state, + DataChunk &chunk) { + VortexGlobalState &global = global_state.Cast(); + VortexLocalState &local = local_state.Cast(); + + duckdb_vx_error error = nullptr; + duckdb_data_chunk ffi_chunk = reinterpret_cast(&chunk); + const void *const ffi_global = global.ffi_global_state->DataPtr(); + void *const ffi_local = local.ffi_local_state->DataPtr(); + const void *const ffi_file_ptr = ffi_file->DataPtr(); + const bool has_more_data = duckdb_reader_scan(ffi_file_ptr, ffi_global, ffi_local, ffi_chunk, &error); + if (error) { + throw InvalidInputException(IntoErrString(error)); + } + return has_more_data ? SourceResultType::HAVE_MORE_OUTPUT : SourceResultType::FINISHED; +} + +void VortexReaderInterface::GetVirtualColumns(ClientContext &, + MultiFileBindData &, + virtual_column_map_t &result) { + // "filename", "file_index" and "empty" come from MultiFileReader + result.insert( + {MultiFileReader::COLUMN_IDENTIFIER_FILE_ROW_NUMBER, {"file_row_number", LogicalType::UBIGINT}}); +} + +bool VortexReaderInterface::FinalizeScan(ClientContext &, + GlobalTableFunctionState &global_state, + DataChunk &output) { + const VortexGlobalState &global = global_state.Cast(); + + duckdb_vx_error error = nullptr; + duckdb_data_chunk ffi_chunk = reinterpret_cast(&output); + const void *const ffi_global = global.ffi_global_state->DataPtr(); + const bool filled = duckdb_reader_finalize_scan(ffi_global, ffi_chunk, &error); + if (error) { + throw InvalidInputException(IntoErrString(error)); + } + return filled; +} + +void VortexReaderInterface::FinishReading(ClientContext &, + GlobalTableFunctionState &global_state, + LocalTableFunctionState &local_state) { + const VortexGlobalState &global = global_state.Cast(); + const VortexLocalState &local = local_state.Cast(); + const void *const ffi_global = global.ffi_global_state->DataPtr(); + void *const ffi_local = local.ffi_local_state->DataPtr(); + duckdb_reader_finish_reading(ffi_global, ffi_local); +} + +static Value &UnwrapValue(duckdb_value value) { + return *(reinterpret_cast(value)); +} + +static unique_ptr numeric_stats(duckdb_column_statistics &stats, LogicalType type) { + BaseStatistics out = NumericStats::CreateUnknown(type); + if (stats.min) { + NumericStats::SetMin(out, UnwrapValue(stats.min)); + duckdb_destroy_value(&stats.min); + } + if (stats.max) { + NumericStats::SetMax(out, UnwrapValue(stats.max)); + duckdb_destroy_value(&stats.max); + } + if (!stats.has_null) { + out.Set(StatsInfo::CANNOT_HAVE_NULL_VALUES); + } + return out.ToUnique(); +} + +static unique_ptr string_stats(duckdb_column_statistics &stats, LogicalType type) { + BaseStatistics out = StringStats::CreateUnknown(type); + if (stats.min) { + StringStats::SetMin(out, StringValue::Get(UnwrapValue(stats.min))); + duckdb_destroy_value(&stats.min); + } + if (stats.max) { + StringStats::SetMax(out, StringValue::Get(UnwrapValue(stats.max))); + duckdb_destroy_value(&stats.max); + } + if (stats.max_string_length >> 63) { + StringStats::SetMaxStringLength(out, uint32_t(stats.max_string_length)); + } + if (!stats.has_null) { + out.Set(StatsInfo::CANNOT_HAVE_NULL_VALUES); + } + + return out.ToUnique(); +} + +static unique_ptr base_stats(duckdb_column_statistics &stats, LogicalType type) { + BaseStatistics out = BaseStatistics::CreateUnknown(type); + if (stats.min) { + duckdb_destroy_value(&stats.min); + } + if (stats.max) { + duckdb_destroy_value(&stats.max); + } + if (!stats.has_null) { + out.Set(StatsInfo::CANNOT_HAVE_NULL_VALUES); + } + return out.ToUnique(); +} + +unique_ptr to_duckdb_statistics(duckdb_column_statistics &statistics) { + using enum LogicalTypeId; + const unique_ptr type(reinterpret_cast(statistics.type)); + switch (type->id()) { + case BOOLEAN: + case TINYINT: + case SMALLINT: + case INTEGER: + case BIGINT: + case FLOAT: + case DOUBLE: + case UTINYINT: + case USMALLINT: + case UINTEGER: + case UBIGINT: + case UHUGEINT: + case HUGEINT: + case DATE: + case TIME: + case TIME_TZ: + case TIMESTAMP_SEC: + case TIMESTAMP_MS: + case TIMESTAMP: + case TIMESTAMP_NS: + case TIMESTAMP_TZ: { + return numeric_stats(statistics, *type); + } + case VARCHAR: + case BLOB: { + return string_stats(statistics, *type); + } + case STRUCT: { + // TODO(myrrc) + // Duckdb's has_null has a different semantics for structs. + // If we propagate our has_null, this breaks Duckdb optimizer. + // You can reproduce it in struct.slt test in vortex-sqllogictests: + if (statistics.min) { + duckdb_destroy_value(&statistics.min); + } + if (statistics.max) { + duckdb_destroy_value(&statistics.max); + } + return {}; + } + default: + return base_stats(statistics, *type); + } +} + +unique_ptr VortexBaseReader::GetStatistics(ClientContext &, const string &name) { + D_ASSERT(ffi_bind); + duckdb_column_statistics statistics = {}; + if (!duckdb_reader_get_statistics(ffi_file->DataPtr(), + ffi_bind, + name.c_str(), + name.size(), + &statistics)) { + return {}; + } + + return to_duckdb_statistics(statistics); +} + +double VortexBaseReader::GetProgressInFile(ClientContext &) { + return duckdb_reader_get_progress_in_file(ffi_file->DataPtr()); +} diff --git a/vortex-duckdb/cpp/optimizer.cpp b/vortex-duckdb/cpp/optimizer.cpp index 87082d57790..9f799fb1e42 100644 --- a/vortex-duckdb/cpp/optimizer.cpp +++ b/vortex-duckdb/cpp/optimizer.cpp @@ -1,17 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors #include "optimizer.hpp" +#include "multi_file_reader.hpp" #include "table_function.hpp" #include "duckdb/planner/expression/bound_columnref_expression.hpp" #include "duckdb/planner/operator/logical_get.hpp" +#include "duckdb/planner/operator/logical_limit.hpp" #include "duckdb/planner/operator/logical_projection.hpp" void FindGetsAndProjections(LogicalOperator &op, Analyses &analyses, Projections &projections) { using enum LogicalOperatorType; switch (op.type) { case LOGICAL_GET: { - if (auto &get = op.Cast(); get.function.bind == duckdb_vx_table_function_bind) { + if (auto &get = op.Cast(); is_vortex_scan(get.function)) { analyses.emplace(get.table_index, GetAnalysis {get, {}}); } break; @@ -20,12 +22,23 @@ void FindGetsAndProjections(LogicalOperator &op, Analyses &analyses, Projections LogicalProjection &projection = op.Cast(); D_ASSERT(projection.children.size() == 1); auto &child = *projection.children[0]; - if (!IsPassthrough(projection) || child.type != LOGICAL_GET) { + if (!IsPassthrough(projection)) { break; } - // The GET itself is recorded when recursion reaches it below. Only - // passthrough projections wrapping a vortex GET act as aliases. - if (auto &get = child.Cast(); get.function.bind == duckdb_vx_table_function_bind) { + LogicalGet *get = nullptr; + + // queries with LIMIT may include a STREAMING_LIMIT between PROJECTION and GET. + // See sqllogictest scalar_function_pushdown_limit.test + if (child.type == LOGICAL_LIMIT) { + if (auto &limit = child.Cast(); + limit.children.size() == 1 && limit.children[0]->type == LOGICAL_GET) { + get = &limit.children[0]->Cast(); + } + } else if (child.type == LOGICAL_GET) { + get = &child.Cast(); + } + + if (get != nullptr && is_vortex_scan(get->function)) { projections.emplace(projection.table_index, projection); } break; diff --git a/vortex-duckdb/cpp/scalar_fn_pushdown.cpp b/vortex-duckdb/cpp/scalar_fn_pushdown.cpp index bcf23496443..3ff85e880df 100644 --- a/vortex-duckdb/cpp/scalar_fn_pushdown.cpp +++ b/vortex-duckdb/cpp/scalar_fn_pushdown.cpp @@ -4,6 +4,7 @@ #include "duckdb/planner/operator/logical_projection.hpp" #include "duckdb/planner/operator/logical_get.hpp" +#include "optimizer.hpp" #include diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index 42c4285c383..6dd0f6e9601 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -4,6 +4,7 @@ #include "data.hpp" #include "error.hpp" #include "table_function.hpp" +#include "multi_file_reader.hpp" #include "expr.h" #include "vortex_duckdb.h" #include "table_function.h" @@ -11,178 +12,35 @@ #include "duckdb.h" #include "duckdb/catalog/catalog.hpp" +#include "duckdb/catalog/default/default_functions.hpp" #include "duckdb/common/insertion_order_preserving_map.hpp" +#include "duckdb/parser/keyword_helper.hpp" #include "duckdb/common/multi_file/multi_file_reader.hpp" #include "duckdb/function/table_function.hpp" #include "duckdb/main/capi/capi_internal.hpp" #include "duckdb/main/connection.hpp" +#include "duckdb/function/partition_stats.hpp" #include "duckdb/parser/parsed_data/create_table_function_info.hpp" #include "duckdb/planner/operator/logical_get.hpp" +#include "duckdb/storage/storage_index.hpp" using namespace std::string_literals; -using namespace duckdb; constexpr column_t COLUMN_IDENTIFIER_FILE_INDEX = MultiFileReader::COLUMN_IDENTIFIER_FILE_INDEX; constexpr column_t COLUMN_IDENTIFIER_FILE_ROW_NUMBER = MultiFileReader::COLUMN_IDENTIFIER_FILE_ROW_NUMBER; -struct CTableBindData final : FunctionData { - CTableBindData(unique_ptr ffi_data_p, const vector &types) - : ffi_data(std::move(ffi_data_p)), types(types) { - } - unique_ptr Copy() const override; - bool Equals(const FunctionData &other_base) const override; - - unique_ptr ffi_data; - vector types; -}; - -unique_ptr CTableBindData::Copy() const { - const auto copied_ffi_data = duckdb_table_function_bind_data_clone(ffi_data->DataPtr()); - auto ffi_data_p = unique_ptr(reinterpret_cast(copied_ffi_data)); - return make_uniq(std::move(ffi_data_p), types); -} - -bool CTableBindData::Equals(const FunctionData &other_base) const { - const CTableBindData &other = other_base.Cast(); - // if "types" are different, "ffi_data" would also be different as it - // contains types inside, so omit "types" from comparison. - return ffi_data.get() == other.ffi_data.get(); -} - -struct CTableGlobalData final : GlobalTableFunctionState { - explicit CTableGlobalData(unique_ptr ffi_data) : ffi_data(std::move(ffi_data)) { - } - - idx_t MaxThreads() const override { - return GlobalTableFunctionState::MAX_THREADS; - } - - unique_ptr ffi_data; -}; - -struct CTableLocalData final : LocalTableFunctionState { - explicit CTableLocalData(unique_ptr ffi_data) : ffi_data(std::move(ffi_data)) { - } - - unique_ptr ffi_data; -}; - -double -table_scan_progress(ClientContext &, const FunctionData *, const GlobalTableFunctionState *global_state) { - void *const c_global_state = global_state->Cast().ffi_data->DataPtr(); - return duckdb_table_function_scan_progress(c_global_state); -} - -static Value &UnwrapValue(duckdb_value value) { - return *(reinterpret_cast(value)); -} - -unique_ptr numeric_stats(duckdb_column_statistics &stats, LogicalType type) { - BaseStatistics out = StringStats::CreateUnknown(type); - if (stats.min) { - NumericStats::SetMin(out, UnwrapValue(stats.min)); - duckdb_destroy_value(&stats.min); - } - if (stats.max) { - NumericStats::SetMax(out, UnwrapValue(stats.max)); - duckdb_destroy_value(&stats.max); - } - if (!stats.has_null) { - out.Set(StatsInfo::CANNOT_HAVE_NULL_VALUES); - } - return out.ToUnique(); -} - -unique_ptr string_stats(duckdb_column_statistics &stats, LogicalType type) { - BaseStatistics out = StringStats::CreateUnknown(type); - if (stats.min) { - StringStats::SetMin(out, StringValue::Get(UnwrapValue(stats.min))); - duckdb_destroy_value(&stats.min); - } - if (stats.max) { - StringStats::SetMax(out, StringValue::Get(UnwrapValue(stats.max))); - duckdb_destroy_value(&stats.max); - } - if (stats.max_string_length >> 63) { - StringStats::SetMaxStringLength(out, uint32_t(stats.max_string_length)); - } - if (!stats.has_null) { - out.Set(StatsInfo::CANNOT_HAVE_NULL_VALUES); - } - - return out.ToUnique(); -} - -unique_ptr base_stats(duckdb_column_statistics &stats, LogicalType type) { - BaseStatistics out = StringStats::CreateUnknown(type); - if (!stats.has_null) { - out.Set(StatsInfo::CANNOT_HAVE_NULL_VALUES); - } - return out.ToUnique(); -} - -unique_ptr statistics(ClientContext &, const FunctionData *bind_data, column_t column_index) { - if (IsVirtualColumn(column_index)) { - return {}; - } - - const auto &bind = bind_data->Cast(); - void *const ffi_bind = bind.ffi_data->DataPtr(); - - duckdb_column_statistics statistics = {}; - if (!duckdb_table_function_statistics(ffi_bind, column_index, &statistics)) { - return {}; - } - - const LogicalType type = bind.types[column_index]; - - switch (type.id()) { - case LogicalTypeId::BOOLEAN: - case LogicalTypeId::TINYINT: - case LogicalTypeId::SMALLINT: - case LogicalTypeId::INTEGER: - case LogicalTypeId::BIGINT: - case LogicalTypeId::FLOAT: - case LogicalTypeId::DOUBLE: - case LogicalTypeId::UTINYINT: - case LogicalTypeId::USMALLINT: - case LogicalTypeId::UINTEGER: - case LogicalTypeId::UBIGINT: - case LogicalTypeId::UHUGEINT: - case LogicalTypeId::HUGEINT: { - return numeric_stats(statistics, type); - } - case LogicalTypeId::VARCHAR: - case LogicalTypeId::BLOB: { - return string_stats(statistics, type); - } - case LogicalTypeId::STRUCT: { - // TODO(myrrc) - // Duckdb's has_null has a different semantics for structs. - // If we propagate our has_null, this breaks Duckdb optimizer. - // You can reproduce it in struct.slt test in vortex-sqllogictests: - return {}; - } - default: - return base_stats(statistics, type); - } -} - -struct CTableBindResult { - vector &return_types; - vector &names; -}; - // This is a flaw of Duckdb API which doesn't allow passing non-const // expressions. We never modify the value on Rust side. static duckdb_vx_expr get_ffi_expr(const Expression &expr) { return reinterpret_cast(const_cast(&expr)); } -bool projection_expression_pushdown(ClientContext &, const TableFunctionProjectionExpressionInput &input) { - const auto &bind = input.get.bind_data->Cast(); +static void *get_ffi_bind(const FunctionData *bind_data) { + return bind_data->Cast().bind_data->Cast().ffi_bind_data->DataPtr(); +} +bool projection_expression_pushdown(ClientContext &, const TableFunctionProjectionExpressionInput &input) { duckdb_vx_expr ffi_expr = get_ffi_expr(input.expression); - void *const ffi_bind = bind.ffi_data->DataPtr(); + void *const ffi_bind = get_ffi_bind(input.get.bind_data.get()); duckdb_vx_error error_out = nullptr; const bool ret = duckdb_table_function_pushdown_projection_expression( // @@ -211,8 +69,7 @@ duckdb_vx_expr duckdb_vx_aggregate_at(duckdb_vx_agg_input ffi_input, idx_t i, id } bool aggregate_pushdown(ClientContext &, const TableFunctionUngroupedAggregateInput &input) { - const auto &bind = input.get.bind_data->Cast(); - void *const ffi_bind = bind.ffi_data->DataPtr(); + void *const ffi_bind = get_ffi_bind(input.get.bind_data.get()); duckdb_vx_error error_out = nullptr; const auto ffi_input = reinterpret_cast(const_cast(&input)); @@ -223,91 +80,10 @@ bool aggregate_pushdown(ClientContext &, const TableFunctionUngroupedAggregateIn return res; } -/** - * Called for every new query. For example, if there is a VIEW over *.vortex, - * and after a query another file is added matching the glob, for second query - * bind() will be called again. - */ -unique_ptr duckdb_vx_table_function_bind(ClientContext &, - TableFunctionBindInput &input, - vector &return_types, - vector &names) { - CTableBindResult result = {return_types, names}; - - duckdb_vx_error error_out = nullptr; - auto ffi_bind_data = duckdb_table_function_bind(reinterpret_cast(&input), - reinterpret_cast(&result), - &error_out); - if (error_out) { - throw BinderException(IntoErrString(error_out)); - } - - auto cdata = unique_ptr(reinterpret_cast(ffi_bind_data)); - return make_uniq(std::move(cdata), return_types); -} - -unique_ptr c_init_global(ClientContext &context, TableFunctionInitInput &input) { - const CTableBindData &bind = input.bind_data->Cast(); - - duckdb_vx_tfunc_init_input ffi_input = { - .bind_data = bind.ffi_data->DataPtr(), - .column_ids = input.column_ids.data(), - .column_ids_count = input.column_ids.size(), - .projection_ids = input.projection_ids.data(), - .projection_ids_count = input.projection_ids.size(), - .filters = reinterpret_cast(input.filters.get()), - .client_context = reinterpret_cast(&context), - }; - - duckdb_vx_error error_out = nullptr; - duckdb_vx_data ffi_global_data = duckdb_table_function_init_global(&ffi_input, &error_out); - if (error_out) { - throw BinderException(IntoErrString(error_out)); - } - - auto cdata = unique_ptr(reinterpret_cast(ffi_global_data)); - return make_uniq(std::move(cdata)); -} - -unique_ptr -init_local(ExecutionContext &, TableFunctionInitInput &input, GlobalTableFunctionState *global_state) { - const void *const ffi_bind = input.bind_data->Cast().ffi_data->DataPtr(); - void *const ffi_global = global_state->Cast().ffi_data->DataPtr(); - - duckdb_vx_data ffi_local_data = duckdb_table_function_init_local(ffi_bind, ffi_global); - auto cdata = unique_ptr(reinterpret_cast(ffi_local_data)); - return make_uniq(std::move(cdata)); -} - -void function(ClientContext &, TableFunctionInput &input, DataChunk &output) { - void *const ffi_global = input.global_state->Cast().ffi_data->DataPtr(); - void *const ffi_local = input.local_state->Cast().ffi_data->DataPtr(); - - duckdb_data_chunk chunk = reinterpret_cast(&output); - duckdb_vx_error error_out = nullptr; - duckdb_table_function_scan(ffi_global, ffi_local, chunk, &error_out); - if (error_out) { - throw InvalidInputException(IntoErrString(error_out)); - } -} - using FilterVec = vector>; -/* - * Table filter pushdown is used for two tasks in duckdb: - * - * 1. Prune files based on filename or hive partitioning, see Parquet - * filter pushdown. We don't use this because we do own file-level pruning in - * FileStatsLayoutReader, and we don't support hive partitioning yet. - * - * 2. Avoid reading unused file data. Filter expressions are pushed to Vortex, - * converted to Vortex expressions and used during the scan. - * Duckdb pushes a subset of expressions i.e. equality operators, and also - * expressions which return true in pushdown_expression. - */ void pushdown_complex_filter(const FunctionData &bind_data, FilterVec &filters) { - const auto &bind = bind_data.Cast(); - void *const ffi_bind = bind.ffi_data->DataPtr(); + void *const ffi_bind = get_ffi_bind(&bind_data); duckdb_vx_error error_out = nullptr; for (auto iter = filters.begin(); iter != filters.end();) { @@ -321,86 +97,20 @@ void pushdown_complex_filter(const FunctionData &bind_data, FilterVec &filters) } } -unique_ptr c_cardinality(ClientContext &, const FunctionData *bind_data) { - auto &bind = bind_data->Cast(); - - duckdb_vx_node_statistics stats = {}; - duckdb_table_function_cardinality(bind.ffi_data->DataPtr(), &stats); - - auto out = make_uniq(); - out->has_estimated_cardinality = stats.has_estimated_cardinality; - out->estimated_cardinality = stats.estimated_cardinality; - out->has_max_cardinality = stats.has_max_cardinality; - out->max_cardinality = stats.max_cardinality; - - return out; -} - -extern "C" duckdb_value duckdb_vx_tfunc_bind_input_get_parameter(duckdb_vx_tfunc_bind_input ffi_input, - size_t index) { - D_ASSERT(ffi_input); - const TableFunctionBindInput &input = *reinterpret_cast(ffi_input); - return reinterpret_cast(new Value(input.inputs[index])); -} - -extern "C" void duckdb_vx_tfunc_bind_result_add_column(duckdb_vx_tfunc_bind_result ffi_result, +extern "C" void duckdb_vx_tfunc_bind_result_add_column(duckdb_bind_result ffi_result, const char *name_str, size_t name_len, duckdb_logical_type ffi_type) { D_ASSERT(ffi_result); D_ASSERT(name_str); D_ASSERT(ffi_type); - const CTableBindResult &result = *reinterpret_cast(ffi_result); + VortexBindResult &result = *reinterpret_cast(ffi_result); const LogicalType logical_type = *reinterpret_cast(ffi_type); result.names.emplace_back(name_str, name_len); result.return_types.emplace_back(logical_type); } -/** - * Called at planning time to determine whether data is partitioned by a - * given set of columns. Requested columns are GROUP BY parameters i.e. columns - * over which the query aggregates. - */ -TablePartitionInfo get_partition_info(ClientContext &, TableFunctionPartitionInput &input) { - const vector &ids = input.partition_ids; - // Our data is partitioned by array exporters. Each exporter processes a - // single Array which belongs to a single file. If data is partitioned only - // by file_index, there is one unique value for an Array. Otherwise there - // may be multiple values. - return (ids.size() == 1 && ids[0] == COLUMN_IDENTIFIER_FILE_INDEX) - ? TablePartitionInfo::SINGLE_VALUE_PARTITIONS - : TablePartitionInfo::NOT_PARTITIONED; -} - -/** - * Duckdb requests this function after exporting the chunk. We answer with - * partition_index we have exported as well as information about constant - * columns in this partition. As data is partitioned by array exporters, in - * each partition ~ exported array file_index is constant. - */ -OperatorPartitionData get_partition_data(ClientContext &, TableFunctionGetPartitionInput &input) { - void *const ffi_global = input.global_state->Cast().ffi_data->DataPtr(); - void *const ffi_local = input.local_state->Cast().ffi_data->DataPtr(); - duckdb_vx_partition_data partition_data; - duckdb_table_function_get_partition_data(ffi_global, ffi_local, &partition_data); - - OperatorPartitionData out(partition_data.partition_index); - - // file_index_column_pos may be INVALID_IDX, but column_index will never - // be INVALID_IDX, so we can compare directly - for (const column_t column_index : input.partition_info.partition_columns) { - if (column_index == partition_data.file_index_column_pos) { - out.partition_data.emplace_back(Value::UBIGINT(partition_data.file_index)); - } else { - throw InternalException(StringUtil::Format( - "get_partition_data: requested column_index %d is not constant for given partition", - column_index)); - } - } - return out; -} - extern "C" void duckdb_vx_string_map_insert(duckdb_vx_string_map map, const char *key, const char *value) { D_ASSERT(map); D_ASSERT(key); @@ -408,60 +118,100 @@ extern "C" void duckdb_vx_string_map_insert(duckdb_vx_string_map map, const char reinterpret_cast *>(map)->insert(key, value); } -InsertionOrderPreservingMap c_to_string(TableFunctionToStringInput &input) { +InsertionOrderPreservingMap to_string(TableFunctionToStringInput &input) { InsertionOrderPreservingMap result; duckdb_vx_string_map ffi_map = reinterpret_cast(&result); - void *const ffi_bind = input.bind_data->Cast().ffi_data->DataPtr(); + const void *const ffi_bind = get_ffi_bind(input.bind_data.get()); duckdb_table_function_to_string(ffi_bind, ffi_map); return result; } +bool is_vortex_scan(const TableFunction &function) { + return function.bind == MultiFileFunction::MultiFileBind; +} + +unique_ptr get_multi_file_reader(const TableFunction &) { + return make_uniq(); +} + +unique_ptr VortexRowGroup::GetColumnStatistics(const StorageIndex &storage_index) { + duckdb_column_statistics statistics = {}; + const idx_t idx = storage_index.GetPrimaryIndex(); + const void *const ffi_footer_ptr = ffi_footer->DataPtr(); + if (!duckdb_footer_get_statistics(ffi_footer_ptr, idx, &statistics)) { + return {}; + } + return to_duckdb_statistics(statistics); +} + +static vector get_partition_stats(ClientContext &, GetPartitionStatsInput &input) { + const MultiFileBindData &bind_data = input.bind_data->Cast(); + VortexBindData &bind = bind_data.bind_data->Cast(); + void *const ffi_bind = bind.ffi_bind_data->DataPtr(); + + if (!duckdb_table_function_can_get_partition_stats(ffi_bind)) { + return {}; + } + + vector files = bind_data.file_list->GetAllFiles(); + vector result(files.size()); + idx_t row_start = 0; + uint64_t count = 0; + duckdb_vx_error error = nullptr; + for (size_t i = 0; i < files.size(); ++i) { + const std::string_view path = files[i].path; + duckdb_vx_data raw = duckdb_footer_get_cached(ffi_bind, path.data(), path.size(), &count, &error); + unique_ptr cdata(reinterpret_cast(raw)); + if (error) { + throw BinderException(IntoErrString(error)); + } + if (!cdata) { + return {}; + } + + PartitionStatistics &stats = result[i]; + stats.row_start = row_start; + stats.count = count; + stats.count_type = CountType::COUNT_EXACT; + row_start += count; + stats.partition_row_group = make_shared_ptr(std::move(cdata)); + } + return result; +} + duckdb_state register_table_function(DatabaseInstance &db, LogicalType parameter, const std::string &name) { - TableFunction tf(name, {}, function, duckdb_vx_table_function_bind, c_init_global, init_local); + MultiFileFunction fn(name); + fn.arguments[0] = parameter; + fn.named_parameters = {{"filename", LogicalType::ANY}, + {"allow_empty", LogicalType::BOOLEAN}, + {"hive_partitioning", LogicalType::BOOLEAN}}; - tf.projection_pushdown = true; - tf.filter_pushdown = true; - tf.filter_prune = true; - tf.sampling_pushdown = false; + fn.filter_pushdown = true; + fn.filter_prune = true; - tf.pushdown_expression = [](auto &, const auto &, Expression &expression) { + fn.pushdown_expression = [](auto &, const auto &, Expression &expression) { return duckdb_table_function_pushdown_expression(reinterpret_cast(&expression)); }; - tf.pushdown_complex_filter = [](auto &, auto &, FunctionData *bind_data, FilterVec &filters) { + fn.pushdown_complex_filter = [](auto &, auto &, FunctionData *bind_data, FilterVec &filters) { pushdown_complex_filter(*bind_data, filters); }; - tf.cardinality = c_cardinality; - tf.get_partition_info = get_partition_info; - tf.get_partition_data = get_partition_data; - tf.to_string = c_to_string; - tf.table_scan_progress = table_scan_progress; - tf.statistics = statistics; - - tf.late_materialization = true; + fn.to_string = to_string; + + fn.late_materialization = true; // Columns that uniquely identify a row for deferred re-fetch in a multi // file scan: (file index, row number in file). - tf.get_row_id_columns = [](auto &, auto) -> vector { + fn.get_row_id_columns = [](auto &, auto) -> vector { return {COLUMN_IDENTIFIER_FILE_INDEX, COLUMN_IDENTIFIER_FILE_ROW_NUMBER}; }; - tf.get_virtual_columns = [](auto &, auto) -> virtual_column_map_t { - return { - {COLUMN_IDENTIFIER_EMPTY, {"", LogicalTypeId::BOOLEAN}}, - {COLUMN_IDENTIFIER_FILE_INDEX, {"file_index", LogicalType::UBIGINT}}, - // MultiFileReader's file_row_number column is BIGINT. - // row_idx() is UBIGINT. Use UBIGINT since there's no difference to - // Duckdb what to compare. - {COLUMN_IDENTIFIER_FILE_ROW_NUMBER, {"file_row_number", LogicalType::UBIGINT}}, - }; - }; - - tf.arguments.resize(1); - tf.arguments[0] = parameter; + fn.statistics = MultiFileFunction::MultiFileScanStats; + fn.get_partition_stats = get_partition_stats; + fn.get_multi_file_reader = get_multi_file_reader; try { auto &system_catalog = Catalog::GetSystemCatalog(db); auto data = CatalogTransaction::GetSystemTransaction(db); - CreateTableFunctionInfo tf_info(tf); + CreateTableFunctionInfo tf_info(fn); tf_info.on_conflict = OnCreateConflict::ALTER_ON_CONFLICT; system_catalog.CreateFunction(data, tf_info); } catch (const std::exception &e) { @@ -472,6 +222,32 @@ duckdb_state register_table_function(DatabaseInstance &db, LogicalType parameter return DuckDBSuccess; } +extern "C" duckdb_state duckdb_vx_register_version_function(duckdb_database ffi_db, const char *version) { + D_ASSERT(ffi_db); + D_ASSERT(version); + const DatabaseWrapper &wrapper = *reinterpret_cast(ffi_db); + DatabaseInstance &db = *wrapper.database->instance; + + const string quoted = KeywordHelper::WriteQuoted(version); + + const DefaultMacro macro {DEFAULT_SCHEMA, + "vortex_version", + {nullptr}, + {{nullptr, nullptr}}, + quoted.c_str()}; + try { + auto info = DefaultFunctionGenerator::CreateInternalMacroInfo(macro); + auto &system_catalog = Catalog::GetSystemCatalog(db); + auto data = CatalogTransaction::GetSystemTransaction(db); + system_catalog.CreateFunction(data, *info); + } catch (const std::exception &e) { + ErrorData data(e); + DUCKDB_LOG_ERROR(db, "Failed to create the vortex_version macro:\t" + data.Message()); + return DuckDBError; + } + return DuckDBSuccess; +} + extern "C" duckdb_state duckdb_vx_register_table_functions(duckdb_database ffi_db) { D_ASSERT(ffi_db); const DatabaseWrapper &wrapper = *reinterpret_cast(ffi_db); diff --git a/vortex-duckdb/cpp/vector.cpp b/vortex-duckdb/cpp/vector.cpp index b6666ff6dcb..47e0492d094 100644 --- a/vortex-duckdb/cpp/vector.cpp +++ b/vortex-duckdb/cpp/vector.cpp @@ -26,11 +26,6 @@ extern "C" void duckdb_vx_vector_dictionary(duckdb_vector ffi_vector, vector->Dictionary(*dict, dictionary_size, *sel_vec, count); } -extern "C" void duckdb_vx_set_dictionary_vector_length(duckdb_vector dict, unsigned int len) { - auto ddict = reinterpret_cast(dict); - ddict->GetBuffer()->Cast().SetDictionarySize(len); -} - extern "C" void duckdb_vx_sequence_vector(duckdb_vector c_vector, int64_t start, int64_t step, idx_t capacity) { auto vector = reinterpret_cast(c_vector); diff --git a/vortex-duckdb/include/vortex.h b/vortex-duckdb/include/vortex.h index e8e4550e108..53b1e0d020c 100644 --- a/vortex-duckdb/include/vortex.h +++ b/vortex-duckdb/include/vortex.h @@ -17,62 +17,93 @@ extern "C" { #endif // __cplusplus -extern void duckdb_table_function_to_string(void *bind_data, duckdb_vx_string_map map); +extern void duckdb_table_function_to_string(const void *bind, duckdb_vx_string_map map); extern -bool duckdb_table_function_statistics(const void *bind_data, - size_t column_index, - duckdb_column_statistics *stats_out); +bool duckdb_table_function_pushdown_complex_filter(void *bind, + duckdb_vx_expr expr, + duckdb_vx_error *error); -extern double duckdb_table_function_scan_progress(void *global_state); +extern +bool duckdb_table_function_pushdown_projection_expression(void *bind, + duckdb_vx_expr expr, + size_t column_id, + duckdb_vx_error *error); extern -void duckdb_table_function_get_partition_data(void *global_init_data, - void *local_init_data, - duckdb_vx_partition_data *partition_data_out); +bool duckdb_table_function_pushdown_projection_aggregates(void *bind, + duckdb_vx_agg_input input, + duckdb_vx_error *error); + +extern bool duckdb_table_function_pushdown_expression(duckdb_vx_expr expr); extern -bool duckdb_table_function_pushdown_complex_filter(void *bind_data, - duckdb_vx_expr expr, - duckdb_vx_error *error_out); +void duckdb_table_function_cardinality(const void *bind, + uint64_t file_count, + duckdb_vx_node_statistics *stats); extern -bool duckdb_table_function_pushdown_projection_expression(void *bind_data, - duckdb_vx_expr expr, - size_t column_id, - duckdb_vx_error *error_out); +duckdb_vx_data duckdb_table_function_init_global(const duckdb_vx_tfunc_init_input *init_input, + duckdb_vx_error *error); + +extern duckdb_vx_data duckdb_table_function_init_local(const void *bind, const void *global); extern -bool duckdb_table_function_pushdown_projection_aggregates(void *bind_data, - duckdb_vx_agg_input input, - duckdb_vx_error *error_out); +duckdb_vx_data duckdb_reader_bind(const void *first_file, + duckdb_bind_result result, + duckdb_vx_error *error_out); extern -void duckdb_table_function_scan(void *global_init_data, - void *local_init_data, - duckdb_data_chunk output, - duckdb_vx_error *error_out); +duckdb_vx_data duckdb_reader_open(const char *file_path, + size_t file_path_len, + duckdb_vx_error *error); -extern bool duckdb_table_function_pushdown_expression(duckdb_vx_expr expr); +extern +bool duckdb_reader_get_statistics(const void *file, + const void *bind, + const char *column_name, + size_t column_name_len, + duckdb_column_statistics *stats_out); + +extern bool duckdb_table_function_can_get_partition_stats(const void *bind); extern -void duckdb_table_function_cardinality(void *bind_data, - duckdb_vx_node_statistics *node_stats_out); +duckdb_vx_data duckdb_footer_get_cached(void *bind, + const char *path, + size_t len, + uint64_t *row_count_out, + duckdb_vx_error *error); extern -duckdb_vx_data duckdb_table_function_init_global(const duckdb_vx_tfunc_init_input *init_input, - duckdb_vx_error *error_out); +bool duckdb_footer_get_statistics(const void *footer, + size_t column_index, + duckdb_column_statistics *stats_out); + +extern bool duckdb_reader_initialize(const void *global, void *file, duckdb_vx_error *error); + +extern duckdb_logical_type duckdb_reader_bind_column_type(const void *bind, size_t index); + +extern bool duckdb_reader_is_aggregate(const void *bind); + +extern bool duckdb_reader_try_initialize_scan(void *local, void *file); extern -duckdb_vx_data duckdb_table_function_init_local(const void *bind_data, - void *global_init_data); +bool duckdb_reader_scan(const void *file, + const void *global, + void *local, + duckdb_data_chunk chunk, + duckdb_vx_error *error); + +extern double duckdb_reader_get_progress_in_file(const void *file); extern -duckdb_vx_data duckdb_table_function_bind(duckdb_vx_tfunc_bind_input bind_input, - duckdb_vx_tfunc_bind_result bind_result, - duckdb_vx_error *error_out); +bool duckdb_reader_finalize_scan(const void *global, + duckdb_data_chunk chunk, + duckdb_vx_error *error); -extern duckdb_vx_data duckdb_table_function_bind_data_clone(const void *bind_data); +extern void duckdb_reader_finish_reading(const void *global, void *local); + +extern duckdb_vx_data duckdb_table_function_bind_data_clone(const void *bind); extern duckdb_vx_data duckdb_copy_function_copy_to_bind(const char *const *column_names, @@ -88,12 +119,35 @@ duckdb_vx_data duckdb_copy_function_copy_to_initialize_global(const void *bind_d extern void duckdb_copy_function_copy_to_sink(const void *bind_data, - void *global_data, + const void *global_data, duckdb_data_chunk data_chunk, duckdb_vx_error *error_out); extern void duckdb_copy_function_copy_to_finalize(void *global_data, duckdb_vx_error *error_out); +extern duckdb_vx_data duckdb_copy_function_prepare_batch_new(void); + +extern +void duckdb_copy_function_prepare_batch_push(const void *bind, + void *batch, + duckdb_data_chunk chunk, + duckdb_vx_error *error); + +extern +void duckdb_copy_function_flush_batch(const void *global, + const void *batch, + duckdb_vx_error *error); + +extern +bool duckdb_copy_function_get_written_file_statistics(const void *global_data, + duckdb_vx_written_file_statistics *out); + +extern +bool duckdb_copy_function_get_written_column_statistics(const void *global_data, + size_t column_index, + duckdb_vx_written_column_statistics *out, + duckdb_vx_error *error_out); + #ifdef __cplusplus } // extern "C" #endif // __cplusplus diff --git a/vortex-duckdb/patches/duckdb-mfr-finalize-scan.diff b/vortex-duckdb/patches/duckdb-mfr-finalize-scan.diff new file mode 100644 index 00000000000..d7214c7493f --- /dev/null +++ b/vortex-duckdb/patches/duckdb-mfr-finalize-scan.diff @@ -0,0 +1,69 @@ +--- a/src/include/duckdb/common/multi_file/multi_file_function.hpp ++++ b/src/include/duckdb/common/multi_file/multi_file_function.hpp +@@ -63,6 +63,9 @@ + virtual void GetVirtualColumns(ClientContext &context, MultiFileBindData &bind_data, virtual_column_map_t &result); + virtual unique_ptr Copy(); + virtual FileGlobInput GetGlobInput(); ++ virtual bool FinalizeScan(ClientContext &context, GlobalTableFunctionState &global_state, DataChunk &output) { ++ return false; ++ } + }; + + template +@@ -590,18 +593,30 @@ + + static OperatorPartitionData MultiFileGetPartitionData(ClientContext &context, + TableFunctionGetPartitionInput &input) { ++ if (!input.local_state) { ++ return OperatorPartitionData(0); ++ } + auto &bind_data = input.bind_data->CastNoConst(); + auto &data = input.local_state->Cast(); + auto &gstate = input.global_state->Cast(); + OperatorPartitionData partition_data(data.batch_index); +- bind_data.multi_file_reader->GetPartitionData(context, bind_data.reader_bind, *data.reader_data, +- gstate.multi_file_reader_state, input.partition_info, +- partition_data); ++ if (data.reader_data) { ++ bind_data.multi_file_reader->GetPartitionData(context, bind_data.reader_bind, *data.reader_data, ++ gstate.multi_file_reader_state, input.partition_info, ++ partition_data); ++ } + return partition_data; + } + + static void MultiFileScan(ClientContext &context, TableFunctionInput &data_p, DataChunk &output) { + if (!data_p.local_state) { ++ auto &finalize_gstate = data_p.global_state->Cast(); ++ auto &finalize_bind_data = data_p.bind_data->CastNoConst(); ++ if (finalize_gstate.global_state && finalize_bind_data.interface && ++ finalize_bind_data.interface->FinalizeScan(context, *finalize_gstate.global_state, output)) { ++ data_p.async_result = SourceResultType::HAVE_MORE_OUTPUT; ++ return; ++ } + data_p.async_result = SourceResultType::FINISHED; + return; + } +@@ -610,6 +625,10 @@ + auto &bind_data = data_p.bind_data->CastNoConst(); + + if (gstate.finished) { ++ if (bind_data.interface->FinalizeScan(context, *gstate.global_state, output)) { ++ data_p.async_result = SourceResultType::HAVE_MORE_OUTPUT; ++ return; ++ } + data_p.async_result = SourceResultType::FINISHED; + return; + } +@@ -660,6 +679,11 @@ + } + + if (!TryInitializeNextBatch(context, bind_data, data, gstate)) { ++ if (output.size() == 0 && ++ bind_data.interface->FinalizeScan(context, *gstate.global_state, output)) { ++ data_p.async_result = SourceResultType::HAVE_MORE_OUTPUT; ++ return; ++ } + if (output.size() > 0 && data_p.results_execution_mode == AsyncResultsExecutionMode::SYNCHRONOUS) { + gstate.finished = true; + data_p.async_result = SourceResultType::HAVE_MORE_OUTPUT; diff --git a/vortex-duckdb/src/column_statistics.rs b/vortex-duckdb/src/column_statistics.rs index ccc71eeade1..9be4ec2fefa 100644 --- a/vortex-duckdb/src/column_statistics.rs +++ b/vortex-duckdb/src/column_statistics.rs @@ -4,37 +4,34 @@ use vortex::array::stats::StatsSet; use vortex::dtype::DType; use vortex::error::VortexExpect as _; +use vortex::error::VortexResult; use vortex::expr::stats::Precision; use vortex::expr::stats::Stat; use vortex::scalar::Scalar; use vortex::scalar::ScalarValue; use crate::convert::ToDuckDBScalar as _; +use crate::duckdb::LogicalType; use crate::duckdb::Value; -#[derive(Debug, Default)] +#[derive(Debug)] pub struct ColumnStatistics { pub min: Option, pub max: Option, pub max_string_length: u64, pub has_null: bool, + pub logical_type: LogicalType, } impl ColumnStatistics { - pub fn from(stats: &ColumnStatisticsAggregate, dtype: DType) -> Self { - let min = stats.min.as_ref().map(|value| { - let value = value.clone(); + pub fn try_from(stats: ColumnStatisticsAggregate, dtype: DType) -> VortexResult { + let to_value = |value: ScalarValue| { Scalar::try_new(dtype.clone(), Some(value)) - .vortex_expect("scalar dtype and value are incompatible") - .try_to_duckdb_scalar() - .vortex_expect("can't convert Scalar to duckdb Value") - }); - let max = stats.max.as_ref().map(|value| { - Scalar::try_new(dtype.clone(), Some(value.clone())) - .vortex_expect("scalar dtype and value are incompatible") - .try_to_duckdb_scalar() - .vortex_expect("can't convert Scalar to duckdb Value") - }); + .and_then(|scalar| scalar.try_to_duckdb_scalar()) + .ok() + }; + let min = stats.min.and_then(to_value); + let max = stats.max.and_then(to_value); let max_string_length = stats .max_string_length @@ -43,12 +40,15 @@ impl ColumnStatistics { // Useful estimate if we didn't get null count stats let has_null = stats.has_null && dtype.is_nullable(); - Self { + let logical_type = LogicalType::try_from(dtype)?; + + Ok(Self { min, max, max_string_length, has_null, - } + logical_type, + }) } } diff --git a/vortex-duckdb/src/convert/dtype.rs b/vortex-duckdb/src/convert/dtype.rs index 6f95cbb8956..c994a5f53a6 100644 --- a/vortex-duckdb/src/convert/dtype.rs +++ b/vortex-duckdb/src/convert/dtype.rs @@ -58,14 +58,17 @@ use vortex::extension::datetime::Time; use vortex::extension::datetime::TimeUnit; use vortex::extension::datetime::Timestamp; use vortex::extension::datetime::is_utc_timezone; -use vortex_geo::extension::GeoMetadata; -use vortex_geo::extension::LineString; -use vortex_geo::extension::MultiLineString; -use vortex_geo::extension::MultiPoint; -use vortex_geo::extension::MultiPolygon; -use vortex_geo::extension::Point; -use vortex_geo::extension::Polygon; -use vortex_geo::extension::WellKnownBinary; +use vortex::extension::uuid::Uuid; +use vortex::extension::uuid::UuidMetadata; +use vortex_spatial::extension::LineString; +use vortex_spatial::extension::MultiLineString; +use vortex_spatial::extension::MultiPoint; +use vortex_spatial::extension::MultiPolygon; +use vortex_spatial::extension::Point; +use vortex_spatial::extension::Polygon; +use vortex_spatial::extension::SpatialMetadata; +use vortex_spatial::extension::WellKnownBinary; +use vortex_utils::aliases::hash_set::HashSet; use crate::cpp::DUCKDB_TYPE; use crate::duckdb::LogicalType; @@ -173,18 +176,25 @@ impl FromLogicalType for DType { let crs = logical_type.geometry_crs().map(|crs| crs.to_string()); DType::Extension( ExtDType::::try_new( - GeoMetadata { crs }, + SpatialMetadata { crs }, DType::Binary(nullability), )? .erased(), ) } DUCKDB_TYPE::DUCKDB_TYPE_VARIANT => DType::Variant(nullability), + DUCKDB_TYPE::DUCKDB_TYPE_UUID => DType::Extension( + ExtDType::::try_with_vtable( + Uuid, + UuidMetadata::default(), + uuid_storage_dtype(nullability), + )? + .erased(), + ), other @ (DUCKDB_TYPE::DUCKDB_TYPE_TIME_TZ | DUCKDB_TYPE::DUCKDB_TYPE_INTERVAL | DUCKDB_TYPE::DUCKDB_TYPE_ENUM | DUCKDB_TYPE::DUCKDB_TYPE_MAP - | DUCKDB_TYPE::DUCKDB_TYPE_UUID | DUCKDB_TYPE::DUCKDB_TYPE_UNION | DUCKDB_TYPE::DUCKDB_TYPE_BIT | DUCKDB_TYPE::DUCKDB_TYPE_ANY @@ -244,6 +254,7 @@ impl TryFrom<&DType> for LogicalType { DType::Struct(struct_type, _) => { return LogicalType::try_from(struct_type); } + DType::Map(..) => vortex_bail!("Vortex Map isn't supported"), // TODO(connor): Union DType::Union(..) => vortex_bail!("Vortex Union isn't supported"), DType::Variant(_) => vortex_bail!("Vortex Variant array aren't supported"), @@ -253,8 +264,12 @@ impl TryFrom<&DType> for LogicalType { return temporal_to_duckdb(temporal); } + if ext_dtype.is::() { + return Ok(LogicalType::new(DUCKDB_TYPE::DUCKDB_TYPE_UUID)); + } + // Native geometry types and WKB all surface to DuckDB as GEOMETRY so `ST_*` bind. - if let Some(geo) = ext_dtype + if let Some(spatial_metadata) = ext_dtype .metadata_opt::() .or_else(|| ext_dtype.metadata_opt::()) .or_else(|| ext_dtype.metadata_opt::()) @@ -263,7 +278,7 @@ impl TryFrom<&DType> for LogicalType { .or_else(|| ext_dtype.metadata_opt::()) .or_else(|| ext_dtype.metadata_opt::()) { - return LogicalType::geometry_type(geo.crs.as_deref()); + return LogicalType::geometry_type(spatial_metadata.crs.as_deref()); } vortex_bail!("Unsupported extension type \"{}\"", ext_dtype.id()); @@ -274,6 +289,14 @@ impl TryFrom<&DType> for LogicalType { } } +fn uuid_storage_dtype(nullability: Nullability) -> DType { + DType::FixedSizeList( + Arc::new(DType::Primitive(U8, Nullability::NonNullable)), + 16, + nullability, + ) +} + fn temporal_to_duckdb(temporal: TemporalMetadata) -> VortexResult { let duckdb_type = match temporal { TemporalMetadata::Timestamp(unit, None) => match unit { @@ -326,10 +349,14 @@ impl TryFrom<&StructFields> for LogicalType { .map(|field_dtype| LogicalType::try_from(&field_dtype)) .collect::>()?; + let mut name_set = HashSet::new(); let child_names: Vec = struct_type .names() .iter() .map(|field_name| { + if name_set.replace(field_name.as_ref()).is_some() { + vortex_bail!("Duplicate field '{field_name}'"); + } CString::new(field_name.as_ref()) .map_err(|e| vortex_err!("Invalid field name '{field_name}': {e}")) }) @@ -379,8 +406,8 @@ mod tests { use vortex::extension::datetime::Time; use vortex::extension::datetime::Timestamp; use vortex::scalar::ScalarValue; - use vortex_geo::extension::GeoMetadata; - use vortex_geo::extension::WellKnownBinary; + use vortex_spatial::extension::SpatialMetadata; + use vortex_spatial::extension::WellKnownBinary; use crate::convert::dtype::FromLogicalType; use crate::cpp; @@ -633,7 +660,7 @@ mod tests { fn test_geometry_roundtrip() -> VortexResult<()> { let vortex_geometry = DType::Extension( ExtDType::::try_new( - GeoMetadata { + SpatialMetadata { crs: Some("EPSG:4326".to_string()), }, DType::Binary(Nullability::NonNullable), @@ -657,6 +684,31 @@ mod tests { Ok(()) } + #[rstest] + #[case(Nullability::NonNullable)] + #[case(Nullability::Nullable)] + fn test_uuid_roundtrip(#[case] nullability: Nullability) -> VortexResult<()> { + use vortex::extension::uuid::Uuid; + use vortex::extension::uuid::UuidMetadata; + + let storage = DType::FixedSizeList( + Arc::new(DType::Primitive(PType::U8, Nullability::NonNullable)), + 16, + nullability, + ); + let vortex_uuid = DType::Extension( + ExtDType::::try_with_vtable(Uuid, UuidMetadata::default(), storage)?.erased(), + ); + + let duckdb_uuid = LogicalType::try_from(&vortex_uuid)?; + assert_eq!(duckdb_uuid.as_type_id(), cpp::DUCKDB_TYPE::DUCKDB_TYPE_UUID); + + let original = DType::from_logical_type(&duckdb_uuid, nullability)?; + assert_eq!(original, vortex_uuid); + + Ok(()) + } + #[test] fn test_unsupported_extension_type() { #[derive(Clone, Debug, Default, PartialEq, Eq, Hash)] diff --git a/vortex-duckdb/src/convert/expr.rs b/vortex-duckdb/src/convert/expr.rs index aee065a2a20..51aeebebb6d 100644 --- a/vortex-duckdb/src/convert/expr.rs +++ b/vortex-duckdb/src/convert/expr.rs @@ -17,7 +17,8 @@ use vortex::aggregate_fn::fns::first::First; use vortex::aggregate_fn::fns::max::Max; use vortex::aggregate_fn::fns::mean::Mean; use vortex::aggregate_fn::fns::min::Min; -use vortex::aggregate_fn::fns::sum::Sum; +use vortex::aggregate_fn::fns::sum_v2::SumV2; +use vortex::arrow::ArrowSessionExt; use vortex::dtype::DType; use vortex::dtype::Nullability; use vortex::dtype::PType; @@ -41,6 +42,7 @@ use vortex::expr::lit; use vortex::expr::not; use vortex::expr::or_collect; use vortex::expr::root; +use vortex::layout::layouts::row_idx::row_idx; use vortex::scalar::Scalar; use vortex::scalar_fn::EmptyOptions as ScalarEmptyOptions; use vortex::scalar_fn::ScalarFnVTableExt; @@ -52,18 +54,19 @@ use vortex::scalar_fn::fns::like::Like; use vortex::scalar_fn::fns::like::LikeOptions; use vortex::scalar_fn::fns::literal::Literal; use vortex::scalar_fn::fns::operators::Operator; -use vortex_geo::extension::LineString; -use vortex_geo::extension::MultiLineString; -use vortex_geo::extension::MultiPoint; -use vortex_geo::extension::MultiPolygon; -use vortex_geo::extension::Point; -use vortex_geo::extension::Polygon; -use vortex_geo::extension::WellKnownBinary; -use vortex_geo::extension::native_geometry_scalar_from_wkb; -use vortex_geo::scalar_fn::contains::GeoContains; -use vortex_geo::scalar_fn::distance::GeoDistance; -use vortex_geo::scalar_fn::intersects::GeoIntersects; - +use vortex_spatial::extension::LineString; +use vortex_spatial::extension::MultiLineString; +use vortex_spatial::extension::MultiPoint; +use vortex_spatial::extension::MultiPolygon; +use vortex_spatial::extension::Point; +use vortex_spatial::extension::Polygon; +use vortex_spatial::extension::WellKnownBinary; +use vortex_spatial::extension::native_geometry_scalar_from_wkb; +use vortex_spatial::scalar_fn::contains::SpatialContains; +use vortex_spatial::scalar_fn::distance::SpatialDistance; +use vortex_spatial::scalar_fn::intersects::SpatialIntersects; + +use crate::SESSION; use crate::convert::dtype::FromLogicalType; use crate::cpp::DUCKDB_TYPE; use crate::cpp::DUCKDB_VX_EXPR_TYPE; @@ -118,10 +121,10 @@ struct ConvertCtx<'a> { fields: Option<&'a [DuckdbField]>, } -/// Whether `name` is a non-nullable native geometry column of the scan. The pushed geo kernels -/// reject nullable operands and cannot evaluate `vortex.geo.wkb` columns, which also surface to +/// Whether `name` is a non-nullable native geometry column of the scan. The pushed spatial kernels +/// reject nullable operands and cannot evaluate `vortex.st.wkb` columns, which also surface to /// DuckDB as `GEOMETRY`. -fn is_native_geo_column(fields: Option<&[DuckdbField]>, name: &str) -> bool { +fn is_native_spatial_column(fields: Option<&[DuckdbField]>, name: &str) -> bool { fields .into_iter() .flatten() @@ -139,9 +142,9 @@ fn is_native_geo_column(fields: Option<&[DuckdbField]>, name: &str) -> bool { }) } -/// Lower a geo operand: a `GEOMETRY` literal arrives as WKB, decoded once to its native type so the -/// pushed `GeoDistance` stays native; a column must be native geometry. `None` skips the push. -fn geo_operand( +/// Lower a spatial operand: a `GEOMETRY` literal arrives as WKB, decoded once to its native type so the +/// pushed `SpatialDistance` stays native; a column must be native geometry. `None` skips the push. +fn spatial_operand( value: &duckdb::ExpressionRef, ctx: ConvertCtx<'_>, ) -> VortexResult> { @@ -158,10 +161,10 @@ fn geo_operand( let Some(buf) = storage.as_binary_opt().and_then(|b| b.value()) else { return Ok(None); }; - Ok(native_geometry_scalar_from_wkb(buf.as_slice())?.map(lit)) + Ok(native_geometry_scalar_from_wkb(buf.as_slice(), &SESSION.arrow())?.map(lit)) } Some(BoundColumnRef(col_ref)) - if is_native_geo_column(ctx.fields, col_ref.name.as_ref()) => + if is_native_spatial_column(ctx.fields, col_ref.name.as_ref()) => { try_from_expression_inner(value, ctx) } @@ -169,72 +172,72 @@ fn geo_operand( } } -/// Lower all geometry operands of a geo function. Returns `None`, skipping the push, when any +/// Lower all geometry operands of a spatial function. Returns `None`, skipping the push, when any /// operand is neither a constant geometry nor a native geometry column. -fn geo_operands( +fn spatial_operands( children: &[&duckdb::ExpressionRef], ctx: ConvertCtx<'_>, ) -> VortexResult>> { children .iter() - .map(|child| geo_operand(child, ctx)) + .map(|child| spatial_operand(child, ctx)) .collect() } -/// Lower geo UDFs to native Vortex geo ops so the work runs in the scan. `None` otherwise. -fn try_from_geo_function( +/// Lower spatial UDFs to native Vortex spatial operations so the work runs in the scan. `None` otherwise. +fn try_from_spatial_function( name: &str, func: &BoundFunction, ctx: ConvertCtx<'_>, ) -> VortexResult> { let children: Vec<_> = func.children().collect(); let expr = match name.to_ascii_lowercase().as_str() { - // Spatial's own st_dwithin folds the radius into bind data; the override + // DuckDB's spatial extension folds the radius of `ST_DWithin` into bind data; the override // (cpp/spatial_overrides.cpp) keeps it visible here as `children[2]`. "st_dwithin" => { if children.len() != 3 { return Ok(None); } - let Some(operands) = geo_operands(&children[..2], ctx)? else { + let Some(operands) = spatial_operands(&children[..2], ctx)? else { return Ok(None); }; // A non-constant radius is left for DuckDB to evaluate. let Some(distance) = from_bound_f64(children[2])? else { return Ok(None); }; - let geo_distance = GeoDistance.new_expr(ScalarEmptyOptions, operands); - Binary.new_expr(Operator::Lte, [geo_distance, lit(distance)]) + let spatial_distance = SpatialDistance.new_expr(ScalarEmptyOptions, operands); + Binary.new_expr(Operator::Lte, [spatial_distance, lit(distance)]) } "st_distance" => { if children.len() != 2 { return Ok(None); } - let Some(operands) = geo_operands(&children, ctx)? else { + let Some(operands) = spatial_operands(&children, ctx)? else { return Ok(None); }; - GeoDistance.new_expr(ScalarEmptyOptions, operands) + SpatialDistance.new_expr(ScalarEmptyOptions, operands) } "st_intersects" => { if children.len() != 2 { return Ok(None); } - let Some(operands) = geo_operands(&children, ctx)? else { + let Some(operands) = spatial_operands(&children, ctx)? else { return Ok(None); }; - GeoIntersects.new_expr(ScalarEmptyOptions, operands) + SpatialIntersects.new_expr(ScalarEmptyOptions, operands) } containment @ ("st_contains" | "st_within") => { if children.len() != 2 { return Ok(None); } - let Some(mut operands) = geo_operands(&children, ctx)? else { + let Some(mut operands) = spatial_operands(&children, ctx)? else { return Ok(None); }; // `st_within(a, b)` is `st_contains(b, a)`; both lower to the contains kernel. if containment == "st_within" { operands.swap(0, 1); } - GeoContains.new_expr(ScalarEmptyOptions, operands) + SpatialContains.new_expr(ScalarEmptyOptions, operands) } _ => return Ok(None), }; @@ -330,8 +333,8 @@ fn try_from_bound_function( return Ok(None); } } - // Geo UDFs are handled here; non-geo names return `None` inside. - name => return try_from_geo_function(name, func, ctx), + // Spatial UDFs are handled here; non-spatial names return `None` inside. + name => return try_from_spatial_function(name, func, ctx), }; Ok(Some(expr)) @@ -354,7 +357,7 @@ pub(super) fn try_from_bound_expression_with_col_sub( value: &duckdb::ExpressionRef, col_sub: &Expression, ) -> VortexResult> { - // No fields: scan-time table filters never carry geo functions, because + // No fields: scan-time table filters never carry spatial functions, because // `can_push_expression` refuses them. try_from_expression_inner( value, @@ -425,7 +428,7 @@ pub fn can_push_expression(value: &duckdb::ExpressionRef) -> bool { || name == "strlen" || name == "array_length" || (matches!(name, "len" | "length") && is_supported_length_alias(&func)) - // Geo functions are absent on purpose: they push only via + // Spatial functions are absent on purpose: they push only via // `pushdown_complex_filter`, which has the scan's fields to verify the geometry // columns are native. } @@ -524,11 +527,16 @@ impl Display for PushedAggregate { impl PushedAggregate { pub fn build(self, dtype: DType) -> VortexResult> { - let opts = NumericalAggregateOpts::default(); + let opts = if dtype.is_float() { + // duckdb treats nan as a real value, vortex defaults skip nans + NumericalAggregateOpts::include_nans() + } else { + NumericalAggregateOpts::default() + }; Ok(match self { Self::Min => Box::new(Accumulator::try_new(Min, opts, dtype)?), Self::Max => Box::new(Accumulator::try_new(Max, opts, dtype)?), - Self::Sum => Box::new(Accumulator::try_new(Sum, opts, dtype)?), + Self::Sum => Box::new(Accumulator::try_new(SumV2, opts, dtype)?), Self::Mean => Box::new(Accumulator::try_new( Mean::combined(), PairOptions(opts, opts), @@ -581,7 +589,22 @@ fn try_from_expression_inner( }; col.clone() } - BoundColumnRef(col_ref) => col(col_ref.name.as_ref()), + BoundColumnRef(col_ref) => { + let name = col_ref.name.as_ref(); + if name == "file_row_number" { + return Ok(Some(row_idx())); + } + + // Duckdb generates some columns (e.g. hive partitions) after we + // load file data, so filters on these columns can't be evaluated + if ctx + .fields + .is_some_and(|fields| !fields.iter().any(|field| field.name == name)) + { + return Ok(None); + } + col(name) + } BoundConstant(const_) => lit(Scalar::try_from(const_.value)?), BoundComparison(compare) => { let operator: Operator = compare.op.try_into()?; diff --git a/vortex-duckdb/src/convert/scalar.rs b/vortex-duckdb/src/convert/scalar.rs index 1549abdb7ee..4b072368041 100644 --- a/vortex-duckdb/src/convert/scalar.rs +++ b/vortex-duckdb/src/convert/scalar.rs @@ -22,6 +22,7 @@ use vortex::array::match_each_native_simd_ptype; use vortex::dtype::DType; use vortex::dtype::DecimalDType; +use vortex::dtype::Nullability::NonNullable; use vortex::dtype::Nullability::Nullable; use vortex::dtype::PType; use vortex::dtype::PType::I32; @@ -30,6 +31,7 @@ use vortex::dtype::half::f16; use vortex::error::VortexError; use vortex::error::VortexResult; use vortex::error::vortex_bail; +use vortex::error::vortex_ensure; use vortex::error::vortex_err; use vortex::extension::datetime::AnyTemporal; use vortex::extension::datetime::Date; @@ -39,6 +41,7 @@ use vortex::extension::datetime::TimeUnit; use vortex::extension::datetime::Timestamp; use vortex::extension::datetime::TimestampOptions; use vortex::extension::datetime::is_utc_timezone; +use vortex::extension::uuid::Uuid; use vortex::scalar::BinaryScalar; use vortex::scalar::BoolScalar; use vortex::scalar::DecimalScalar; @@ -48,7 +51,7 @@ use vortex::scalar::PrimitiveScalar; use vortex::scalar::Scalar; use vortex::scalar::ScalarValue; use vortex::scalar::Utf8Scalar; -use vortex_geo::extension::WellKnownBinary; +use vortex_spatial::extension::WellKnownBinary; use crate::convert::dtype::FromLogicalType; use crate::duckdb::LogicalType; @@ -83,6 +86,7 @@ impl ToDuckDBScalar for Scalar { DType::FixedSizeList(..) => { vortex_bail!("Vortex FixedSizeList scalars aren't supported") } + DType::Map(..) => vortex_bail!("Vortex Map scalars aren't supported"), DType::Variant(_) => vortex_bail!("Vortex Variant scalars aren't supported"), DType::Struct(..) => vortex_bail!("Vortex Struct scalars aren't supported"), // TODO(connor): Union @@ -178,8 +182,6 @@ impl ToDuckDBScalar for ExtScalar<'_> { /// Converts an extension scalar (temporal types or `WellKnownBinary` geometries) to a DuckDB /// value. fn try_to_duckdb_scalar(&self) -> VortexResult { - let logical_type = LogicalType::try_from(&DType::Extension(self.ext_dtype().clone()))?; - if let Some(wkb) = self.ext_dtype().metadata_opt::() { let storage = self.to_storage_scalar(); let binary = storage @@ -187,7 +189,7 @@ impl ToDuckDBScalar for ExtScalar<'_> { .ok_or_else(|| vortex_err!("WellKnownBinary storage must be a binary scalar"))?; return Ok(match binary.value() { Some(bytes) => Value::new_geometry(bytes.as_slice(), wkb.crs.as_deref())?, - None => Value::null(&logical_type), + None => Value::null(&*ext_logical_type(self)?), }); } @@ -195,12 +197,9 @@ impl ToDuckDBScalar for ExtScalar<'_> { vortex_bail!("Cannot convert non-temporal extension scalar to duckdb value"); }; + let storage = PrimitiveScalar::try_new(self.ext_dtype().storage_dtype(), self.value())?; let value = || { - self.to_storage_scalar() - .as_primitive_opt() - .ok_or_else(|| { - vortex_err!("Cannot have a temporal time type not packed by a primitive scalar") - })? + storage .as_::() .ok_or_else(|| vortex_err!("temporal types must be convertible to i64")) }; @@ -215,7 +214,11 @@ impl ToDuckDBScalar for ExtScalar<'_> { "Currently only UTC timezone is supported for duckdb timestamp(tz) conversion" ); } - return Ok(Value::new_timestamp_tz(value()?)); + // TIMESTAMP_TZ stores time in UTC microseconds, tz is a display sign. + return Ok(Value::new_timestamp_tz(timestamp_tz_micros( + *unit, + value()?, + )?)); } match unit { TimeUnit::Nanoseconds => Value::new_timestamp_ns(value()?), @@ -228,29 +231,40 @@ impl ToDuckDBScalar for ExtScalar<'_> { } } TemporalMetadata::Date(unit) => match unit { - TimeUnit::Days => self - .to_storage_scalar() - .as_primitive_opt() - .ok_or_else(|| { - vortex_err!("temporal types must be backed by primitive scalars") - })? - .as_::() - .map(Value::new_date) - .unwrap_or_else(|| Value::null(&logical_type)), + TimeUnit::Days => match storage.as_::() { + Some(days) => Value::new_date(days), + None => Value::null(&*ext_logical_type(self)?), + }, _ => vortex_bail!("cannot have TimeUnit {unit}, so represent a day"), }, TemporalMetadata::Time(unit) => match unit { TimeUnit::Microseconds => Value::new_time(value()?), TimeUnit::Milliseconds => Value::new_time(value()? * 1000), TimeUnit::Seconds => Value::new_time(value()? * 1000 * 1000), - TimeUnit::Nanoseconds | TimeUnit::Days => { - vortex_bail!("cannot convert timeunit {unit} to a duckdb MS time") + TimeUnit::Nanoseconds => Value::new_time_ns(value()?), + TimeUnit::Days => { + vortex_bail!("cannot convert timeunit {unit} to a duckdb time") } }, }) } } +fn ext_logical_type(scalar: &ExtScalar<'_>) -> VortexResult { + LogicalType::try_from(&DType::Extension(scalar.ext_dtype().clone())) +} + +fn timestamp_tz_micros(unit: TimeUnit, raw: i64) -> VortexResult { + let overflow = || vortex_err!("timestamp_tz overflow rescaling {raw}{unit} to micros"); + match unit { + TimeUnit::Seconds => raw.checked_mul(1_000_000).ok_or_else(overflow), + TimeUnit::Milliseconds => raw.checked_mul(1_000).ok_or_else(overflow), + TimeUnit::Microseconds => Ok(raw), + TimeUnit::Nanoseconds => Ok(raw / 1_000), + TimeUnit::Days => vortex_bail!("timestamp_tz cannot have a day time unit"), + } +} + impl TryFrom for Scalar { type Error = VortexError; @@ -299,6 +313,19 @@ impl<'a> TryFrom<&'a ValueRef> for Scalar { ext.clone(), Scalar::binary(b, Nullable), )), + DType::Extension(ext) if ext.is::() => { + vortex_ensure!(b.len() == 16, "UUID blob must be 16 bytes, got {}", b.len()); + let children = b + .iter() + .map(|&byte| Scalar::primitive(byte, NonNullable)) + .collect(); + let storage = Scalar::fixed_size_list( + DType::Primitive(PType::U8, NonNullable), + children, + Nullable, + ); + Ok(Scalar::extension_ref(ext.clone(), storage)) + } _ => vortex_bail!("Cannot convert DuckDB blob to Vortex scalar of dtype {dtype}"), }, ExtractedValue::Date(days) => Ok(Scalar::extension::( @@ -315,6 +342,13 @@ impl<'a> TryFrom<&'a ValueRef> for Scalar { Some(ScalarValue::from(micros)), )?, )), + ExtractedValue::TimeNs(nanos) => Ok(Scalar::extension::

(array_len)?; let array_len_u64 = array_len as u64; // Load kernel function @@ -115,7 +114,7 @@ where let cuda_function = ctx.load_function("constant_numeric", &kernel_ptypes)?; ctx.launch_kernel(&cuda_function, array_len, |args| { - args.arg(&output_view); + args.arg(&mut output_buffer); args.arg(&value); args.arg(&array_len_u64); })?; @@ -163,8 +162,7 @@ where .ok_or_else(|| vortex_err!("Failed to cast decimal value to native type"))?; // Allocate output buffer on device - let output_buffer = ctx.device_alloc::(array_len)?; - let output_view = output_buffer.as_view(); + let mut output_buffer = ctx.device_alloc::(array_len)?; let array_len_u64 = array_len as u64; // Load kernel function @@ -172,7 +170,7 @@ where ctx.load_function_with_suffixes("constant_numeric", &[&D::DECIMAL_TYPE.to_string()])?; ctx.launch_kernel(&cuda_function, array_len, |args| { - args.arg(&output_view); + args.arg(&mut output_buffer); args.arg(&value); args.arg(&array_len_u64); })?; diff --git a/vortex-cuda/src/kernel/arrays/dict.rs b/vortex-cuda/src/kernel/arrays/dict.rs index 840a3cdbbac..dea04a0cd0e 100644 --- a/vortex-cuda/src/kernel/arrays/dict.rs +++ b/vortex-cuda/src/kernel/arrays/dict.rs @@ -138,12 +138,10 @@ async fn execute_dict_bool_typed( // Each CUDA thread owns complete output bytes, avoiding races between threads writing // different bits in the same byte. The kernel handles the final partial byte explicitly. let output_bytes = codes_len.div_ceil(8); - let output_slice = ctx.device_alloc::(output_bytes)?; - let output_device = CudaDeviceBuffer::new(output_slice); + let mut output_slice = ctx.device_alloc::(output_bytes)?; let values_view = values_device.cuda_view::()?; let codes_view = codes_device.cuda_view::()?; - let output_view = output_device.as_view::(); let codes_len_u64 = codes_len as u64; let values_offset_u64 = values_meta.offset() as u64; @@ -154,9 +152,10 @@ async fn execute_dict_bool_typed( .arg(&codes_len_u64) .arg(&values_view) .arg(&values_offset_u64) - .arg(&output_view); + .arg(&mut output_slice); })?; + let output_device = CudaDeviceBuffer::new(output_slice); Ok(Canonical::Bool(BoolArray::new_handle( BufferHandle::new_device(Arc::new(output_device)), 0, @@ -190,13 +189,11 @@ async fn execute_dict_prim_typed(codes_len)?; - let output_device = CudaDeviceBuffer::new(output_slice); + let mut output_slice = ctx.device_alloc::(codes_len)?; // Get views for kernel launch let values_view = values_device.cuda_view::()?; let codes_view = codes_device.cuda_view::()?; - let output_view = output_device.as_view::(); let codes_len_u64 = codes_len as u64; @@ -205,9 +202,10 @@ async fn execute_dict_prim_typed(codes_len)?; - let output_device = CudaDeviceBuffer::new(output_slice); + let mut output_slice = ctx.device_alloc::(codes_len)?; // Get views for kernel launch let values_view = values_device.cuda_view::()?; let codes_view = codes_device.cuda_view::()?; - let output_view = output_device.as_view::(); // Load kernel function using string suffixes let cuda_function = ctx.load_function_with_suffixes( @@ -291,9 +287,10 @@ async fn execute_dict_decimal_typed< args.arg(&codes_view) .arg(&codes_len_u64) .arg(&values_view) - .arg(&output_view); + .arg(&mut output_slice); })?; + let output_device = CudaDeviceBuffer::new(output_slice); Ok(Canonical::Decimal(DecimalArray::new_handle( BufferHandle::new_device(Arc::new(output_device)), V::DECIMAL_TYPE, @@ -338,15 +335,13 @@ async fn execute_dict_varbinview( let codes_device = ctx.ensure_on_device(codes_buffer).await?; // Allocate output: one i128 per code. - let output_slice = ctx.device_alloc::(codes_len)?; - let output_device = CudaDeviceBuffer::new(output_slice); + let mut output_slice = ctx.device_alloc::(codes_len)?; // Dispatch by code type, reusing the existing dict_i128_ kernel. // BinaryView is repr(C, align(16)) and 16 bytes — identical layout to i128. match_each_integer_ptype!(codes_ptype, |C| { let values_view = values_device.cuda_view::()?; let codes_view = codes_device.cuda_view::()?; - let output_view = output_device.as_view::(); let codes_ptype_str = C::PTYPE.to_string(); let cuda_function = ctx.load_function_with_suffixes("dict", &["i128", &codes_ptype_str])?; @@ -357,10 +352,12 @@ async fn execute_dict_varbinview( args.arg(&codes_view); args.arg(&codes_len_u64); args.arg(&values_view); - args.arg(&output_view); + args.arg(&mut output_slice); })?; }); + let output_device = CudaDeviceBuffer::new(output_slice); + // Output views gathered by the kernel share the values' data buffers. // Outlined views reference into these buffers via buffer_index + offset, // and inlined views are self-contained within the 16-byte view. diff --git a/vortex-cuda/src/kernel/arrays/list.rs b/vortex-cuda/src/kernel/arrays/list.rs new file mode 100644 index 00000000000..d6aab480e4e --- /dev/null +++ b/vortex-cuda/src/kernel/arrays/list.rs @@ -0,0 +1,214 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use async_trait::async_trait; +use cudarc::driver::DeviceRepr; +use cudarc::driver::PushKernelArg; +use tracing::instrument; +use vortex::array::ArrayRef; +use vortex::array::Canonical; +use vortex::array::IntoArray; +use vortex::array::arrays::List; +use vortex::array::arrays::ListViewArray; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::list::ListArrayExt; +use vortex::array::arrays::list::ListArraySlotsExt; +use vortex::array::arrays::primitive::PrimitiveDataParts; +use vortex::array::buffer::BufferHandle; +use vortex::array::match_each_integer_ptype; +use vortex::array::validity::Validity; +use vortex::dtype::NativePType; +use vortex::dtype::Nullability; +use vortex::error::VortexResult; +use vortex::error::vortex_ensure; +use vortex::error::vortex_err; + +use crate::CudaBufferExt; +use crate::CudaDeviceBuffer; +use crate::executor::CudaArrayExt; +use crate::executor::CudaExecute; +use crate::executor::CudaExecutionCtx; +use crate::executor::execute_validity_cuda; + +/// CUDA executor for `ListArray`. +/// +/// `List` stores `len + 1` Arrow-style offsets; its canonical form, `ListView`, stores one +/// offset and one size per list. Decode the elements on the GPU and derive the view pair from +/// the offsets with a single kernel. +#[derive(Debug)] +pub(crate) struct ListExecutor; + +#[async_trait] +impl CudaExecute for ListExecutor { + #[instrument(level = "trace", skip_all, fields(executor = ?self))] + async fn execute( + &self, + array: ArrayRef, + ctx: &mut CudaExecutionCtx, + ) -> VortexResult { + let list = array + .try_downcast::() + .map_err(|_| vortex_err!("Expected ListArray"))?; + + let list_len = list.len(); + let validity = execute_validity_cuda(list.list_validity(), list_len, ctx).await?; + let elements = list + .elements() + .clone() + .execute_cuda(ctx) + .await? + .into_array(); + + if list_len == 0 { + let empty = PrimitiveArray::empty::(Nullability::NonNullable); + return Ok(Canonical::List(ListViewArray::try_new( + elements, + empty.clone().into_array(), + empty.into_array(), + validity, + )?)); + } + + let offsets = list + .offsets() + .clone() + .execute_cuda(ctx) + .await? + .into_primitive(); + vortex_ensure!( + offsets.len() == list_len + 1, + "ListArray must have {} offsets, got {}", + list_len + 1, + offsets.len() + ); + + let offsets_ptype = offsets.ptype(); + match_each_integer_ptype!(offsets_ptype, |O| { + list_views_typed::(offsets, elements, validity, list_len, ctx).await + }) + } +} + +async fn list_views_typed( + offsets: PrimitiveArray, + elements: ArrayRef, + validity: Validity, + list_len: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + let PrimitiveDataParts { + buffer: offsets_buffer, + .. + } = offsets.into_data_parts(); + + let offsets_device = ctx.ensure_on_device(offsets_buffer).await?; + let offsets_view = offsets_device.cuda_view::()?; + + let mut view_offsets = ctx.device_alloc::(list_len)?; + let mut view_sizes = ctx.device_alloc::(list_len)?; + let list_len_u64 = list_len as u64; + + let offsets_ptype = O::PTYPE.to_string(); + let cuda_function = ctx.load_function_with_suffixes("list", &["views", &offsets_ptype])?; + ctx.launch_kernel(&cuda_function, list_len, |args| { + args.arg(&offsets_view) + .arg(&mut view_offsets) + .arg(&mut view_sizes) + .arg(&list_len_u64); + })?; + + let view_offsets = PrimitiveArray::from_buffer_handle( + BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(view_offsets))), + O::PTYPE, + Validity::NonNullable, + ); + let view_sizes = PrimitiveArray::from_buffer_handle( + BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(view_sizes))), + O::PTYPE, + Validity::NonNullable, + ); + + Ok(Canonical::List(ListViewArray::try_new( + elements, + view_offsets.into_array(), + view_sizes.into_array(), + validity, + )?)) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex::array::arrays::BoolArray; + use vortex::array::arrays::ListArray; + use vortex::array::assert_arrays_eq; + use vortex::buffer::Buffer; + use vortex::buffer::buffer; + use vortex::error::VortexExpect; + use vortex_array::VortexSessionExecute; + + use super::*; + use crate::CanonicalCudaExt; + use crate::session::CudaSession; + + #[rstest] + #[case::single_run(vec![0i32, 2, 5, 9])] + #[case::empty_lists(vec![0i32, 0, 3, 3])] + #[crate::test] + async fn test_cuda_list_decompression(#[case] offsets: Vec) -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let element_count = *offsets.last().vortex_expect("offsets are non-empty"); + let elements = PrimitiveArray::new( + (0..element_count).collect::>(), + Validity::NonNullable, + ) + .into_array(); + let offsets_array = + PrimitiveArray::new(Buffer::from(offsets), Validity::NonNullable).into_array(); + let list = ListArray::try_new(elements, offsets_array, Validity::NonNullable)?; + + let gpu_result = ListExecutor + .execute(list.clone().into_array(), &mut cuda_ctx) + .await + .vortex_expect("GPU decompression failed") + .into_host() + .await? + .into_array(); + + assert_arrays_eq!(list, gpu_result, &mut ctx); + + Ok(()) + } + + #[crate::test] + async fn test_cuda_list_with_nulls() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let elements = + PrimitiveArray::new(buffer![10i32, 20, 30, 40], Validity::NonNullable).into_array(); + let offsets = + PrimitiveArray::new(buffer![0i32, 2, 2, 4], Validity::NonNullable).into_array(); + let validity = + Validity::Array(BoolArray::from_iter([true, false, true].into_iter()).into_array()); + let list = ListArray::try_new(elements, offsets, validity)?; + + let gpu_result = ListExecutor + .execute(list.clone().into_array(), &mut cuda_ctx) + .await + .vortex_expect("GPU decompression failed") + .into_host() + .await? + .into_array(); + + assert_arrays_eq!(list, gpu_result, &mut ctx); + + Ok(()) + } +} diff --git a/vortex-cuda/src/kernel/arrays/masked.rs b/vortex-cuda/src/kernel/arrays/masked.rs new file mode 100644 index 00000000000..b739ed2bf60 --- /dev/null +++ b/vortex-cuda/src/kernel/arrays/masked.rs @@ -0,0 +1,99 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use async_trait::async_trait; +use tracing::instrument; +use vortex::array::ArrayRef; +use vortex::array::Canonical; +use vortex::array::arrays::Masked; +use vortex::array::arrays::masked::MaskedArrayExt; +use vortex::array::arrays::masked::MaskedArraySlotsExt; +use vortex::array::arrays::masked::mask_validity_canonical; +use vortex::array::validity::Validity; +use vortex::error::VortexResult; +use vortex::error::vortex_bail; +use vortex::error::vortex_err; + +use crate::executor::CudaArrayExt; +use crate::executor::CudaExecute; +use crate::executor::CudaExecutionCtx; +use crate::executor::execute_validity_cuda; + +/// CUDA executor for MaskedArray. +/// +/// A `MaskedArray` is a child array that carries no nulls of its own, plus the validity +/// bitmap that supplies them. Decode the child on the GPU, decode the mask on the GPU, and +/// attach the mask to the result. +#[derive(Debug)] +pub(crate) struct MaskedExecutor; + +#[async_trait] +impl CudaExecute for MaskedExecutor { + #[instrument(level = "trace", skip_all, fields(executor = ?self))] + async fn execute( + &self, + array: ArrayRef, + ctx: &mut CudaExecutionCtx, + ) -> VortexResult { + let masked = array + .try_downcast::() + .map_err(|_| vortex_err!("Expected MaskedArray"))?; + + let len = masked.len(); + let validity = masked.masked_validity(); + + // `MaskedArray` guarantees its child holds no nulls, so the mask alone determines the + // output validity. Combining two device-resident bitmaps would need a CPU compute pass. + if matches!(masked.child().validity()?, Validity::Array(_)) { + vortex_bail!( + "MaskedArray child carries a per-element validity bitmap, which cannot be combined with the mask on the GPU" + ); + } + + let child = masked.child().clone().execute_cuda(ctx).await?; + + let validity = execute_validity_cuda(validity, len, ctx).await?; + mask_validity_canonical(child, validity, ctx.execution_ctx()) + } +} + +#[cfg(test)] +mod tests { + use vortex::array::IntoArray; + use vortex::array::arrays::BoolArray; + use vortex::array::arrays::MaskedArray; + use vortex::array::arrays::PrimitiveArray; + use vortex::array::assert_arrays_eq; + use vortex::buffer::buffer; + use vortex::error::VortexExpect; + use vortex_array::VortexSessionExecute; + + use super::*; + use crate::CanonicalCudaExt; + use crate::session::CudaSession; + + #[crate::test] + async fn test_cuda_masked_applies_validity() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let child = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); + let validity = Validity::Array( + BoolArray::from_iter([true, false, true, true].into_iter()).into_array(), + ); + let masked = MaskedArray::try_new(child, validity)?; + + let gpu_result = MaskedExecutor + .execute(masked.clone().into_array(), &mut cuda_ctx) + .await + .vortex_expect("GPU decompression failed") + .into_host() + .await? + .into_array(); + + assert_arrays_eq!(masked, gpu_result, &mut ctx); + + Ok(()) + } +} diff --git a/vortex-cuda/src/kernel/arrays/mod.rs b/vortex-cuda/src/kernel/arrays/mod.rs index ab81934bb27..56d4c483480 100644 --- a/vortex-cuda/src/kernel/arrays/mod.rs +++ b/vortex-cuda/src/kernel/arrays/mod.rs @@ -3,8 +3,12 @@ mod constant; mod dict; +mod list; +mod masked; mod shared; pub(crate) use constant::ConstantNumericExecutor; pub(crate) use dict::DictExecutor; +pub(crate) use list::ListExecutor; +pub(crate) use masked::MaskedExecutor; pub(crate) use shared::SharedExecutor; diff --git a/vortex-cuda/src/kernel/encodings/alp.rs b/vortex-cuda/src/kernel/encodings/alp.rs index d027aa0bbd6..832d67725e4 100644 --- a/vortex-cuda/src/kernel/encodings/alp.rs +++ b/vortex-cuda/src/kernel/encodings/alp.rs @@ -85,9 +85,7 @@ where // Allocate output rounded up to a full chunk: the fused kernel writes a // whole 1024-element chunk per block, and we slice off any padding below. - let output_slice = ctx.device_alloc::(array_len.next_multiple_of(1024))?; - let output_buf = CudaDeviceBuffer::new(output_slice); - let output_view = output_buf.as_view::(); + let mut output_slice = ctx.device_alloc::(array_len.next_multiple_of(1024))?; // Patch validity does not need to be scattered: the ALP encoder strips null // positions from the exception list, so patches only exist at valid @@ -116,7 +114,7 @@ where let array_len_u64 = array_len as u64; ctx.launch_kernel_config(&cuda_function, config, array_len, |args| { args.arg(&input_view) - .arg(&output_view) + .arg(&mut output_slice) .arg(&f) .arg(&e) .arg(&array_len_u64) @@ -127,6 +125,7 @@ where ctx.synchronize_stream()?; drop(device_patches); + let output_buf = CudaDeviceBuffer::new(output_slice); let output_handle = BufferHandle::new_device(output_buf.slice_typed::(0..array_len)); Ok(Canonical::Primitive(PrimitiveArray::from_buffer_handle( output_handle, diff --git a/vortex-cuda/src/kernel/encodings/bitpacked.rs b/vortex-cuda/src/kernel/encodings/bitpacked.rs index 0f00156058b..86b7a88b276 100644 --- a/vortex-cuda/src/kernel/encodings/bitpacked.rs +++ b/vortex-cuda/src/kernel/encodings/bitpacked.rs @@ -16,7 +16,7 @@ use vortex::array::ArrayView; use vortex::array::Canonical; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::Slice; -use vortex::array::arrays::slice::SliceArrayExt; +use vortex::array::arrays::slice::SliceArraySlotsExt; use vortex::array::buffer::BufferHandle; use vortex::array::buffer::DeviceBufferExt; use vortex::array::match_each_integer_ptype; @@ -180,9 +180,7 @@ where let output_len = offset + len; // Allocate output buffer - let output_slice = ctx.device_alloc::(output_len.next_multiple_of(1024))?; - let output_buf = CudaDeviceBuffer::new(output_slice); - let output_view = output_buf.as_view::(); + let mut output_slice = ctx.device_alloc::(output_len.next_multiple_of(1024))?; let output_width = size_of::() * 8; let cuda_function = bitpacked_cuda_kernel(bit_width, output_width, ctx)?; @@ -203,14 +201,17 @@ where ctx.launch_kernel_config(&cuda_function, config, output_len, |args| { args.arg(&input_view) - .arg(&output_view) + .arg(&mut output_slice) .arg(&reference) .arg(&patches_arg); })?; - // NOTE: we must synchronize here, as the device patches are only alive for this call. - ctx.synchronize_stream()?; + // Patch-free decodes need no host synchronization. + if device_patches.is_some() { + ctx.synchronize_stream()?; + } + let output_buf = CudaDeviceBuffer::new(output_slice); let output_handle = BufferHandle::new_device(output_buf.slice_typed::(offset..(offset + len))); diff --git a/vortex-cuda/src/kernel/encodings/date_time_parts.rs b/vortex-cuda/src/kernel/encodings/date_time_parts.rs index dcb1fd1f69e..5bd8b019ee2 100644 --- a/vortex-cuda/src/kernel/encodings/date_time_parts.rs +++ b/vortex-cuda/src/kernel/encodings/date_time_parts.rs @@ -15,7 +15,7 @@ use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::TemporalArray; use vortex::array::arrays::primitive::PrimitiveDataParts; use vortex::array::buffer::BufferHandle; -use vortex::array::match_each_signed_integer_ptype; +use vortex::array::match_each_integer_ptype; use vortex::array::validity::Validity; use vortex::dtype::DType; use vortex::dtype::NativePType; @@ -110,9 +110,9 @@ impl CudaExecute for DateTimePartsExecutor { let seconds_ptype = seconds_prim.ptype(); let subseconds_ptype = subseconds_prim.ptype(); - match_each_signed_integer_ptype!(days_ptype, |DaysT| { - match_each_signed_integer_ptype!(seconds_ptype, |SecondsT| { - match_each_signed_integer_ptype!(subseconds_ptype, |SubsecondsT| { + match_each_integer_ptype!(days_ptype, |DaysT| { + match_each_integer_ptype!(seconds_ptype, |SecondsT| { + match_each_integer_ptype!(subseconds_ptype, |SubsecondsT| { decode_datetimeparts_typed::( days_prim, seconds_prim, @@ -167,13 +167,11 @@ where let subseconds_device = ctx.ensure_on_device(subseconds_buffer).await?; // Allocate output buffer - let output_slice = ctx.device_alloc::(output_len)?; - let output_device = CudaDeviceBuffer::new(output_slice); + let mut output_slice = ctx.device_alloc::(output_len)?; let days_view = days_device.cuda_view::()?; let seconds_view = seconds_device.cuda_view::()?; let subseconds_view = subseconds_device.cuda_view::()?; - let output_view = output_device.as_view::(); let cuda_function = ctx.load_function( "date_time_parts", @@ -187,10 +185,11 @@ where .arg(&seconds_view) .arg(&subseconds_view) .arg(&divisor) - .arg(&output_view) + .arg(&mut output_slice) .arg(&array_len_u64); })?; + let output_device = CudaDeviceBuffer::new(output_slice); let output_buffer = BufferHandle::new_device(Arc::new(output_device)); let output_primitive = PrimitiveArray::from_buffer_handle(output_buffer, PType::I64, validity); @@ -298,6 +297,88 @@ mod tests { Ok(()) } + /// Compression narrows each component to its smallest ptype, choosing unsigned whenever the + /// component is non-negative. Every post-epoch timestamp column therefore decomposes into + /// entirely unsigned components, which is what panicked on `taxi`. + #[crate::test] + async fn test_cuda_datetimeparts_unsigned_components() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let len = 3; + let days_arr = PrimitiveArray::new(buffer![1u32, 2, 3], Validity::NonNullable).into_array(); + let seconds_arr = + PrimitiveArray::new(buffer![3600u16, 0, 60], Validity::NonNullable).into_array(); + let subseconds_arr = + PrimitiveArray::new(buffer![250u8, 0, 99], Validity::NonNullable).into_array(); + + let temporal = TemporalArray::new_timestamp( + PrimitiveArray::new(buffer![0i64; len], Validity::NonNullable).into_array(), + TimeUnit::Milliseconds, + None, + ); + let dtp_array = DateTimeParts::try_new( + temporal.dtype().clone(), + days_arr, + seconds_arr, + subseconds_arr, + )?; + + let gpu_result = DateTimePartsExecutor + .execute(dtp_array.clone().into_array(), &mut cuda_ctx) + .await + .vortex_expect("GPU decompression failed") + .into_host() + .await? + .into_array(); + + assert_arrays_eq!(dtp_array, gpu_result, &mut ctx); + + Ok(()) + } + + /// Components are narrowed independently, so their signedness can differ within one array. + /// A timestamp column straddling the epoch narrows to unsigned days (all zero) with signed + /// seconds, which is only decoded correctly if each component keeps its own signedness. + #[crate::test] + async fn test_cuda_datetimeparts_mixed_signedness() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + // Milliseconds -32_000, 0 and 31_000 split into these components. + let days_arr = PrimitiveArray::new(buffer![0u8, 0, 0], Validity::NonNullable).into_array(); + let seconds_arr = + PrimitiveArray::new(buffer![-32i8, 0, 31], Validity::NonNullable).into_array(); + let subseconds_arr = + PrimitiveArray::new(buffer![0u8, 0, 0], Validity::NonNullable).into_array(); + + let temporal = TemporalArray::new_timestamp( + PrimitiveArray::new(buffer![0i64; 3], Validity::NonNullable).into_array(), + TimeUnit::Milliseconds, + None, + ); + let dtp_array = DateTimeParts::try_new( + temporal.dtype().clone(), + days_arr, + seconds_arr, + subseconds_arr, + )?; + + let gpu_result = DateTimePartsExecutor + .execute(dtp_array.clone().into_array(), &mut cuda_ctx) + .await + .vortex_expect("GPU decompression failed") + .into_host() + .await? + .into_array(); + + assert_arrays_eq!(dtp_array, gpu_result, &mut ctx); + + Ok(()) + } + #[crate::test] async fn test_cuda_datetimeparts_large_array() -> VortexResult<()> { let mut ctx = vortex_array::array_session().create_execution_ctx(); diff --git a/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs b/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs index 5f7e639b8a1..3475f26a175 100644 --- a/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs +++ b/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs @@ -10,7 +10,7 @@ use vortex::array::Canonical; use vortex::array::arrays::DecimalArray; use vortex::array::arrays::primitive::PrimitiveDataParts; use vortex::encodings::decimal_byte_parts::DecimalByteParts; -use vortex::encodings::decimal_byte_parts::DecimalBytePartsArrayExt; +use vortex::encodings::decimal_byte_parts::DecimalBytePartsArraySlotsExt; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_bail; diff --git a/vortex-cuda/src/kernel/encodings/delta.rs b/vortex-cuda/src/kernel/encodings/delta.rs new file mode 100644 index 00000000000..5dde85cfd65 --- /dev/null +++ b/vortex-cuda/src/kernel/encodings/delta.rs @@ -0,0 +1,285 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! CUDA executor for FastLanes delta. +//! +//! Delta stores its values in the FastLanes transposed layout: a 1024-element chunk is +//! `LANES` independent columns, each carrying its own running total seeded from that lane's +//! base. A chunk therefore decodes as `LANES` independent scans, not one scan of 1024 +//! elements, so the whole array is data-parallel across both chunks and lanes. +//! +//! This mirrors the CPU decoder in `vortex-fastlanes`: `undelta` into the transposed layout, +//! then untranspose into natural order, then apply the array's logical slice. + +use std::fmt::Debug; +use std::sync::Arc; + +use async_trait::async_trait; +use cudarc::driver::LaunchConfig; +use cudarc::driver::PushKernelArg; +use tracing::instrument; +use vortex::array::ArrayRef; +use vortex::array::Canonical; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::primitive::PrimitiveDataParts; +use vortex::array::buffer::BufferHandle; +use vortex::array::match_each_unsigned_integer_ptype; +use vortex::dtype::NativePType; +use vortex::encodings::fastlanes::Delta; +use vortex::encodings::fastlanes::DeltaArray; +use vortex::encodings::fastlanes::DeltaArrayExt; +use vortex::encodings::fastlanes::DeltaArraySlotsExt; +use vortex::error::VortexResult; +use vortex::error::vortex_ensure; +use vortex::error::vortex_err; + +use crate::CudaBufferExt; +use crate::CudaDeviceBuffer; +use crate::executor::CudaArrayExt; +use crate::executor::CudaExecute; +use crate::executor::CudaExecutionCtx; +use crate::executor::execute_validity_cuda; + +/// Elements per FastLanes chunk. Must match `FL_CHUNK` in `kernels/src/fastlanes_common.cuh`. +const FL_CHUNK: usize = 1024; +/// Threads per block: covers the widest lane count (128, for 8-bit values) in a single pass. +const BLOCK_THREADS: u32 = 128; + +/// CUDA decoder for FastLanes delta. +#[derive(Debug)] +pub(crate) struct DeltaExecutor; + +#[async_trait] +impl CudaExecute for DeltaExecutor { + #[instrument(level = "trace", skip_all, fields(executor = ?self))] + async fn execute( + &self, + array: ArrayRef, + ctx: &mut CudaExecutionCtx, + ) -> VortexResult { + let delta = array + .try_downcast::() + .map_err(|_| vortex_err!("Expected DeltaArray"))?; + decode_delta(delta, ctx).await + } +} + +#[instrument(skip_all)] +async fn decode_delta(array: DeltaArray, ctx: &mut CudaExecutionCtx) -> VortexResult { + let dtype = array.dtype().clone(); + let len = array.len(); + if len == 0 { + return Ok(Canonical::empty(&dtype)); + } + + // The vtable already narrows validity to the logical slice. + let validity = execute_validity_cuda(array.validity()?, len, ctx).await?; + + let deltas = array + .deltas() + .clone() + .execute_cuda(ctx) + .await? + .into_primitive(); + let bases = array + .bases() + .clone() + .execute_cuda(ctx) + .await? + .into_primitive(); + + // Signed values decode through their unsigned counterpart: the kernel's wrapping add + // inverts the wrapping subtract applied at compress time regardless of signedness. The + // buffer is untyped, so only the kernel and the device view need the unsigned type. + let ptype = deltas.ptype(); + let deltas_len = deltas.len(); + let offset = array.offset(); + vortex_ensure!( + deltas_len % FL_CHUNK == 0, + "Delta deltas child must be padded to a multiple of {FL_CHUNK}, got {deltas_len}" + ); + vortex_ensure!( + offset + len <= deltas_len, + "Delta slice {offset}..{} exceeds its {deltas_len} decoded values", + offset + len + ); + let num_chunks = deltas_len / FL_CHUNK; + let lanes = FL_CHUNK / (ptype.byte_width() * 8); + let required_bases = num_chunks * lanes; + vortex_ensure!( + bases.len() >= required_bases, + "Delta needs {required_bases} bases for {num_chunks} chunks, got {}", + bases.len() + ); + + let PrimitiveDataParts { + buffer: deltas_buffer, + .. + } = deltas.into_data_parts(); + let PrimitiveDataParts { + buffer: bases_buffer, + .. + } = bases.into_data_parts(); + let deltas_device = ctx.ensure_on_device(deltas_buffer).await?; + let bases_device = ctx.ensure_on_device(bases_buffer).await?; + + let num_chunks_u64 = num_chunks as u64; + let config = LaunchConfig { + grid_dim: (u32::try_from(num_chunks)?, 1, 1), + block_dim: (BLOCK_THREADS, 1, 1), + shared_mem_bytes: 0, + }; + + let decoded: BufferHandle = match_each_unsigned_integer_ptype!(ptype.to_unsigned(), |U| { + let deltas_view = deltas_device.cuda_view::()?; + let bases_view = bases_device.cuda_view::()?; + let mut output = ctx.device_alloc::(deltas_len)?; + let function = ctx.load_function("delta", &[U::PTYPE])?; + ctx.launch_kernel_config(&function, config, deltas_len, |args| { + args.arg(&deltas_view) + .arg(&bases_view) + .arg(&mut output) + .arg(&num_chunks_u64); + })?; + BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(output))) + }); + + // Chunks are decoded whole; the logical slice is applied to the result. + let width = ptype.byte_width(); + let sliced = decoded.slice(offset * width..(offset + len) * width); + + Ok(Canonical::Primitive(PrimitiveArray::from_buffer_handle( + sliced, ptype, validity, + ))) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex::array::IntoArray; + use vortex::array::assert_arrays_eq; + use vortex::array::validity::Validity; + use vortex::buffer::Buffer; + use vortex::error::VortexExpect; + use vortex_array::VortexSessionExecute; + + use super::*; + use crate::CanonicalCudaExt; + use crate::session::CudaSession; + + /// Decodes `array` on the GPU and asserts it matches the CPU canonical form. + async fn assert_gpu_matches_cpu(delta: DeltaArray) -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let gpu = DeltaExecutor + .execute(delta.clone().into_array(), &mut cuda_ctx) + .await + .vortex_expect("GPU decompression failed") + .into_host() + .await? + .into_array(); + + assert_arrays_eq!(delta, gpu, &mut ctx); + Ok(()) + } + + /// Every element width is worth covering: lane count is `1024 / bit-width`, so each width + /// splits the 1024-element chunk differently — 128 lanes of 8 rows for `u8` through 16 + /// lanes of 64 rows for `u64` — and each has its own kernel instantiation. Every case + /// spans several chunks. + #[crate::test] + async fn test_cuda_delta_u8() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let primitive = PrimitiveArray::new( + Buffer::from_iter((0u8..=255).cycle().take(3000)), + Validity::NonNullable, + ); + + let delta = Delta::try_from_primitive_array(&primitive, &mut ctx)?; + assert_gpu_matches_cpu(delta).await + } + + #[crate::test] + async fn test_cuda_delta_u16() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let primitive = PrimitiveArray::new(Buffer::from_iter(0u16..3000), Validity::NonNullable); + + let delta = Delta::try_from_primitive_array(&primitive, &mut ctx)?; + assert_gpu_matches_cpu(delta).await + } + + #[crate::test] + async fn test_cuda_delta_u32() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let primitive = PrimitiveArray::new( + Buffer::from_iter((0u32..3000).map(|i| i * 7)), + Validity::NonNullable, + ); + + let delta = Delta::try_from_primitive_array(&primitive, &mut ctx)?; + assert_gpu_matches_cpu(delta).await + } + + #[crate::test] + async fn test_cuda_delta_u64() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let primitive = PrimitiveArray::new( + Buffer::from_iter((0u64..3000).map(|i| i * 1_000_003)), + Validity::NonNullable, + ); + + let delta = Delta::try_from_primitive_array(&primitive, &mut ctx)?; + assert_gpu_matches_cpu(delta).await + } + + /// Deltas across negative values wrap at compress time, so signed input must decode + /// through the unsigned kernel unchanged. + #[crate::test] + async fn test_cuda_delta_signed_values() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let primitive = PrimitiveArray::new( + Buffer::from_iter((0..3000i32).map(|i| i - 1500)), + Validity::NonNullable, + ); + + let delta = Delta::try_from_primitive_array(&primitive, &mut ctx)?; + assert_gpu_matches_cpu(delta).await + } + + /// A sliced Delta keeps whole chunks and carries a nonzero offset, which the decode + /// applies only after the chunks are decoded. + #[rstest] + #[case::within_first_chunk(5, 100)] + #[case::across_chunks(1000, 1500)] + #[crate::test] + async fn test_cuda_delta_sliced(#[case] start: usize, #[case] end: usize) -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let primitive = PrimitiveArray::new( + Buffer::from_iter((0..3000u32).map(|i| i * 3)), + Validity::NonNullable, + ); + + let delta = Delta::try_from_primitive_array(&primitive, &mut ctx)? + .into_array() + .slice(start..end)?; + let Ok(delta) = delta.try_downcast::() else { + // A slice that the encoding chose to canonicalise is not this test's concern. + return Ok(()); + }; + assert_gpu_matches_cpu(delta).await + } + + /// Nulls ride alongside the values, and the validity is narrowed to the logical slice. + #[crate::test] + async fn test_cuda_delta_nullable() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let primitive = PrimitiveArray::from_option_iter( + (0..3000u32).map(|value| (value % 7 != 0).then_some(value)), + ); + + let delta = Delta::try_from_primitive_array(&primitive, &mut ctx)?; + assert_gpu_matches_cpu(delta).await + } +} diff --git a/vortex-cuda/src/kernel/encodings/for_.rs b/vortex-cuda/src/kernel/encodings/for_.rs index 6348f8479b7..cc3ceed2869 100644 --- a/vortex-cuda/src/kernel/encodings/for_.rs +++ b/vortex-cuda/src/kernel/encodings/for_.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::fmt::Debug; +use std::sync::Arc; use async_trait::async_trait; use cudarc::driver::DeviceRepr; @@ -13,7 +14,8 @@ use vortex::array::IntoArray; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::Slice; use vortex::array::arrays::primitive::PrimitiveDataParts; -use vortex::array::arrays::slice::SliceArrayExt; +use vortex::array::arrays::slice::SliceArraySlotsExt; +use vortex::array::buffer::BufferHandle; use vortex::array::match_each_integer_ptype; use vortex::array::match_each_native_simd_ptype; use vortex::dtype::NativePType; @@ -21,12 +23,15 @@ use vortex::encodings::fastlanes::BitPacked; use vortex::encodings::fastlanes::FoR; use vortex::encodings::fastlanes::FoRArray; use vortex::encodings::fastlanes::FoRArrayExt; +use vortex::encodings::fastlanes::FoRArraySlotsExt; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_ensure; use vortex::error::vortex_err; use crate::CudaBufferExt; +use crate::CudaDeviceBuffer; +use crate::device_buffer::with_cuda_view_mut; use crate::executor::CudaArrayExt; use crate::executor::CudaExecute; use crate::executor::CudaExecutionCtx; @@ -102,23 +107,42 @@ where buffer, validity, .. } = primitive.into_data_parts(); - let device_buffer = ctx.ensure_on_device(buffer).await?; - - // Get CUDA view of the buffer - let cuda_view = device_buffer.cuda_view::